Prevent Voxtral NaNs in CUDA split-K attention

[digantdesai]

Voxtral attention scores can exceed the fixed-phi exponent range, turning partial softmax values into infinities and decoder logits into NaNs.

Use stable normalization because a fixed offset cannot cover both large positive and negative score ranges. No perf regression.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 2dd38cf with merge base d0fe35d ():
💚 Looks good so far! There are no failures yet. 💚

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

[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.

[Copilot]

Pull request overview

This PR fixes numerical instability in the CUDA Triton split-K decode SDPA path by replacing the prior fixed-offset (“phi”) softmax approximation with a stable online-softmax that tracks per-split maxima and rescales partials during the cross-split reduction. This directly targets Voxtral decode cases where very large positive or negative attention scores previously caused overflow/underflow leading to NaNs or silent zero outputs.

Changes:

• Implement stable per-split online softmax in _sdpa_decode_splitk_kernel, and stable global rescaling in _sdpa_decode_reduce_kernel via a new M_partial buffer.
• Remove the fixed _DEFAULT_SPLITK_PHI usage from the split-K decode implementation (while keeping phi in the operator signature for schema compatibility and explicitly documenting it as ignored).
• Add regression tests covering large positive logits (overflow case), large negative logits (underflow-to-zero case), and correctness with kv_len excluding empty trailing splits.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
backends/cuda/triton/kernels/sdpa.py	Reworks split-K decode softmax to be numerically stable using per-split max tracking + global rescaling; introduces M_partial and updates launch plumbing.
backends/cuda/tests/test_triton_sdpa_splitk.py	Adds targeted regression tests to prevent NaNs/Infs and validate correctness under extreme logits and kv_len-bounded decode.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[shoumikhin]

The algorithm change looks correct to me, and it matches what tq4_sdpa.py
already does. I reproduced both the bug and the fix on an H100 (torch 2.13,
triton 3.7.1) using the inputs from your new test_voxtral_large_logits_stable:
on the pre-change kernel sdpa() returns 4096 NaNs at L_q=1, on this branch it
returns finite output with max abs error 0.0074. Accuracy against a float64
reference is also slightly better than the fixed-phi path, out to L_kv=32768.

Three things before it lands.

1. The rebase has to cover both split-K kernels.

Since this branched, main added _sdpa_small_query_splitk_kernel (#21628), and
sdpa() routes L_q 2 to 4 into it. It still computes exp(qk - phi) with a
fixed phi and sums partials with no rescaling, so it has the same bug. On current
main, with the same inputs your new test uses:

sdpa() L_q=1  nan=4096
sdpa() L_q=2  nan=8192
sdpa() L_q=4  nan=16384

The merge is also not safe to resolve mechanically. This PR deletes
_DEFAULT_SPLITK_PHI at sdpa.py:54, and that deletion sits outside the conflict
hunk while main's two remaining uses of it sit inside. Taking main's side of the
conflict leaves the constant referenced but not defined, so it fails at runtime.

2. "No perf regression" does not hold at D=256.

Alternating base/head processes on an idle H100, do_bench medians, best of three
rounds:

B, H_q, H_kv, D, L_kv	before	after	change
1, 16, 2, 256, 16384	36.7 us	48.6 us	+32%
1, 16, 2, 256, 8192	25.6 us	29.9 us	+17%
1, 32, 8, 128, 32768	98.0 us	91.9 us	-6%
1, 32, 8, 128, 8192	33.6 us	33.3 us	-1%

Profiling attributes essentially all of it to the reduce kernel, which gets 3.2x
to 3.4x slower (5.1 to 17.1 us at L_kv=16384). grid_reduce is (B * H_q,), so
that shape launches 16 blocks on 132 SMs and the kernel is a latency-bound serial
loop over splits. The change puts two tl.exp on the loop-carried dependency
through m_global, plus a HEAD_DIM-wide multiply-add where there used to be an
add. Voxtral shapes have more query heads and a narrower accumulator, so they do
not see it.

For what it is worth, the torch.zeros to torch.empty change does not offset
this: the memsets it removes measure about 1 us at that shape.

Happy either way on the fix, but the claim in the description should probably be
scoped to the Voxtral shapes, or the reduce restructured (a cheap option is
parallelizing the reduce over D, or a separate max pass so the exponentials come
off the serial chain).

3. test_kv_len_ignores_empty_trailing_splits is not a regression test.

It passes on the pre-change kernel too, so it cannot fail for the reason this PR
exists. The other two do fail there, so those two are doing real work.

It is still a useful test of the new uninitialized-buffer contract, just a weak
one, since dirty allocator memory rarely shows up in a fresh process. Pre-filling
M_partial and L_partial with NaN before the launch would pin the invariant
properly. I NaN-poisoned every float32 CUDA allocation and swept batch sizes,
group counts, and L_kv from 1 to 40000, and every partial slot did get written,
so the invariant does hold today.

Comment nit, sdpa.py:1389:

# The split grid unconditionally writes every partial, including empty
# splits, so these buffers do not need initialization kernels.

The stores are not unconditional, they are masked by g_valid, and the kernel
comment three lines above says so. What actually makes it safe is that the grid
partitions the h_q range exactly. Worth rewording, because as written it invites
someone to add an early return for empty splits later, which would silently
reintroduce garbage.

Also, since phi is now ignored, the docstring line reads better as "deprecated,
accepted for schema compatibility and ignored" rather than implying it still does
something.