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 function

mtp_eh_norm_kernel

def mtp_eh_norm_kernel[dtype: DType, OutLayoutType: TensorLayout, out_origin: MutOrigin, EmbedLayoutType: TensorLayout, embed_origin: ImmOrigin, PrevLayoutType: TensorLayout, prev_origin: ImmOrigin, EWLayoutType: TensorLayout, ew_origin: ImmOrigin, HWLayoutType: TensorLayout, hw_origin: ImmOrigin, hidden_size: Int, max_warps_per_block: Int, OutEngine: TensorEngine, EmbedEngine: TensorEngine, PrevEngine: TensorEngine, EWEngine: TensorEngine, HWEngine: TensorEngine](out_buf: TileTensor[dtype, OutLayoutType, out_origin, Engine=OutEngine], embed: TileTensor[dtype, EmbedLayoutType, embed_origin, Engine=EmbedEngine], prev: TileTensor[dtype, PrevLayoutType, prev_origin, Engine=PrevEngine], enorm_weight: TileTensor[dtype, EWLayoutType, ew_origin, Engine=EWEngine], hnorm_weight: TileTensor[dtype, HWLayoutType, hw_origin, Engine=HWEngine], epsilon: Float32, num_tokens: Int32)

Normalizes both inputs for one token and writes them side by side.

Parameters:

• ​dtype (DType): Element type of the inputs, the weights and the output.
• ​OutLayoutType (TensorLayout): Layout of out_buf.
• ​out_origin (MutOrigin): Origin of out_buf.
• ​EmbedLayoutType (TensorLayout): Layout of embed.
• ​embed_origin (ImmOrigin): Origin of embed.
• ​PrevLayoutType (TensorLayout): Layout of prev.
• ​prev_origin (ImmOrigin): Origin of prev.
• ​EWLayoutType (TensorLayout): Layout of enorm_weight.
• ​ew_origin (ImmOrigin): Origin of enorm_weight.
• ​HWLayoutType (TensorLayout): Layout of hnorm_weight.
• ​hw_origin (ImmOrigin): Origin of hnorm_weight.
• ​hidden_size (Int): Channels per input; the output row is twice this.
• ​max_warps_per_block (Int): Power-of-two ceiling on the block's warp count, not the launch's actual warp count. block_reduce_dual_sum reduces over a power-of-two lane group; slots past the real warp count hold zero.
• ​OutEngine (TensorEngine): Engine policy of out_buf.
• ​EmbedEngine (TensorEngine): Engine policy of embed.
• ​PrevEngine (TensorEngine): Engine policy of prev.
• ​EWEngine (TensorEngine): Engine policy of enorm_weight.
• ​HWEngine (TensorEngine): Engine policy of hnorm_weight.

Args:

• ​out_buf (TileTensor[dtype, OutLayoutType, out_origin, Engine=OutEngine]): Output [num_tokens, 2 * hidden_size]. Columns [0, hidden_size) hold the normalized embedding, the rest the normalized hidden state.
• ​embed (TileTensor[dtype, EmbedLayoutType, embed_origin, Engine=EmbedEngine]): Token embeddings [num_tokens, hidden_size].
• ​prev (TileTensor[dtype, PrevLayoutType, prev_origin, Engine=PrevEngine]): Target hidden states [num_tokens, hidden_size].
• ​enorm_weight (TileTensor[dtype, EWLayoutType, ew_origin, Engine=EWEngine]): Embedding norm weight [hidden_size].
• ​hnorm_weight (TileTensor[dtype, HWLayoutType, hw_origin, Engine=HWEngine]): Hidden-state norm weight [hidden_size].
• ​epsilon (Float32): Added inside the square root, matching ops.rms_norm.
• ​num_tokens (Int32): Rows to process.