ONNX Attention opset-24 external KV cache: Flash split-KV sized by buffer length, not valid KV length (CUDA)

[titaiwangms]

Summary

The ONNX Attention opset-24 external-KV-cache decode path (nonpad_kv_seqlen + in-place TensorScatter, CUDA EP) sizes its Flash split-KV launch from the KV cache buffer length, not from the valid KV length. For a fixed valid length, enlarging the cache buffer makes the same Flash kernel dramatically slower — worst exactly for long-max-context deployments, where a large pre-allocated buffer is the whole point.

This was surfaced by the baseline decode benchmark in #28352 (script #29684).

Evidence (RTX PRO 6000, SM120, fp16, S_past=2048, per decode step)

Shrinking only the cache buffer (--max-seq-len) at a fixed valid length collapses the identical flash_splitkv kernel:

attn_scatter config	flash_splitkv GPU µs
buffer 8192	34.1
buffer 2560	13.4

For comparison, attn_past (opset-23 past/present) — whose K/V shape is the valid length — runs the same-tier Flash at 13.2 µs. Cache layout is ruled out (4D BNSH cache: same ~34 µs). Raw artifacts (CSV / nsys): https://github.com/namgyu-youn/onnxruntime/releases/tag/bench-28352-artifacts

Root cause

onnxruntime/core/providers/cuda/llm/attention.cc:257-260 computes num_splits from parameters.total_sequence_length:

auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes] =
    onnxruntime::flash::get_num_splits_and_buffer_sizes(
        parameters.batch_size, parameters.sequence_length,
        parameters.total_sequence_length, parameters.q_num_heads, ...);

In the nonpad_kv_seqlen path, K/V are the full cache, so total_sequence_length = buffer length. The valid length only exists as the nonpad_kv_seqlen device tensor, which isn't consumed until LaunchConvertNonpadKvSeqlenToFlashSeqlensK at attention.cc:311-320 — after the split count is already fixed. So the split-KV launch (and its partitioning) is sized for the whole padded buffer and does wasted work over the padding region.

By contrast, contrib GQA sizes its split from a CPU-input valid length (onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc:595), so it never over-sizes.

Impact

• The "TensorScatter → near-parity with GQA" expectation in the ONNX Attention CUDA: Performance Optimization Opportunities vs Contrib GQA #28352 gap table does not hold today: measured +66% e2e at buffer 8192, and it gets worse as the buffer grows.
• This is architecture-independent host-side sizing logic, not an SM120 artifact.
• It's the single largest term separating the ONNX Attention external-cache decode path from GQA.

Proposed fix (for evaluation)

Size the Flash split from the valid KV length instead of the buffer shape. The valid length is device-side (nonpad_kv_seqlen), which host-side split sizing can't read — mirror GQA and accept the valid length (or a max-valid-length hint) as a CPU input to the opset-24 path, then feed it into get_num_splits_and_buffer_sizes.

Open questions for the reviewing team:

• Cleanest way to plumb a CPU-side length/hint through the opset-24 Attention schema + kernel without regressing the general (non-decode) path.
• Whether a conservative "max valid length" hint is enough (avoids a device→host sync) vs. needing the exact per-batch length.
• Correctness interaction with the existing device-side seqlens_k masking (the split sizing is a perf-only knob; the mask still bounds the math).

Where this sits in the bigger picture

This is step 1 of "close the ONNX-Attention-vs-GQA decode gap" from #28352. It only brings attn_scatter up to the Flash-tier attn_past level (~13-20 µs GPU); it does not by itself reach GQA's XQA/cuDNN decode tier (~10-11 µs GPU). The remaining gap needs (2) fused KV-append prep to cut launch overhead (see the existing TODO at attention.cc:346), and (3) wiring ONNX Attention decode onto a decode-specialized kernel tier — ORT already has a cuDNN SDPA integration used by GQA. Filing this one separately because it's the smallest, most self-contained, highest-ROI item.

References

• Benchmark write-up: ONNX Attention CUDA: Performance Optimization Opportunities vs Contrib GQA #28352 (comment)
• Benchmark script: Add decode benchmark comparing ONNX Attention CUDA with contrib GQA #29684
• attention.cc:257-260 (split sizing), :311-320 (nonpad→seqlens_k conversion)
• group_query_attention.cc:595 (GQA CPU-input length)

[titaiwangms]

Evaluation: won't fix for now (backlog)

After a design evaluation we've decided not to pursue this for now. Recording the conclusion and a technical correction for anyone who picks it up later.

An important technical correction

The original description implies over-partitioning wastes attention compute. Reading the flash split-KV kernel confirms this is not the case:

• The flash split-KV kernel already early-exits any split that falls entirely in the padding region (onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h:490): when n_block_min >= n_block_max, it writes -inf to gLSEaccum / 0 to gOaccum as placeholders and returns — doing no QK / gmem load.
• The valid length is already available device-side (binfo.actual_seqlen_k, from device seqlens_k), and n_block_max is already clamped to it (:485).

So sizing num_splits from the buffer length over-partitions, but the waste is not attention FLOPs — it's (a) extra threadblock launches plus -inf writes to the gmem accumulator, (b) the combine kernel reducing over num_splits lanes (padding splits are -inf, contributing 0 after exp, but the load/exp/compare still runs), and (c) suboptimal occupancy/reuse when each valid split only covers 1–4 blocks. The upside is therefore smaller than originally estimated.

Scope of benefit: eager-only

The benchmark (#28352) was measured with enable_cuda_graph=False. Under CUDA Graph, num_splits determines gridDim and is frozen at capture; replay never re-enters the host heuristic, so dynamically choosing splits from the valid length cannot help under graph. Production decode often uses CUDA Graph, which weakens the practical value of this optimization.

Trade-off across the three options

Option	Verdict
A. Add an early-exit in the kernel	Already exists (:490) — no-op
B. Tweak n_blocks_per_split to redistribute blocks by valid length	Limited benefit: removes the empty-split gmem writes, but num_splits / grid launch / combine lane count are all set by the host-chosen num_splits, which the kernel can't change
C. Have the host pick a smaller num_splits from valid length	On-target, but num_splits must be decided host-side before launch → runs back into the valid-vs-buffer input/device problem (D2H sync breaks graphs; H2D workaround adds launch cost; a second input changes the ONNX standard-op spec). This is a launch-configuration problem, not a kernel problem
D. Persistent / device-scheduled kernel	The only approach that fixes it including under graph and is graph-compatible (fixed grid, valid read purely device-side, no sync) — i.e. the FlashInfer / Flash-Decoding v2 approach. But this is a standalone kernel project, maintaining an ORT-owned decode-scheduled flash variant, far beyond the scope of this issue

Why we're not fixing it

• The small fixes (B/C) have limited benefit and graph-compatibility risk;
• The large fix (D) is out of proportion to the actual benefit;
• And GQA's decode lead comes mainly from decode-specialized tiers (XQA / cuDNN SDPA), not from split precision — even with split sizing fixed, ONNX Attention (dispatch is only Flash→MEA→Unfused) still won't close the fundamental gap.

Keeping this analysis here as a starting point should anyone later build a device-scheduled decode kernel (option D). The associated PR #29689 has been closed.