Add paged / block-table KV cache export (--features paged-cache) - #395
Closed
justinchuby wants to merge 5 commits into
Closed
Add paged / block-table KV cache export (--features paged-cache)#395justinchuby wants to merge 5 commits into
justinchuby wants to merge 5 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. |
# Conflicts: # src/mobius/__main__.py
- attention: reject paged_cache combined with a non-None attention_bias
instead of silently dropping it, so ALiBi / sliding-window models fail
fast rather than producing wrong results.
- _causal_lm: parametrize _validate_static_cache_support with a `mode`
label and pass "Paged cache" from the paged path so the raised message
names the layout the user actually enabled.
- __main__: correct the --paged-cache help text ("Gather-based page
assembly" — the graph uses ONNX Gather, not GatherElements).
- paged_cache_test: drop the redundant sys.path manipulation and import
_test_configs directly like the rest of the suite; add coverage for the
attention_bias fail-fast and the paged validation message wording.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
🟡 Not ready to approve
There are a few user-facing documentation/error-message inconsistencies and a test input-shape mismatch that should be corrected before merge.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
Suppressed comments (5)
src/mobius/tasks/_causal_lm.py:585
- _validate_static_cache_support() is now used for both static and paged cache validation, but the first docstring line still claims it only checks StaticCacheState. This is confusing when the caller passes mode="Paged cache".
def _validate_static_cache_support(module: nn.Module, mode: str = "Static cache") -> None:
"""Check that the module's decoder layers support StaticCacheState.
src/mobius/tasks/_causal_lm.py:652
- When validating paged cache support (mode="Paged cache"), the raised TypeError still tells users to "add StaticCacheState dispatch", which is misleading. Custom layers need to dispatch both StaticCacheState and PagedCacheState as appropriate.
raise TypeError(
f"{mode} mode requires decoder layers that "
f"inherit from DecoderLayer or MoEDecoderLayer (or set "
f"_supports_static_cache=True), but "
f"{name}[{i}] is {type(layer).__name__}. Either use a "
f"compatible model or add StaticCacheState dispatch to "
f"{type(layer).__name__}.forward()."
src/mobius/tasks/_causal_lm.py:80
- Paged-cache docstring says key_pool/value_pool and updated pools are FLOAT, but the implementation uses the model's dtype (config.dtype) for these tensors. This is misleading for fp16/bf16 exports and for callers building the expected I/O contract.
This issue also appears in the following locations of the same file:
- line 584
- line 646
- key_pool.{i}: [num_pages, page_size, kv_hidden] FLOAT per layer
- value_pool.{i}: [num_pages, page_size, kv_hidden] FLOAT per layer
src/mobius/components/_attention.py:151
- The _apply_paged_attention() docstring states that query/key/value projections all have shape [1, seq_len, heads*head_dim], but key/value use num_key_value_heads (kv_num_heads) rather than num_attention_heads. This is an externally visible contract and should be accurate.
query/key/value: 3-D ``[1, seq_len, heads * head_dim]`` projections.
tests/paged_cache_test.py:278
- This test builds a PagedCacheState, but the dummy inputs don't match the paged-cache I/O contract (block_table should be 1-D [num_blocks], and pool last dim should be kv_hidden=num_key_value_heads*head_dim). Using contract-shaped inputs makes the test clearer and avoids accidentally normalizing a mismatched shape convention.
bt = builder.input("block_table", dtype=ir.DataType.INT64, shape=[1, 4])
sm = builder.input("slot_mapping", dtype=ir.DataType.INT64, shape=[1])
kp = builder.input("kpool", dtype=config.dtype, shape=[16, 8, config.head_dim])
vp = builder.input("vpool", dtype=config.dtype, shape=[16, 8, config.head_dim])
seqlen = builder.input("nonpad_kv_seqlen", dtype=ir.DataType.INT64, shape=[1])
- Files reviewed: 8/8 changed files
- Comments generated: 0 new
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
Member
Author
|
Not chasing in favor of the VMM solution. |
…atures Update from main (brings in the cargo-style --features build option) and fold the paged / block-table KV cache toggle into it: - Add 'paged-cache' to _BUILD_FEATURES so 'mobius build --features paged-cache' enables the paged KV cache, matching static-cache / fp8-kv-cache / text-only. - Remove the standalone --paged-cache boolean flag (main dropped the other boolean feature flags); --page-size / --num-pages remain as tuning params. - Reword paged validation errors and docs (README, cli_reference, CHANGELOG) to reference --features paged-cache. - Add a cli_test covering --features paged-cache. Conflicts resolved in __main__.py, tasks/_causal_lm.py (keep both paged params and prune_lm_head), and CHANGELOG.md. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
justinchuby
changed the title
Member
Author
|
Updated from
Merge resolved conflicts in |
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).