[kernel] QAD 5090: Add Attn-QAT training Triton kernels (11/12)

[alexzms]

Purpose

Continues the Attn-QAT upstreaming stack (#1225). Adds the fake-quantized
attention Triton kernels imported by AttnQatTrainBackend, completing the
training-side attention path. With these present, is_attn_qat_train_available()
(from #1459) returns True and ATTN_QAT_TRAIN performs real fake-quant instead
of falling back to Flash Attention.

Changes (all new, under fastvideo-kernel/)

• triton_kernels/attn_qat_train.py: QAT attention fwd/bwd — fake NVFP4 (E2M1)
  quantization of Q/K/V/P with a straight-through estimator; exports attention
  (the symbol fastvideo/attention/backends/attn_qat_train.py imports).
• triton_kernels/quant_utils.py + triton_kernels/nvfp4_utils.py: NVFP4
  fake-quant helpers.
• triton_kernels/fused_attention.py: non-QAT Triton reference used by the test.
• tests/test_attn_qat_train.py.

Test Plan / Test Results

Verified on a Blackwell node (GB200, sm_100) — the full test_attn_qat_train.py
suite passes: forward/backward, causal & non-causal, varied sequence lengths, and
a Wan (non-divisible-by-64) shape. QAT output and gradients match the FP32 / naive
reference at cos ≈ 0.98 (the gap is the modeled FP4 attention error — exactly
what QAT trains the model to absorb).

This is a Triton kernel, so unlike the sm_120-only inference CUDA kernel (#1455)
it runs and validates on datacenter Blackwell too.

Notes

• tests/test_fake_quant.py from the source branch was left out: it imports
  flashinfer at module top and would break collection where flashinfer isn't
  installed. Can add it later behind an import guard.
• These kernel files live under fastvideo-kernel/, which is excluded from the
  repo's python pre-commit hooks (yapf/ruff/mypy), consistent with the other
  kernels there.

Depends on #1459 (backend wiring). Part of #1225.

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces Triton kernels for Quantization-Aware Training (QAT) fused attention, including core attention passes, legacy compatibility shims, NVFP4 quantization utilities, and comprehensive tests. The code review highlights several critical issues and optimizations, such as a runtime AttributeError when smooth_q is enabled, potential undefined values from masked tl.load calls without default values, a type mismatch in the sticky variable, and inefficiencies in boolean mask creation and FP8 sum reductions.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When smooth_q is enabled, q_m is initialized to None and the line that computes it (triton_group_mean) is commented out. However, q_m.repeat_interleave is still called on the next line, which will result in a guaranteed AttributeError: 'NoneType' object has no attribute 'repeat_interleave' at runtime.

If smooth_q is not currently supported, please raise a NotImplementedError to prevent silent failures or unexpected crashes.

Suggested change

	q_m = None
	if ctx.smooth_q:
	# _, q_m = triton_group_mean(q)
	q_m = q_m.repeat_interleave(q.shape[2] // q_m.shape[2], dim=2) # B,H,L,D
	q_m = None
	if ctx.smooth_q:
	raise NotImplementedError("smooth_q is not currently supported in the backward pass.")

[gemini-code-assist]

In Triton, calling tl.load with a mask but without specifying the other parameter results in undefined values (garbage or NaNs) for the masked-out elements. These undefined values will propagate through subsequent operations (like tl.sum or tl.dot), corrupting the output and gradients.

Please specify other=0.0 (or another appropriate default) for all masked tl.load calls in this file (e.g., lines 306-307, 347, 358, 427-428, etc.), similar to how it is correctly done in quant_utils.py.

Suggested change

	o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None])
	do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None]).to(tl.float32)
	o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None], other=0.0)
	do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None], other=0.0).to(tl.float32)

[gemini-code-assist]

There is a type mismatch for the sticky variable:

1. In the if IS_SRC_FP32 branch, sticky is initialized as a boolean tensor (is_subnormal & ...), but then updated with |= using a uint32 tensor.
2. In the else branch, sticky is initialized directly as a uint32 tensor.

This type mismatch can cause Triton compilation failures or silent bugs depending on the compiler version. We should consistently use uint32 for sticky in both branches.

Suggested change

	sticky = is_subnormal & (bit0_dropped | dropped_post)
	sticky |= ((mantissas & 0x1FFFFF) != 0).to(tl.uint32)
	else:
	sticky = ((mantissas & 0x1FFFFF) != 0).to(tl.uint32)
	sticky = (is_subnormal & (bit0_dropped | dropped_post)).to(tl.uint32)
	sticky |= ((mantissas & 0x1FFFFF) != 0).to(tl.uint32)
	else:
	sticky = ((mantissas & 0x1FFFFF) != 0).to(tl.uint32)

[gemini-code-assist]

Using tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0 to create a boolean mask of all True values is inefficient because it allocates a float tensor, fills it, and then performs an elementwise comparison.

Instead, you can directly create a boolean mask using tl.full(p.shape, True, dtype=tl.int1). This also applies to lines 365 and 597.

Suggested change

	src_tensor=p,
	valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0,
	valid_src_mask=tl.full(p.shape, True, dtype=tl.int1),

[gemini-code-assist]

When dtype is tl.float8e5 (FP8), performing tl.sum directly on high_prec_p or p can lead to significant precision loss or overflow during reduction due to the limited dynamic range and precision of FP8.

It is highly recommended to upcast the tensor to tl.float32 before performing the sum reduction to ensure numerical stability.

Suggested change

	l_ij = tl.sum(high_prec_p, 1)
	else:
	l_ij = tl.sum(p, 1)
	l_ij = tl.sum(high_prec_p.to(tl.float32), 1)
	else:
	l_ij = tl.sum(p.to(tl.float32), 1)

[gemini-code-assist]

The variable scale is reshaped here, but it is never used anywhere else in the remainder of the _compute_dequant function. This line can be safely removed to avoid redundant operations.

[gemini-code-assist]

Calling triton.set_allocator inside the forward method of an autograd function is a global state mutation that runs on every forward pass. This can introduce overhead and potential race conditions in multi-threaded environments.

It is cleaner to either set the allocator once at the module level or rely on Triton's default integration with PyTorch's caching allocator (which is automatically used when PyTorch is imported).