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

FA4Config

struct FA4Config[qkv_dtype: DType, *, rope_dtype_: Optional[DType] = None, scale_dtype_: Optional[DType] = None]

Fields​

  • ​MMA_M (Int):
  • ​BM (Int):
  • ​BN (Int):
  • ​BK0 (Int):
  • ​BK1 (Int):
  • ​qk_depth (Int):
  • ​padded_qk_depth (Int):
  • ​ov_depth (Int):
  • ​padded_ov_depth (Int):
  • ​nope_depth (Int):
  • ​padded_nope_depth (Int):
  • ​group (Int):
  • ​num_q_heads (Int):
  • ​num_kv_heads (Int):
  • ​TMEM_S1 (Int):
  • ​TMEM_O0 (Int):
  • ​TMEM_O1 (Int):
  • ​TMEM_P0 (Int):
  • ​TMEM_P1 (Int):
  • ​tmem_used (Int):
  • ​num_kv_stages (Int):
  • ​num_qk_stages (Int):
  • ​num_pv_stages (Int):
  • ​smem_used (Int):
  • ​fuse_gqa (Bool):
  • ​swizzle_mode (TensorMapSwizzle):
  • ​use_shared_kv (Bool):
  • ​pair_cta (Bool):
  • ​num_q (Int):
  • ​single_o (Bool):
  • ​page_size (Int):
  • ​is_mla (Bool):
  • ​splitk_partitions (Int):
  • ​dynamic_cluster_dim (Bool):
  • ​row_major_v_atoms (Bool):
  • ​row_major_k_atoms (Bool):
  • ​use_ws (Bool):
  • ​m_pack (Int):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

mbar_size​

comptime mbar_size = size_of[DType.int64]()

MMA_K​

comptime MMA_K = Int(16) if qkv_dtype.is_half_float() else Int(32)

num_correction_cols​

comptime num_correction_cols = 1

num_threads​

comptime num_threads = Int(512)

qkv_dtype_size​

comptime qkv_dtype_size = size_of[qkv_dtype]()

rope_dtype​

comptime rope_dtype = rope_dtype_.or_else(qkv_dtype)

rope_dtype_size​

comptime rope_dtype_size = size_of[(load_from_mem rope_dtype_.value())]() if rope_dtype_.__bool__() else Int(0)

scale_dtype​

comptime scale_dtype = scale_dtype_.or_else(qkv_dtype)

scale_dtype_size​

comptime scale_dtype_size = size_of[(load_from_mem scale_dtype_.value())]() if scale_dtype_.__bool__() else Int(0)

sm100_smem_carveout​

comptime sm100_smem_carveout = (GPUInfo.from_family(AcceleratorArchitectureFamily(Int(32), Int(2048), Int(233472), Int(65536), Int(1024)), StringSlice("B200"), Vendor(Int8(2)), StringSlice("cuda"), StringSlice("blackwell"), SIMD(10), StringSlice("sm_100a"), Int(148)) - Int(1024))

sm100_tmem_cols​

comptime sm100_tmem_cols = 512

TMEM_S0​

comptime TMEM_S0 = Int(0)

Methods​

__init__​

def __init__(*, num_q_heads: Int, group: Int, qk_depth: Int, ov_depth: Int, swizzle_mode: TensorMapSwizzle, page_size: Int, is_mla: Bool, pair_cta: Bool = False, BM: Int = Int(256), num_qk_stages: Int = Int(0), splitk_partitions: Int = Int(1), dynamic_cluster_dim: Bool = False, nope_depth: Int = Int(-1), single_o: Bool = False, bn_cap: Int = Int(0)) -> Self

BM_eff​

def BM_eff(self) -> Int

Number of distinct sequence positions per full tile. When fuse_gqa, each tile covers BM // group seq positions x group heads.

Returns:

Int

cta_group​

def cta_group(self) -> Int

Returns:

Int

cluster_size​

def cluster_size(self) -> Int

CTAs per launch cluster.

Unifies the two cluster uses: cta_group (pair-CTA shares one MMA across 2 CTAs) and splitk_partitions (num_q==1 split-K groups CTAs that independently attend over K/V partitions and combine via DSMEM). These are mutually exclusive today (split-K is single-CTA only), so the product is cta_group for pair-CTA and splitk_partitions for split-K. Drives the static nvvm.cluster_dim metadata and the launch cluster_dim.

Returns:

Int

PairBM_eff​

def PairBM_eff(self) -> Int

