Add paged / block-table KV cache export (--features paged-cache)

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

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/ \

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

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",

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)