> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/cudnn/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/cudnn/_mcp/server. # Gated Attention Block (SM107) **This is an experimental API and subject to change.** ## Overview The gated attention block is the first **model-level** FE-OSS API: a set of FROST CuTe-DSL kernels behind one Python class, one workspace and one `execute()` call. It implements the gated attention sub-layer used by Qwen3.5-style models (the API is named by op geometry, the model is provenance only): ```text h [B, S, d_model] (post input-layernorm) | | (1) QKV + GATE projection h @ W_qkvg^T FROST GEMM +----------------------------------------------------------------------------+ | Q [B,S,H_q,D] | GATE [B,S,H_q,D] | K [B,S,H_kv,D] | V [B,S,H_kv,D] | | (2)+(3) QK-RMSNorm per head (optional) + partial RoPE on the first rope_dim ONE kernel (Q and K only) | | (4) SDPA O = softmax(Q K^T * scale + mask) V, GQA H_q / H_kv FROST SDPA (Rubin d256 rows) | | (5) gate O_gated = O * sigmoid(GATE) elementwise, after the SDPA epilogue | | (6) out projection O_gated @ W_o^T FROST GEMM v out [B, S, d_model] ``` Every stage is a FROST kernel: the two projections drive the shipped FROST GEMM engine (pinned by name in the graph's ranked plan list), stage (4) drives the shipped FROST SDPA through its standalone adapter, and stages (2)+(3) and (5) are this block's own kernels. There are no cuBLAS or cuDNN-backend call-outs. **Target: NVIDIA Rubin (SM107, compute capability 10.7) only.** Every other architecture is declined with a typed error (`NotImplementedError`), never served slowly. Three precisions share the signature, and two fp4 modes ride the MXFP8 one (both are fields on `MxQuantSpec`, so they are unspellable on the bf16 and per-tensor FP8 pipelines rather than declined): | mode | how it is selected | pipeline | |---|---|---| | bf16 / fp16 | `h`, weights and `cos`/`sin` in bf16 or fp16, `quant=None` | 5 stages, 5 launches (4 with `inplace_qkv`) | | FP8, per-tensor static scales | e4m3 `h` / `W_qkvg` / `W_o` and a `QuantSpec` | unfused: FP8 projections with the descale folded into the epilogue, bf16 norm+RoPE, two quantize passes (Q/K/V, gated O), the Rubin per-tensor FP8 SDPA and an FP8 out projection (9 launches); fully fused: 3 launches | | MXFP8, per-32-element E8M0 block scales | e4m3 `h` / `W_qkvg` codes, the two scale-factor blobs (`h_sf`, `w_qkvg_sf`) and an `MxQuantSpec` | unfused: block-scale projection GEMM, bf16 norm+RoPE, MXFP8 quantize (Q/K rowwise, V columnwise), the Rubin d256 MXFP8 SDPA, per-tensor quantize of O, FP8 out projection (9 stages); fully fused: 3 launches | | MXFP8 with **MXFP4 weights** (unfused only) | `MxQuantSpec(w_qkvg_dtype=torch.float4_e2m1fn_x2)` and an e2m1 `W_qkvg [N, d_model // 2]`; `h`, `h_sf` and `w_qkvg_sf` unchanged | the MXFP8 unfused pipeline with stage (1) on the catalog's mixed MXFP8 x MXFP4 block-scale row (E8M0 scales per 32 on both sides) -- the same 9 launches; `fuse_norm_rope` is a typed decline (the fused projection fork is rendered for an e4m3 B) | | MXFP8 with **fp4 O** (`NVFP4` or `MXFP4`) | `MxQuantSpec(o_fp4=Fp4Format.NVFP4 \| Fp4Format.MXFP4)`, an e2m1 `W_o [d_model, H_q * D // 2]` of the SAME format and its scale blob (`sample_w_o_sf` / `w_o_sf`) | the per-tensor tail is replaced: one `quantize_fp4` launch writes the gated `O` as e2m1 codes plus the out projection's scale blob, and stage (6) becomes the fp4 x fp4 block-scale GEMM (no per-tensor scale on either side). Unfused: 9 launches (`quantize_fp4` takes the per-tensor quantize's place); fully fused: **4 launches** (the gated MXFP8 SDPA writes bf16 `O`, then `quantize_fp4`, then the fp4 out projection). Composes with the MXFP4 weights on the unfused pipeline | ### Fusion knobs Two optional fusions are constructor flags; each is a different compiled specialization behind the same `execute()` signature, and both default to off. - `fuse_norm_rope=True` folds stages (2)+(3) into stage (1)'s epilogue: Q/K tiles are normed and rotated on the fp32 accumulator and written once (needs `inplace_qkv`; inference only, no pre-norm Q/K is kept). Under FP8 / MXFP8 the same epilogue also quantizes: compact e4m3 `q8` / `k8` / `v8` (and the MXFP8 scale factors) come straight out of the GEMM, so the quantize passes disappear. - `fuse_gate=True` folds stage (5) into the SDPA epilogue: `O *= sigmoid(GATE)` after the dead-row select, with the gate tile TMA-staged by the load warp (inference only: no pre-gate `O` for the backward). With both on, the block is **three launches**: `proj(+norm+RoPE[+quant]) -> sdpa(+gate) -> out_proj`. `geometry.qk_norm=False` runs RoPE-only Q/K on every path (norm weights are passed as `None`, no `rstd` is produced, the backward drops the norm-weight gradients). ## API ```python import cudnn from cudnn.gated_attention_block import ( GatedAttentionBlockFwd, GatedAttentionBlockGeometry, QuantSpec, MxQuantSpec, Fp4Format, build_fused_qkvg_weight, GatedAttentionBlockBwd, RecomputePolicy, SavedForBackward, ) ``` ### Geometry `GatedAttentionBlockGeometry` (frozen dataclass; `validate()` raises `ValueError`): | field | meaning | |---|---| | `d_model` | model width | | `h_q`, `h_kv` | query heads and KV heads (GQA); `h_q % h_kv == 0` | | `d_head` | head dim of Q, K, V and O (`d_qk == d_v`); selects the SDPA flavor (d256 today) | | `rope_dim` | the LEADING dims of each head that rotate; `[rope_dim, d_head)` pass through | | `qk_norm` (`True`) | apply QK-RMSNorm before RoPE; `False` = RoPE only, no norm weights, no `rstd` | | `qk_norm_eps` (`1e-6`) | RMSNorm epsilon | | `attn_scale` (`None`) | softmax scale; `None` = `d_head ** -0.5` | | `is_causal` (`True`), `causal_bottom_right` (`False`) | causal mask and its diagonal alignment | | `window_left`, `window_right` (`-1`) | sliding window bounds; `-1` = unbounded | ### Weights and tables - `W_qkvg [N, d_model]` with `N = (2*H_q + 2*H_kv) * D`, column blocks `Q | GATE | K | V`. Build it from separate projection weights with `build_fused_qkvg_weight(w_q_gate, w_k, w_v, geometry, q_gate_layout="flat")`; the block's tile alignment is `QKVG_TILE_ALIGN = 64` columns. - `W_o [d_model, H_q * D]`. - `cos`, `sin` `[B, S, rope_dim]` rotary tables in the activation dtype (rotate-half convention on the first `rope_dim` dims of every head). - `w_q_norm`, `w_k_norm` `[D]` (both `None` iff `geometry.qk_norm` is `False`). - MXFP8 only: `h_sf` and `w_qkvg_sf`, the E8M0 scale factors of `h` and `W_qkvg` in cuDNN's F8_128x4 order (`uint8` or `float8_e8m0fnu`; byte counts from `cudnn.gated_attention_block.kernels.proj_gemm.sf_blob_bytes`). - fp4 weights are **packed e2m1**, dtype `torch.float4_e2m1fn_x2`, stored `[N, K // 2]` -- two codes per byte along the contraction axis, LOW nibble = even `k`. The block checks the STORAGE shape (`[N, d_model // 2]` for `W_qkvg`, `[d_model, H_q * D // 2]` for `W_o`); a logical `[N, K]` fp4 tensor, or `uint8` storage, is a typed `ValueError` (torch can `.view(torch.float4_e2m1fn_x2)` packed bytes but cannot cast to fp4). - MXFP4 `W_qkvg` (`MxQuantSpec.w_qkvg_dtype=torch.float4_e2m1fn_x2`): `w_qkvg_sf` is UNCHANGED -- the same E8M0 / 32 F8_128x4 blob over `n_qkvg x d_model` as for e4m3 codes. - fp4 `W_o` (`MxQuantSpec.o_fp4`): its scale blob `w_o_sf` is in the SAME format as `O` -- `float8_e4m3fn` scales per 16 for `Fp4Format.NVFP4`, `float8_e8m0fnu` per 32 for `Fp4Format.MXFP4` (either as `uint8` or as that dtype), `sf_blob_bytes(d_model, H_q * D, block)` bytes in F8_128x4 order (padded to whole 128-row x 4-block atoms; pad bytes, if any, `0x00`). ### Forward ```python blk = GatedAttentionBlockFwd( sample_h, sample_w_qkvg, sample_w_q_norm, sample_w_k_norm, sample_cos, sample_sin, sample_w_o, sample_out, geometry, return_lse=False, # also write the softmax stats [B, H_q, S] fp32 save_for_backward=False, # keep what GatedAttentionBlockBwd needs (implies return_lse) seq_lens_present=False, # per-batch KV lengths (padding) at execute time inplace_qkv=None, # None -> not save_for_backward: norm/RoPE Q,K in place in the projection slab fuse_norm_rope=False, # stages (2)+(3) inside the projection epilogue fuse_gate=False, # stage (5) inside the SDPA epilogue quant=None, # QuantSpec (FP8) | MxQuantSpec (MXFP8, + fp4 weights / fp4 O) | None (bf16 / fp16) sample_h_sf=None, sample_w_qkvg_sf=None, # MXFP8 scale-factor blobs sample_w_o_sf=None, # fp4 O only (MxQuantSpec.o_fp4): the e2m1 W_o's scale blob, in O's format ) workspace = torch.empty(blk.get_workspace_size(), dtype=torch.uint8, device=h.device) blk.execute(h, w_qkvg, w_q_norm, w_k_norm, cos, sin, w_o, out, workspace, seq_lens=None, lse=None, saved=None, current_stream=None, h_sf=None, w_qkvg_sf=None, w_o_sf=None) # required iff MxQuantSpec.o_fp4, refused otherwise ``` Every appended argument (`sample_w_o_sf`, `w_o_sf`) sits at the end with a `None` default, so positional callers of the bf16, FP8 and MXFP8 pipelines are unchanged; `MxQuantSpec.o_fp4` and `sample_w_o_sf` must be given together (a typed `ValueError` names the missing half). `execute()` allocates nothing, reads nothing back to the host and converts nothing: every intermediate is a strided view of the caller's workspace, sized honestly by `get_workspace_size()`, so the call is CUDA-graph friendly. All stages run on one launch stream (torch's current stream, or `current_stream`), on both the FROST GEMM route and the DSL stages. Dtype, layout and shape mismatches, an unsupported architecture, or a feature the selected specialization cannot serve raise typed errors at construction (`check_support`) rather than at launch. Quantization specs: - `QuantSpec(descale_h, descale_w_qkvg, descale_w_o, scale_q, scale_k, scale_v, scale_o, dtype=torch.float8_e4m3fn)` — static per-tensor scales for the FP8 pipeline. - `MxQuantSpec(descale_w_o, scale_o=1.0, dtype=torch.float8_e4m3fn, block_size=32, w_qkvg_dtype=torch.float8_e4m3fn, o_fp4=None)` — the MXFP8 pipeline; the MXFP8 SDPA writes e4m3 `O` unscaled, so the fully fused path needs `scale_o == 1.0`. The two appended fields select the fp4 modes: - `w_qkvg_dtype`: `torch.float8_e4m3fn` (default, MXFP8 x MXFP8) or `torch.float4_e2m1fn_x2` (MXFP4 `W_qkvg`, the mixed row); any other dtype is a typed `NotImplementedError`. - `o_fp4`: `None` (default, per-tensor e4m3 `O`) or an `Fp4Format` member (anything else is a `TypeError`). Under `o_fp4` neither side of the out projection has a per-tensor scale -- the quantizer writes block scales only and `W_o` dequantizes through `w_o_sf` in the MMA -- so **`scale_o` and `descale_w_o` must both be `1.0`** (typed `ValueError` otherwise; a non-unit value would be silently dropped by the block-scale GEMM). - `Fp4Format` (enum; one member = codes x scale dtype x block, so an illegal pairing cannot be spelled): | member | codes | scale dtype | scale block along K | |---|---|---|---| | `Fp4Format.NVFP4` | e2m1 | `torch.float8_e4m3fn` | 16 | | `Fp4Format.MXFP4` | e2m1 | `torch.float8_e8m0fnu` | 32 | Properties: `block_size`, `sf_torch_dtype`, `sf_cudnn_dtype` (and `fmt_name`, the kernel's format key). Neither format carries a global (per-tensor) scale: an NVFP4 block's e4m3 scale is `max(amax * fp32(1/6), 2^-9)` (the e4m3 min-subnormal floor keeps an all-zero block's scale NONZERO, so the encode `x / scale` stays finite), an MXFP4 block's E8M0 scale is the power of two at or above `amax * fp32(1/6)` -- `fp32(1/6)`, not an exact `/ 6`: the kernel and the reference multiply by the fp32 constant, and the two differ by one ulp; a dead (fully masked or zero-length) row quantizes to codes `0` exactly. The `O` codes are round-to-nearest-even on the e2m1 grid, saturating at 6. ### Backward `GatedAttentionBlockBwd(sample_dy, sample_saved, sample_w_qkvg, sample_w_q_norm, sample_w_k_norm, sample_cos, sample_sin, sample_w_o, geometry, *, recompute=RecomputePolicy.RECOMPUTE_QK_PRE, need_dh=True, need_dw_qkvg=True, need_dw_o=True, need_dw_norms=None)` consumes the forward's `SavedForBackward(h, gate, o, lse, rstd_q, rstd_k, q_pre=None, k_pre=None)`. `RecomputePolicy` chooses between re-running stage (1) for the pre-norm Q/K (`RECOMPUTE_QK_PRE`, the default) and reading them from the save set (`SAVE_ALL`); which input gradients are wanted is fixed at build time because it decides which GEMMs exist. `need_dw_norms=None` follows `geometry.qk_norm`; asking for norm-weight gradients under `qk_norm=False` is a typed decline. The backward is bf16 / fp16 only. ## Requirements and limits - Rubin (SM107) only; cuDNN 9.x, `nvidia-cutlass-dsl >= 4.8.0.dev0` (the Rubin arch names), torch. - `d_head = 256` (the Rubin d256 SDPA flavor with the fused gate); `d_model % 128 == 0` under MXFP8. - FP8 / MXFP8 are inference only; the backward is bf16 / fp16. - FP8: a dense (no-mask) sequence length must be a multiple of 128 unless the causal mask or a padding mask covers the KV tail (the Rubin per-tensor FP8 SDPA contract); MXFP8: e4m3 codes only (e5m2 is a typed decline); the fully fused MXFP8 path needs `scale_o == 1.0` and, at `B > 1`, `S % 128 == 0` (a scale-factor atom is per sequence). - `fuse_gate` and `fuse_norm_rope` are inference-only specializations (no pre-gate `O`, no pre-norm Q/K). - fp4 (`MxQuantSpec.w_qkvg_dtype` / `o_fp4`): MXFP8 pipeline only (unrepresentable on `QuantSpec` / bf16); inference only (`save_for_backward` is a typed decline, as for every quantized pipeline); no global per-tensor scale in either fp4 format (`scale_o == descale_w_o == 1.0` under `o_fp4`); `d_head % (4 * block) == 0` under `o_fp4` (whole 4-block scale words per head: 64 for NVFP4, 128 for MXFP4; `d_head = 256` passes both); the MXFP4 `W_qkvg` runs on the unfused pipeline only -- `fuse_norm_rope` with an e2m1 `W_qkvg` is a typed `NotImplementedError` (the fused MXFP8 projection fork is rendered for an e4m3 B). `h` stays e4m3 (an fp4 `h` is not served), and an fp4 `W_o` with an e4m3 `O` is not a served pairing. ## Performance Whole block, B=1, `h_q=32 h_kv=2 d=256 d_model=5120` (the 397B geometry), Rubin perf node (212 SMs, SM clock locked at 2376 MHz), speedup over the same bf16 torch chain (median of 5 launch-interleaved rounds x 30 launches; the bf16 FROST control pair stayed within 0.6 %). The fp4 modes (MXFP4 weights, NVFP4 / MXFP4 `O`) are not in these tables: their perf-node measurement is pending, and no number is quoted until it exists. Causal: | S | bf16 FROST | bf16 fully fused | FP8 unfused | FP8 fully fused | MXFP8 unfused | MXFP8 fully fused | |---|---|---|---|---|---|---| | 2048 | 2.96x | 3.16x | 4.22x | 4.90x | 3.84x | 4.30x | | 4096 | 3.06x | 3.27x | 4.52x | 5.28x | 4.10x | 4.81x | | 8192 | 2.79x | 2.95x | 4.33x | 4.90x | 4.06x | 4.68x | | 16384 | 2.37x | 2.44x | 3.93x | 4.27x | 3.75x | 4.15x | | 32768 | 1.90x | 1.96x | 3.41x | 3.61x | 3.22x | 3.44x | Dense (no mask): | S | bf16 FROST | bf16 fully fused | FP8 unfused | FP8 fully fused | MXFP8 unfused | MXFP8 fully fused | |---|---|---|---|---|---|---| | 2048 | 2.78x | 2.98x | 4.07x | 4.82x | 3.85x | 4.45x | | 4096 | 2.74x | 2.93x | 4.24x | 4.87x | 4.00x | 4.58x | | 8192 | 2.33x | 2.44x | 3.84x | 4.27x | 3.71x | 4.13x | | 16384 | 1.85x | 1.90x | 3.37x | 3.63x | 3.30x | 3.51x | | 32768 | 1.50x | 1.53x | 2.92x | 3.03x | 2.81x | 2.86x | ## Related - The SDPA epilogue gate is also reachable through the **graph API**: an `sdpa` node followed by `sigmoid` and `mul` pointwise nodes on `O` is served fused by the Rubin d256 FROST SDPA engines — see [Attention](/cudnn/operations/Attention), "Fused epilogue gate". - How the block is composed (workspace, streams, fusion knobs, typed declines) and how to build the next one: [Composing multi-kernel blocks in Python](/cudnn/utilities/composing_kernel_blocks). - The MLA sibling of the fused projection epilogue: [GEMM + RoPE + MXFP8 Projection](/cudnn/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8). - Tests: `test/python/fe_api/gated_attention_block/` (layout contract, reference oracle, end to end, FP8, MXFP8, fp4 weights / fp4 O (`test_block_fp4.py`, `test_proj_gemm_fp4.py`, `test_quantize_fp4.py`), per-stage kernels, stream ordering).