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

grouped_matmul_rowwise_dynamic_scaled_fp8

def grouped_matmul_rowwise_dynamic_scaled_fp8[c_type: DType, a_type: DType, b_type: DType, a_scales_type: DType, b_scales_type: DType, a_offsets_type: DType, expert_ids_type: DType, //, transpose_b: Bool = True, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu"), elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[c_type, Storage=c.Storage, linear_idx_type=c.linear_idx_type], a: TileTensor[a_type, Storage=a.Storage, linear_idx_type=a.linear_idx_type], b: TileTensor[b_type, Storage=b.Storage, linear_idx_type=b.linear_idx_type], a_scales: TileTensor[a_scales_type, Storage=a_scales.Storage, linear_idx_type=a_scales.linear_idx_type], b_scales: TileTensor[b_scales_type, Storage=b_scales.Storage, linear_idx_type=b_scales.linear_idx_type], a_offsets: TileTensor[a_offsets_type, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type], expert_ids: TileTensor[expert_ids_type, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type], max_num_tokens_per_expert: Int, num_active_experts: Int, ctx: DeviceContext)

Grouped (ragged MoE) FP8 matmul with rowwise weight + per-token act scales.

Target: NVIDIA SM100 (B200). Correctness-first naive grouped kernel; no persistent / TMA path. Computes, for each token t in group g's row range and each output channel n::

out[t, n] = (sum_k a[t, k] * b[expert_ids[g], n, k])
            * a_scale[t] * b_scale[expert_ids[g], n]

accumulated in fp32 with a single post-reduction scale.

Parameters:

c_type (DType): Output dtype (typically bfloat16).
a_type (DType): Activation dtype (float8_e4m3fn).
b_type (DType): Weight dtype (float8_e4m3fn).
a_scales_type (DType): Per-token activation scale dtype (float32).
b_scales_type (DType): Per-channel weight scale dtype (float32).
a_offsets_type (DType): Ragged-offset dtype (uint32).
expert_ids_type (DType): Expert-id dtype (int32).
transpose_b (Bool): Must be True; b is [E, N, K].
target (StringSlice[ImmStaticOrigin]): Compilation target string.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional output epilogue applied with the (global_row, n) index.

Args:

c (TileTensor[c_type, Storage=c.Storage, linear_idx_type=c.linear_idx_type]): Output [total_tokens, N].
a (TileTensor[a_type, Storage=a.Storage, linear_idx_type=a.linear_idx_type]): Activations [total_tokens, K].
b (TileTensor[b_type, Storage=b.Storage, linear_idx_type=b.linear_idx_type]): Weights [num_experts, N, K] (already transposed; K innermost).
a_scales (TileTensor[a_scales_type, Storage=a_scales.Storage, linear_idx_type=a_scales.linear_idx_type]): Per-token activation scales [total_tokens, 1].
b_scales (TileTensor[b_scales_type, Storage=b_scales.Storage, linear_idx_type=b_scales.linear_idx_type]): Per-channel weight scales [num_experts, N, 1].
a_offsets (TileTensor[a_offsets_type, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type]): Ragged row offsets [num_active_experts + 1].
expert_ids (TileTensor[expert_ids_type, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type]): Real expert ids [num_active_experts].
max_num_tokens_per_expert (Int): Max tokens routed to any active expert.
num_active_experts (Int): Number of active experts (groups).
ctx (DeviceContext): Device context.