Sequence positions covered by both CTAs in a pair.

Returns:

Int

v_cols_per_cta​

def v_cols_per_cta(self) -> Int

V columns stored in this CTA's SMEM.

Returns:

Int

v_box_cols​

def v_box_cols(self) -> Int

V TMA box depth (columns) per issued V load.

WS shared sub-tile ring: V is depth-split into num_qk_stages 256x64 sub-tiles, so each V TMA loads v_cols_per_cta() // num_qk_stages columns (mirrors K's BK0). Non-WS loads V whole (v_cols_per_cta()). MUST be used identically at EVERY V TMA type-param site (fa4_load signature, dispatch create_tma_tile, kernel launch param); a bare inline ... if use_ws else ... does NOT fold to v_cols_per_cta() at parse time, so the type-param expressions would mismatch across sites. Routing all sites through this single method keeps them syntactically identical.

Returns:

Int

v_e_box_rows​

def v_e_box_rows(self) -> Int

Layout-E (m_pack == 2) V TMA box KEY-row count per issued sub-tile.

Layout-E KEY-splits V into m_pack * num_qk_stages per-partition reduction-chunk sub-tiles instead of Layout-G's DEPTH-split (v_box_cols()): each sub-tile covers (BN // m_pack) // num_qk_stages keys (this partition's OWN reduction chunk) and the FULL depth (v_e_box_cols()). Sibling of v_box_cols(), kept as a SEPARATE method (not folded into v_box_cols()) so Layout-G's v_row_major() / kv_tma_fold_chunks derivation -- which reads v_box_cols() -- stays byte-identical; see docs/plans/sm100-fa4-layout-e-mma64.md "the reduction-split geometry". Meaningful only for m_pack == 2; a harmless (unused) value otherwise.

Returns:

Int

v_e_box_cols​

def v_e_box_cols(self) -> Int

Layout-E (m_pack == 2) V TMA box depth (columns) per issued sub-tile: the FULL padded_ov_depth, since Layout-E splits V by KEY (reduction), not by depth. Sibling of v_e_box_rows().

Returns:

Int

v_tma_box_rows​

def v_tma_box_rows(self, page_size: Int) -> Int

V TMA box KEY-row count for this config's layout.

The single selector every V TMA type-param site routes through (dispatch create_tma_tile, kernel VTMAOpType, the fa4_load signature), so the box shape stays syntactically identical across sites -- the same single-source-of-truth rule v_box_cols() documents. Layout-E (m_pack == 2) uses the KEY-split v_e_box_rows(); Layout-G / non-WS use kv_sub_tile_rows(BN, page_size) (byte-identical to the historical inline expression).

Returns:

Int

v_tma_box_cols​

def v_tma_box_cols(self) -> Int

V TMA box depth (columns) for this config's layout.

Sibling of v_tma_box_rows(): Layout-E (m_pack == 2) uses the full-depth v_e_box_cols(), Layout-G / non-WS the depth-split v_box_cols().

Returns:

Int

nope_cols_per_cta​

def nope_cols_per_cta(self) -> Int

K_nope columns stored in this CTA's SMEM (per-CTA padded nope width).

Sibling of v_cols_per_cta() on padded_nope_depth. Equals v_cols_per_cta() for MHA / DeepSeek (nope == ov).

Returns:

Int

shared_kv_cols​

def shared_kv_cols(self) -> Int

Un-halved width of one shared K_nope/V SMEM stage.

K_nope (padded_nope_depth) and V (padded_ov_depth) share one buffer, so a stage fits the wider of the two. This is the full (non-pair-halved) column count; pair-CTA halving is applied at the call site where needed. Equals padded_ov_depth for MHA / DeepSeek (nope == ov).

Returns:

Int

k_rows_per_cta​

def k_rows_per_cta(self) -> Int

K rows stored in this CTA's SMEM.

Returns:

Int

v_row_major​

def v_row_major(self) -> Bool

Effective row-major (page-dense, chunk-inner) V layout selector.

Drives BOTH the V TMA producer fold (tma_copy_v[row_major=...]) and the P@V MMA consumer descriptor (smem_descriptor[page_dense=...] / SM100TensorAccumulator[b_page_dense=...]); a single accessor keeps the producer's page-dense SMEM and the consumer's descriptor in agreement.

Returns True only when the page-dense layout is actually applicable — i.e. when the V-side kv_tma_fold_chunks[row_major=True] would fold (>= 2). The base_ok geometry comes from the SHARED _kv_fold_base_ok gate (the single source of truth the predicate also uses), so the two cannot drift; a comptime assert at the dispatch site still cross-checks this accessor against the real fold result (it spans the box_rows == page_size bridge the shared gate does not).

Restricted to genuine multi-page paging (0 < page_size < BN):

  • This is the only regime where the fold helps: with page_size >= BN or page_size == 0 the tile is a single page (pages_per_iter == 1) that the chunk-outer rank-4 fold already loads in one TMA.
  • It is also the only regime where the rank-5 atom-row coordinate (gmem_row // _SWIZZLE_ATOM_ROWS) is safe: a block-indirected PagedKVCache gives gmem_row = block_idx * stride with stride a multiple of page_size, so page_size % 8 == 0 (checked below) makes every page row 8-aligned. Continuous / ragged operands (page_size >= BN) place tiles at arbitrary token offsets that are NOT 8-aligned, which would corrupt the atom-row coordinate.

Gated to single-CTA: under pair_cta the descriptor (BMN = v_cols_per_cta() = ov/2) and the accumulator advance (b_BMN = MMA_N = ov) disagree on mn_dim for the native layout — a fast-follow change. SWIZZLE_128B only (the native layout is defined there).

Returns:

Bool

k_row_major​

def k_row_major(self) -> Bool

Effective row-major (page-dense, chunk-inner) K layout selector.

K-side analog of v_row_major(): drives BOTH the K TMA producer fold (tma_copy_k[row_major=...]) and the Q@K' MMA consumer descriptor (smem_descriptor[is_k_major=True, page_dense=...] / SM100TensorAccumulator[b_page_dense=...] for the k-major B operand). A single accessor keeps the producer's page-dense SMEM and the consumer's descriptor in agreement (disagreement is silent wrong output, not a crash).

Gated identically to V: SWIZZLE_128B only, single-CTA only (not pair_cta), and genuine multi-page paging (0 < page_size < k_rows_per_cta() with page_size % 8 == 0, so a block-indirected PagedKVCache gives 8-aligned page rows for the rank-5 atom-row coordinate — see v_row_major() for the full rationale). The base_ok geometry comes from the SHARED _kv_fold_base_ok gate (the K-side kv_tma_fold_chunks uses the same gate: BK0 % gran == 0, num_chunks >= 2, qk_depth % BK0 == 0), so the two cannot drift; a comptime assert at the dispatch site still cross-checks this against the real fold result.

K and V share the same constructor-computed default (row_major_{v,k}_atoms = not is_mla and 0 < page_size < BN); the two accessors then apply their own feasibility gates independently.

Returns:

Bool

q_nope_bytes​

def q_nope_bytes(self) -> Int

Q nope region bytes: BM * padded_nope_depth * dtype_size.

The Q_nope tile feeds the Q@K_nope' contraction, so its width is the non-rope Q depth (padded_nope_depth), not the V/output depth. They coincide for MHA / DeepSeek MLA.

Returns:

Int

q_rope_bytes​

def q_rope_bytes(self) -> Int

Q rope region bytes. Uses rope_dtype_size when set, else dtype_size.

Returns:

Int

rope_depth​

def rope_depth(self) -> Int

Depth of the rope part. Calculated as: padded_qk_depth - padded_nope_depth (0 for MHA where qk_depth == nope_depth). Uses the non-rope Q/K width (padded_nope_depth), NOT the V/output depth — the two differ when v_head_dim != qk_nope_head_dim.

Returns:

Int

num_rope_buffers​

def num_rope_buffers(self) -> Int

Number of separate rope smem buffers (shared mode only).

In shared mode K tiles alternate with V tiles in the pipeline. At most ceildiv(num_kv_stages, 2) K tiles can be in-flight simultaneously, so we only need that many rope buffers. For MHA (rope_depth=0), no rope buffers are needed.

Returns:

Int

num_k_scale_bufs​

def num_k_scale_bufs(self) -> Int

Number of staged k_scale smem buffers.

In shared mode, K tiles alternate with V tiles so at most ceildiv(num_kv_stages, 2) K tiles are in-flight simultaneously. In non-shared mode, each KV stage has its own K buffer. Returns 0 when scale_dtype_size == 0 (no per-token scaling).

Returns:

Int

supported​

def supported(self) -> Bool

Returns:

Bool

with_num_q​

def with_num_q(self, num_q: Int, *, num_qk_stages: Int = Int(0)) -> Self

Reconstruct this config with a different num_q (single-CTA).

num_qk_stages == 0 (default) lets the constructor derive the optimal staging for the new shape — appropriate for the dispatch-time 1Q/2Q selection, where each launch config is free-standing. A nonzero value pins the staging (see switch_1q_config).

pair_cta is forced False because num_q == 1 is single-CTA only (see supported()). Re-passing the stored swizzle_mode is faithful: the constructor re-derives it (FP8 re-forces 64B), and it is already the post-override value here. The row_major_{v,k}_atoms fields are not re-passed: the constructor recomputes them from page_size/BN/is_mla, and BN is num_q-independent, so the value is identical to self's.

with_splitk​

def with_splitk(self, splitk_partitions: Int) -> Self

Reconstruct this config with a split-K cluster size (num_q==1).

Split-K groups splitk_partitions single-CTA kernels in a launch cluster that partition the K/V sequence and (from M4) combine via DSMEM. pair_cta is forced False — split-K is single-CTA only: each CTA runs its own cta_group::1 MMA over its own TMEM/SMEM, and the cluster exists purely to group the split-K partitions. num_q and the derived num_qk_stages are preserved, so with_splitk(1) is a no-op (identical config) and the split-K plumbing folds away.

nope_depth (the Q@K'/Q_nope width) and single_o (the wide-V 1Q TMEM mode) are re-passed so a GLM-style config (v_head_dim != qk_nope) or a single-O config survives the reconstruction; both are byte-identical for the DeepSeek/MHA shapes (nope == ov, single_o == False).

with_bm​

def with_bm(self, bm: Int) -> Self

Reconstruct this config with an explicit BM (single-CTA).

Used by dispatch to force a warp-specialized packed-TMEM datapath (BM=32 -> MMA_M=32, m_pack=4, Layout-G; BM=64 -> MMA_M=64, m_pack=2, Layout-E), with use_ws=True, for short prompts. pair_cta is forced False (BM=32/64 are single-CTA only, per supported()). num_qk_stages/splitk_partitions/single_o are left at their constructor defaults (derive staging, no split-K, 2-O) so the reconstruction is byte-identical to a direct FA4Config(..., BM=bm) build (see test_fa4_config_ws_bm64_probe); use_ws/m_pack are derived from the new BM. nope_depth is re-passed so a GLM-style shape survives (byte-identical for MHA where nope == ov).

switch_1q_config​

def switch_1q_config(self) -> Self

The 1Q variant used by the in-kernel per-sequence 1Q/2Q switch.

Unlike the dispatch-time conversion (with_num_q(1)), which is free to pick the optimal staging, this pins num_qk_stages to this (2Q) config's value: the switch feeds the 2Q-built TMA ops to the 1Q body, so the per-stage K split (QTMATile's smem-tile last dim and k_tma's BK = padded_qk_depth // num_qk_stages) must match. The pinned value is always arithmetically valid here because padded_qk_depth and swizzle_mode are identical across the two configs; if its extra barriers do not fit in 1Q smem, the constructor falls back to 1 stage and can_switch_to_1q() rejects the switch.

can_switch_to_1q​

def can_switch_to_1q(self) -> Bool

Whether a 2Q-launched kernel may dispatch to the 1Q body at runtime.

True only when this is a 2Q single-CTA config AND a valid 1Q variant exists whose TMA-op types match the 2Q ones. switch_1q_config() pins num_qk_stages (the one TMA-op parameter that could otherwise diverge — BN, the per-half Q BM (128), v_tma_op, and ragged_tma_store already match), so the equality check below only fails when the pinned staging could not be honored (smem fallback to 1 stage). When this returns False the kernel runs pure 2Q.

Returns:

Bool

launch_smem_used​

def launch_smem_used(self) -> Int

Dynamic smem to reserve when launching this config's kernel.

When the launched kernel may dispatch to the 1Q body at runtime (can_switch_to_1q()), it constructs the 1Q SM100AttentionSMem over the same dynamic smem region, so the launch must reserve the max of both footprints. Otherwise this is just smem_used.

Returns:

Int

description​

def description(self) -> String

Returns:

String

correction_smem_elements​

def correction_smem_elements(self) -> Int

Returns:

Int

num_active_warps_per_group​

def num_active_warps_per_group(self) -> Int

Returns:

Int

num_active_threads_per_group​

def num_active_threads_per_group(self) -> Int

Returns:

Int

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact