Enable split-KV for paged FlashAttention decode

[Baiju Meswani (baijumeswani)]

Summary

EnableFlashAttention's existing split-KV path for paged decode.

Split-KV divides long KV sequences across multiple CUDA thread blocks and combines their partial results. This improves GPU utilization during low-batch, long-context decoding.

The optimization is limited to safe cases:

• decode with one new token per request;
• global attention without attention sinks;
• KV sequences longer than 512 tokens.

Short contexts, local-window attention, prefill, and mixed-length query batches retain their existing paths.

[Copilot]

Pull request overview

Enables FlashAttention split-KV parallelism for eligible paged decode workloads.

Changes:

• Adds split-KV workspace allocation and dispatch plumbing.
• Extends the FlashAttention varlen API with split accumulators.
• Adds heuristic and CUDA end-to-end coverage.

Reviewed changes

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

Show a summary per file

File	Description
attention_split_heuristic_test.cc	Tests split selection.
paged_attention_op_test.cc	Tests paged split-KV decode.
paged_attention.cc	Selects splits and allocates workspaces.
paged_attention_impl.cu	Passes split configuration to FlashAttention.
flash_api.h	Extends the varlen API.
flash_api.cc	Configures split-KV dispatch.
attention_data.h	Stores split workspace metadata.

---

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

[Copilot]

Pull request overview

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

Suppressed comments (1)

docs/contrib_ops/cuda/paged_attention.md:387

• This row is missing the > blockquote prefix used by the surrounding table, so Markdown renders it outside the table and splits the documented bounds table. Keep the row inside the blockquote.

| split-KV eligibility | FlashAttention decode dispatch | `min_max_kv_len_for_split`, else disabled unless exact lengths were read back |

[Tianlei Wu (tianleiwu)]

Reviewed the split-KV dispatch, accumulator sizing, packed decode strides, metadata compatibility, and CUDA graph replay behavior. The follow-up cleanly separates the replay-wide upper bound used for fixed workspace sizing from the lower bound used for split eligibility; I found no active CUDA correctness blocker. One documentation formatting issue is noted inline.

[Tianlei Wu (tianleiwu)]

No correctness blocker found in the current head. The replay-wide lower bound addresses the earlier short-context dispatch concern, and the varlen API guards the split-combine assumptions. I left two non-blocking coverage suggestions below; the provider-neutral schema concern is already tracked in the existing attention_metadata thread.

[Tianlei Wu (tianleiwu)]

Review of PR #32102: Enable split-KV for paged FlashAttention decode

PR: #32102

Reviewed head: dcc407b3ea3fed57dc83ca794b73299189930354

Overall assessment

The split/combine plumbing is consistent with FlashAttention's existing API, and the restrictions to pure decode, global attention, and no attention sinks are appropriate. The added ordinary, INT8-cache, and CUDA Graph tests cover the most important correctness paths. I did not find an active output-correctness blocker in the current head.

I recommend addressing the replay-range split partitioning issue below before merging because it can defeat the purpose of the optimization for a valid and realistic metadata range. The schema/session-option question is not a correctness blocker, but the workload bound and CUDA tuning policy should remain distinct concepts.

Findings

1. High: A loose replay upper bound can leave almost all selected splits empty

Locations:

• onnxruntime/contrib_ops/cuda/bert/paged_attention.cc:613-625
• onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h:481-485

PagedAttention enables splitting using min_max_kv_len_bound, but computes num_splits and passes max_kv_len using the replay-wide upper bound:

if (min_max_kv_len_bound > kFlashSplitKvMinSequenceLength) {
  get_num_splits_and_buffer_sizes(..., max_kv_len, ...);
}

The split kernel then partitions the N tiles using params.seqlen_k, which is that same upper bound:

const int n_blocks_per_split =
    ceil_div(ceil_div(params.seqlen_k, kBlockN), num_n_splits);

The live length is used only to cap n_block_max. Therefore, when the upper bound is much larger than the current live length, the useful tiles are concentrated in the first split(s); the remaining CTAs only write empty partials, and the combine kernel still reduces every configured split.

For example, with B=2, H=2, D=128, and 108 SMs, the current heuristic selects roughly:

KV length supplied to heuristic	Splits
513	5
2,048	16
32,768	24

For a valid replay range [513, 32768], params.seqlen_k=32768 gives 256 N tiles and 11 tiles per split. At the 513-token replay there are only 5 live tiles, so only split 0 performs useful attention work while all 24 partials are combined. The lower bound proves that at least one sequence is longer than 512; it does not prove that an upper-bound-derived partition provides parallelism at the lower end of the replay range.

Suggested fix: for varlen split-KV, derive n_blocks_per_split from binfo.actual_seqlen_k on device. The number of splits and workspace extents can remain replay-fixed, while the tile distribution adapts on each replay:

