For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo struct
MTPEhNorm
struct MTPEhNorm
Registers the mo.mtp.eh_norm graph op with the graph compiler.
Implemented traits
Methods
execute
static def execute[dtype: DType, target: StringSpan[ImmStaticOrigin], //, hidden_size: Int, block_threads: Int](out_buf: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_buf.static_spec], embed: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=embed.static_spec], prev: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=prev.static_spec], enorm_weight: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=enorm_weight.static_spec], hnorm_weight: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=hnorm_weight.static_spec], epsilon: Float32, ctx: DeviceContext)
Normalizes both draft inputs and writes them side by side.
Parameters:
- dtype (
DType): Element type of the inputs, weights and output. - target (
StringSpan[ImmStaticOrigin]): Compilation target. - hidden_size (
Int): Channels per input; the output row is twice this. - block_threads (
Int): Threads per block.
Args:
- out_buf (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_buf.static_spec]): Output[num_tokens, 2 * hidden_size]. - embed (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=embed.static_spec]): Token embeddings[num_tokens, hidden_size]. - prev (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=prev.static_spec]): Target hidden states[num_tokens, hidden_size]. - enorm_weight (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=enorm_weight.static_spec]): Embedding norm weight[hidden_size]. - hnorm_weight (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=hnorm_weight.static_spec]): Hidden-state norm weight[hidden_size]. - epsilon (
Float32): Added inside the square root. - ctx (
DeviceContext): Device context.
Raises:
Error: If the target is not a GPU.