Prevent Voxtral NaNs in CUDA split-K attention by digantdesai · Pull Request #22133 · pytorch/executorch

digantdesai · 2026-08-25T04:06:44Z

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.

 ### A100 performance vs fixed phi=5 (quick and dirty benchmarking)
Measured the complete graph-captured split-K SDPA operator, not full-model latency. Results are two-run averages using PyTorch 2.13 and Triton 3.7.1.
Shape phi=5 this PR Change
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━
Voxtral, D128, Lq1, Lkv4096 16.9 µs 16.6 µs −1.7%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv8192 28.2 µs 26.8 µs −4.9%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv32768 107.5 µs 93.8 µs −12.8%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv4096 18.5 µs 18.9 µs +2.2%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv32768 116.7 µs 103.5 µs −11.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv4096 15.0 µs 15.7 µs +5.0%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv8192 19.4 µs 17.6 µs −9.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv16384 34.2 µs 31.7 µs −7.5%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv32768 70.7 µs 59.4 µs −15.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv16384 45.3 µs 44.9 µs −0.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv32768 76.7 µs 69.7 µs −9.0%
At long context, the stable implementation is 5–16% faster than phi=5. Short-context results range from approximately neutral to 5–7% slower, with more run-to-run noise.

pytorch-bot · 2026-08-25T04:06:47Z

🔗 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 238a723 with merge base 469debd ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-08-25T04:07:29Z

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 · 2026-08-25T05:57:38Z

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.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

A fixed softmax offset is only valid while every attention score remains within the fp32 exponent range. Voxtral can exceed that range, producing infinities and NaNs. Sufficiently negative scores can instead underflow every partial weight and silently return zeros. Both split-K paths have this risk. Use each split's observed maximum and rescale against the global maximum. Each split processes only a few KV tiles, so online softmax removes the model-specific score assumption without materially extending the split-kernel recurrence. Reduce all partials together so cross-split normalization does not become a serial bottleneck. This change was authored with Codex.

Copilot

Pull request overview

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

Gasoonjia

LGTM, it would be great if we can have a perf comparsion between static phi version and online softmax

Gasoonjia · 2026-08-26T16:53:02Z

 (filled) positions (O(context) instead of O(max_seq_len)). Read
 on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+ phi is deprecated, accepted for operator-schema compatibility, and ignored.


maybe we no longer need to mention phi here; there's no phi in the op

will do it later, don't want to run CI again for this :p

digantdesai · 2026-08-27T14:47:32Z

phi is deprecated, accepted for operator-schema compatibility, and ignored.

Updated the PR summary.

digantdesai · 2026-08-25T04:06:44Z

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.

 ### A100 performance vs fixed phi=5 (quick and dirty benchmarking)
Measured the complete graph-captured split-K SDPA operator, not full-model latency. Results are two-run averages using PyTorch 2.13 and Triton 3.7.1.
Shape phi=5 this PR Change
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━
Voxtral, D128, Lq1, Lkv4096 16.9 µs 16.6 µs −1.7%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv8192 28.2 µs 26.8 µs −4.9%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv32768 107.5 µs 93.8 µs −12.8%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv4096 18.5 µs 18.9 µs +2.2%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv32768 116.7 µs 103.5 µs −11.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv4096 15.0 µs 15.7 µs +5.0%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv8192 19.4 µs 17.6 µs −9.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv16384 34.2 µs 31.7 µs −7.5%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv32768 70.7 µs 59.4 µs −15.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv16384 45.3 µs 44.9 µs −0.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv32768 76.7 µs 69.7 µs −9.0%
At long context, the stable implementation is 5–16% faster than phi=5. Short-context results range from approximately neutral to 5–7% slower, with more run-to-run noise.

pytorch-bot · 2026-08-25T04:06:47Z

🔗 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 238a723 with merge base 469debd ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-08-25T04:07:29Z

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 · 2026-08-25T05:57:38Z

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.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

A fixed softmax offset is only valid while every attention score remains within the fp32 exponent range. Voxtral can exceed that range, producing infinities and NaNs. Sufficiently negative scores can instead underflow every partial weight and silently return zeros. Both split-K paths have this risk. Use each split's observed maximum and rescale against the global maximum. Each split processes only a few KV tiles, so online softmax removes the model-specific score assumption without materially extending the split-kernel recurrence. Reduce all partials together so cross-split normalization does not become a serial bottleneck. This change was authored with Codex.

Copilot

