[gemma4_31b][cuda] length-aware bf16 global attention

[Gasoonjia]

• Global (full-attention) bf16 layers: bound SDPA to a runtime kv_len scalar (CUDA-graph-safe) instead of the full max_seq_len KV buffer, reduce complexity of decode from O(kvcache_length) -> O(context); restores decode scaling (was flat ~36.5 t/s at all depths -> 46.5@512, 34.9@127K).
• Prefill global full-attention: replace fixed m32/m64 BLOCK_M selection with a head_dim-keyed autotuned _sdpa_fwd_kernel + register-budget prune (BLOCK_MHEAD_DIM <= 4096num_warps), fixing acc[64,512] fp32 register spill at head_dim=512. Prefill +24% @8K, +63% @32k, +117% @127k; head_dim-agnostic (no split-D needed for D<=512). (sdpa.py)
• Validated: output bitwise-identical to prior kernel (cos=1.0, D=64/128/256/512), no decode regression; non-tq prefill now beats llama.cpp at all 5 cells and turbo TQ4 at 4/5. Op-level autotune profiling (A100) confirms the config set is near-optimal (in-set optimum at every regime; only <=1.3% marginal candidates).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20506

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: Gasoonjia / name: Songhao Jia (74c2a9d)

[metascroy]

Is the right threshold device specific?

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 49s —— View job

---

Code Review

• Read the changed files
• Review sdpa.py (kv_len bound, split-K dispatch, autotune/prune)
• Review cuda_source_transformations.py (length-aware global attention)
• Post review feedback

Reviewed both files against origin/main. This is clean, well-documented work — the kv_len bound is correctly plumbed through both the standard and split-K paths, kept on-device (no .item()) so it survives CUDA-graph replay, and clamped to L_kv for OOB safety. The HAS_KV_LEN=False fallback preserves the original behavior exactly. Output bitwise-identity claim is plausible given the loop bound only skips positions that were empty/future anyway.

A few notes below — none blocking.

---

Device-specific thresholds (answers @metascroy's question)

The reviewer asked whether the threshold is device-specific. It is — and there are now three A100-tuned magic constants in this kernel:

• sdpa.py:582 — BLOCK_M * HEAD_DIM <= 4096 * num_warps, explicitly documented as "the measured A100 no-spill boundary". The register-file budget per SM differs across archs (e.g. consumer/Ada vs. A100/H100), so this prune boundary is genuinely device-specific.
• sdpa.py _SPLITK_LKV_THRESHOLD = 256 — the split-K cutover point.
• sdpa.py:1047 — threshold = 4 * 84 (84 = A100 SM count) for the pack-GQA block_m heuristic.

These are safe rather than correct-only-on-A100: the prune guarantees a non-empty config set (the BLOCK_M=16 configs satisfy the budget for any HEAD_DIM<=1024, plus the explicit fallback), and the autotuner benchmarks the survivors, so a suboptimal-but-valid config is always chosen on other GPUs. Worst case is leaving perf on the table off-A100, not incorrectness.

If you want to make these portable without a re-tune, consider deriving from torch.cuda.get_device_properties(dev).multi_processor_count (replacing the hardcoded 84) and the arch's regs/SM. Not required for this PR, but worth a TODO so the A100 assumption is discoverable.

Stale comment / heuristic mismatch at the pack-GQA decision (sdpa.py:1039-1047)

# Use the actual BLOCK_M that the launched kernel will use:
# - non-pow2 path always uses BLOCK_M=32
# - pow2 path selects M32 or M64 based on CTA occupancy
...
block_m = 32 if total_ctas_m64 < 4 * 84 else 64

This comment and block_m guess are now stale. _sdpa_fwd_kernel is autotuned over BLOCK_M ∈ {16..128}, so the kernel's actual BLOCK_M is no longer the 32/64 value computed here. block_m only feeds the _should_pack_gqa heuristic (which is purely a perf decision, not correctness), so this isn't a bug — but the comment's "actual BLOCK_M that the launched kernel will use" is no longer true and is misleading to the next reader. Suggest updating the comment to note it's an approximation for the pack-GQA heuristic only. Fix this →

_sdpa_prefill_prune HEAD_DIM lookup robustness (sdpa.py:582)

The prune falls back to returning all configs if HEAD_DIM is found in neither kwargs nor nargs. In that (unexpected) case the BLOCK_M=128 config could be selected for HEAD_DIM=512 — the exact acc[128,512] fp32 spill this prune exists to prevent. In practice HEAD_DIM is always passed as a constexpr kwarg to the launch, so this is defensive-only, but since the whole point of the prune is spill-avoidance, consider asserting HEAD_DIM is not None rather than silently returning the unpruned (potentially-spilling) set.

Minor / confirmed-correct

• Split-K empty splits (sdpa.py:1199-1208): splits with start_n >= kv_len get end_n <= start_n, skip the loop, and store the in-kernel zero-initialized acc/l_i. Since all H_q = H_kv * NUM_GROUPS slots are written, the reduce sees valid zeros — correct. O_partial being torch.empty (vs. zeros) is fine because every (split, head) slot is unconditionally stored.
• kv_len = input_pos[0] + input_pos.shape[0] (cuda_source_transformations.py): correct only if input_pos is a contiguous range; that's the invariant for decode (T=1) and prefill chunks here, so fine — just flagging the assumption.
• The new import ...kernels.sdpa # noqa: F401 for op registration mirrors the existing tq4_sdpa import — consistent.

Nice work — the O(context) decode scaling restoration and the head_dim-agnostic prefill autotune are well-motivated and the validation (cos=1.0 across D, no decode regression) is reassuring.
• branch gemma4_31b-cuda-attn-perf-git