IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

PreShuffledBGroupedGEMM

struct PreShuffledBGroupedGEMM[cu_count: Int, wg_per_cu: Int = Int(2)]

Grouped GEMM for MXFP4 on AMD CDNA4 with pre-shuffled weights.

This grouped GEMM operates on weights B that have been pre-shuffled into a layout enabling coalesced reads from shared memory and direct MFMA usage. It offers a persistent kernel (grid-stride over work tiles with XCD-aware work-group swizzling) and a direct kernel (one block per output tile, expert dispatched via block_idx.z), selected at launch time by the persistent comptime flag.

Parameters​

  • ​cu_count (Int): Number of compute units on the target device.
  • ​wg_per_cu (Int): Work groups per compute unit (default 2).

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

num_xcd​

comptime num_xcd = 8

total_wg​

comptime total_wg = (cu_count * wg_per_cu)

wg_per_xcd​

comptime wg_per_xcd = (PreShuffledBGroupedGEMM[cu_count, wg_per_cu].total_wg // Int(8))

Methods​

to_swizzled_idx​

static def to_swizzled_idx(linear_idx: Int) -> Int

Returns:

Int

persistent_kernel​

static def persistent_kernel[BM: Int, BN: Int, BK_ELEMS: Int, WN: Int, out_dtype: DType, LayoutC: TensorLayout, LayoutA: TensorLayout, LayoutBPre: TensorLayout, LayoutSFA: TensorLayout, LayoutSFB: TensorLayout, AOffsetsLayout: TensorLayout, ExpertIdsLayout: TensorLayout, N: Int, K_BYTES: Int, b_cache_policy: CacheOperation = CacheOperation.ALWAYS, dram_to_lds: Bool = False, cluster_drain_sched: Bool = False, mfma_cluster: Int = Int(4), deep_prime: Bool = False, pipeline_depth: Int = Int(2), waves_per_eu: Int = Int(0)](c_tensor: TileTensor[out_dtype, LayoutC, MutAnyOrigin], a_tensor: TileTensor[DType.uint8, LayoutA, ImmutAnyOrigin], b_pre_tensor: TileTensor[DType.uint8, LayoutBPre, ImmutAnyOrigin], sfa_tensor: TileTensor[DType.float8_e8m0fnu, LayoutSFA, ImmutAnyOrigin], sfb_tensor: TileTensor[DType.float8_e8m0fnu, LayoutSFB, ImmutAnyOrigin], a_offsets: TileTensor[DType.uint32, AOffsetsLayout, ImmutAnyOrigin], expert_ids: TileTensor[DType.int32, ExpertIdsLayout, ImmutAnyOrigin], num_active_experts: Int, max_padded_M: Int)

kernel​

static def kernel[BM: Int, BN: Int, BK_ELEMS: Int, WN: Int, out_dtype: DType, LayoutC: TensorLayout, LayoutA: TensorLayout, LayoutBPre: TensorLayout, LayoutSFA: TensorLayout, LayoutSFB: TensorLayout, AOffsetsLayout: TensorLayout, ExpertIdsLayout: TensorLayout, N: Int, K_BYTES: Int, b_cache_policy: CacheOperation = CacheOperation.ALWAYS, dram_to_lds: Bool = False, cluster_drain_sched: Bool = False, mfma_cluster: Int = Int(4), deep_prime: Bool = False, pipeline_depth: Int = Int(2), waves_per_eu: Int = Int(0)](c_tensor: TileTensor[out_dtype, LayoutC, MutAnyOrigin], a_tensor: TileTensor[DType.uint8, LayoutA, ImmutAnyOrigin], b_pre_tensor: TileTensor[DType.uint8, LayoutBPre, ImmutAnyOrigin], sfa_tensor: TileTensor[DType.float8_e8m0fnu, LayoutSFA, ImmutAnyOrigin], sfb_tensor: TileTensor[DType.float8_e8m0fnu, LayoutSFB, ImmutAnyOrigin], a_offsets: TileTensor[DType.uint32, AOffsetsLayout, ImmutAnyOrigin], expert_ids: TileTensor[DType.int32, ExpertIdsLayout, ImmutAnyOrigin], num_active_experts: Int, max_padded_M: Int)

launch​

static def launch[BM: Int, BN: Int, BK_ELEMS: Int, WN: Int, persistent: Bool, b_cache_policy: CacheOperation = CacheOperation.ALWAYS, dram_to_lds: Bool = False, cluster_drain_sched: Bool = False, mfma_cluster: Int = Int(4), deep_prime: Bool = False, pipeline_depth: Int = Int(2), waves_per_eu: Int = Int(0), static_grid_z: Bool = False](c: TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[DType.uint8, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b_pre: TileTensor[DType.uint8, Storage=b_pre.Storage, address_space=b_pre.address_space, linear_idx_type=b_pre.linear_idx_type], a_scales: TileTensor[DType.float8_e8m0fnu, Storage=a_scales.Storage, address_space=a_scales.address_space, linear_idx_type=a_scales.linear_idx_type], b_scales: TileTensor[DType.float8_e8m0fnu, Storage=b_scales.Storage, address_space=b_scales.address_space, linear_idx_type=b_scales.linear_idx_type], a_offsets: TileTensor[DType.uint32, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type], expert_ids: TileTensor[DType.int32, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type], max_num_tokens_per_expert: Int, num_active_experts: Int, ctx: DeviceContext, grid_m_cap: Int = Int(-1))

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact