[None][feat] Run Gemma4 head_dim 512 attention on Triton

[yuzisun]

Description

Gemma 4's full-attention layers use head_dim=512, which no SM90 paged attention kernel serves:

• trtllm-gen has the cubins, but ships them for datacenter Blackwell only (is_sm_100f).
• TRTLLM paged FMHA/MMHA does not cover 512.
• FlashInfer's fa2/fa3 paged kernels stop at 256.

PR #17557 routed all non-SM100 architectures to fa2, but that branch was validated only through unit tests with is_sm_100f mocked, and the same PR removed the comment # head_dim>256 needs trtllm-gen (fa2 JIT doesn't support it). So Gemma 4 does not currently run on Hopper through the PyTorch backend at all. Text-only inference is available today only via AutoDeploy, whose configs select attn_backend: triton for exactly this reason.

This PR routes those layers through Triton in the PyTorch backend.

Approach

Prefill already had a suitable kernel in-tree: attention_backend/triton_prefill.py handles head_dim 512, reads FlashInfer's HND paged layout, gives causal attention when custom_mask is None, and its _get_block_sizes already has a Hopper branch for Lq > 256. Only the call-site gate needed widening.

Decode is ported from the AutoDeploy Triton attention backend (auto_deploy/custom_ops/attention/triton_attention.py), which uses the identical combined HND cache layout [num_pages, 2, num_kv_heads, page_size, head_dim]. The AutoDeploy file imports flashinfer only for a metadata helper, not for the attention math, so the kernels port cleanly. AutoDeploy registry wiring and the context-phase kernels were stripped.

The plumbing lines up with what the backend already maintains:

triton_decode arg	Source
kv_cache	kv_cache_manager.get_buffers(layer_idx, kv_layout="HND") returns exactly this shape; its docstring already covers Gemma 4's 256/512 multi-pool case
kv_indptr / kv_indices / kv_last_page_len	Same slicing the FlashInfer decode plan uses
VSWA per-layer pools	metadata.swap_paged_kv_indices_for_layer(), already called in forward_impl

FP8 KV cache keeps working: the kernels cast on load, so the cache is dequantized in-kernel.

Why both phases are handled in one branch

PR #17557 noted that "uniform backend avoids workspace corruption between different wrapper types under CUDA graphs". Doing prefill and decode in Triton lets the branch return before metadata.plan(), so head_dim 512 layers never create a FlashInfer wrapper or touch workspace_buffer. The sliding layers (head_dim 256) all stay on fa2, leaving exactly one wrapper type in play. Sliding-layer behaviour is unchanged.

The trigger is head_dim > 256 and not is_sm_100f() — a general statement about FlashInfer paged kernel capability, not a Gemma 4 special case — so no model code changes.

Not covered

These raise NotImplementedError rather than returning incorrect results:

• KV-shared layers (k is None). Their current tokens' KV is already paged, so feeding them as Triton's "prefix" would drop causal masking between them. Models with num_kv_shared_layers > 0 will fail loudly on this path.
• Speculative-decoding draft views, which require the trtllm-gen decode backend.
• Multi-token generation steps — triton_decode derives batch size from q's leading dim.

CUDA graph capture and perf tuning are deliberately deferred; triton_decode still allocates its split-K workspaces per call. Triton FlashDecoding will not match trtllm-gen. Multimodal is out of scope (the vision tower forces attn_backend="TRTLLM" with sm100a-derived head-dim padding).

Test Coverage

• tests/unittest/_torch/attention/test_triton_decode.py — new kernel tests: head_dims 64–512, the real Gemma 4 E2B/31B shapes, GQA ratios including non-power-of-2 (12/4, 24/4) that exercise head padding, page sizes 1–64, partial and exact-multiple last pages, split-K, sliding windows, FP8 KV.
• backend_capability.py — head_dim 512 is no longer skipped for FLASHINFER below sm100, activating the existing gemma4_e2b_mqa_hd512 / gemma4_26b_gqa_hd512 / gemma4_31b_gqa_hd512 cases against the VanillaAttention golden. It remains skipped on sm100+, where the harness builds FlashInfer with the default fa2 and this path is off.
• test_modeling_gemma4.py — four tests running the existing HF-comparison harness with the non-Blackwell dispatch forced on, across the E2B/31B/26B real-dims configs, plus test_triton_path_is_actually_taken, which spies on both Triton entry points so a dispatch regression cannot masquerade as passing coverage. Both is_sm_100f call sites are patched together so the simulation is faithful on any GPU.
• l0_h100.yml — individual node ids rather than the whole file, since the rest of test_modeling_gemma4.py assumes trtllm-gen. test_triton_decode.py and the attention-backend sweep are already picked up by the existing unittest/_torch/attention directory entry.

This has not been run on a GPU — hence draft. Opening it to get CI signal on H100.

PR Checklist

• PR title follows [JIRA/NVBUG/None][type] Summary
• Commit is signed off (DCO)
• Tests pass (unvalidated — awaiting CI)
• New file carries the NVIDIA copyright header

[yuzisun]

/bot run

[Hudayday]

Hi Dan, thanks for your contribution. We’re currently preparing PR #18002, which adds SM90 support to the FlashInfer FA2 backend and will unblock Gemma 4 on Hopper. Could you please verify whether this official implementation works for your use case?