diff --git a/CHANGELOG.md b/CHANGELOG.md index 29b8d268d..2587cb891 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Paged / block-table KV cache export (`--features paged-cache`) + +#### Added + +- `CausalLMTask(paged_cache=True)` and `mobius build --features paged-cache` + export a **paged / block-table KV cache** (onnx-genai `docs/DESIGN.md` §39.4 + Option C — vLLM PagedAttention / SGLang RadixAttention layout). KV lives in a + shared per-layer **page pool** `key_pool.{i}` / `value_pool.{i}` + `[num_pages, page_size, kv_hidden]`; a per-sequence `block_table` maps logical + page slots to physical pages and a `slot_mapping` gives the flat physical slot + for each newly written token. Attention writes new K/V into the pool via + `ScatterND`, assembles the sequence's pages contiguously via + `Gather(pool, block_table)`, then runs the opset-24 `Attention` op with + `nonpad_kv_seqlen` (input #6) — identical op contract to `--features + static-cache`, but over non-contiguous pages. Because sequences can list the + *same* physical page in their `block_table`, the same graph expresses + RadixAttention shared-prefix pages with no change. Paging uses only standard + ONNX ops (no custom op). New tuning flags `--page-size` (default 16) and + `--num-pages` (dynamic when omitted). Requires `DecoderLayer` / + `MoEDecoderLayer` models; mutually exclusive with `--features static-cache`. + Targets a single active sequence (`batch == 1`); multi-sequence batching is a + documented TODO. + ### NVIDIA Cosmos 3 Edge vision-language model (`cosmos3_edge`) #### Added diff --git a/README.md b/README.md index f9f2320b6..d5e47283e 100644 --- a/README.md +++ b/README.md @@ -113,8 +113,8 @@ mobius build --model openai/whisper-tiny output_dir/ ``` Build-mode toggles use the cargo-style `--features` option. Available features -are `static-cache`, `fp8-kv-cache`, `prune-lm-head`, and `text-only`. Pass them -as a comma-separated list or repeat the option: +are `static-cache`, `fp8-kv-cache`, `paged-cache`, `prune-lm-head`, and +`text-only`. Pass them as a comma-separated list or repeat the option: ```sh mobius build --model meta-llama/Llama-3.2-1B output_dir/ \ diff --git a/docs/cli_reference.md b/docs/cli_reference.md index 989608091..75ded0288 100644 --- a/docs/cli_reference.md +++ b/docs/cli_reference.md @@ -185,6 +185,7 @@ Available features: |---------|--------| | `static-cache` | Pre-allocate fixed-size KV cache buffers using `TensorScatter` (pair with `--max-seq-len N`). Requires `DecoderLayer` / `MoEDecoderLayer` models. Cannot combine with `--task`. | | `fp8-kv-cache` | Store the `GroupQueryAttention` KV cache as `FLOAT8E4M3FN` (per-tensor E4M3), halving KV-cache memory. Requires a GQA build (e.g. `--ep cuda --dtype f16`) and an ORT runtime with the FP8 KV-cache kernel (SM89+). Pair with `--kv-cache-scale-file` for calibrated scales. | +| `paged-cache` | Export a paged / block-table KV cache (vLLM PagedAttention / SGLang RadixAttention layout): a shared page pool plus `block_table` / `slot_mapping`, using only standard ONNX ops. Tune with `--page-size N` (default 16) and `--num-pages N` (dynamic when omitted). Requires `DecoderLayer` / `MoEDecoderLayer` models. Cannot combine with `--task` or the `static-cache` feature. | | `prune-lm-head` | Select the final hidden-state position before the LM-head projection and emit logits shaped `[B, 1, vocab]`. Supported by models using the base `CausalLMModel.forward()` path; unsupported custom forwards fail explicitly. Use only when the downstream workflow does not need per-token logits. | | `text-only` | Export the text backbone of a multimodal checkpoint as a standalone decoder-only LLM (see below). | diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 821490e8f..a3039972a 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -39,6 +39,7 @@ _BUILD_FEATURES: dict[str, str] = { "static-cache": "static_cache", "fp8-kv-cache": "fp8_kv_cache", + "paged-cache": "paged_cache", "prune-lm-head": "prune_lm_head", "text-only": "text_only", } @@ -199,6 +200,26 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: "Remove --task to use --features static-cache." ) + # Validate --features paged-cache combinations + if args.paged_cache and args.static_cache: + raise SystemExit( + "Error: --features paged-cache cannot be combined with " + "--features static-cache; choose one KV cache layout." + ) + if args.paged_cache and args.task is not None: + raise SystemExit( + "Error: --features paged-cache cannot be combined with --task. " + "Remove --task to use --features paged-cache." + ) + if args.page_size is not None and not args.paged_cache: + raise SystemExit("Error: --page-size can only be used with --features paged-cache.") + if args.page_size is not None and args.page_size <= 0: + raise SystemExit("Error: --page-size must be a positive integer.") + if args.num_pages is not None and not args.paged_cache: + raise SystemExit("Error: --num-pages can only be used with --features paged-cache.") + if args.num_pages is not None and args.num_pages <= 0: + raise SystemExit("Error: --num-pages must be a positive integer.") + # text-only resolution lives in build() (model_type remap + config # stripping), which is only reached on the HuggingFace model-ID path. if args.text_only and args.config: @@ -242,6 +263,13 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: "static_cache": True, "max_seq_len": args.max_seq_len, } + elif args.paged_cache: + task = CausalLMTask( + paged_cache=True, + page_size=args.page_size if args.page_size is not None else 16, + num_pages=args.num_pages, + ) + static_cache_params = None else: static_cache_params = None trust_remote_code = args.trust_remote_code @@ -733,6 +761,22 @@ def main(argv: list[str] | None = None) -> None: "Only used with --features static-cache. " "Defaults to max_position_embeddings from config.", ) + build_parser.add_argument( + "--page-size", + type=int, + default=None, + metavar="N", + help="Tokens per page for the paged KV cache. Only used with " + "--features paged-cache. Defaults to 16.", + ) + build_parser.add_argument( + "--num-pages", + type=int, + default=None, + metavar="N", + help="Number of physical pages in the pool for the paged KV cache. " + "Only used with --features paged-cache. Left dynamic (symbolic) when omitted.", + ) build_parser.add_argument( "--ep", "--execution-provider", diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index e33fe4247..f45cd5c29 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -74,6 +74,139 @@ class StaticCacheState(NamedTuple): nonpad_kv_seqlen: ir.Value +class PagedCacheState(NamedTuple): + """Paged (block-table) KV cache state for opset-24 attention. + + Implements the paged / block-table KV cache described in onnx-genai + ``docs/DESIGN.md`` §39.4 Option C ("ONNX Scatter/GatherElements in + Graph"). The KV for the whole batch lives in a shared *page pool* of + fixed-size pages that are NOT contiguous per sequence. A per-sequence + ``block_table`` maps logical page slots to physical page indices, and a + ``slot_mapping`` gives the flat physical slot for each newly written + token. This is the format used by vLLM PagedAttention and, because + multiple sequences can list the *same* physical page in their + ``block_table``, it also supports SGLang RadixAttention (shared prefix + pages) with no graph change — the sharing lives entirely in the runtime's + ``block_table`` / ``slot_mapping`` bookkeeping. + + The paged attention body is (per layer): + + 1. ``ScatterND(pool_flat, slot_mapping, new_kv)`` — write the new + tokens' K/V into their physical slots in the pool. + 2. ``Gather(updated_pool, block_table, axis=0)`` — assemble this + sequence's pages into a contiguous ``[num_blocks, page_size, ...]`` + tensor, reshaped to ``[1, num_blocks * page_size, kv_hidden]``. + 3. ``Attention(query, K_gathered, V_gathered, nonpad_kv_seqlen, + is_causal=1)`` — attend over the contiguous KV, bounded to the + valid prefix by ``nonpad_kv_seqlen`` (same op contract as the + static cache). + + Only standard ONNX ops are used (``Reshape``/``Shape``/``Unsqueeze``/ + ``ScatterND``/``Gather``/``Attention``); no custom op is required. + + .. note:: + The current implementation targets a single active sequence per + forward (``batch == 1``), so ``block_table`` and ``slot_mapping`` are + 1-D. Multi-sequence batching (2-D block tables + per-row gather) is a + documented TODO. + + Fields: + key_pool: Physical key page pool ``[num_pages, page_size, kv_hidden]``. + value_pool: Physical value page pool ``[num_pages, page_size, kv_hidden]``. + block_table: Physical page indices for the active sequence in logical + order ``[num_blocks]`` int64. + slot_mapping: Flat physical slot (``page_id * page_size + offset``) for + each newly written token ``[seq_len]`` int64. + nonpad_kv_seqlen: Valid KV length for the active sequence ``[batch]`` + int64 (``write_start + valid_token_count``). + """ + + key_pool: ir.Value + value_pool: ir.Value + block_table: ir.Value + slot_mapping: ir.Value + nonpad_kv_seqlen: ir.Value + + +def _apply_paged_attention( + op: OpBuilder, + query: ir.Value, + key: ir.Value, + value: ir.Value, + paged_cache: PagedCacheState, + *, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + scale: float, + softcap: float = 0.0, +) -> tuple[ir.Value, ir.Value, ir.Value]: + """Apply attention against a paged (block-table) KV cache. + + See :class:`PagedCacheState` for the I/O contract and op sequence. RoPE + must already be baked into ``key`` (applied by the caller) so cached page + entries carry the rotated keys, exactly as in the static cache path. + + Args: + query/key/value: 3-D ``[1, seq_len, heads * head_dim]`` projections. + head_dim: Per-head dim; used to recover ``kv_hidden`` for the pool + reshape. + + Returns: + ``(attn_output, updated_key_pool, updated_value_pool)`` where the + updated pools have the same ``[num_pages, page_size, kv_hidden]`` shape + as the inputs and are registered as graph outputs by the task. + """ + kv_hidden = num_key_value_heads * head_dim + + # Original 3-D pool shapes [num_pages, page_size, kv_hidden]; we scatter on + # a 2-D flattened view then restore the pool shape for the graph output. + key_pool_shape = op.Shape(paged_cache.key_pool) + value_pool_shape = op.Shape(paged_cache.value_pool) + + key_pool_flat = op.Reshape(paged_cache.key_pool, [-1, kv_hidden]) + value_pool_flat = op.Reshape(paged_cache.value_pool, [-1, kv_hidden]) + + # New tokens' K/V as [seq_len, kv_hidden] (batch == 1). + key_rows = op.Reshape(key, [-1, kv_hidden]) + value_rows = op.Reshape(value, [-1, kv_hidden]) + + # ScatterND row-writes: pool_flat[slot_mapping[t]] = new_kv[t]. + slot_indices = op.Unsqueeze(paged_cache.slot_mapping, [-1]) # [seq_len, 1] + updated_key_flat = op.ScatterND(key_pool_flat, slot_indices, key_rows) + updated_value_flat = op.ScatterND(value_pool_flat, slot_indices, value_rows) + + updated_key_pool = op.Reshape(updated_key_flat, key_pool_shape) + updated_value_pool = op.Reshape(updated_value_flat, value_pool_shape) + + # Gather this sequence's physical pages into logical order, then flatten the + # (num_blocks, page_size) axes into a contiguous KV sequence for Attention. + gathered_key = op.Gather(updated_key_pool, paged_cache.block_table, axis=0) + gathered_value = op.Gather(updated_value_pool, paged_cache.block_table, axis=0) + gathered_key = op.Reshape(gathered_key, [1, -1, kv_hidden]) + gathered_value = op.Reshape(gathered_value, [1, -1, kv_hidden]) + + # Maskless is_causal=1 + nonpad_kv_seqlen (Attention input #6): identical op + # contract to the static cache path — causal + padding masking is derived + # internally, bounding attention to the valid prefix of the gathered pages. + attn_output, _, _ = op.Attention( + query, + gathered_key, + gathered_value, + None, # no attn_mask — is_causal handles masking + None, # no past_key (full gathered KV is provided) + None, # no past_value + paged_cache.nonpad_kv_seqlen, + q_num_heads=num_attention_heads, + kv_num_heads=num_key_value_heads, + scale=scale, + softcap=softcap, + is_causal=1, + _outputs=3, + ) + return attn_output, updated_key_pool, updated_value_pool + + def _apply_attention( op: OpBuilder, query: ir.Value, @@ -321,6 +454,7 @@ def forward( position_embeddings: tuple | None = None, past_key_value: tuple | None = None, static_cache: StaticCacheState | None = None, + paged_cache: PagedCacheState | None = None, ): query_states = self.q_proj(op, hidden_states) key_states = self.k_proj(op, hidden_states) @@ -381,20 +515,48 @@ def forward( interleaved=self._rope_interleave, ) - attn_output, present_key, present_value = _apply_attention( - op, - query_states, - key_states, - value_states, - attention_bias, - past_key_value[0] if past_key_value is not None else None, - past_key_value[1] if past_key_value is not None else None, - num_attention_heads=self.num_attention_heads, - num_key_value_heads=self.num_key_value_heads, - scale=self.scaling, - softcap=self._softcap, - static_cache=static_cache, - ) + if paged_cache is not None: + # Paged attention attends over gathered pages with a causal mask + # only; it has no path for an additive/float attention_bias. Models + # that need bias masking (ALiBi, sliding-window/block overlays) + # would be silently miscomputed, so reject the combination up front. + if attention_bias is not None: + raise ValueError( + "paged_cache is incompatible with a non-None attention_bias " + "(e.g. ALiBi or sliding-window additive masks); the paged " + "attention path applies causal masking only and cannot honor " + "an additive bias." + ) + # 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, + ) + else: + attn_output, present_key, present_value = _apply_attention( + op, + query_states, + key_states, + value_states, + attention_bias, + past_key_value[0] if past_key_value is not None else None, + past_key_value[1] if past_key_value is not None else None, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + scale=self.scaling, + softcap=self._softcap, + static_cache=static_cache, + ) attn_output = self.o_proj(op, attn_output) return attn_output, (present_key, present_value) diff --git a/src/mobius/components/_decoder.py b/src/mobius/components/_decoder.py index 67ca345e4..bceea91f7 100644 --- a/src/mobius/components/_decoder.py +++ b/src/mobius/components/_decoder.py @@ -9,7 +9,7 @@ from onnxscript import OpBuilder, nn from mobius._configs import ArchitectureConfig -from mobius.components._attention import Attention, StaticCacheState +from mobius.components._attention import Attention, PagedCacheState, StaticCacheState from mobius.components._mlp import MLP from mobius.components._rms_norm import RMSNorm @@ -93,15 +93,18 @@ def forward( hidden_states: ir.Value, attention_bias: ir.Value | None, position_embeddings: tuple, - past_key_value: tuple | StaticCacheState | None, + past_key_value: tuple | StaticCacheState | PagedCacheState | None, ): - # Dispatch StaticCacheState to the static_cache parameter; - # custom DecoderLayer subclasses must add this check themselves. - if isinstance(past_key_value, StaticCacheState): + # Dispatch StaticCacheState / PagedCacheState to their attention + # parameters; custom DecoderLayer subclasses must add this themselves. + static_cache = None + paged_cache = None + if isinstance(past_key_value, PagedCacheState): + paged_cache = past_key_value + past_key_value = None + elif isinstance(past_key_value, StaticCacheState): static_cache = past_key_value past_key_value = None - else: - static_cache = None if self._post_norm: return self._forward_post_norm( @@ -111,6 +114,7 @@ def forward( position_embeddings, past_key_value, static_cache, + paged_cache, ) return self._forward_pre_norm( op, @@ -119,6 +123,7 @@ def forward( position_embeddings, past_key_value, static_cache, + paged_cache, ) def _forward_pre_norm( @@ -129,6 +134,7 @@ def _forward_pre_norm( position_embeddings: tuple, past_key_value: tuple | None, static_cache: StaticCacheState | None, + paged_cache: PagedCacheState | None = None, ): residual = hidden_states hidden_states = self.input_layernorm(op, hidden_states) @@ -140,6 +146,7 @@ def _forward_pre_norm( position_embeddings=position_embeddings, past_key_value=past_key_value, static_cache=static_cache, + **({"paged_cache": paged_cache} if paged_cache is not None else {}), ) if not math.isclose(self._residual_multiplier, 1.0): @@ -164,6 +171,7 @@ def _forward_post_norm( position_embeddings: tuple, past_key_value: tuple | None, static_cache: StaticCacheState | None, + paged_cache: PagedCacheState | None = None, ): residual = hidden_states attn_output, present_key_value = self.self_attn( @@ -173,6 +181,7 @@ def _forward_post_norm( position_embeddings=position_embeddings, past_key_value=past_key_value, static_cache=static_cache, + **({"paged_cache": paged_cache} if paged_cache is not None else {}), ) hidden_states = self.post_attention_layernorm(op, attn_output) hidden_states = op.Add(residual, hidden_states) diff --git a/src/mobius/models/moe.py b/src/mobius/models/moe.py index 7d153b5bc..77a002a63 100644 --- a/src/mobius/models/moe.py +++ b/src/mobius/models/moe.py @@ -23,7 +23,7 @@ create_attention_bias, initialize_rope, ) -from mobius.components._attention import StaticCacheState +from mobius.components._attention import PagedCacheState, StaticCacheState from mobius.components._moe import MLP, SigmoidTopKGate, SoftmaxTopKGate from mobius.models.base import CausalLMModel @@ -76,15 +76,18 @@ def forward( hidden_states: ir.Value, attention_bias: ir.Value | None, position_embeddings: tuple, - past_key_value: tuple | StaticCacheState | None, + past_key_value: tuple | StaticCacheState | PagedCacheState | None, ): - # Dispatch StaticCacheState to the static_cache parameter; - # custom MoEDecoderLayer subclasses must add this check themselves. - if isinstance(past_key_value, StaticCacheState): + # Dispatch StaticCacheState / PagedCacheState to their attention + # parameters; custom MoEDecoderLayer subclasses must add this too. + static_cache = None + paged_cache = None + if isinstance(past_key_value, PagedCacheState): + paged_cache = past_key_value + past_key_value = None + elif isinstance(past_key_value, StaticCacheState): static_cache = past_key_value past_key_value = None - else: - static_cache = None if self._post_feedforward_norm: # Post-norm style (FlexOLMo): norm is applied after each sub-layer output @@ -97,6 +100,7 @@ def forward( position_embeddings=position_embeddings, past_key_value=past_key_value, static_cache=static_cache, + **({"paged_cache": paged_cache} if paged_cache is not None else {}), ) attn_output = self.post_attention_layernorm(op, attn_output) if not math.isclose(self._residual_multiplier, 1.0): @@ -121,6 +125,7 @@ def forward( position_embeddings=position_embeddings, past_key_value=past_key_value, static_cache=static_cache, + **({"paged_cache": paged_cache} if paged_cache is not None else {}), ) if not math.isclose(self._residual_multiplier, 1.0): attn_output = op.Mul(attn_output, self._residual_multiplier) diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index cc665c9fb..65ac972b7 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -11,7 +11,7 @@ from mobius._build_context import lm_head_pruning from mobius._configs import ArchitectureConfig from mobius._model_package import ModelPackage -from mobius.components._attention import StaticCacheState +from mobius.components._attention import PagedCacheState, StaticCacheState from mobius.tasks._base import ( ModelTask, _make_graph, @@ -64,18 +64,53 @@ class CausalLMTask(ModelTask): No ``attention_mask`` input — causal masking uses ``is_causal=1``. + **Paged cache** (``paged_cache=True``): + Block-table / paged KV cache (onnx-genai ``docs/DESIGN.md`` §39.4 + Option C). KV lives in a shared *page pool* of fixed-size pages that + are non-contiguous per sequence. New tokens are written to physical + slots via ``ScatterND`` and the sequence's pages are assembled + contiguously via ``Gather(pool, block_table)`` before attention. This + is the vLLM PagedAttention layout; because sequences can share physical + pages through their ``block_table``, it also supports SGLang + RadixAttention (shared prefix pages) with no graph change. + + Inputs: + - input_ids: [batch, seq_len] INT64 + - position_ids: [batch, seq_len] INT64 + - key_pool.{i}: [num_pages, page_size, kv_hidden] FLOAT per layer + - value_pool.{i}: [num_pages, page_size, kv_hidden] FLOAT per layer + - block_table: [num_blocks] INT64 (physical page ids, logical order) + - slot_mapping: [seq_len] INT64 (flat slot per new token) + - nonpad_kv_seqlen: [batch] INT64 + Outputs: + - logits: FLOAT + - updated_key_pool.{i} / updated_value_pool.{i}: FLOAT + + No ``attention_mask`` input — causal masking uses ``is_causal=1``. + Targets a single active sequence (``batch == 1``); multi-sequence + batching is a documented TODO. + The module's ``forward()`` must accept ``(op, input_ids, attention_mask, position_ids, past_key_values)`` - and return ``(logits, list_of_(key, value)_tuples)``. In static cache - mode, ``attention_mask`` will be ``None`` and ``past_key_values`` - entries will be :class:`StaticCacheState` tuples. + and return ``(logits, list_of_(key, value)_tuples)``. In static/paged + cache mode, ``attention_mask`` will be ``None`` and ``past_key_values`` + entries will be :class:`StaticCacheState` / :class:`PagedCacheState` + tuples. Args: static_cache: If ``True``, use pre-allocated static KV cache buffers instead of dynamic concatenation. + paged_cache: If ``True``, use a paged / block-table KV cache (page + pool + block_table + slot_mapping). Mutually exclusive with + ``static_cache``. max_seq_len: Maximum sequence length for static cache buffers. Only used when ``static_cache=True``. Defaults to ``config.max_position_embeddings``. + page_size: Number of tokens per page for the paged cache. Only used + when ``paged_cache=True``. Defaults to 16. + num_pages: Number of physical pages in the pool. Only used when + ``paged_cache=True``. Left symbolic (dynamic) when ``None`` so the + runtime can size the pool; pass an int to stamp a fixed pool size. prune_lm_head: If ``True``, insert ``Gather(axis=1, index=-1)`` before the LM head so only the last token's hidden state is projected to logits. Output logits shape becomes ``[B, 1, @@ -92,11 +127,21 @@ def __init__( self, *, static_cache: bool = False, + paged_cache: bool = False, max_seq_len: int | None = None, + page_size: int = 16, + num_pages: int | None = None, prune_lm_head: bool = False, ): + if static_cache and paged_cache: + raise ValueError( + "static_cache and paged_cache are mutually exclusive; enable at most one." + ) self._static_cache = static_cache + self._paged_cache = paged_cache self._max_seq_len = max_seq_len + self._page_size = page_size + self._num_pages = num_pages self._prune_lm_head = prune_lm_head def build( @@ -105,6 +150,7 @@ def build( config: ArchitectureConfig, ) -> ModelPackage: static = self._static_cache + paged = self._paged_cache # --- Static-cache pre-validation --- if static: @@ -119,6 +165,16 @@ def build( ) _validate_static_cache_support(module) + # --- Paged-cache pre-validation --- + if paged: + if self._page_size is None or self._page_size <= 0: + raise ValueError("page_size must be a positive integer for paged cache.") + if self._num_pages is not None and self._num_pages <= 0: + raise ValueError("num_pages must be a positive integer when provided.") + # 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, mode="Paged cache") + # --- Graph input dims --- batch = ir.SymbolicDim("batch") seq_len = ir.SymbolicDim("sequence_len") @@ -130,8 +186,22 @@ def build( # --- Inputs common to both modes --- input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) - # --- Cache setup (static vs dynamic) --- - if static: + # --- Cache setup (paged vs static vs dynamic) --- + if paged: + attention_mask = None + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] + ) + past_key_values = _make_paged_cache_inputs( + builder, + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + config.dtype, + page_size=self._page_size, + num_pages=self._num_pages, + ) + elif static: attention_mask = None position_ids = builder.input( "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] @@ -207,8 +277,13 @@ def build( builder.add_output(logits, "logits") - # --- Output registration (static vs dynamic) --- - if static: + # --- Output registration (paged vs static vs dynamic) --- + if paged: + _register_paged_cache_outputs( + builder, + present_key_values, + ) + elif static: _register_static_cache_outputs( builder, present_key_values, @@ -459,9 +534,104 @@ def _register_static_cache_outputs( builder.add_output(updated_value, f"updated_value_cache.{i}") -def _validate_static_cache_support(module: nn.Module) -> None: +def _make_paged_cache_inputs( + builder: GraphBuilder, + num_layers: int, + num_key_value_heads: int, + head_dim: int, + dtype: ir.DataType, + *, + page_size: int, + num_pages: int | None, +) -> list[PagedCacheState]: + """Create paged (block-table) KV cache inputs for ``num_layers`` layers. + + Emits a per-layer page pool ``key_pool.{i}`` / ``value_pool.{i}`` of shape + ``[num_pages, page_size, kv_hidden]`` plus the shared ``block_table``, + ``slot_mapping`` and ``nonpad_kv_seqlen`` inputs, and packs them into one + :class:`PagedCacheState` per layer (shared block/slot/nonpad tensors). + + ``num_pages`` is left symbolic (dynamic dimension ``num_pages``) when + ``None`` so the runtime can size the pool; passing an int stamps a fixed + pool size. ``block_table`` (``[num_blocks]``) and ``slot_mapping`` + (``[seq_len]``) are 1-D — the current implementation targets a single + active sequence (``batch == 1``). + + Returns: + A list of :class:`PagedCacheState` tuples for passing to the module + via ``past_key_values``. + """ + kv_hidden = num_key_value_heads * head_dim + pages_dim: int | str = num_pages if num_pages is not None else "num_pages" + + pool_pairs: list[tuple[ir.Value, ir.Value]] = [] + for i in range(num_layers): + key_pool = builder.input( + f"key_pool.{i}", + dtype=dtype, + shape=[pages_dim, page_size, kv_hidden], + ) + value_pool = builder.input( + f"value_pool.{i}", + dtype=dtype, + shape=[pages_dim, page_size, kv_hidden], + ) + pool_pairs.append((key_pool, value_pool)) + + # Shared inputs across all layers. + block_table = builder.input( + "block_table", + dtype=ir.DataType.INT64, + shape=["num_blocks"], + ) + slot_mapping = builder.input( + "slot_mapping", + dtype=ir.DataType.INT64, + shape=["sequence_len"], + ) + nonpad_kv_seqlen = builder.input( + "nonpad_kv_seqlen", + dtype=ir.DataType.INT64, + shape=["batch"], + ) + + paged_caches: list[PagedCacheState] = [] + for key_pool, value_pool in pool_pairs: + paged_caches.append( + PagedCacheState( + key_pool=key_pool, + value_pool=value_pool, + block_table=block_table, + slot_mapping=slot_mapping, + nonpad_kv_seqlen=nonpad_kv_seqlen, + ) + ) + return paged_caches + + +def _register_paged_cache_outputs( + builder: GraphBuilder, + present_key_values: list[tuple[ir.Value, ir.Value]], +) -> None: + """Name and register paged (block-table) cache outputs on the graph. + + Each layer produces the UPDATED page pools (the ``ScatterND`` result), + registered as ``updated_key_pool.{i}`` / ``updated_value_pool.{i}``. + Shapes/dtypes are inferred by the shape inference pass. + """ + for i, (updated_key_pool, updated_value_pool) in enumerate(present_key_values): + builder.add_output(updated_key_pool, f"updated_key_pool.{i}") + builder.add_output(updated_value_pool, f"updated_value_pool.{i}") + + +def _validate_static_cache_support(module: nn.Module, mode: str = "Static cache") -> None: """Check that the module's decoder layers support StaticCacheState. + The same layer dispatch backs both the static and paged KV caches (both + flow their state through ``past_key_value``), so this helper serves both; + pass ``mode`` (e.g. ``"Paged cache"``) so the raised message names the + layout the caller actually enabled. + Shared decoder layers have the ``isinstance(StaticCacheState)`` dispatch in ``forward()``. Custom decoder layers must opt in with the ``_supports_static_cache`` marker after implementing equivalent handling; @@ -517,7 +687,7 @@ def _validate_static_cache_support(module: nn.Module) -> None: if getattr(type(layer), "_supports_static_cache", False): continue raise TypeError( - f"Static cache mode requires decoder layers that " + 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 " diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 8c986cdca..f2a0c8d04 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -5616,9 +5616,175 @@ def test_gemma4_static_cache_input_ordering(self): ) +class TestBuildPagedCacheGraph: + """Verify CausalLMTask(paged_cache=True) builds a valid paged/block-table graph. + + Paged cache = the vLLM PagedAttention / SGLang RadixAttention layout from + onnx-genai ``docs/DESIGN.md`` §39.4 Option C: a per-layer page pool + + shared ``block_table`` / ``slot_mapping`` / ``nonpad_kv_seqlen``, with + ``ScatterND`` writes into the pool and ``Gather`` page assembly before the + ``Attention`` op. + """ + + PAGE_SIZE = 8 + NUM_PAGES = 32 + + def _build_paged_cache_model( + self, model_type: str = "qwen2", num_pages=NUM_PAGES, **config_overrides + ): + """Build a model with CausalLMTask(paged_cache=True); return (model, config).""" + from mobius.tasks import CausalLMTask + + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = CausalLMTask(paged_cache=True, page_size=self.PAGE_SIZE, num_pages=num_pages) + pkg = task.build(module, config) + return pkg["model"], config + + def test_paged_cache_graph_builds(self): + """Build a Qwen2 model with a paged cache.""" + model, _ = self._build_paged_cache_model() + assert model.graph is not None + assert len(model.graph.inputs) > 0 + assert len(model.graph.outputs) > 0 + + def test_paged_cache_graph_inputs(self): + """Verify expected inputs: standard + per-layer pools + shared paging.""" + model, config = self._build_paged_cache_model() + input_names = {inp.name for inp in model.graph.inputs} + num_layers = config.num_hidden_layers + + assert "input_ids" in input_names + assert "position_ids" in input_names + # No attention_mask — causal masking is is_causal=1 on Attention. + assert "attention_mask" not in input_names + + for i in range(num_layers): + assert f"key_pool.{i}" in input_names, f"Missing key_pool.{i}" + assert f"value_pool.{i}" in input_names, f"Missing value_pool.{i}" + + # Shared paging inputs + assert "block_table" in input_names + assert "slot_mapping" in input_names + assert "nonpad_kv_seqlen" in input_names + + # Exact count: 2 standard + 2*num_layers pools + 3 shared paging tensors + expected_count = 2 + 2 * num_layers + 3 + assert len(model.graph.inputs) == expected_count, ( + f"Expected {expected_count} inputs, got {len(model.graph.inputs)}" + ) + + def test_paged_cache_pool_shapes(self): + """Verify page pools are [num_pages, page_size, kv_hidden].""" + model, config = self._build_paged_cache_model() + kv_hidden = config.num_key_value_heads * config.head_dim + pools = { + inp.name: inp for inp in model.graph.inputs if inp.name.startswith("key_pool") + } + assert pools, "No key_pool inputs found" + for inp in pools.values(): + dims = list(inp.shape) + assert dims[0] == self.NUM_PAGES + assert dims[1] == self.PAGE_SIZE + assert dims[2] == kv_hidden + + def test_paged_cache_num_pages_dynamic_by_default(self): + """Omitting num_pages leaves the pool's first dim symbolic.""" + model, _ = self._build_paged_cache_model(num_pages=None) + key_pool0 = next(inp for inp in model.graph.inputs if inp.name == "key_pool.0") + first_dim = next(iter(key_pool0.shape)) + assert not isinstance(first_dim, int), ( + f"Expected symbolic num_pages dim, got {first_dim!r}" + ) + + def test_paged_cache_graph_outputs(self): + """Verify outputs: logits + updated page pools per layer.""" + model, config = self._build_paged_cache_model() + output_names = {out.name for out in model.graph.outputs} + num_layers = config.num_hidden_layers + + assert "logits" in output_names + for i in range(num_layers): + assert f"updated_key_pool.{i}" in output_names, f"Missing updated_key_pool.{i}" + assert f"updated_value_pool.{i}" in output_names, f"Missing updated_value_pool.{i}" + # No dynamic/static cache outputs + assert not any(n.startswith("present.") for n in output_names) + assert not any(n.startswith("updated_key_cache.") for n in output_names) + + # Exact count: 1 logits + 2*num_layers updated pools + expected_count = 1 + 2 * num_layers + assert len(model.graph.outputs) == expected_count + + def test_paged_cache_has_scatternd_gather_and_attention(self): + """Verify graph uses ScatterND (write), Gather (assemble) and Attention.""" + model, _ = self._build_paged_cache_model() + op_types = {n.op_type for n in model.graph} + assert "ScatterND" in op_types, "Paged cache should write via ScatterND" + assert "Gather" in op_types, "Paged cache should assemble pages via Gather" + assert "Attention" in op_types, "Paged cache should use Attention" + # Paging uses standard ONNX ops — no custom paged-attention op. + assert "TensorScatter" not in op_types + + def test_paged_cache_attention_is_causal(self): + """Verify Attention ops use is_causal=1 in paged cache mode.""" + model, config = self._build_paged_cache_model() + attention_nodes = [n for n in model.graph if n.op_type == "Attention"] + assert len(attention_nodes) == config.num_hidden_layers + for node in attention_nodes: + is_causal = node.attributes.get("is_causal") + assert is_causal is not None and is_causal.as_int() == 1 + + def test_paged_cache_attention_consumes_nonpad_kv_seqlen(self): + """Attention input #6 must be nonpad_kv_seqlen (opset-24 external cache).""" + model, config = self._build_paged_cache_model() + attention_nodes = [n for n in model.graph if n.op_type == "Attention"] + assert len(attention_nodes) == config.num_hidden_layers + for node in attention_nodes: + assert len(node.inputs) > 6 + nonpad = node.inputs[6] + assert nonpad is not None and nonpad.name == "nonpad_kv_seqlen" + # attn_mask (input #3) is not connected — masking is is_causal. + attn_mask = node.inputs[3] + assert attn_mask is None or attn_mask.name == "" + + def test_paged_cache_graph_validates(self): + """Verify the paged graph survives a serialization round-trip.""" + model, _ = self._build_paged_cache_model() + proto = ir.serde.serialize_model(model) + assert len(proto.SerializeToString()) > 0 + + def test_paged_cache_moe_graph_builds(self): + """Build a MoE model (qwen2_moe) with a paged cache.""" + model, _ = self._build_paged_cache_model( + model_type="qwen2_moe", + num_local_experts=4, + num_experts_per_tok=2, + attn_qkv_bias=True, + shared_expert_intermediate_size=64, + ) + op_types = {n.op_type for n in model.graph} + assert "ScatterND" in op_types + assert "Gather" in op_types + assert "Attention" in op_types + + def test_paged_and_static_cache_mutually_exclusive(self): + """CausalLMTask rejects enabling both cache layouts.""" + from mobius.tasks import CausalLMTask + + with pytest.raises(ValueError, match="mutually exclusive"): + CausalLMTask(static_cache=True, paged_cache=True) + + def test_outputs_have_shapes_and_dtypes(self): + """Verify shape inference populates all output shapes and dtypes.""" + model, _ = self._build_paged_cache_model() + _assert_outputs_have_shapes_and_dtypes({"model": model}, "qwen2-paged") + + # === Parametrized Vision-Language configs (imported from _test_configs) === _VL_MODEL_PARAMS = _make_params(VL_CONFIGS) + # VL models that produce a single "model" key instead of 3-model split _VL_SINGLE_MODEL_TASKS = {"qwen3-vl-vision-language"} diff --git a/tests/cli_test.py b/tests/cli_test.py index 7a3e7cf6e..c8d3e6d60 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -268,6 +268,31 @@ def test_features_prune_lm_head_passed_through(self): ) assert mock_build.call_args.kwargs.get("prune_lm_head") is True + def test_features_paged_cache_builds_paged_task(self): + """--features paged-cache passes a paged CausalLMTask to build().""" + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius._diffusers_builder._load_diffusers_pipeline_index", + return_value=None, + ), + mock.patch("mobius.__main__.build", return_value=mock.MagicMock()) as mock_build, + mock.patch("mobius.__main__._save_package"), + ): + main( + [ + "build", + "--model", + "some/model", + tmpdir, + "--no-weights", + "--features", + "paged-cache", + ] + ) + task = mock_build.call_args.kwargs.get("task") + assert getattr(task, "_paged_cache", False) is True + def test_features_comma_separated_multiple(self): """A single --features accepts a comma-separated list.""" with ( diff --git a/tests/paged_cache_test.py b/tests/paged_cache_test.py new file mode 100644 index 000000000..d99f658c2 --- /dev/null +++ b/tests/paged_cache_test.py @@ -0,0 +1,293 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Paged / block-table KV cache tests (onnx-genai DESIGN §39.4 Option C). + +These tests guard the paged-attention KV cache layout mobius emits with +``CausalLMTask(paged_cache=True)``. KV lives in a shared *page pool* of +fixed-size pages that are non-contiguous per sequence; a per-sequence +``block_table`` maps logical page slots to physical pages, and a +``slot_mapping`` gives the flat physical slot for each newly written token. +This is the vLLM PagedAttention layout, and — because sequences can list the +*same* physical page in their ``block_table`` — it also expresses SGLang +RadixAttention (shared prefix pages) with no graph change. + +The in-graph paging is built from standard ONNX ops +(``Reshape``/``Shape``/``Unsqueeze``/``ScatterND``/``Gather``/``Attention``); +see :class:`mobius.components._attention.PagedCacheState`. + +Two levels of coverage live here: + +* :class:`TestPagedPagingOpsCpuParity` — runs the paging ops + (``ScatterND`` write + ``Gather`` page assembly) on the **CPU** EP and + checks them against a NumPy reference, including a RadixAttention + shared-page case. These ops are EP-agnostic, so this runs anywhere. + +* The full-model ``Attention`` op with ``nonpad_kv_seqlen`` (Attention input + #6) is CUDA-only until onnxruntime#28958 ships (same constraint as the + static cache), so end-to-end paged *execution* parity is intentionally not + asserted here — the graph-construction guarantees are covered by + ``build_graph_test.py::TestBuildPagedCacheGraph``. + +Run:: + + pytest tests/paged_cache_test.py -v +""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +from onnxscript import GraphBuilder + +from mobius._constants import OPSET_VERSION + + +def _build_paging_probe_graph( + *, + num_pages: int, + page_size: int, + kv_hidden: int, + seq_len: int, + num_blocks: int, +) -> bytes: + """Build a standalone graph mirroring the paged write+assemble sub-graph. + + Replicates :func:`mobius.components._attention._apply_paged_attention` up + to (but excluding) the ``Attention`` op: ``ScatterND`` the new K rows into + the flattened pool at ``slot_mapping``, restore the pool shape, then + ``Gather`` the sequence's physical pages via ``block_table`` and flatten to + a contiguous ``[1, num_blocks * page_size, kv_hidden]`` KV sequence. + + Inputs: ``key_pool`` ``[num_pages, page_size, kv_hidden]``, ``key`` + ``[1, seq_len, kv_hidden]``, ``block_table`` ``[num_blocks]`` and + ``slot_mapping`` ``[seq_len]`` int64. + Outputs: ``updated_pool`` (same shape as ``key_pool``) and ``gathered`` + ``[1, num_blocks * page_size, kv_hidden]``. + """ + + def _value(name: str, dims: list[int], dt: ir.DataType) -> ir.Value: + return ir.Value(name=name, shape=ir.Shape(dims), type=ir.TensorType(dt)) + + key_pool = _value("key_pool", [num_pages, page_size, kv_hidden], ir.DataType.FLOAT) + key = _value("key", [1, seq_len, kv_hidden], ir.DataType.FLOAT) + block_table = _value("block_table", [num_blocks], ir.DataType.INT64) + slot_mapping = _value("slot_mapping", [seq_len], ir.DataType.INT64) + + graph = ir.Graph( + inputs=[key_pool, key, block_table, slot_mapping], + outputs=[], + nodes=[], + name="paged_paging_probe", + opset_imports={"": OPSET_VERSION}, + ) + op = GraphBuilder(graph).op + + pool_shape = op.Shape(key_pool) + pool_flat = op.Reshape(key_pool, [-1, kv_hidden]) + key_rows = op.Reshape(key, [-1, kv_hidden]) + slot_idx = op.Unsqueeze(slot_mapping, [-1]) + updated_flat = op.ScatterND(pool_flat, slot_idx, key_rows) + updated_pool = op.Reshape(updated_flat, pool_shape) + gathered = op.Gather(updated_pool, block_table, axis=0) + gathered = op.Reshape(gathered, [1, -1, kv_hidden]) + + updated_pool.name = "updated_pool" + gathered.name = "gathered" + graph.outputs.extend([updated_pool, gathered]) + + model = ir.Model(graph, ir_version=10) + return ir.serde.serialize_model(model).SerializeToString() + + +class TestPagedPagingOpsCpuParity: + """The paged write+assemble sub-graph matches a NumPy reference on CPU.""" + + def _run(self, proto: bytes, feeds: dict[str, np.ndarray]): + sess = ort.InferenceSession(proto, providers=["CPUExecutionProvider"]) + return sess.run(None, feeds) + + def test_scatter_then_gather_matches_numpy(self): + """Write new tokens to physical slots, then gather contiguous pages.""" + num_pages, page_size, kv_hidden = 6, 4, 8 + seq_len, num_blocks = 3, 2 + proto = _build_paging_probe_graph( + num_pages=num_pages, + page_size=page_size, + kv_hidden=kv_hidden, + seq_len=seq_len, + num_blocks=num_blocks, + ) + + rng = np.random.default_rng(0) + pool0 = rng.standard_normal((num_pages, page_size, kv_hidden)).astype(np.float32) + new_k = rng.standard_normal((1, seq_len, kv_hidden)).astype(np.float32) + # Sequence uses physical pages 4 then 1; write 3 tokens into page 4. + block = np.array([4, 1], dtype=np.int64) + slots = np.array( + [4 * page_size + 0, 4 * page_size + 1, 4 * page_size + 2], dtype=np.int64 + ) + + out_pool, out_gathered = self._run( + proto, + {"key_pool": pool0, "key": new_k, "block_table": block, "slot_mapping": slots}, + ) + + ref_flat = pool0.copy().reshape(-1, kv_hidden) + ref_flat[slots] = new_k.reshape(-1, kv_hidden) + ref_pool = ref_flat.reshape(num_pages, page_size, kv_hidden) + ref_gathered = ref_pool[block].reshape(1, -1, kv_hidden) + + np.testing.assert_allclose(out_pool, ref_pool, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(out_gathered, ref_gathered, rtol=1e-6, atol=1e-6) + # The written tokens are visible at the front of the first gathered page. + np.testing.assert_allclose(out_gathered[0, 0:seq_len], new_k[0], rtol=1e-6, atol=1e-6) + assert out_gathered.shape == (1, num_blocks * page_size, kv_hidden) + + def test_radix_shared_page_is_read_by_two_block_tables(self): + """A physical page shared by two sequences' block_tables reads back once. + + This is the SGLang RadixAttention property: two sequences that share a + common prefix point their ``block_table`` at the *same* physical page, + and both gather identical KV for that page — no duplication, no graph + change. + """ + num_pages, page_size, kv_hidden = 5, 2, 4 + proto = _build_paging_probe_graph( + num_pages=num_pages, + page_size=page_size, + kv_hidden=kv_hidden, + seq_len=2, + num_blocks=2, + ) + rng = np.random.default_rng(7) + pool0 = rng.standard_normal((num_pages, page_size, kv_hidden)).astype(np.float32) + # Write a fresh "prefix" page into physical page 3. + prefix_kv = rng.standard_normal((1, 2, kv_hidden)).astype(np.float32) + slots = np.array([3 * page_size + 0, 3 * page_size + 1], dtype=np.int64) + + # Sequence A: [shared page 3, own page 0]; Sequence B: [shared page 3, own page 1]. + block_a = np.array([3, 0], dtype=np.int64) + block_b = np.array([3, 1], dtype=np.int64) + + pool_a, gathered_a = self._run( + proto, + { + "key_pool": pool0, + "key": prefix_kv, + "block_table": block_a, + "slot_mapping": slots, + }, + ) + # Sequence B reads the already-written shared page (no new write needed); + # feed a no-op write that rewrites the same slots with the same values. + _, gathered_b = self._run( + proto, + { + "key_pool": pool_a, + "key": prefix_kv, + "block_table": block_b, + "slot_mapping": slots, + }, + ) + + # Both sequences see identical KV for the shared prefix page (first page). + np.testing.assert_allclose( + gathered_a[0, 0:page_size], gathered_b[0, 0:page_size], rtol=1e-6, atol=1e-6 + ) + # And that shared KV equals what was written. + np.testing.assert_allclose( + gathered_a[0, 0:page_size], prefix_kv[0], rtol=1e-6, atol=1e-6 + ) + + +def test_paged_cache_task_builds_valid_onnx(): + """End-to-end: CausalLMTask(paged_cache=True) builds a checker-valid graph.""" + from _test_configs import _base_config + from onnx_ir.passes.common import CheckerPass + + from mobius._optimizations import SymbolicShapeInferencePass + from mobius._registry import registry + from mobius.tasks import CausalLMTask + + config = _base_config() + module = registry.get("qwen2")(config) + task = CausalLMTask(paged_cache=True, page_size=8, num_pages=16) + pkg = task.build(module, config) + model = pkg["model"] + + # Fill dummy weights so the checker can serialize. + for init in model.graph.initializers.values(): + if init.const_value is None: + dims = [d if isinstance(d, int) else 1 for d in (init.shape or [1])] + dtype = init.dtype or ir.DataType.FLOAT + init.const_value = ir.Tensor(np.zeros(dims, dtype=dtype.numpy())) + + CheckerPass(True)(model) + SymbolicShapeInferencePass()(model) + + kv_hidden = config.num_key_value_heads * config.head_dim + out_shapes = {o.name: list(o.shape) for o in model.graph.outputs} + for i in range(config.num_hidden_layers): + assert out_shapes[f"updated_key_pool.{i}"] == [16, 8, kv_hidden] + assert out_shapes[f"updated_value_pool.{i}"] == [16, 8, kv_hidden] + + +def test_paged_validation_message_names_paged_mode(): + """An unsupported decoder layer reports the *paged* mode in the error.""" + import pytest + from onnxscript import nn + + from mobius.tasks._causal_lm import _validate_static_cache_support + + class _BadLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + + class _BadModel(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_BadLayer()]) + + with pytest.raises(TypeError, match="Paged cache mode"): + _validate_static_cache_support(_BadModel(), mode="Paged cache") + + +def test_paged_cache_rejects_attention_bias(): + """Paged attention has no additive-bias path, so it must fail fast.""" + import pytest + from _test_configs import _base_config + + from mobius._registry import registry + from mobius.components._attention import PagedCacheState + from mobius.tasks._base import _make_graph + + config = _base_config() + module = registry.get("qwen2")(config) + attn = module.model.layers[0].self_attn + + _graph, builder = _make_graph() + op = builder.op + hidden = builder.input("hidden", dtype=config.dtype, shape=[1, 1, config.hidden_size]) + 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]) + paged = PagedCacheState( + key_pool=kp, + value_pool=vp, + block_table=bt, + slot_mapping=sm, + nonpad_kv_seqlen=seqlen, + ) + + with pytest.raises(ValueError, match="attention_bias"): + attn( + op, + hidden, + attention_bias=op.Identity(hidden), + paged_cache=paged, + )