Pull request overview

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

Gasoonjia

LGTM, it would be great if we can have a perf comparsion between static phi version and online softmax

Gasoonjia · 2026-08-26T16:53:02Z

 (filled) positions (O(context) instead of O(max_seq_len)). Read
 on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+ phi is deprecated, accepted for operator-schema compatibility, and ignored.


maybe we no longer need to mention phi here; there's no phi in the op

will do it later, don't want to run CI again for this :p

digantdesai · 2026-08-27T14:47:32Z

phi is deprecated, accepted for operator-schema compatibility, and ignored.

Updated the PR summary.

digantdesai · 2026-08-25T04:06:44Z

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.

 ### A100 performance vs fixed phi=5 (quick and dirty benchmarking)
Measured the complete graph-captured split-K SDPA operator, not full-model latency. Results are two-run averages using PyTorch 2.13 and Triton 3.7.1.
Shape phi=5 this PR Change
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━
Voxtral, D128, Lq1, Lkv4096 16.9 µs 16.6 µs −1.7%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv8192 28.2 µs 26.8 µs −4.9%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv32768 107.5 µs 93.8 µs −12.8%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv4096 18.5 µs 18.9 µs +2.2%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv32768 116.7 µs 103.5 µs −11.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv4096 15.0 µs 15.7 µs +5.0%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv8192 19.4 µs 17.6 µs −9.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv16384 34.2 µs 31.7 µs −7.5%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv32768 70.7 µs 59.4 µs −15.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv16384 45.3 µs 44.9 µs −0.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv32768 76.7 µs 69.7 µs −9.0%
At long context, the stable implementation is 5–16% faster than phi=5. Short-context results range from approximately neutral to 5–7% slower, with more run-to-run noise.

pytorch-bot · 2026-08-25T04:06:47Z

🔗 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 238a723 with merge base 469debd ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-08-25T04:07:29Z

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 · 2026-08-25T05:57:38Z

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.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

A fixed softmax offset is only valid while every attention score remains within the fp32 exponent range. Voxtral can exceed that range, producing infinities and NaNs. Sufficiently negative scores can instead underflow every partial weight and silently return zeros. Both split-K paths have this risk. Use each split's observed maximum and rescale against the global maximum. Each split processes only a few KV tiles, so online softmax removes the model-specific score assumption without materially extending the split-kernel recurrence. Reduce all partials together so cross-split normalization does not become a serial bottleneck. This change was authored with Codex.

Copilot

Pull request overview

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

Gasoonjia

LGTM, it would be great if we can have a perf comparsion between static phi version and online softmax

Gasoonjia · 2026-08-26T16:53:02Z

 (filled) positions (O(context) instead of O(max_seq_len)). Read
 on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+ phi is deprecated, accepted for operator-schema compatibility, and ignored.


maybe we no longer need to mention phi here; there's no phi in the op

will do it later, don't want to run CI again for this :p

digantdesai · 2026-08-27T14:47:32Z

phi is deprecated, accepted for operator-schema compatibility, and ignored.

Updated the PR summary.

digantdesai · 2026-08-25T04:06:44Z

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.

 ### A100 performance vs fixed phi=5 (quick and dirty benchmarking)
Measured the complete graph-captured split-K SDPA operator, not full-model latency. Results are two-run averages using PyTorch 2.13 and Triton 3.7.1.
Shape phi=5 this PR Change
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━
Voxtral, D128, Lq1, Lkv4096 16.9 µs 16.6 µs −1.7%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv8192 28.2 µs 26.8 µs −4.9%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv32768 107.5 µs 93.8 µs −12.8%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv4096 18.5 µs 18.9 µs +2.2%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv32768 116.7 µs 103.5 µs −11.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv4096 15.0 µs 15.7 µs +5.0%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv8192 19.4 µs 17.6 µs −9.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv16384 34.2 µs 31.7 µs −7.5%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv32768 70.7 µs 59.4 µs −15.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv16384 45.3 µs 44.9 µs −0.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv32768 76.7 µs 69.7 µs −9.0%
At long context, the stable implementation is 5–16% faster than phi=5. Short-context results range from approximately neutral to 5–7% slower, with more run-to-run noise.

pytorch-bot · 2026-08-25T04:06:47Z

🔗 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 238a723 with merge base 469debd ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-08-25T04:07:29Z

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 · 2026-08-25T05:57:38Z

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.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

A fixed softmax offset is only valid while every attention score remains within the fp32 exponent range. Voxtral can exceed that range, producing infinities and NaNs. Sufficiently negative scores can instead underflow every partial weight and silently return zeros. Both split-K paths have this risk. Use each split's observed maximum and rescale against the global maximum. Each split processes only a few KV tiles, so online softmax removes the model-specific score assumption without materially extending the split-kernel recurrence. Reduce all partials together so cross-split normalization does not become a serial bottleneck. This change was authored with Codex.

Copilot

Pull request overview

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

Gasoonjia

LGTM, it would be great if we can have a perf comparsion between static phi version and online softmax

Gasoonjia · 2026-08-26T16:53:02Z

 (filled) positions (O(context) instead of O(max_seq_len)). Read
 on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+ phi is deprecated, accepted for operator-schema compatibility, and ignored.


maybe we no longer need to mention phi here; there's no phi in the op

will do it later, don't want to run CI again for this :p

digantdesai · 2026-08-27T14:47:32Z

phi is deprecated, accepted for operator-schema compatibility, and ignored.

Updated the PR summary.

digantdesai · 2026-08-25T04:06:44Z

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.

 ### A100 performance vs fixed phi=5 (quick and dirty benchmarking)
Measured the complete graph-captured split-K SDPA operator, not full-model latency. Results are two-run averages using PyTorch 2.13 and Triton 3.7.1.
Shape phi=5 this PR Change
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━
Voxtral, D128, Lq1, Lkv4096 16.9 µs 16.6 µs −1.7%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv8192 28.2 µs 26.8 µs −4.9%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv32768 107.5 µs 93.8 µs −12.8%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv4096 18.5 µs 18.9 µs +2.2%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv32768 116.7 µs 103.5 µs −11.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv4096 15.0 µs 15.7 µs +5.0%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv8192 19.4 µs 17.6 µs −9.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv16384 34.2 µs 31.7 µs −7.5%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv32768 70.7 µs 59.4 µs −15.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv16384 45.3 µs 44.9 µs −0.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv32768 76.7 µs 69.7 µs −9.0%
At long context, the stable implementation is 5–16% faster than phi=5. Short-context results range from approximately neutral to 5–7% slower, with more run-to-run noise.

pytorch-bot · 2026-08-25T04:06:47Z

🔗 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 238a723 with merge base 469debd ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-08-25T04:07:29Z

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 · 2026-08-25T05:57:38Z

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.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

A fixed softmax offset is only valid while every attention score remains within the fp32 exponent range. Voxtral can exceed that range, producing infinities and NaNs. Sufficiently negative scores can instead underflow every partial weight and silently return zeros. Both split-K paths have this risk. Use each split's observed maximum and rescale against the global maximum. Each split processes only a few KV tiles, so online softmax removes the model-specific score assumption without materially extending the split-kernel recurrence. Reduce all partials together so cross-split normalization does not become a serial bottleneck. This change was authored with Codex.

Copilot

Pull request overview

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

Gasoonjia

LGTM, it would be great if we can have a perf comparsion between static phi version and online softmax

Gasoonjia · 2026-08-26T16:53:02Z

 (filled) positions (O(context) instead of O(max_seq_len)). Read
 on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+ phi is deprecated, accepted for operator-schema compatibility, and ignored.


maybe we no longer need to mention phi here; there's no phi in the op

will do it later, don't want to run CI again for this :p

digantdesai · 2026-08-27T14:47:32Z

phi is deprecated, accepted for operator-schema compatibility, and ignored.

Updated the PR summary.

digantdesai · 2026-08-25T04:06:44Z

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.

 ### A100 performance vs fixed phi=5 (quick and dirty benchmarking)
Measured the complete graph-captured split-K SDPA operator, not full-model latency. Results are two-run averages using PyTorch 2.13 and Triton 3.7.1.
Shape phi=5 this PR Change
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━
Voxtral, D128, Lq1, Lkv4096 16.9 µs 16.6 µs −1.7%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv8192 28.2 µs 26.8 µs −4.9%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv32768 107.5 µs 93.8 µs −12.8%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv4096 18.5 µs 18.9 µs +2.2%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv32768 116.7 µs 103.5 µs −11.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv4096 15.0 µs 15.7 µs +5.0%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv8192 19.4 µs 17.6 µs −9.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv16384 34.2 µs 31.7 µs −7.5%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv32768 70.7 µs 59.4 µs −15.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv16384 45.3 µs 44.9 µs −0.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv32768 76.7 µs 69.7 µs −9.0%
At long context, the stable implementation is 5–16% faster than phi=5. Short-context results range from approximately neutral to 5–7% slower, with more run-to-run noise.

