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 module

mtp_eh_norm

Input projection for a multi-token-prediction draft layer.

An MTP draft layer predicts token t+2 from the embedding of the token the target just produced and the target's hidden state at t:

eh_proj( concat( enorm(embedding), hnorm(hidden_state) ) )

This kernel does everything before the projection. Both halves reduce over the same hidden_size, so one block computes both row sums and writes the whole [tokens, 2 * hidden_size] row.

Normalization matches ops.rms_norm in its Llama-style configuration, which is what every MTP draft in the tree uses:

out = cast(x * rsqrt(mean(x^2) + eps)) * weight

The cast comes before the multiply (multiply_before_cast=False). Swapping that order changes the low bits, so the kernel would no longer match ops.rms_norm.

Every tensor must be contiguous and row-major with a last dimension of exactly hidden_size. The kernel addresses elements by linear offset through raw_load and raw_store, which bypass the layout, so a transposed or strided view would read the wrong elements.

Callers own position handling. MAX shifts a draft's token stream left per request and appends a bonus token (eagle_prefill_shift_tokens), so every row arriving here already holds the embedding it needs. Masking the first position here as well would blank a correct row.

Functions

• ​mtp_eh_norm_kernel: Normalizes both inputs for one token and writes them side by side.