Repository navigation
[CUDA] Enable non-causal PagedAttention with small cache pages - #32619
Merged
Tianlei Wu (tianleiwu) merged 2 commits intoSep 21, 2026
Merged
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Tianlei Wu (tianleiwu)
marked this pull request as ready for review
September 15, 2026 23:38
Tianlei Wu (tianleiwu)
changed the title
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
The changes affect several specialized CUDA attention kernels and require hardware-backed validation across backend configurations.
Pull request overview
Enables non-causal CUDA PagedAttention across MEA, paged decode, XQA, and latent-attention backends.
Changes:
- Propagates causal mode through CUDA kernels and adjusts masking.
- Supports non-causal left-local windows in CUTLASS FMHA.
- Adds backend, page-size, ragged-batch, and masking regression coverage.
No actionable findings identified.
File summaries
| File | Description |
|---|---|
test_paged_attention.py |
Adds non-causal backend regression tests. |
paged_attention.cc |
Removes Flash-only non-causal dispatch restriction. |
paged_attention_impl.cu |
Implements backend-specific non-causal masking. |
kernel_forward.h |
Supports non-causal left-window masking. |
fmha_launch_template.h |
Configures CUTLASS masking offsets. |
paged_attention.md |
Documents CUDA non-causal support. |
Review details
- Files reviewed: 6/6 changed files
- Comments generated: 0
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Tianlei Wu (tianleiwu)
force-pushed
the
copilot/support-causal-0-in-paged-attention
branch
from
September 17, 2026 07:41
f5c8354 to
6a3edab
Compare
Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>
Tianlei Wu (tianleiwu)
force-pushed
the
copilot/support-causal-0-in-paged-attention
branch
from
September 17, 2026 07:42
6a3edab to
b269161
Compare
Tianlei Wu (tianleiwu)
requested review from
Akshay Sonawane (apsonawane),
Baiju Meswani (baijumeswani) and
kunal-vaishnavi
September 17, 2026 18:42
Baiju Meswani (baijumeswani)
left a comment
Baiju Meswani (baijumeswani)
left a comment
Contributor
There was a problem hiding this comment.
I found one blocking dispatch issue. The kernel changes and test coverage otherwise look good.
Tianlei Wu (tianleiwu)
added a commit
to microsoft/onnxruntime-genai
that referenced
this pull request
Sep 17, 2026
### Description The builder required `paged_block_size` to be a multiple of 256, documented as "required by the ONNX Runtime PagedAttention CUDA kernel". That is not the op's contract, and the rule is wrong in both directions. `CheckInputs` in [`paged_attention_helper.h`](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h) requires a **power of two of at least 16**. The multiple-of-256 figure comes from one FlashAttention tile size; `paged_attention.cc` treats that as a backend-selection input, not a validity check, and falls back to PagedDecode or MemoryEfficientAttention when the page is smaller. A quantized KV cache is exempt from the alignment check alone (`flash_block_size_ok = kIsQuantizedCache || block_size % tile == 0`); it still reaches FlashAttention, through a dense dequantized path that has no page to align. So the old rule: * **rejected valid values.** 16, 32, 64 and 128 all build and run. * **accepted an invalid one.** 768 passes `% 256 == 0`, exports without complaint, and then fails on the first forward pass: `block_size must be a power of two and at least 16. Got block_size == 768`. (`CheckInputs` runs from `ComputeInternal`, not at load.) This replaces it with the op's actual contract and corrects the documentation. ### Motivation and Context Smaller pages cut the internal fragmentation of the KV pool, which matters most for short sequences and high batch counts, and a quantized cache pays nothing for them. They were simply unavailable. ### Testing New `test/python/builder/test_paged_block_size.py`: 16 through 2048, the string form Olive and the CLI pass, sub-16, non-powers-of-two (including 768), and non-positive values. 17 cases. `test_precision.py` asserted that 128 was rejected; it now asserts that 768 is, which is the case the old rule got wrong. Full builder suite: 717 passed, 3 skipped. Measured end to end on Qwen3.8-27B (H200, bf16, unquantized KV, `head_size` 256, DFlash2 drafter), 200 MMLU-Pro questions at concurrency 8, with `num_blocks` scaled to hold the KV pool at 262144 tokens: | `paged_block_size` | tok/s | drafts/round | output vs 256 | | --- | --- | --- | --- | | 64 | 161.0 | 0.00 | 84/200 identical | | 128 | 374.6 | 7.00 | 200/200 identical | | 256 (default) | 381.9 | 7.00 | reference | | 512 | 369.8 | 7.00 | 200/200 identical | 128, 256 and 512 are byte-identical and their throughput spread is inside this harness's noise, so the default stays the right choice for this model. ### A note on the 64 column, and why it is not in the code The 64 row is not a block-size effect. On today's `main`, PagedAttention serves `is_causal=0` from FlashAttention only, and the DFlash2 drafter has `head_size` 128, so a 64-token page makes the drafter's node fail and the engine falls back to target-only decoding for the rest of the run: ``` PagedAttention: is_causal=0 requires the FlashAttention backend (sm>=80, fp16/bf16, head_size 128, block_size 64). ``` microsoft/onnxruntime#32619 removes exactly that restriction: it gives `is_causal=0` non-causal MEA, paged decode, XQA and latent backends, and adds head-size-128 regression cases at 16/32/64-token pages. Once it lands, a small page stops being fatal for a block drafter and becomes what it already is for the target, a backend-selection input. Because of that, this PR deliberately keeps the runtime detail **out** of the source. The comments in `builder.py`, `block_drafter.py` and the README state only what stays true either way: the tile is a per-head-size throughput choice, not a validity rule, and a drafter usually has the smaller head size and therefore the larger tile, so the block that is fastest for the target is not necessarily the one that is fastest for the drafter. The `TODO` in `block_drafter.py` is scoped to that, giving the drafter its own block size instead of inheriting the target's, which is what a vLLM-style 16-token target page would need. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Baiju Meswani (baijumeswani)
approved these changes
Sep 17, 2026
Tianlei Wu (tianleiwu)
enabled auto-merge (squash)
September 17, 2026 21:52
Tianlei Wu (tianleiwu)
deleted the
copilot/support-causal-0-in-paged-attention
branch
September 21, 2026 00:57
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
is_causal=0; support non-causal MEA, paged decode, XQA speculative decoding, and latent attention.Motivation and Context
DFlash2’s non-causal drafter uses head size 128, for which FlashAttention requires pages of at least 128 tokens. A valid 64-token page therefore fails the drafter’s PagedAttention node and forces target-only decoding for the rest of the run.