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):

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):

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

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):

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

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:

Dense (no mask):

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, “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.
• The MLA sibling of the fused projection epilogue: GEMM + RoPE + MXFP8 Projection.
• 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).