For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
Struct_router_gate_mixed_gemv
struct Struct_router_gate_mixed_gemv
MOGG wrapper for the mixed-input router-gate GEMV.
Computes c = a @ b^T for a bf16 activation a and an fp32 router weight
b, returning fp32 c. M is runtime-dynamic (one graph serves both
decode and prefill), so the op branches on runtime M:
M == 0(graph-capture warmup): no launch.- tiny
M(M <= 16, decode): the fused mixed GEMV —ais loaded as bf16 and widened to fp32 in registers, then dotted against the unchanged fp32bin a SINGLE launch, fusing away the standalone bf16->fp32 activation cast that otherwise precedes the fp32 router GEMV. - large
M(prefill): the fused GEMV is catastrophically slow, so the op falls back to the baseline chain — castato fp32, then an ordinary fp32 matmul against the fp32b(two launches, matching baseline).
Because bf16->fp32 widening is lossless, every route is numerically
equivalent to casting a to fp32 first and running the fp32 GEMM.
This is not part of generic matmul/GEMV dispatch. The MiniMax-M3 router emits this explicit graph op on MI355X; other architectures retain the standard router path.
N (expert count) and K (hidden size) are read from the weight b's
static [N, K] layout at compile time — the op carries no redundant N/K
custom-op parameters. N still selects the kernel's check_bounds_n guard.
Implemented traits
Methods
execute
static def execute[target: StringSlice[ImmStaticOrigin]](c: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=c.static_spec], a: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a.static_spec], b: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b.static_spec], context: DeviceContext)
Executes the mixed bf16-A × fp32-B router-gate GEMV with an M-based route to the baseline cast + fp32 matmul for large (prefill) M.
Constraints:
The weight b must have a static [N, K] shape (both dims known at
compile time).
Args:
- c (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=c.static_spec]): Output[M, N]fp32 tensor. - a (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a.static_spec]): Activation[M, K]bf16 tensor. - b (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b.static_spec]): Weight[N, K]fp32 tensor (transpose_b layout). - context (
DeviceContext): The device context.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!