Size Flash split-KV from valid KV length in ONNX Attention opset-24 external-cache path (CUDA)

[Copilot]

Description

In the opset-24 external-KV-cache decode path (nonpad_kv_seqlen, CUDA EP), RunFlashAttention sized the Flash split-KV num_splits from parameters.total_sequence_length. On this path K/V are the full cache, so total_sequence_length is the buffer length, not the valid KV length — over-partitioning the split-KV launch and doing wasted work over the padding region. The same kernel gets slower as the cache buffer grows, exactly the opposite of what long-max-context deployments want.

• Split sizing (onnxruntime/core/providers/cuda/llm/attention.cc): read the per-batch max valid length from the device-side nonpad_kv_seqlen onto the host (one batch_size-element copy) and feed it into get_num_splits_and_buffer_sizes. Clamped to the buffer length and floored at 1.
• Scope: only the nonpad_kv_seqlen Flash path; the non-decode paths keep total_sequence_length sizing.
• Correctness invariant: the kernel's seqlen_k stays = total_sequence_length — it drives the KV-cache strides (k_batch_stride = seqlen_k*h*d) and n_blocks_per_split. Device-side seqlens_k masking still bounds the math per batch, so num_splits is a perf-only knob and outputs are unchanged.
• Tests (test_tensorscatter_attention.py): added CUDA fp16 large-buffer / small-valid cases (buffer 2048, valid 64/96/128) to exercise the new sizing branch across num_splits values.

int split_kv_seqlen = parameters.total_sequence_length;
if (nonpad_kv_seqlen != nullptr && parameters.batch_size > 0) {
  // DtoH copy of nonpad_kv_seqlen -> host max, clamped to buffer, floored at 1
  split_kv_seqlen = /* max valid KV length across batch */;
}
auto [num_splits, ...] = get_num_splits_and_buffer_sizes(
    batch_size, q_sequence_length, split_kv_seqlen, q_num_heads, head_size, num_SMs);

Tradeoff: reading the valid length host-side costs one small device→host copy + stream sync, scoped to the external-cache Flash path. Valid lengths change every decode step, so cross-call host caching does not apply. A nonpad_kv_seqlen-as-CPU-input design (mirroring GQA) would avoid the sync but is a larger, cross-cutting change touching all three dispatch paths and input memory placement — left for reviewer input per the issue's open questions.

Motivation and Context

The nonpad_kv_seqlen + TensorScatter decode path is meant to reach near-parity with contrib GQA, but buffer-sized split launches make it measurably slower (~+66% e2e at buffer 8192, worse as the buffer grows). This is architecture-independent host-side sizing logic, and the single largest term separating the ONNX Attention external-cache decode path from GQA. It brings attn_scatter up to the Flash-tier attn_past level; the remaining GQA decode-tier gap (fused KV-append prep, decode-specialized kernel) is out of scope here.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

This PR optimizes the CUDA EP implementation of the ONNX-domain Attention op (opset 24) in the external KV-cache decode path by sizing FlashAttention split-KV work (num_splits) using the valid KV length (from nonpad_kv_seqlen) rather than the buffer length (total_sequence_length). This avoids wasted work over padded cache regions as the preallocated cache grows.

Changes:

• Read back nonpad_kv_seqlen (device → host) to compute a per-batch max valid KV length and use it to size Flash split-KV launches.
• Keep seqlen_k unchanged (= buffer length) to preserve KV stride correctness while adjusting only the split heuristic.
• Add CUDA fp16 test cases with large cache buffers but small valid lengths to exercise the new sizing branch.

Reviewed changes

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

File	Description
onnxruntime/core/providers/cuda/llm/attention.cc	Size FlashAttention split-KV from max valid KV length (from nonpad_kv_seqlen) instead of cache buffer length.
onnxruntime/test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py	Add fp16 CUDA tests for large-buffer/small-valid external-cache decode scenarios to validate correctness in the new sizing regime.

[titaiwangms]

Review synthesis (multi-model review team)

Thanks for the fix — it targets the right root cause from #29686, and the numerical correctness is solid (see below). But there's one Critical integration problem that should block merge as-is: the per-step host sync breaks CUDA Graph capture.

Critical — the unconditional cudaStreamSynchronize breaks CUDA Graph capture

attention.cc (new block, ~L275-280) does a device→host copy of nonpad_kv_seqlen + cudaStreamSynchronize(cuda_stream) on every external-cache Flash call, with no capture guard.

• CUDA forbids synchronizing a stream that is being captured — the call returns cudaErrorStreamCaptureUnsupported.
• Capture begins in CUDAExecutionProvider::OnRunStart (cuda_execution_provider.cc:497-503, CaptureBegin), i.e. before any op's ComputeInternal. So under enable_cuda_graph=True, the whole model — including this op — runs on the capturing stream, and this sync aborts capture.
• Decode with static shapes is the primary CUDA Graph use case, and GenAI pipelines wrap the decoder in enable_cuda_graph. This op (opset-24 Attention) has no graph-capture exclusion, so those models would fail at capture time.
• Regression vs GQA: contrib GQA keeps its decode length as a CPU input (group_query_attention.cc:64,600) and never syncs per step; its only cudaStreamSynchronize (group_query_attention.cc:188) is one-time init (head-sink conversion, cached). So GQA stays graph-capturable and this PR gives that property up.
• Note also: under CUDA Graph replay the launch config is baked at capture, so a per-step-varying num_splits couldn't take effect during replay anyway — another reason the dynamic sizing must fall back under capture.

Major — per-step host sync serializes the decode hot path

Even without CUDA graphs, cudaStreamSynchronize forces all prior GPU work to finish before the host can compute num_splits and enqueue Attention, killing CPU/GPU overlap every token. The InlinedVector destination is pageable, so the "async" copy gives no real overlap before the immediate sync. For launch-bound B=1 decode (~20-30 µs host floor) this can offset or exceed the ~20 µs Flash-kernel saving the PR is chasing. Recommend measuring end-to-end token latency (PR vs base) before merge.
[needs-run: per-token DtoH sync offsets the split-KV kernel gain for B=1 decode; repro=benchmark_onnx_attention_vs_gqa.py attn_scatter fp16 B=1 H=32/8 D=128 max-seq-len 8192, valid 64/128/256, nsys, PR vs base; expect=host/API sync gap ≈ or > the Flash time saved; cost=expensive]

Recommended fix

The clean, precedent-matching fix (also the direction noted in #29686) is to make nonpad_kv_seqlen a CPU input for the opset-24 path — register it with .InputMemoryType(OrtMemTypeCPUInput, 6) like GQA does for its length input. That removes both the graph break and the host stall, and matches the established ORT contract.

If a schema/memory-placement change is too large for this PR, the minimum band-aid is to gate the sync on capture status and fall back to buffer-length sizing during capture:

cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(cuda_stream, &capture_status));
if (capture_status == cudaStreamCaptureStatusNone && nonpad_kv_seqlen != nullptr && batch_size > 0) {
  // DtoH copy + sync + max-over-batch sizing
} // else: split_kv_seqlen = parameters.total_sequence_length

This stops the crash, but keeps the Major host-stall cost on the non-capture path and forgoes the optimization under capture — so the CPU-input approach is preferred.

Confirmed correct (deep invariant review)

The load-bearing claim — num_splits is a perf-only knob; outputs cannot change — was traced end-to-end (attention.cc → mha_fwd_kvcache → set_params_fprop → split-KV kernel → combine) and holds:

