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)]

It is cuda flash attention specified parameter so it might not fit other EPs.

There are a few options: (1) session option if the value is same for all paged attention node in a model. (2) run some offline experiment then use simple rules (like if...else...) to determine the value given input shapes.

[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)]

Suggestion: Since this optimization is designed around replay-wide bounds, could this test capture and replay a CUDA Graph while updating device-resident past_seqlens and cumulative_sequence_length at stable addresses? Varying unequal sequence lengths across replays and comparing each result against the reference would catch regressions that freeze capture-time lengths or scratch state; a single ordinary Run cannot.

[Tianlei Wu (tianleiwu)]

Suggestion: This changed launch also covers quantized caches after dense gather/dequantization, but the current split execution test uses only unquantized paged KV. Please add an INT8 or FP8 case with decoder attention disabled, GQA or unequal long lengths, and a verified NumSplits > 1 to exercise packed offsets/strides and the accumulator layout.