pytorch-bot · 2026-08-25T04:06:47Z

🔗 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 238a723 with merge base 469debd ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-08-25T04:07:29Z

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 · 2026-08-25T05:57:38Z

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.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

A fixed softmax offset is only valid while every attention score remains within the fp32 exponent range. Voxtral can exceed that range, producing infinities and NaNs. Sufficiently negative scores can instead underflow every partial weight and silently return zeros. Both split-K paths have this risk. Use each split's observed maximum and rescale against the global maximum. Each split processes only a few KV tiles, so online softmax removes the model-specific score assumption without materially extending the split-kernel recurrence. Reduce all partials together so cross-split normalization does not become a serial bottleneck. This change was authored with Codex.

Copilot

Pull request overview

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

Gasoonjia

LGTM, it would be great if we can have a perf comparsion between static phi version and online softmax

Gasoonjia · 2026-08-26T16:53:02Z

 (filled) positions (O(context) instead of O(max_seq_len)). Read
 on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+ phi is deprecated, accepted for operator-schema compatibility, and ignored.


maybe we no longer need to mention phi here; there's no phi in the op

will do it later, don't want to run CI again for this :p

digantdesai · 2026-08-27T14:47:32Z

phi is deprecated, accepted for operator-schema compatibility, and ignored.

Updated the PR summary.

digantdesai · 2026-08-25T04:06:44Z

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.

 ### A100 performance vs fixed phi=5 (quick and dirty benchmarking)
Measured the complete graph-captured split-K SDPA operator, not full-model latency. Results are two-run averages using PyTorch 2.13 and Triton 3.7.1.
Shape phi=5 this PR Change
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━
Voxtral, D128, Lq1, Lkv4096 16.9 µs 16.6 µs −1.7%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv8192 28.2 µs 26.8 µs −4.9%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv32768 107.5 µs 93.8 µs −12.8%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv4096 18.5 µs 18.9 µs +2.2%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv32768 116.7 µs 103.5 µs −11.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv4096 15.0 µs 15.7 µs +5.0%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv8192 19.4 µs 17.6 µs −9.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv16384 34.2 µs 31.7 µs −7.5%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv32768 70.7 µs 59.4 µs −15.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv16384 45.3 µs 44.9 µs −0.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv32768 76.7 µs 69.7 µs −9.0%
At long context, the stable implementation is 5–16% faster than phi=5. Short-context results range from approximately neutral to 5–7% slower, with more run-to-run noise.

pytorch-bot · 2026-08-25T04:06:47Z

🔗 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 238a723 with merge base 469debd ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-08-25T04:07:29Z

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 · 2026-08-25T05:57:38Z

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.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

A fixed softmax offset is only valid while every attention score remains within the fp32 exponent range. Voxtral can exceed that range, producing infinities and NaNs. Sufficiently negative scores can instead underflow every partial weight and silently return zeros. Both split-K paths have this risk. Use each split's observed maximum and rescale against the global maximum. Each split processes only a few KV tiles, so online softmax removes the model-specific score assumption without materially extending the split-kernel recurrence. Reduce all partials together so cross-split normalization does not become a serial bottleneck. This change was authored with Codex.

Copilot

Pull request overview

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

Gasoonjia

LGTM, it would be great if we can have a perf comparsion between static phi version and online softmax

Gasoonjia · 2026-08-26T16:53:02Z

 (filled) positions (O(context) instead of O(max_seq_len)). Read
 on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+ phi is deprecated, accepted for operator-schema compatibility, and ignored.


maybe we no longer need to mention phi here; there's no phi in the op

will do it later, don't want to run CI again for this :p

digantdesai · 2026-08-27T14:47:32Z

phi is deprecated, accepted for operator-schema compatibility, and ignored.

Updated the PR summary.

digantdesai · 2026-08-25T04:06:44Z

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.

 ### A100 performance vs fixed phi=5 (quick and dirty benchmarking)
Measured the complete graph-captured split-K SDPA operator, not full-model latency. Results are two-run averages using PyTorch 2.13 and Triton 3.7.1.
Shape phi=5 this PR Change
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━
Voxtral, D128, Lq1, Lkv4096 16.9 µs 16.6 µs −1.7%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv8192 28.2 µs 26.8 µs −4.9%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq1, Lkv32768 107.5 µs 93.8 µs −12.8%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv4096 18.5 µs 18.9 µs +2.2%
────────────────────────────── ────────── ─────────────────── ────────
Voxtral, D128, Lq4, Lkv32768 116.7 µs 103.5 µs −11.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv4096 15.0 µs 15.7 µs +5.0%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv8192 19.4 µs 17.6 µs −9.4%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv16384 34.2 µs 31.7 µs −7.5%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq1, Lkv32768 70.7 µs 59.4 µs −15.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv16384 45.3 µs 44.9 µs −0.9%
────────────────────────────── ────────── ─────────────────── ────────
Qwen3.5, D256, Lq4, Lkv32768 76.7 µs 69.7 µs −9.0%
At long context, the stable implementation is 5–16% faster than phi=5. Short-context results range from approximately neutral to 5–7% slower, with more run-to-run noise.

pytorch-bot · 2026-08-25T04:06:47Z

🔗 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 238a723 with merge base 469debd ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-08-25T04:07:29Z

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 · 2026-08-25T05:57:38Z

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.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

A fixed softmax offset is only valid while every attention score remains within the fp32 exponent range. Voxtral can exceed that range, producing infinities and NaNs. Sufficiently negative scores can instead underflow every partial weight and silently return zeros. Both split-K paths have this risk. Use each split's observed maximum and rescale against the global maximum. Each split processes only a few KV tiles, so online softmax removes the model-specific score assumption without materially extending the split-kernel recurrence. Reduce all partials together so cross-split normalization does not become a serial bottleneck. This change was authored with Codex.

Copilot

Pull request overview

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

Gasoonjia

LGTM, it would be great if we can have a perf comparsion between static phi version and online softmax

Gasoonjia · 2026-08-26T16:53:02Z

 (filled) positions (O(context) instead of O(max_seq_len)). Read
 on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+ phi is deprecated, accepted for operator-schema compatibility, and ignored.


maybe we no longer need to mention phi here; there's no phi in the op

will do it later, don't want to run CI again for this :p

digantdesai · 2026-08-27T14:47:32Z

phi is deprecated, accepted for operator-schema compatibility, and ignored.

Updated the PR summary.

digantdesai requested review from Gasoonjia and a lite review from Copilot August 25, 2026 04:06

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 25, 2026

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 04:07 View session

CopilotAI reviewed Aug 25, 2026
View reviewed changes

shoumikhin mentioned this pull request Aug 25, 2026
Enable fast ROCm Voxtral decode #22005
Merged

CopilotAI review requested due to automatic review settings August 25, 2026 21:45

digantdesai force-pushed the split_k_nan branch from 2dd38cf to 02f7a00Compare August 25, 2026 21:45

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 21:46 View session

digantdesai temporarily deployed to cadence August 25, 2026 21:47 — with GitHub Actions Inactive

CopilotAI reviewed Aug 25, 2026
View reviewed changes

Comment threadbackends/cuda/triton/kernels/sdpa.py

digantdesai temporarily deployed to cadence August 25, 2026 22:46 — with GitHub Actions Inactive

CopilotAI review requested due to automatic review settings August 26, 2026 14:56

digantdesai force-pushed the split_k_nan branch from 02f7a00 to 238a723Compare August 26, 2026 14:56

digantdesai temporarily deployed to cadence August 26, 2026 14:56 — with GitHub Actions Inactive

Copilotstarted reviewing on behalf of digantdesai August 26, 2026 14:57 View session

CopilotAI reviewed Aug 26, 2026
View reviewed changes

digantdesai temporarily deployed to cadence August 26, 2026 15:27 — with GitHub Actions Inactive

Gasoonjia approved these changes Aug 26, 2026
View reviewed changes

Gasoonjia reviewed Aug 26, 2026
View reviewed changes

digantdesai merged commit d58fc25 into mainAug 27, 2026
259 checks passed

digantdesai deleted the split_k_nan branch August 27, 2026 14:47

digantdesai requested review from Gasoonjia and a lite review from Copilot August 25, 2026 04:06

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 25, 2026

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 04:07 View session

CopilotAI reviewed Aug 25, 2026
View reviewed changes

shoumikhin mentioned this pull request Aug 25, 2026
Enable fast ROCm Voxtral decode #22005
Merged

CopilotAI review requested due to automatic review settings August 25, 2026 21:45

digantdesai force-pushed the split_k_nan branch from 2dd38cf to 02f7a00Compare August 25, 2026 21:45

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 21:46 View session

digantdesai temporarily deployed to cadence August 25, 2026 21:47 — with GitHub Actions Inactive

CopilotAI reviewed Aug 25, 2026
View reviewed changes

Comment threadbackends/cuda/triton/kernels/sdpa.py

digantdesai temporarily deployed to cadence August 25, 2026 22:46 — with GitHub Actions Inactive

CopilotAI review requested due to automatic review settings August 26, 2026 14:56

digantdesai force-pushed the split_k_nan branch from 02f7a00 to 238a723Compare August 26, 2026 14:56

digantdesai temporarily deployed to cadence August 26, 2026 14:56 — with GitHub Actions Inactive

Copilotstarted reviewing on behalf of digantdesai August 26, 2026 14:57 View session

CopilotAI reviewed Aug 26, 2026
View reviewed changes

digantdesai temporarily deployed to cadence August 26, 2026 15:27 — with GitHub Actions Inactive

Gasoonjia approved these changes Aug 26, 2026
View reviewed changes

Gasoonjia reviewed Aug 26, 2026
View reviewed changes

digantdesai merged commit d58fc25 into mainAug 27, 2026
259 checks passed

digantdesai deleted the split_k_nan branch August 27, 2026 14:47

digantdesai requested review from Gasoonjia and a lite review from Copilot August 25, 2026 04:06

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 25, 2026

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 04:07 View session

CopilotAI reviewed Aug 25, 2026
View reviewed changes

shoumikhin mentioned this pull request Aug 25, 2026
Enable fast ROCm Voxtral decode #22005
Merged

CopilotAI review requested due to automatic review settings August 25, 2026 21:45

digantdesai force-pushed the split_k_nan branch from 2dd38cf to 02f7a00Compare August 25, 2026 21:45

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 21:46 View session

digantdesai temporarily deployed to cadence August 25, 2026 21:47 — with GitHub Actions Inactive

CopilotAI reviewed Aug 25, 2026
View reviewed changes

Comment threadbackends/cuda/triton/kernels/sdpa.py

digantdesai temporarily deployed to cadence August 25, 2026 22:46 — with GitHub Actions Inactive

CopilotAI review requested due to automatic review settings August 26, 2026 14:56

digantdesai force-pushed the split_k_nan branch from 02f7a00 to 238a723Compare August 26, 2026 14:56

digantdesai temporarily deployed to cadence August 26, 2026 14:56 — with GitHub Actions Inactive

Copilotstarted reviewing on behalf of digantdesai August 26, 2026 14:57 View session

CopilotAI reviewed Aug 26, 2026
View reviewed changes

digantdesai temporarily deployed to cadence August 26, 2026 15:27 — with GitHub Actions Inactive

Gasoonjia approved these changes Aug 26, 2026
View reviewed changes

Gasoonjia reviewed Aug 26, 2026
View reviewed changes

digantdesai merged commit d58fc25 into mainAug 27, 2026
259 checks passed

digantdesai deleted the split_k_nan branch August 27, 2026 14:47

digantdesai requested review from Gasoonjia and a lite review from Copilot August 25, 2026 04:06

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 25, 2026

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 04:07 View session

CopilotAI reviewed Aug 25, 2026
View reviewed changes

shoumikhin mentioned this pull request Aug 25, 2026
Enable fast ROCm Voxtral decode #22005
Merged

CopilotAI review requested due to automatic review settings August 25, 2026 21:45

digantdesai force-pushed the split_k_nan branch from 2dd38cf to 02f7a00Compare August 25, 2026 21:45

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 21:46 View session

digantdesai temporarily deployed to cadence August 25, 2026 21:47 — with GitHub Actions Inactive

CopilotAI reviewed Aug 25, 2026
View reviewed changes

Comment threadbackends/cuda/triton/kernels/sdpa.py

digantdesai temporarily deployed to cadence August 25, 2026 22:46 — with GitHub Actions Inactive

CopilotAI review requested due to automatic review settings August 26, 2026 14:56

digantdesai force-pushed the split_k_nan branch from 02f7a00 to 238a723Compare August 26, 2026 14:56

digantdesai temporarily deployed to cadence August 26, 2026 14:56 — with GitHub Actions Inactive

Copilotstarted reviewing on behalf of digantdesai August 26, 2026 14:57 View session

CopilotAI reviewed Aug 26, 2026
View reviewed changes

digantdesai temporarily deployed to cadence August 26, 2026 15:27 — with GitHub Actions Inactive

Gasoonjia approved these changes Aug 26, 2026
View reviewed changes

Gasoonjia reviewed Aug 26, 2026
View reviewed changes

digantdesai merged commit d58fc25 into mainAug 27, 2026
259 checks passed

digantdesai deleted the split_k_nan branch August 27, 2026 14:47

digantdesai requested review from Gasoonjia and a lite review from Copilot August 25, 2026 04:06

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 25, 2026

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 04:07 View session

CopilotAI reviewed Aug 25, 2026
View reviewed changes

shoumikhin mentioned this pull request Aug 25, 2026
Enable fast ROCm Voxtral decode #22005
Merged

CopilotAI review requested due to automatic review settings August 25, 2026 21:45

digantdesai force-pushed the split_k_nan branch from 2dd38cf to 02f7a00Compare August 25, 2026 21:45

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 21:46 View session

digantdesai temporarily deployed to cadence August 25, 2026 21:47 — with GitHub Actions Inactive

CopilotAI reviewed Aug 25, 2026
View reviewed changes

Comment threadbackends/cuda/triton/kernels/sdpa.py

digantdesai temporarily deployed to cadence August 25, 2026 22:46 — with GitHub Actions Inactive

CopilotAI review requested due to automatic review settings August 26, 2026 14:56

digantdesai force-pushed the split_k_nan branch from 02f7a00 to 238a723Compare August 26, 2026 14:56

digantdesai temporarily deployed to cadence August 26, 2026 14:56 — with GitHub Actions Inactive

Copilotstarted reviewing on behalf of digantdesai August 26, 2026 14:57 View session

CopilotAI reviewed Aug 26, 2026
View reviewed changes

digantdesai temporarily deployed to cadence August 26, 2026 15:27 — with GitHub Actions Inactive

Gasoonjia approved these changes Aug 26, 2026
View reviewed changes

Gasoonjia reviewed Aug 26, 2026
View reviewed changes

digantdesai merged commit d58fc25 into mainAug 27, 2026
259 checks passed

digantdesai deleted the split_k_nan branch August 27, 2026 14:47

digantdesai requested review from Gasoonjia and a lite review from Copilot August 25, 2026 04:06

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 25, 2026

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 04:07 View session

CopilotAI reviewed Aug 25, 2026
View reviewed changes

shoumikhin mentioned this pull request Aug 25, 2026
Enable fast ROCm Voxtral decode #22005
Merged

CopilotAI review requested due to automatic review settings August 25, 2026 21:45

digantdesai force-pushed the split_k_nan branch from 2dd38cf to 02f7a00Compare August 25, 2026 21:45

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 21:46 View session

digantdesai temporarily deployed to cadence August 25, 2026 21:47 — with GitHub Actions Inactive

CopilotAI reviewed Aug 25, 2026
View reviewed changes

Comment threadbackends/cuda/triton/kernels/sdpa.py

digantdesai temporarily deployed to cadence August 25, 2026 22:46 — with GitHub Actions Inactive

CopilotAI review requested due to automatic review settings August 26, 2026 14:56

digantdesai force-pushed the split_k_nan branch from 02f7a00 to 238a723Compare August 26, 2026 14:56

digantdesai temporarily deployed to cadence August 26, 2026 14:56 — with GitHub Actions Inactive

Copilotstarted reviewing on behalf of digantdesai August 26, 2026 14:57 View session

CopilotAI reviewed Aug 26, 2026
View reviewed changes

digantdesai temporarily deployed to cadence August 26, 2026 15:27 — with GitHub Actions Inactive

Gasoonjia approved these changes Aug 26, 2026
View reviewed changes

Gasoonjia reviewed Aug 26, 2026
View reviewed changes

digantdesai merged commit d58fc25 into mainAug 27, 2026
259 checks passed

digantdesai deleted the split_k_nan branch August 27, 2026 14:47

digantdesai requested review from Gasoonjia and a lite review from Copilot August 25, 2026 04:06

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 25, 2026

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 04:07 View session

CopilotAI reviewed Aug 25, 2026
View reviewed changes

shoumikhin mentioned this pull request Aug 25, 2026
Enable fast ROCm Voxtral decode #22005
Merged

CopilotAI review requested due to automatic review settings August 25, 2026 21:45

digantdesai force-pushed the split_k_nan branch from 2dd38cf to 02f7a00Compare August 25, 2026 21:45

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 21:46 View session

digantdesai temporarily deployed to cadence August 25, 2026 21:47 — with GitHub Actions Inactive

CopilotAI reviewed Aug 25, 2026
View reviewed changes

Comment threadbackends/cuda/triton/kernels/sdpa.py

digantdesai temporarily deployed to cadence August 25, 2026 22:46 — with GitHub Actions Inactive

CopilotAI review requested due to automatic review settings August 26, 2026 14:56

digantdesai force-pushed the split_k_nan branch from 02f7a00 to 238a723Compare August 26, 2026 14:56

digantdesai temporarily deployed to cadence August 26, 2026 14:56 — with GitHub Actions Inactive

Copilotstarted reviewing on behalf of digantdesai August 26, 2026 14:57 View session

CopilotAI reviewed Aug 26, 2026
View reviewed changes

digantdesai temporarily deployed to cadence August 26, 2026 15:27 — with GitHub Actions Inactive

Gasoonjia approved these changes Aug 26, 2026
View reviewed changes

Gasoonjia reviewed Aug 26, 2026
View reviewed changes

digantdesai merged commit d58fc25 into mainAug 27, 2026
259 checks passed

digantdesai deleted the split_k_nan branch August 27, 2026 14:47

digantdesai requested review from Gasoonjia and a lite review from Copilot August 25, 2026 04:06

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 25, 2026

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 04:07 View session

CopilotAI reviewed Aug 25, 2026
View reviewed changes

shoumikhin mentioned this pull request Aug 25, 2026
Enable fast ROCm Voxtral decode #22005
Merged

CopilotAI review requested due to automatic review settings August 25, 2026 21:45

digantdesai force-pushed the split_k_nan branch from 2dd38cf to 02f7a00Compare August 25, 2026 21:45

Copilotstarted reviewing on behalf of digantdesai August 25, 2026 21:46 View session

digantdesai temporarily deployed to cadence August 25, 2026 21:47 — with GitHub Actions Inactive

CopilotAI reviewed Aug 25, 2026
View reviewed changes

Comment threadbackends/cuda/triton/kernels/sdpa.py

digantdesai temporarily deployed to cadence August 25, 2026 22:46 — with GitHub Actions Inactive

CopilotAI review requested due to automatic review settings August 26, 2026 14:56

digantdesai force-pushed the split_k_nan branch from 02f7a00 to 238a723Compare August 26, 2026 14:56

digantdesai temporarily deployed to cadence August 26, 2026 14:56 — with GitHub Actions Inactive

Copilotstarted reviewing on behalf of digantdesai August 26, 2026 14:57 View session

CopilotAI reviewed Aug 26, 2026
View reviewed changes

digantdesai temporarily deployed to cadence August 26, 2026 15:27 — with GitHub Actions Inactive

Gasoonjia approved these changes Aug 26, 2026
View reviewed changes

Gasoonjia reviewed Aug 26, 2026
View reviewed changes

digantdesai merged commit d58fc25 into mainAug 27, 2026
259 checks passed

digantdesai deleted the split_k_nan branch August 27, 2026 14:47

Uh oh!

Conversation

digantdesai commented Aug 25, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-botBot commented Aug 25, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

github-actionsBot commented Aug 25, 2026

This PR needs a release notes: label

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

shoumikhin commented Aug 25, 2026

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Gasoonjia left a comment

Choose a reason for hiding this comment

Uh oh!

GasoonjiaAug 26, 2026

Choose a reason for hiding this comment

Uh oh!

digantdesaiAug 27, 2026

Choose a reason for hiding this comment

Uh oh!

digantdesai commented Aug 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

digantdesai commented Aug 25, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-botBot commented Aug 25, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

github-actionsBot commented Aug 25, 2026

This PR needs a release notes: label

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

shoumikhin commented Aug 25, 2026

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

digantdesai commented Aug 25, 2026•
edited
Loading

pytorch-botBot commented Aug 25, 2026•
edited
Loading

This PR needs a `release notes:` label

digantdesai commented Aug 25, 2026•
edited
Loading

pytorch-botBot commented Aug 25, 2026•
edited
Loading

This PR needs a `release notes:` label

digantdesai commented Aug 25, 2026•
edited
Loading

pytorch-botBot commented Aug 25, 2026•
edited
Loading

This PR needs a `release notes:` label