Kimi K3 stores routed expert matrices in grouped MXFP4 while the shared-expert output remains in higher precision. Let there be T tokens, E experts, input width I=32G, output width O, and K selected experts per token. Latent tokens have shape (T,I), selected expert indices and mixture weights have shape (T,K), and the supplied shared output has shape (T,O).
Packed weights have shape (E,O,G,16) and scale bytes have shape (E,O,G). Each group of sixteen bytes stores 32 E2M1 values. Decode the low nibble before the high nibble in every byte, preserving this interleaved order. Map nibble codes through
[0,21,1,23,2,3,4,6,−0,−21,−1,−23,−2,−3,−4,−6].For scale byte s, multiply all 32 decoded values in that group by
2s−127.Concatenate the G decoded groups to reconstruct selected expert matrix We∈RO×I. For token xt, selected set Tt, mixture weight pt,e, and shared output ytshared, compute
yt=ytshared+e∈Tt∑pt,eWext.Reconstruct only the selected expert matrices. Return the combined output as one tensor of shape (T,O).
Examples
Input: latent_tokens = [[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]], packed_weights = [[[[34,34,34,34,34,34,34,34,34,34,34,34,34,34,34,34]]]], scale_bytes = [[[127]]], selected_experts = [[0]], mixture_weights = [[1]], shared_output = [[0]]
Output: tensor([[32.0]])
Explanation: Every packed nibble decodes to 1 at unit scale, so the dot product over 32 values is 32.
Input: latent tokens shape (1, 32), packed weights shape (2, 1, 1, 16), selected experts shape (1, 2), shared output shape (1, 1)
Output: tensor of shape (1, 1)
Input: latent tokens shape (1, 64), packed weights shape (2, 2, 2, 16), selected experts shape (1, 1), shared output shape (1, 2)
Output: tensor of shape (1, 2)
Hint 1
Extract each low nibble with a bit mask and each high nibble with a right shift followed by the same mask.
Hint 2
Interleave the decoded low and high values before flattening the 32-value groups.
Hint 3
Convert each scale byte to a power of two and apply it only to its matching group.
Requirements
- Return one floating-point tensor shaped like the supplied shared output.
- Decode each weight block with its own E8M0 scale.
- Preserve the latent token dtype and device in the result.
- Do not mutate packed weights, routing data, or shared output.
Constraints
- Input width is a positive multiple of 32.
- Packed weights have sixteen bytes per 32-value group.
- Selected expert indices are valid, and mixture weights sum to one per token.
- Scale bytes and packed values use unsigned 8-bit storage.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
MXFP4 Experts
Kimi K3 stores routed expert matrices in grouped MXFP4 while the shared-expert output remains in higher precision. Let there be T tokens, E experts, input width I=32G, output width O, and K selected experts per token. Latent tokens have shape (T,I), selected expert indices and mixture weights have shape (T,K), and the supplied shared output has shape (T,O).
Packed weights have shape (E,O,G,16) and scale bytes have shape (E,O,G). Each group of sixteen bytes stores 32 E2M1 values. Decode the low nibble before the high nibble in every byte, preserving this interleaved order. Map nibble codes through
[0,21,1,23,2,3,4,6,−0,−21,−1,−23,−2,−3,−4,−6].For scale byte s, multiply all 32 decoded values in that group by
2s−127.Concatenate the G decoded groups to reconstruct selected expert matrix We∈RO×I. For token xt, selected set Tt, mixture weight pt,e, and shared output ytshared, compute
yt=ytshared+e∈Tt∑pt,eWext.Reconstruct only the selected expert matrices. Return the combined output as one tensor of shape (T,O).
Examples
Input: latent_tokens = [[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]], packed_weights = [[[[34,34,34,34,34,34,34,34,34,34,34,34,34,34,34,34]]]], scale_bytes = [[[127]]], selected_experts = [[0]], mixture_weights = [[1]], shared_output = [[0]]
Output: tensor([[32.0]])
Explanation: Every packed nibble decodes to 1 at unit scale, so the dot product over 32 values is 32.
Input: latent tokens shape (1, 32), packed weights shape (2, 1, 1, 16), selected experts shape (1, 2), shared output shape (1, 1)
Output: tensor of shape (1, 1)
Input: latent tokens shape (1, 64), packed weights shape (2, 2, 2, 16), selected experts shape (1, 1), shared output shape (1, 2)
Output: tensor of shape (1, 2)
Hint 1
Extract each low nibble with a bit mask and each high nibble with a right shift followed by the same mask.
Hint 2
Interleave the decoded low and high values before flattening the 32-value groups.
Hint 3
Convert each scale byte to a power of two and apply it only to its matching group.
Requirements
- Return one floating-point tensor shaped like the supplied shared output.
- Decode each weight block with its own E8M0 scale.
- Preserve the latent token dtype and device in the result.
- Do not mutate packed weights, routing data, or shared output.
Constraints
- Input width is a positive multiple of 32.
- Packed weights have sixteen bytes per 32-value group.
- Selected expert indices are valid, and mixture weights sum to one per token.
- Scale bytes and packed values use unsigned 8-bit storage.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array