flashinfer.gemm.mm_M1_16_K7168_N384_bf16

flashinfer.gemm.mm_M1_16_K7168_N384_bf16(mat_a: Tensor, mat_b: Tensor, out: Tensor, launch_with_pdl: bool = True) → None

Optimized GEMM for the router operation in Kimi-K2.

Performs a highly optimized matrix multiplication specifically tailored for the expert routing GEMM in the Kimi-K2 family’s (K2/K2.5/K2.6) Mixture-of-Experts (MoE) architecture. Computes out = mat_a @ mat_b where mat_a is a small batch of token embeddings (1-16 rows) and mat_b is the expert routing weight matrix. Specialized for the dimensions used in Kimi-K2 MoE (K = 7168, N = 384).

Produces bfloat16 router logits (as opposed to the float32 produced by mm_M1_16_K7168_N384()).

Parameters:

  • mat_a (torch.Tensor) – Input token embeddings of shape (M, K) where M is the number of tokens (1-16) and K is the hidden dimension (7168). Must be bfloat16, row-major (contiguous).

  • mat_b (torch.Tensor) – Expert routing weights of shape (K, N) where N is the number of experts (384). Must be bfloat16, column-major (transposed layout).

  • out (torch.Tensor) – Pre-allocated output tensor of shape (M, N) containing the routing scores. Must be bfloat16, row-major (contiguous). Mutated in place.

  • launch_with_pdl (bool) – Whether to launch the kernel using Programmatic Dependent Launch. Defaults to True.

Notes

Requires SM90 (Hopper) or newer. The specialized problem-size optimization makes this significantly faster than general-purpose GEMM implementations for the router op, for small token counts; past roughly 4-8 tokens (shape dependent) a general-purpose GEMM is faster. Raises ValueError if tensor dimensions, strides, or dtypes do not match the expected Kimi-K2 router configuration.