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

MXFP4MatmulAMD_PreB

struct MXFP4MatmulAMD_PreB[BM: Int = Int(64), BN: Int = Int(128), BK_ELEMS: Int = Int(512), WN: Int = Int(64), b_prefetch: Bool = False, 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)]

Preshuffled-B variant of MXFP4MatmulAMD.

The preb path requires num_warps_m == 1 (no LDS staging for B = no cross-warp M-direction B reuse), so WM is structurally fixed to BM.

When b_prefetch=True, runs a depth-2 outer-K software pipeline: while the current iter's MFMAs execute, the next iter's B fragments stream from DRAM into the alternate b_reg slot. Doubles _b_reg size (extra VGPRs) but hides DRAM B latency across the inner MFMA chain. Targets K-heavy shapes (e.g. gate/up, K=7168) where outer-iter serialization dominates.

cluster_drain_sched (b_prefetch only) switches the 1-deep steady loop to an interleaved B-issue schedule: the next tile's B fragments are issued per-k between the current tile's MFMA phases (not front-loaded), each phase pinned by sched_barrier(0) + bracketed by s_setprio, and the end-of-tile sync is a bare s_barrier + lgkmcnt-only drain so in-flight B DMAs cross it. (deep_prime / the epilogue still use the per-cluster vmcnt staircase, mma_chain_scheduled.) Default off: callers bit-identical unless opted in.

deep_prime (b_prefetch only, num_tiles >= 2) deepens the A pipeline to 2-tiles-ahead: the prologue stages BOTH tile0 -> slot0 and tile1 -> slot1 into LDS so each steady iter reads an A tile that has had a full extra iteration of MFMA shadow to land. Iter i reads slot[i%2] and issues the A DMA for tile i+2 into that same (just-freed) slot. Reuses the existing num_a_slots=2 LDS buffers: no extra LDS/VGPR. Composes with cluster_drain_sched/mfma_cluster (the MFMA chain is unchanged). Falls back to the 1-deep path when num_tiles < 2. Default off: existing callers are bit-identical.

pipeline_depth sizes the B-fragment register ring (num_b_slots) and, when > 2, switches the b_prefetch steady loop's end-of-iter sync to the non-draining s_waitcnt[lgkmcnt=0] + bare s_barrier so in-flight B DMAs are NOT drained every iteration. At the default (2) the ring is the same 2 slots and the draining barrier() is kept — bit-identical to before. The deeper prefetch schedule that consumes slots >= 2 is a follow-up; this param only plumbs the depth + the non-draining-barrier seam.

MFMA consumption order is B-major (n-outer / m-inner): the B fragment is held resident across the m-loop for better MFMA ILP. See mma.

Parameters​

  • ​BM (Int): CTA tile size along M in elements, either 16 or a multiple of 32. WM is locked to BM (single warp along M).
  • ​BN (Int): CTA tile size along N in elements, split across num_warps_n = BN // WN warps.
  • ​BK_ELEMS (Int): K tile size in MXFP4 elements per outer-K iteration; must be a multiple of 256 so num_k_mmas is even. BK_BYTES = BK_ELEMS // 2.
  • ​WN (Int): Per-warp tile size along N in elements, either 16 or a multiple of 32.
  • ​b_prefetch (Bool): Enables a depth-2 outer-K software pipeline that double-buffers B fragments across iterations (defaults to False).
  • ​b_cache_policy (CacheOperation): CacheOperation hint applied to preshuffled B DRAM loads (defaults to CacheOperation.ALWAYS).
  • ​dram_to_lds (Bool): Routes A loads through the shared swizzled TileLoaderLDS DRAM-to-LDS path instead of a register bounce (defaults to False).
  • ​cluster_drain_sched (Bool): Switches the prefetch steady loop to an interleaved B-issue schedule with per-cluster s_setprio and a partial-vmcnt staircase (defaults to False).
  • ​mfma_cluster (Int): Number of MFMAs per cluster in the scheduled MFMA chain used by cluster_drain_sched (defaults to 4).
  • ​deep_prime (Bool): Deepens the A pipeline to 2-tiles-ahead on the prefetch path, staging two A tiles in the prologue (defaults to False).
  • ​pipeline_depth (Int): Depth of the B-fragment register ring (num_b_slots) on the prefetch path (defaults to 2). When > 2 it also switches the steady loop's end-of-iter sync to the non-draining s_waitcnt[lgkmcnt=0] + bare s_barrier and co-deepens the A/scale rings. At the default (2) the behavior is bit-identical to before.

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

BK_BYTES​

comptime BK_BYTES = (BK_ELEMS // Int(2))

c_frag_size​

comptime c_frag_size = MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].MmaOpType.c_frag_size

MMA_K​

comptime MMA_K = 128

MMA_K_BYTES​

comptime MMA_K_BYTES = MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].MmaOpType.MMA_K_BYTES

MMA_M​

comptime MMA_M = 16

MMA_N​

comptime MMA_N = 16

MmaOpType​

comptime MmaOpType = BlockScaledMmaOp_PreB[IndexList(Int(16), Int(16), Int(128), __list_literal__=NoneType(None)), IndexList(MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].WM, WN, BK_ELEMS, __list_literal__=NoneType(None)), MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].num_b_slots, MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].num_scale_slots]

num_a_load_slots​

comptime num_a_load_slots = pipeline_depth if (pipeline_depth > Int(2)) else Int(1)

num_a_slots​

comptime num_a_slots = Int(2) if b_prefetch else Int(1)

num_b_slots​

comptime num_b_slots = pipeline_depth if b_prefetch else Int(1)

num_k_mmas​

comptime num_k_mmas = MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].MmaOpType.num_k_mmas

num_m_mmas​

comptime num_m_mmas = MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].MmaOpType.num_m_mmas

num_n_mmas​

comptime num_n_mmas = MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].MmaOpType.num_n_mmas

num_scale_slots​

comptime num_scale_slots = pipeline_depth if (pipeline_depth > Int(2)) else Int(1)

num_threads​

comptime num_threads = (MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].num_warps * _resolve_warp_size())

num_warps​

comptime num_warps = MXFP4MatmulAMD_PreB[BM, BN, BK_ELEMS, WN, b_prefetch, b_cache_policy, dram_to_lds, cluster_drain_sched, mfma_cluster, deep_prime, pipeline_depth].num_warps_n

num_warps_m​

comptime num_warps_m = 1

num_warps_n​

comptime num_warps_n = (BN // WN)

simd_width​

comptime simd_width = simd_width_of[DType.uint8]()

WM​

comptime WM = BM

Methods​

run​

static def run[out_dtype: DType, c_layout: TensorLayout, a_layout: TensorLayout, b_pre_layout: TensorLayout, sfa_layout: TensorLayout, sfb_layout: TensorLayout, N: Int, K_BYTES: Int](c: TileTensor[out_dtype, c_layout, MutAnyOrigin], a: TileTensor[DType.uint8, a_layout, ImmutAnyOrigin], b_pre: TileTensor[DType.uint8, b_pre_layout, ImmutAnyOrigin], sfa: TileTensor[DType.float8_e8m0fnu, sfa_layout, ImmutAnyOrigin], sfb: TileTensor[DType.float8_e8m0fnu, sfb_layout, ImmutAnyOrigin], n_tile_idx: Int, m_tile_idx: Int)

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact