Add paged / block-table KV cache export (--paged-cache) - #395
Draft
justinchuby wants to merge 2 commits into
Draft
Add paged / block-table KV cache export (--paged-cache)#395justinchuby wants to merge 2 commits into
justinchuby wants to merge 2 commits into
Conversation
Add a paged (block-table) KV cache variant of the static cache, emitting models whose attention reads KV from non-contiguous pages via a page pool + block_table + slot_mapping — the vLLM PagedAttention layout, which also expresses SGLang RadixAttention (shared prefix pages) with no graph change. Implements onnx-genai DESIGN §39.4 Option C using only standard ONNX ops. Attention writes new K/V into the per-layer page pool with ScatterND, gathers the sequence's physical pages contiguously with Gather(pool, block_table), then runs the opset-24 Attention op with nonpad_kv_seqlen (input #6) — the same op contract as --static-cache, but over paged KV. - PagedCacheState + _apply_paged_attention in components/_attention.py - PagedCacheState dispatch in DecoderLayer / MoEDecoderLayer (backward compatible: paged_cache only forwarded to self_attn when set) - CausalLMTask(paged_cache=True, page_size=, num_pages=) with _make_paged_cache_inputs / _register_paged_cache_outputs - CLI flags --paged-cache / --page-size / --num-pages with validation - Graph-level tests (TestBuildPagedCacheGraph) + CPU numerical parity for the scatter/gather paging ops incl. a RadixAttention shared-page case I/O contract (batch == 1): inputs: input_ids, position_ids, key_pool.{i}/value_pool.{i} [num_pages, page_size, kv_hidden], block_table [num_blocks] i64, slot_mapping [seq_len] i64, nonpad_kv_seqlen [batch] i64 outputs: logits, updated_key_pool.{i}/updated_value_pool.{i} Multi-sequence batching is a documented TODO. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Performance Comparison
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new causal-LM export mode that emits a paged / block-table KV cache graph (vLLM PagedAttention / SGLang RadixAttention layout) driven by block_table + slot_mapping, using standard ONNX ScatterND + Gather + opset-24 Attention with nonpad_kv_seqlen.
Changes:
- Introduces
PagedCacheStateand paged-cache attention path (ScatterNDwrite +Gatherassemble + opset-24Attention) in the attention component. - Extends
CausalLMTaskand the CLI withpaged_cache/--paged-cachepluspage_size/--page-sizeandnum_pages/--num-pages, including validation and IO wiring. - Adds graph-level build tests and CPU parity tests for the paging subgraph.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/paged_cache_test.py | New CPU parity tests for ScatterND+Gather paging ops and a task build validity check. |
| tests/build_graph_test.py | New TestBuildPagedCacheGraph to assert IO contract and key ops presence for paged-cache builds. |
| src/mobius/tasks/_causal_lm.py | Adds paged_cache option to CausalLMTask, paged-cache IO creation, and output registration. |
| src/mobius/models/moe.py | Extends MoE decoder layer dispatch to route PagedCacheState into attention. |
| src/mobius/components/_decoder.py | Extends decoder layer dispatch to route PagedCacheState into attention. |
| src/mobius/components/_attention.py | Implements PagedCacheState and paged cache attention body; wires into Attention.forward. |
| src/mobius/main.py | Adds --paged-cache, --page-size, --num-pages CLI args and validation. |
| CHANGELOG.md | Documents the new paged/block-table cache export feature and CLI flags. |
Comment on lines
+207
to
+212
| import sys | ||
| from pathlib import Path | ||
|
|
||
| sys.path.insert(0, str(Path(__file__).parent)) | ||
| from _test_configs import _base_config | ||
| from onnx_ir.passes.common import CheckerPass |
Comment on lines
+518
to
+533
| if paged_cache is not None: | ||
| # Paged (block-table) KV cache: scatter new K/V into the physical | ||
| # page pool, gather this sequence's pages contiguously, then attend. | ||
| # present_* here are the UPDATED page pools (registered as outputs). | ||
| attn_output, present_key, present_value = _apply_paged_attention( | ||
| op, | ||
| query_states, | ||
| key_states, | ||
| value_states, | ||
| paged_cache, | ||
| num_attention_heads=self.num_attention_heads, | ||
| num_key_value_heads=self.num_key_value_heads, | ||
| head_dim=self.head_dim, | ||
| scale=self.scaling, | ||
| softcap=self._softcap, | ||
| ) |
Comment on lines
+568
to
+571
| help="Use a paged / block-table KV cache (vLLM PagedAttention / SGLang " | ||
| "RadixAttention layout): a page pool + block_table + slot_mapping, with " | ||
| "GatherElements-style page assembly and ScatterND writes (onnx-genai " | ||
| "DESIGN §39.4 Option C). Requires DecoderLayer or MoEDecoderLayer models.", |
Comment on lines
+161
to
+163
| # Paged cache reuses the same DecoderLayer/MoEDecoderLayer dispatch | ||
| # as the static cache (both flow their state through past_key_value). | ||
| _validate_static_cache_support(module) |
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.
What this adds
A paged / block-table KV cache export mode for causal LMs — a variant of the
existing
--static-cachethat emits models whose attention reads KV fromnon-contiguous pages via a page pool +
block_table+slot_mapping.This is the vLLM PagedAttention layout, and because sequences can list the
same physical page in their
block_table, the identical graph also expressesSGLang RadixAttention (shared-prefix pages) with no change. It implements
onnx-genai
docs/DESIGN.md§39.4 Option C ("ONNX Scatter/GatherElements inGraph") using only standard ONNX ops (no custom op).
Why
The onnx-genai runtime needs models that operate directly on a paged KV pool so
it can do vLLM-style paged attention and SGLang-style prefix sharing without
gathering pages into a contiguous buffer in the runtime on every step.
How it works (per attention layer)
Step 3 is the same op contract as
--static-cache(Attentionwithis_causal=1+nonpad_kv_seqleninput #6, noattention_mask), so it reusesthe existing opset-24 retention logic (
_graph_requires_opset24already detectsAttentionwith a non-empty input #6) and the same DecoderLayer/MoEDecoderLayerdispatch pattern as the static cache.
Ops used:
Reshape,Shape,Unsqueeze,ScatterND(write),Gather(page assemble),
Attention. All opset-24 compatible.New flag
Also available programmatically:
CausalLMTask(paged_cache=True, page_size=16, num_pages=None).Mutually exclusive with
--static-cache; requiresDecoderLayer/MoEDecoderLayermodels.ONNX I/O contract (what onnx-genai must drive)
Single active sequence (
batch == 1), per layeri,kv_hidden = num_kv_heads * head_dim:Inputs
input_ids[batch, seq_len]position_ids[batch, seq_len]key_pool.{i}/value_pool.{i}[num_pages, page_size, kv_hidden]block_table[num_blocks]slot_mapping[seq_len]page_id*page_size + offsetper new token)nonpad_kv_seqlen[batch]write_start + valid_token_count)Outputs
logits[batch, seq_len, vocab]updated_key_pool.{i}/updated_value_pool.{i}[num_pages, page_size, kv_hidden]No
attention_mask— causal + padding masking is derived fromis_causal=1+nonpad_kv_seqlen. RoPE is baked into the keys before they are written to thepool. RadixAttention: point multiple sequences'
block_tableat the samephysical page — the graph gathers it identically for each.
Complete vs TODO
Complete
components/_attention.pyPagedCacheStatedispatch inDecoderLayerandMoEDecoderLayer(backward-compatible:paged_cacheonly forwarded to customself_attnmodules when set)CausalLMTask(paged_cache=…)+ input/output wiring, mutual-exclusion with static cache--paged-cache/--page-size/--num-pageswith validationonnx.checkerpasses, shape inference yields correct pool shapesTODO (documented in code)
block_table/ per-row gather); current impl targetsbatch == 1Attentionop withnonpad_kv_seqlenis CUDA-only until onnxruntime#28958 ships (same constraint as--static-cache), so full-model run parity is deferred to the shared static-cache Flash probe. The paging ops themselves are validated on CPU.Validation
tests/build_graph_test.py::TestBuildPagedCacheGraph— 12 graph-level tests(inputs/outputs, pool shapes, dynamic
num_pages,ScatterND/Gather/Attentionpresence,
is_causal=1,nonpad_kv_seqlenwiring, MoE, mutual exclusion, shape inference).tests/paged_cache_test.py— CPU numerical parity for the scatter+gather pagingops vs a NumPy reference, including a RadixAttention shared-page case, plus an
end-to-end
onnx.checker+ shape-inference build check.TestBuildGraph/TestBuildStaticCacheGraph/cli_testsuites pass (no regressions).