const int live_n_blocks = cute::ceil_div(binfo.actual_seqlen_k, kBlockN);
const int n_blocks_per_split = cute::ceil_div(live_n_blocks, num_n_splits);

This also adapts the partition to unequal sequence lengths within the batch. Please verify that this change is safe for the other callers of the shared split kernel, or gate it to the varlen/cumulative-length case.

At minimum, add a benchmark or focused test with a much wider valid replay interval. The existing CUDA Graph test uses [2048, 4096], where the upper bound is close enough to the live length that this behavior is easy to miss.

2. Medium: Keep workload metadata separate from CUDA tuning policy

Locations:

• onnxruntime/core/graph/contrib_ops/bert_defs.cc:1721-1740
• onnxruntime/contrib_ops/cuda/bert/paged_attention.cc:27,396-423,616

The third metadata value and the hard-coded threshold serve different roles:

• min_max_kv_len_bound is a replay/workload fact: a lower bound on the largest live per-sequence KV length.
• kFlashSplitKvMinSequenceLength = 512 is CUDA FlashAttention tuning policy.

I would not move the workload bound to a session option. A session option applies to every PagedAttention node and every capture in the session, whereas the valid lower bound can differ by node, request bucket, and CUDA Graph. A session-wide value would generally need to be so conservative that splitting remains disabled.

Keeping the optional third value in attention_metadata is reasonable now that its contract no longer names split-KV or CUDA. Consider renaming it to max_kv_len_lower_bound; that is easier to parse than min_max_kv_len_bound while preserving the same semantics.

If benchmarks later show a need for a user override of the CUDA threshold, expose the policy separately. The repository's current session-config convention suggests a name such as:

ep.cuda.paged_attention.split_kv_min_length

rather than cuda.op.paged_attention.min_max_kv_len_for_split. I would keep the threshold internal until there is evidence that a stable public override is needed; ideally it should be part of a device- and shape-aware heuristic rather than a universal constant.

3. Medium: The split correctness tests use layouts that are too regular

Location: onnxruntime/test/contrib_ops/paged_attention_op_test.cc:735-882

The new tests distinguish batches, which is good, but they use:

• an identity/contiguous block table;
• one KV head and uniform values within every head dimension;
• mostly equal or near-equal long sequence lengths.

Those inputs can allow page-table, KV-head, or head/dimension stride mistakes to produce the expected scalar-like output. Add one split-KV correctness case with a permuted/non-contiguous block table, at least two KV heads, and distinct values by KV head and dimension. Unequal live lengths should span different numbers of N tiles. This would more directly validate the packed varlen offsets and accumulator layout introduced by the PR.

4. Low: Add contract-validation coverage for the third metadata entry

Locations:

• onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h:571-581
• onnxruntime/contrib_ops/cuda/bert/paged_attention.cc:397-423

The implementation accepts shape (2) or (3), rejects negative entries, and rejects a lower bound greater than the effective upper bound. Add small validation tests for:

• shape (2) retaining the old behavior;
• shape (3) with a negative lower bound;
• lower bound greater than the KV upper bound;
• lower bound greater than static cache capacity when the upper bound is unknown.

These are inexpensive tests for a schema contract that producers will need to implement precisely.

What looks good

• Split-KV is limited to exactly one query token per sequence by the combination of token_count == batch_size and max_query_len == 1.
• Local-window attention and smooth-softmax/head-sink cases remain unsplit.
• The varlen API validates accumulator pointers and pure-decode assumptions before dispatch.
• The output batch stride required by the split-combine kernel is explicitly set for packed pure decode.
• Quantized-cache coverage exercises the gather/dequantize path.
• The CUDA Graph test updates device-resident lengths across replays and verifies outputs rather than only checking successful capture.
• Shape (2) metadata remains accepted, preserving the previous contract.

Validation performed

• Inspected the full diff from origin/main to PR head.
• Traced split selection, workspace sizing, varlen parameter setup, split tile partitioning, and combine output indexing.
• Reviewed the ordinary, INT8-cache, short-range, and CUDA Graph replay tests.
• Ran git diff --check origin/main...HEAD; it passed.
• Checked PR CI: Linux/Windows CUDA builds and tests, CUDA plugin EP builds/tests, minimal builds, lint, and other reported checks were passing at review time. One documentation-validation check was still pending.

No local CUDA test binary was available in this checkout, so no additional GPU execution was run locally.

[Baiju Meswani (baijumeswani)]

Updated the varlen path to divide the KV blocks across the fixed split count.

Renamed the input to max_kv_len_lower_bound to make the meaning clearer.

Also, updated tests.

[Baiju Meswani (baijumeswani)]

Thank you for the review.