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 function

router_gate_mixed_gemv

def router_gate_mixed_gemv[static_N: Int, c_layout: TensorLayout, a_layout: TensorLayout, b_layout: TensorLayout, c_storage: TensorStorage, a_storage: TensorStorage, b_storage: TensorStorage](c: TileTensor[DType.float32, c_layout, MutAnyOrigin, Storage=c_storage], a: TileTensor[DType.bfloat16, a_layout, ImmutAnyOrigin, Storage=a_storage], b: TileTensor[DType.float32, b_layout, ImmutAnyOrigin, Storage=b_storage], m: Int, n: Int, k: Int, ctx: DeviceContext)

Launches the mixed bf16-activation × fp32-weight router-gate GEMV.

Fuses the standalone bf16→fp32 activation cast into the router GEMV: a is loaded as bf16 and widened to fp32 in registers, then dotted against the unchanged fp32 weight b (c = a @ b^T) in a single launch. MiniMax-M3 uses this path on MI355X; other architectures retain the standard router.

Because bf16→fp32 widening is lossless and the reduction structure matches gemv_split_k, the result is numerically identical to casting a to fp32 first and running the fp32 GEMV.

The launch config is the MI355X (gfx950 / CDNA4) cache-busting sweep winner for the N=128, K=6144, M<=16 router-gate shape: simd_width=4, tile_m=1, tile_n=2, 128 threads, unroll=2, with the weight kept cache-resident (weight_non_temporal=False) so the per-row-block weight rereads hit L2 — the same cache policy the merged KERN-3219 fp32 router path selects.

Parameters:

static_N (Int): Static output width (weight rows / expert count). Selects the check_bounds_n guard at compile time.
c_layout (TensorLayout): Layout of the fp32 output tensor.
a_layout (TensorLayout): Layout of the bf16 activation tensor.
b_layout (TensorLayout): Layout of the fp32 weight tensor.
c_storage (TensorStorage): Storage of the fp32 output tensor.
a_storage (TensorStorage): Storage of the bf16 activation tensor.
b_storage (TensorStorage): Storage of the fp32 weight tensor.

Args:

c (TileTensor[DType.float32, c_layout, MutAnyOrigin, Storage=c_storage]): Output [M, N] fp32 tensor.
a (TileTensor[DType.bfloat16, a_layout, ImmutAnyOrigin, Storage=a_storage]): Activation [M, K] bf16 tensor.
b (TileTensor[DType.float32, b_layout, ImmutAnyOrigin, Storage=b_storage]): Weight [N, K] fp32 tensor (transpose_b layout).
m (Int): Runtime row count (M). Optimal for tiny M (router: M<=16).
n (Int): Output width (N).
k (Int): Contraction dim (K).
ctx (DeviceContext): The device context.