Implement FlashDecoding++ async softmax for split-K SDPA

[Gasoonjia]

Replace online softmax (per-tile max tracking + cross-split rescaling) with a unified maximum value (phi=5.0) approach from FlashDecoding++.

Key changes:

• Split kernel: subtract fixed phi instead of tracking running max m_i, eliminating alpha rescaling between tiles
• Reduce kernel: simple summation of partial outputs instead of max-aware weighted combination; removes M_partial buffer
• ~12.9% average kernel-level speedup (6.8%-20.1% range) by saving HBM bandwidth (no M_partial reads/writes) and reducing ALU ops

The unified phi works because exp(qk - phi) is numerically stable for typical attention score ranges, and the fixed constant allows all splits to compute independently without synchronization.

Also used KernelAgent(https://github.com/meta-pytorch/KernelAgent) to further optimized the kernel.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 6 New Failures, 1 Pending, 4 Unrelated Failures

As of commit 1a79d9d with merge base c48ea12 ():

NEW FAILURES - The following jobs have failed:

• Cadence Build & Test / cpu-test / test-aot / test-aot (gh)
  backends/cadence/aot/tests/test_replace_ops_passes.py::TestReplaceOpsPasses::test_replace_conv2d_with_linear
• Test CUDA Builds / unittest-cuda / linux-job (gh)
  examples/models/qwen3_5_moe/test_turboquant.py::TestTurboQuant::test_kv_cache_state_matters
• Test CUDA Windows Export and E2E / export-model-cuda-windows-artifact (facebook, dinov2-small-imagenet1k-1-layer, non-quantized) / linux-job (gh)
  RuntimeError: Command docker exec -t 506200b9a24ca4734272fda19ce92d8d9796949b25fe548827d3e0951c8ef257 /exec failed with exit code 1
• Test CUDA Windows Export and E2E / export-model-cuda-windows-artifact (nvidia, parakeet-tdt, non-quantized) / linux-job (gh)
  RuntimeError: Command docker exec -t 5787cc7a7e33240a38d2327ca60eb1268e14e22e29dcaaec825c79aeafa78cf1 /exec failed with exit code 1
• Test CUDA Windows Export and E2E / export-model-cuda-windows-artifact (nvidia, parakeet-tdt, quantized-int4-weight-only) / linux-job (gh)
  RuntimeError: Command docker exec -t 13bf3f36575557e48f932a31332ea5fe8cc34db6190d752865c7b10c0f7bb852 /exec failed with exit code 1
• trunk / unittest-release / linux / linux-job (gh)
  backends/xnnpack/test/ops/test_conv2d.py::TestConv2d::test_qs8_conv2d_relu_multi_users

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / unittest / windows / windows-job (gh) (matched win rule in flaky-rules.json)
  ##[error]The operation was canceled.
• trunk / unittest-release / windows / windows-job (gh) (matched win rule in flaky-rules.json)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• trunk / test-models-macos-coreml (edsr) / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

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

[github-actions]

This PR needs a release notes: label

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

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

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

[digantdesai]

TIL FlashDecoding++, thanks.

Couple of things, seems like this is still online softmax (i.e you need to do reduce later). And can you add performance numbers with a sweep in the PR summary, thanks.

Also ask claude to do a quality comparisons using various prompts just to spotcheck quality.

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 1m 17s —— View job

---

Review of PR #18867

• Gather context and read changed files
• Analyze the diff and implementation
• Post review feedback

---

This PR makes several simplifying changes to the split-K decode kernel. Here's my analysis:

Summary of Changes

1. Reverted exp2/log2-space optimization back to tl.exp — the previous version pre-scaled Q by sm_scale * log2(e) and used tl.math.exp2 (single PTX ex2 instruction). This PR reverts to tl.exp with per-tile qk * sm_scale multiply.

2. Removed BATCH_ONE constexpr specialization — previously avoided a division for B=1 decode; now always computes b = pid_bh // H_kv.

3. Removed _compute_num_splits heuristic — replaced the SM-aware occupancy heuristic (which queried get_device_properties) with a simpler min(max(cdiv(L_kv, 256), 1), 128).

4. Added _splitk_buf_cache — caches O_partial and L_partial buffers to avoid repeated CUDA allocations. Previous version allocated fresh each call (and used torch.zeros for L_partial vs. torch.empty now).

5. Removed num_stages=2 from reduce kernel loop — changed from tl.range(0, num_splits, num_stages=2) to tl.range(0, num_splits).

6. Removed two autotune configs from the split kernel (256/8/3 and 128/8/3).

Issues & Questions

1. L_partial changed from torch.zeros to torch.empty — potential correctness issue

In the previous version, L_partial was initialized with torch.zeros. Now it uses torch.empty (via the buffer cache). This is safe only if every split fully writes its entry before the reduce kernel reads it. However, when L_kv is not evenly divisible by num_splits, the last split may process fewer than chunk_size tokens, but it still writes its l_i (initialized to zero in the kernel via tl.zeros). So this should be correct — but what about splits where start_n >= Lk? If num_splits * chunk_size > Lk, some splits will have start_n >= end_n (since end_n = tl.minimum(start_n + chunk_size, Lk)), meaning the loop body never executes. In that case, l_i stays at tl.zeros (0.0) and acc stays at tl.zeros, and both get stored — writing zeros. So this is fine. But the comment should acknowledge that correctness relies on the kernel writing zeros for empty splits.

2. tl.exp vs tl.math.exp2 regression

The previous version used tl.math.exp2 which maps to a single PTX ex2.approx.f32 instruction. Reverting to tl.exp (which uses ex2 internally but with a multiply by log2(e)) adds ALU work. The PR description says "KernelAgent" was used to further optimize — but this particular change goes in the opposite direction. Was this an intentional revert for readability/simplicity, or an oversight? The exp2 path with pre-scaled Q is strictly fewer FLOPs.

3. _compute_num_splits removal trades adaptivity for simplicity

The old heuristic targeted 2 full SM waves and adapted to the actual GPU (A100 vs RTX 4090 have different SM counts). The new cdiv(L_kv, 256) is simpler but can over-split on small L_kv or under-utilize on GPUs with many SMs. For example, with L_kv=512 you get 2 splits regardless of GPU, whereas the old heuristic would target more splits on a 128-SM GPU. Was there a measurable regression from the old heuristic that motivated this change?

4. BATCH_ONE removal

This was a minor optimization (avoiding an integer division in the kernel). Since B=1 is extremely common in decode, this may have had measurable benefit. Was it benchmarked before removal?

5. Reduce kernel num_stages change

Removing num_stages=2 from the reduce loop means no software pipelining. For large num_splits (up to 128), this could regress the reduce kernel. Was this benchmarked?

Minor Observations

• The PR title says "Implement FlashDecoding++ async softmax" but the async-softmax (fixed phi) approach was already present in the previous version. This commit is more accurately described as "simplify/tune the FlashDecoding++ split-K kernel" — reverting some micro-optimizations and adding buffer caching.
• The buffer cache _splitk_buf_cache is a module-level dict that grows unboundedly. In long-running inference servers with varying shapes, this could leak memory. Consider using an LRU cache or at least bounding the cache size.

Overall

The core FlashDecoding++ logic (fixed phi, simple sum in reduce) is correct and unchanged from the prior commit. The main changes are simplification/tuning tradeoffs. The buffer caching is a good optimization for decode latency. The exp2 -> exp revert and _compute_num_splits removal warrant explanation — if the benchmarks in the PR description already account for these changes, then the net speedup speaks for itself, but it would be good to confirm.

[Gasoonjia]

@digantdesai here we replace the online softmax with async softmax by using a unified phi.