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_sm100_persistent

def grouped_matmul_sm100_persistent[c_type: DType, a_type: DType, b_type: DType, transpose_b: Bool, *, config: MatmulConfig[a_type, b_type, c_type, transpose_b], cta_group: Int = Int(1), num_pipeline_stages: Optional[Int] = None, a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, a_plane_splits: IndexList[Int(2)] = Index[Int, Int](Int(0), Int(0))](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], a_offsets: TileTensor[DType.uint32, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type], b: TileTensor[b_type, Storage=b.Storage, linear_idx_type=b.linear_idx_type], expert_ids: TileTensor[DType.int32, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type], expert_usage_stats: TileTensor[DType.uint32, Storage=expert_usage_stats.Storage, linear_idx_type=expert_usage_stats.linear_idx_type], ctx: DeviceContext)

Launches the persistent grouped GEMM kernel for SM100 from host tensors.

Swaps A and B to match the kernel's transposed-B convention, delegates to _grouped_matmul_sm100_persistent which builds TMA descriptors and shared-memory layouts from the matmul config, selects the pipeline depth from available shared memory, and enqueues the warp-specialized kernel on the device context.

Parameters:

c_type (DType): Element type of the C output matrix (inferred).
a_type (DType): Element type of the A operand matrix (inferred).
b_type (DType): Element type of the B operand matrix (inferred).
transpose_b (Bool): Whether B is stored transposed in global memory (inferred).
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Matmul config carrying the block tile shape, MMA shape, and cluster shape.
cta_group (Int): Number of CTAs cooperating per MMA along the M dimension (defaults to 1).
num_pipeline_stages (Optional[Int]): Number of TMA producer and consumer pipeline stages; when None, auto-selected from available shared memory (defaults to None).
a_swizzle (TensorMapSwizzle): TMA swizzle mode applied to A shared-memory tiles (defaults to SWIZZLE_128B).
b_swizzle (TensorMapSwizzle): TMA swizzle mode applied to B shared-memory tiles (defaults to SWIZZLE_128B).
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional elementwise epilogue applied to stored C fragments (defaults to None).
a_plane_splits (IndexList[Int(2)]): Per-plane split sizes for fused LoRA QKV A-plane row offsetting; (0, 0) disables it (defaults to (0, 0)).

Args:

c (TileTensor[c_type, Storage=c.Storage, linear_idx_type=c.linear_idx_type]): Output C tile tensor in generic address space.
a (TileTensor[a_type, Storage=a.Storage, linear_idx_type=a.linear_idx_type]): A operand tile tensor (activations); swapped to B by the kernel's transposed-B convention.
a_offsets (TileTensor[DType.uint32, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type]): Per-expert A offset prefix sum tile tensor.
b (TileTensor[b_type, Storage=b.Storage, linear_idx_type=b.linear_idx_type]): B operand tile tensor (weights); swapped to A by the kernel's transposed-B convention.
expert_ids (TileTensor[DType.int32, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type]): Per-group expert id tile tensor; negative entries skip the tile.
expert_usage_stats (TileTensor[DType.uint32, Storage=expert_usage_stats.Storage, linear_idx_type=expert_usage_stats.linear_idx_type]): Per-expert usage stats tile tensor; index 1 holds the active expert count.
ctx (DeviceContext): Device context used to enqueue the kernel.