• mha_fwd_kvcache still receives kv_sequence_length (= total_sequence_length = buffer) as seqlen_k, so k_batch_stride and n_blocks_per_split stay buffer-sized (attention.cc:368, flash_api.cc:84,115).
• split_kv_seqlen only feeds the split heuristic + accum buffer sizes; for any num_splits ≥ 1, num_splits * ceil(B/num_splits) ≥ B, so the split ranges still cover every valid block exactly once — coverage is independent of num_splits (flash_fwd_kernel.h:481-510). Empty/over-partitioned splits write -inf/0 and are absorbed by the combine reduction's all--inf guard.
• Ragged batch (max-over-batch length) and the empty-cache boundary (max_valid=0 → floored to 1) are both safe.

So the sizing math and the clamp/floor are correct; the clamp/floor comment and the "perf-only knob" comment are accurate and genuinely helpful.

Minor / test coverage

• The new tests assert output parity only, but this is a perf-only change — they would still pass if the old buffer-length sizing were reintroduced, so they don't actually protect the fix (test_tensorscatter_attention.py:570-606). Consider a unit assertion on the split heuristic (e.g. extend attention_split_heuristic_test.cc to check num_splits from valid vs buffer length), plus a CUDA-graph capture/replay regression test for this path.
• With q_seq_len=1, head_size=64, valid 64/96/128, the heuristic likely collapses to num_splits=1 on typical SM counts, so both split/no-split branches may not be exercised. Add a case with a valid length large enough to yield num_splits>1.
• The class is gated on has_cuda_device(53), but the changed path is Flash-only; gate on has_flash_attention() (or assert Flash usage) so coverage is deterministic.

Nits

• split_kv_seqlen sits two lines above num_splits and reads like "number of splits"; kv_seqlen_for_split_sizing / effective_kv_seqlen_for_splits would disambiguate (attention.cc:264).
• The ~15-line comment mixes bug rationale + fix + clamp rationale; splitting the clamp note onto its one line would ease skimming.

Follow-up (out of scope for this PR)

The measurement note in benchmark_onnx_attention_vs_gqa.py:26 (from #29684) now says the host "sizes the Flash split-KV launch from the buffer length" — this PR makes that stale. Worth a one-line update in that file's next touch.

---

GPU-dependent items are labeled [needs-run] — no CUDA-built ORT in the review environment. The Critical/CUDA-graph claim is grounded in source (cuda_execution_provider.cc:497-503, CUDA capture semantics); the correctness invariant is a static proof over the flash kernel source.

[titaiwangms]

Evaluation: not adopting, closing

Thanks for the submission. After a design evaluation we've decided not to take this PR, for the following reasons:

1. Benefit is eager-only. The benchmark cited in ONNX Attention opset-24 external KV cache: Flash split-KV sized by buffer length, not valid KV length (CUDA) #29686 was measured with enable_cuda_graph=False. Production decode typically runs under CUDA Graph, where num_splits is frozen at capture time — the per-step dynamic split sizing cannot pay off there.

2. This approach breaks CUDA Graph capture (Critical). The PR adds an unconditional per-step cudaMemcpyAsync(DtoH) + cudaStreamSynchronize in RunFlashAttention. CUDA forbids synchronizing a stream that is being captured (cudaErrorStreamCaptureUnsupported), and attention.cc has no IsGraphCaptureEnabled gate — so graph capture fails, a regression versus the current (capturable) behavior.

3. The waste ceiling is smaller than it looks. The flash split-KV kernel already early-exits any split that falls entirely in the padding region (flash_fwd_kernel.h:490: when n_block_min >= n_block_max it writes -inf/0 placeholders and returns, doing no QK / gmem load). The valid length is also already available device-side (binfo.actual_seqlen_k). So over-partitioning does not waste attention FLOPs — only extra threadblock launches, combine-reduction lanes, and suboptimal occupancy. The upside is limited.

4. The correct fix has poor ROI. Fixing this properly requires the host to pick a smaller num_splits from the valid length (it determines gridDim, so it must be decided host-side before launch). That runs back into the valid-vs-buffer input/device problem (a D2H sync breaks graphs; an H2D workaround adds launch cost; a second input would change the ONNX standard-op spec) — not worth it.

Full analysis is captured in #29686. Closing this PR.