From 1ba1ad48de197ff2b850f8225bde6663045f2235 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Tue, 11 Aug 2026 23:00:03 +0000 Subject: [PATCH 01/21] feat(inference): add the neutral vocabulary for a paged KV runtime Introduces maxtext/inference/kv_common/, the kernel-neutral types a paged KV control plane is built on: KvStorageLayoutV1 for pool geometry and KvPageTableV1 for one step's page bookkeeping. The layer is deliberately narrow. It carries no strides, no packing and no vendor shapes, because those belong to whichever kernel backend is in use, and a control plane that knew them could not host a second backend. It is also pure host numpy, since allocation, free lists, refcounts and prefix matching are data-dependent irregular logic that cannot live inside a traced computation. An import rule follows from that and is enforced statically rather than by convention: these modules may import only the standard library and numpy, never jax, jax_aiter or the rest of maxtext. The test parses the sources and checks it, because a rule nobody can verify erodes, and this one is what keeps the layer CPU-testable today and mechanically extractable later. Two details worth knowing: Element sizes are tabulated rather than taken from numpy. numpy has neither bfloat16 nor the fp8 variants, and those are precisely the dtypes that matter for KV, so np.dtype("bfloat16") raises. Sizing a pool is too load-bearing to depend on ml_dtypes being installed. Sharding distinguishes three regimes rather than assuming divisibility. When the shard count exceeds the KV head count, which is where GQA at high TP and MQA land, heads are replicated rather than partitioned, so the footprint is multiplied by TP / num_kv_heads instead of divided. replication_factor() and total_pool_bytes() expose that, and a configuration where neither divides the other is rejected at construction rather than mis-sharded silently. Tests run without an accelerator and without MaxText's own dependencies. They load the modules directly from their files, which is what actually demonstrates the isolation property: importing through the package path executes maxtext/__init__.py and pulls in the full config stack, so the modules are clean but the package path is not. Co-authored-by: Cursor --- src/maxtext/inference/kv_common/__init__.py | 39 +++ src/maxtext/inference/kv_common/page_table.py | 159 +++++++++ .../inference/kv_common/storage_layout.py | 159 +++++++++ tests/unit/kv_common_test.py | 309 ++++++++++++++++++ 4 files changed, 666 insertions(+) create mode 100644 src/maxtext/inference/kv_common/__init__.py create mode 100644 src/maxtext/inference/kv_common/page_table.py create mode 100644 src/maxtext/inference/kv_common/storage_layout.py create mode 100644 tests/unit/kv_common_test.py diff --git a/src/maxtext/inference/kv_common/__init__.py b/src/maxtext/inference/kv_common/__init__.py new file mode 100644 index 0000000000..08570bd315 --- /dev/null +++ b/src/maxtext/inference/kv_common/__init__.py @@ -0,0 +1,39 @@ +"""Neutral vocabulary for the paged KV runtime. + +This package is the durable interoperability surface. It describes pool geometry +and per-step page tables in kernel-neutral terms: no strides, no packing, no +vendor shapes, and no accelerator framework. + +Import rule, enforced by CI: this package may import only the standard library +and ``numpy``. It must never import ``jax``, ``jax_aiter``, ``maxtext.layers``, or +``maxtext.models``. A second frontend or a second kernel vendor meets these types, +not a vendor ABI, which is why the restriction is worth policing rather than +merely documenting. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from maxtext.inference.kv_common.page_table import KV_PAGE_TABLE_VERSION, KvPageTableV1 +from maxtext.inference.kv_common.storage_layout import ( + KV_STORAGE_LAYOUT_VERSION, + KvStorageLayoutV1, +) + +__all__ = [ + "KV_PAGE_TABLE_VERSION", + "KV_STORAGE_LAYOUT_VERSION", + "KvPageTableV1", + "KvStorageLayoutV1", +] diff --git a/src/maxtext/inference/kv_common/page_table.py b/src/maxtext/inference/kv_common/page_table.py new file mode 100644 index 0000000000..f030ef8117 --- /dev/null +++ b/src/maxtext/inference/kv_common/page_table.py @@ -0,0 +1,159 @@ +"""Per-step page table: which pages each request holds, and where to write. + +Purely semantic. No strides, no packing, no vendor shapes, and no device arrays -- +everything here is host numpy, because every field is produced by data-dependent +irregular host logic that cannot live inside a traced computation. + +A backend converts this into whatever flat arrays its kernels take. That +conversion is the only place a vendor contract appears. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses + +import numpy as np + +KV_PAGE_TABLE_VERSION = 1 + + +@dataclasses.dataclass +class KvPageTableV1: + """One step's worth of page bookkeeping. + + Attributes: + page_ids: per request, the pages holding its context in sequence order. + seq_lens: int32 [num_reqs] total context length after this step. + query_lens: int32 [num_reqs] new tokens contributed this step. All ones + for decode; the uncached suffix length for prefill. + write_positions: int32 [num_tokens] absolute token index within its own + sequence for each new token, flattened in request order. + request_order: int32 [num_reqs] the order requests appear in the batch. + """ + + page_ids: list[list[int]] = dataclasses.field(default_factory=list) + seq_lens: np.ndarray = dataclasses.field( + default_factory=lambda: np.zeros((0,), dtype=np.int32) + ) + query_lens: np.ndarray = dataclasses.field( + default_factory=lambda: np.zeros((0,), dtype=np.int32) + ) + write_positions: np.ndarray = dataclasses.field( + default_factory=lambda: np.zeros((0,), dtype=np.int32) + ) + request_order: np.ndarray = dataclasses.field( + default_factory=lambda: np.zeros((0,), dtype=np.int32) + ) + version: int = KV_PAGE_TABLE_VERSION + + @property + def num_requests(self) -> int: + return len(self.page_ids) + + @property + def num_tokens(self) -> int: + return int(self.write_positions.shape[0]) + + def validate(self, tokens_per_page: int) -> None: + """Check internal consistency before anything reaches a kernel. + + Cheap here, and the alternative is a silent out-of-bounds read inside an + attention kernel. + """ + n = self.num_requests + for name, arr in ( + ("seq_lens", self.seq_lens), + ("query_lens", self.query_lens), + ("request_order", self.request_order), + ): + if arr.shape != (n,): + raise ValueError(f"{name} must have shape ({n},), got {arr.shape}") + if arr.dtype != np.int32: + raise ValueError(f"{name} must be int32, got {arr.dtype}") + + if int(self.query_lens.sum()) != self.num_tokens: + raise ValueError( + f"query_lens sums to {int(self.query_lens.sum())} but there are " + f"{self.num_tokens} write positions" + ) + + for i, pages in enumerate(self.page_ids): + needed = -(-int(self.seq_lens[i]) // tokens_per_page) # ceil + if len(pages) < needed: + raise ValueError( + f"request {i} holds {len(pages)} pages but seq_len " + f"{int(self.seq_lens[i])} needs {needed} at " + f"{tokens_per_page} tokens per page" + ) + + def last_page_lens(self, tokens_per_page: int) -> np.ndarray: + """Occupancy of each request's final page, in tokens. + + Kept exact rather than rounded up: an over-stated last-page length is how + a kernel reads bytes belonging to a previous occupant of a recycled page. + """ + lens = np.empty((self.num_requests,), dtype=np.int32) + for i in range(self.num_requests): + seq_len = int(self.seq_lens[i]) + if seq_len == 0: + lens[i] = 0 + continue + rem = seq_len % tokens_per_page + lens[i] = rem if rem else tokens_per_page + return lens + + def indptr(self) -> np.ndarray: + """Exclusive prefix sum over per-request page counts, int32 [num_reqs+1].""" + counts = np.array([len(p) for p in self.page_ids], dtype=np.int32) + out = np.zeros((self.num_requests + 1,), dtype=np.int32) + if self.num_requests: + np.cumsum(counts, out=out[1:]) + return out + + def flat_page_indices(self) -> np.ndarray: + """All page ids concatenated in request order, int32.""" + if not self.page_ids: + return np.zeros((0,), dtype=np.int32) + return np.concatenate( + [np.asarray(p, dtype=np.int32) for p in self.page_ids] + ).astype(np.int32) + + def slot_mapping(self, tokens_per_page: int, padding_page_id: int = 0) -> np.ndarray: + """Absolute pool slot for each new token, int32 [num_tokens]. + + A slot is ``page_id * tokens_per_page + offset_within_page``, which is + what an append kernel scatters on. Tokens whose page is the padding + sentinel map to -1 so the kernel skips them. + """ + slots = np.empty((self.num_tokens,), dtype=np.int32) + t = 0 + for i, pages in enumerate(self.page_ids): + for _ in range(int(self.query_lens[i])): + pos = int(self.write_positions[t]) + page_slot = pos // tokens_per_page + if page_slot >= len(pages): + raise ValueError( + f"request {i} token at position {pos} needs page slot " + f"{page_slot} but only {len(pages)} pages are held" + ) + page_id = pages[page_slot] + if page_id == padding_page_id: + slots[t] = -1 + else: + slots[t] = page_id * tokens_per_page + (pos % tokens_per_page) + t += 1 + return slots diff --git a/src/maxtext/inference/kv_common/storage_layout.py b/src/maxtext/inference/kv_common/storage_layout.py new file mode 100644 index 0000000000..6142245fa4 --- /dev/null +++ b/src/maxtext/inference/kv_common/storage_layout.py @@ -0,0 +1,159 @@ +"""Logical geometry of a paged KV pool. + +This type exists to produce two things: the pool byte count that the allocator +needs, and enough shape information for a backend to derive whatever physical +layout its kernels want. It deliberately stops short of that physical layout -- +strides, packing and vendor shapes belong to the backend's own ABI. + +The byte count is also the entire accelerator-facing surface of the pool. XLA's +total knowledge of a KV pool is a size in and a pointer out; it never learns what +a page means. That is the design, not an omission: fp8 KV, MLA's fused tensor and +x-packing each change this class or a vendor ABI, and none of them should require +rebuilding a runtime plugin. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses + +import numpy as np + +KV_STORAGE_LAYOUT_VERSION = 1 + +# Element sizes are tabulated rather than taken from numpy because the dtypes +# that matter most for KV are ones numpy does not have: bfloat16 and the fp8 +# variants live in ml_dtypes, which this layer may not import. Sizing a pool is +# too load-bearing to depend on an optional package. +_ITEMSIZE_BY_NAME = { + "bfloat16": 2, + "float16": 2, + "float32": 4, + "float64": 8, + "float8_e4m3fn": 1, + "float8_e4m3fnuz": 1, + "float8_e5m2": 1, + "float8_e5m2fnuz": 1, + "int8": 1, + "uint8": 1, +} + + +@dataclasses.dataclass(frozen=True) +class KvStorageLayoutV1: + """Pool geometry, in logical terms. + + Attributes: + tokens_per_page: page size in tokens; "block_size" in vLLM vocabulary. + num_pages: total pages in the pool, per layer, per shard. + num_layers: decoder layers backed by this pool. + num_kv_heads: global KV head count, before any tensor-parallel sharding. + head_dim: per-head dimension. + dtype: numpy dtype name of the stored K/V. + kv_head_shards: tensor-parallel width the pool is sharded over. + padding_page_id: page reserved as a padding target and never allocated. + This buys safe targets for padded rows and simpler invalid-page + handling. It is emphatically *not* a security mechanism: stale KV in a + recycled real page is a separate obligation. + """ + + version: int = KV_STORAGE_LAYOUT_VERSION + tokens_per_page: int = 16 + num_pages: int = 0 + num_layers: int = 0 + num_kv_heads: int = 0 + head_dim: int = 0 + dtype: str = "bfloat16" + kv_head_shards: int = 1 + padding_page_id: int = 0 + + def __post_init__(self): + if self.tokens_per_page <= 0: + raise ValueError(f"tokens_per_page must be positive, got {self.tokens_per_page}") + if self.kv_head_shards <= 0: + raise ValueError(f"kv_head_shards must be positive, got {self.kv_head_shards}") + if self.num_kv_heads and self.kv_head_shards > self.num_kv_heads: + if self.kv_head_shards % self.num_kv_heads != 0: + raise ValueError( + f"kv_head_shards {self.kv_head_shards} exceeds num_kv_heads " + f"{self.num_kv_heads} without dividing evenly: no clean " + f"replication factor exists. Reject this configuration at " + f"startup rather than mis-sharding." + ) + elif self.num_kv_heads and self.num_kv_heads % self.kv_head_shards != 0: + raise ValueError( + f"num_kv_heads {self.num_kv_heads} is not divisible by " + f"kv_head_shards {self.kv_head_shards}" + ) + + def itemsize(self) -> int: + """Bytes per stored element.""" + if self.dtype in _ITEMSIZE_BY_NAME: + return _ITEMSIZE_BY_NAME[self.dtype] + try: + return int(np.dtype(self.dtype).itemsize) + except TypeError as exc: + raise ValueError( + f"unknown KV dtype {self.dtype!r}; add it to _ITEMSIZE_BY_NAME" + ) from exc + + def heads_per_shard(self) -> int: + """KV heads held by one shard. + + Two regimes, and conflating them mis-sizes the pool. When the head count + divides the shard count the heads are partitioned. When the shard count + exceeds the head count -- GQA at high TP, and MQA always -- each head is + *replicated*, so a shard still holds at least one whole head and the + total footprint is multiplied rather than divided. + """ + if self.kv_head_shards >= self.num_kv_heads: + return max(self.num_kv_heads, 1) if self.num_kv_heads else 0 + return self.num_kv_heads // self.kv_head_shards + + def replication_factor(self) -> int: + """How many shards hold a copy of the same KV head.""" + if self.num_kv_heads and self.kv_head_shards > self.num_kv_heads: + return self.kv_head_shards // self.num_kv_heads + return 1 + + def bytes_per_page(self) -> int: + """Bytes of one page of one of K or V, for one layer, on one shard. + + This is the only layout-derived quantity a page transfer needs, which is + what keeps a transfer ABI from becoming a second copy of the attention + ABI. + """ + return self.tokens_per_page * self.heads_per_shard() * self.head_dim * self.itemsize() + + def bytes_per_token(self) -> int: + """Bytes of K and V across all layers for one token, on one shard.""" + return 2 * self.num_layers * self.heads_per_shard() * self.head_dim * self.itemsize() + + def pool_bytes_per_shard(self) -> int: + """Total pool bytes on one shard: K and V, all layers, all pages. + + This is the value handed to the allocator, and the entire + accelerator-facing surface of this class. + """ + return 2 * self.num_layers * self.num_pages * self.bytes_per_page() + + def total_pool_bytes(self) -> int: + """Pool bytes summed over shards, including any replication.""" + return self.pool_bytes_per_shard() * self.kv_head_shards + + def max_tokens(self) -> int: + """Live tokens the pool can hold, excluding the padding page.""" + return max(self.num_pages - 1, 0) * self.tokens_per_page diff --git a/tests/unit/kv_common_test.py b/tests/unit/kv_common_test.py new file mode 100644 index 0000000000..6553d45333 --- /dev/null +++ b/tests/unit/kv_common_test.py @@ -0,0 +1,309 @@ +"""CPU-only tests for the neutral paged-KV vocabulary. + +These run with no accelerator, which is the point of keeping the layer pure +numpy: pool sizing, slot arithmetic and last-page occupancy are all exactly the +sort of data-dependent host logic that cannot be traced, and all of it is +testable in ordinary CI. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import ast +import importlib.util +import pathlib +import sys +import unittest + +import numpy as np + +# Loaded straight from their files rather than through `maxtext.inference.kv_common` +# on purpose. The neutral layer is meant to run with no accelerator and no +# framework, but importing it by package path executes `maxtext/__init__.py`, +# which pulls in the full config stack. Loading the modules directly is what +# actually demonstrates the property, and it keeps these tests runnable in plain +# CI. `_ImportRuleTest` below pins the rule statically. +_KV_COMMON_DIR = ( + pathlib.Path(__file__).resolve().parents[2] + / "src" + / "maxtext" + / "inference" + / "kv_common" +) + + +def _load(module_name: str): + path = _KV_COMMON_DIR / f"{module_name}.py" + spec = importlib.util.spec_from_file_location(f"_kv_common_{module_name}", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +KvStorageLayoutV1 = _load("storage_layout").KvStorageLayoutV1 +KvPageTableV1 = _load("page_table").KvPageTableV1 + +_ALLOWED_IMPORTS = { + "numpy", + "dataclasses", + "annotations", # from __future__ + "__future__", + "ast", + "typing", + "math", + "enum", + "collections", + "hashlib", +} + + +class KvStorageLayoutTest(unittest.TestCase): + """Pool geometry and sizing.""" + + def _layout(self, **kw): + base = dict( + tokens_per_page=16, + num_pages=1024, + num_layers=32, + num_kv_heads=8, + head_dim=128, + dtype="bfloat16", + ) + base.update(kw) + return KvStorageLayoutV1(**base) + + def test_bytes_per_page(self): + layout = self._layout() + # 16 tokens * 8 heads * 128 dim * 2 bytes + self.assertEqual(layout.bytes_per_page(), 16 * 8 * 128 * 2) + + def test_pool_bytes_matches_hand_calculation(self): + layout = self._layout() + expected = 2 * 32 * 1024 * (16 * 8 * 128 * 2) + self.assertEqual(layout.pool_bytes_per_shard(), expected) + + def test_bytes_per_token_is_the_capacity_unit(self): + """Per-token cost is what turns pool size into concurrency.""" + layout = self._layout(num_layers=80) + # 8 heads * 128 dim * 2 bytes * 2 (K,V) * 80 layers = 320 KiB + self.assertEqual(layout.bytes_per_token(), 320 * 1024) + + def test_clean_partition_when_heads_divide_shards(self): + layout = self._layout(num_kv_heads=8, kv_head_shards=4) + self.assertEqual(layout.heads_per_shard(), 2) + self.assertEqual(layout.replication_factor(), 1) + + def test_replication_when_shards_exceed_heads(self): + """GQA above its KV-head count replicates, multiplying the footprint. + + This is the case a naive num_kv_heads // shards would get silently wrong. + """ + layout = self._layout(num_kv_heads=2, kv_head_shards=8) + self.assertEqual(layout.heads_per_shard(), 2) + self.assertEqual(layout.replication_factor(), 4) + self.assertEqual(layout.total_pool_bytes(), layout.pool_bytes_per_shard() * 8) + + def test_mqa_replicates_everywhere(self): + layout = self._layout(num_kv_heads=1, kv_head_shards=8) + self.assertEqual(layout.heads_per_shard(), 1) + self.assertEqual(layout.replication_factor(), 8) + + def test_indivisible_sharding_is_rejected_at_construction(self): + with self.assertRaises(ValueError): + self._layout(num_kv_heads=6, kv_head_shards=4) + + def test_max_tokens_excludes_padding_page(self): + layout = self._layout(num_pages=100, tokens_per_page=16) + self.assertEqual(layout.max_tokens(), 99 * 16) + + +class KvPageTableTest(unittest.TestCase): + """Per-step page bookkeeping.""" + + def _decode_table(self): + """Two requests, one new token each, mid-sequence.""" + return KvPageTableV1( + page_ids=[[1, 2, 3], [4, 5]], + seq_lens=np.array([33, 20], dtype=np.int32), + query_lens=np.array([1, 1], dtype=np.int32), + write_positions=np.array([32, 19], dtype=np.int32), + request_order=np.array([0, 1], dtype=np.int32), + ) + + def test_validate_accepts_a_consistent_table(self): + self._decode_table().validate(tokens_per_page=16) + + def test_validate_rejects_too_few_pages(self): + table = KvPageTableV1( + page_ids=[[1]], + seq_lens=np.array([33], dtype=np.int32), + query_lens=np.array([1], dtype=np.int32), + write_positions=np.array([32], dtype=np.int32), + request_order=np.array([0], dtype=np.int32), + ) + with self.assertRaises(ValueError): + table.validate(tokens_per_page=16) + + def test_validate_rejects_query_len_mismatch(self): + table = self._decode_table() + table.query_lens = np.array([2, 1], dtype=np.int32) + with self.assertRaises(ValueError): + table.validate(tokens_per_page=16) + + def test_indptr_is_exclusive_prefix_sum(self): + np.testing.assert_array_equal( + self._decode_table().indptr(), np.array([0, 3, 5], dtype=np.int32) + ) + + def test_flat_page_indices_is_request_ordered(self): + np.testing.assert_array_equal( + self._decode_table().flat_page_indices(), + np.array([1, 2, 3, 4, 5], dtype=np.int32), + ) + + def test_last_page_lens_are_exact(self): + """Exact occupancy is what stops a kernel reading a recycled page's tail.""" + table = self._decode_table() + # seq_len 33 -> 33 % 16 == 1; seq_len 20 -> 20 % 16 == 4 + np.testing.assert_array_equal( + table.last_page_lens(tokens_per_page=16), np.array([1, 4], dtype=np.int32) + ) + + def test_last_page_len_of_a_full_page_is_the_page_size(self): + table = KvPageTableV1( + page_ids=[[1, 2]], + seq_lens=np.array([32], dtype=np.int32), + query_lens=np.array([1], dtype=np.int32), + write_positions=np.array([31], dtype=np.int32), + request_order=np.array([0], dtype=np.int32), + ) + np.testing.assert_array_equal( + table.last_page_lens(tokens_per_page=16), np.array([16], dtype=np.int32) + ) + + def test_slot_mapping_decode(self): + table = self._decode_table() + slots = table.slot_mapping(tokens_per_page=16) + # req 0: position 32 -> page slot 2 -> page 3, offset 0 -> 3*16 + 0 + # req 1: position 19 -> page slot 1 -> page 5, offset 3 -> 5*16 + 3 + np.testing.assert_array_equal(slots, np.array([48, 83], dtype=np.int32)) + + def test_slot_mapping_prefill_spans_pages(self): + """A prefill writes several tokens per request, crossing a page boundary.""" + table = KvPageTableV1( + page_ids=[[7, 9]], + seq_lens=np.array([18], dtype=np.int32), + query_lens=np.array([18], dtype=np.int32), + write_positions=np.arange(18, dtype=np.int32), + request_order=np.array([0], dtype=np.int32), + ) + table.validate(tokens_per_page=16) + slots = table.slot_mapping(tokens_per_page=16) + expected = [7 * 16 + i for i in range(16)] + [9 * 16 + 0, 9 * 16 + 1] + np.testing.assert_array_equal(slots, np.array(expected, dtype=np.int32)) + + def test_padding_page_maps_to_skip_sentinel(self): + """Tokens landing on the padding page become -1 so kernels drop them.""" + table = KvPageTableV1( + page_ids=[[0]], + seq_lens=np.array([4], dtype=np.int32), + query_lens=np.array([4], dtype=np.int32), + write_positions=np.arange(4, dtype=np.int32), + request_order=np.array([0], dtype=np.int32), + ) + slots = table.slot_mapping(tokens_per_page=16, padding_page_id=0) + np.testing.assert_array_equal(slots, np.full((4,), -1, dtype=np.int32)) + + def test_slot_mapping_rejects_position_beyond_held_pages(self): + table = KvPageTableV1( + page_ids=[[1]], + seq_lens=np.array([20], dtype=np.int32), + query_lens=np.array([1], dtype=np.int32), + write_positions=np.array([19], dtype=np.int32), + request_order=np.array([0], dtype=np.int32), + ) + with self.assertRaises(ValueError): + table.slot_mapping(tokens_per_page=16) + + +class ImportRuleTest(unittest.TestCase): + """The neutral layer must stay free of framework and vendor dependencies. + + Checked statically, so it holds without importing anything. A convention no + one can verify erodes, and this one is load-bearing: it is what makes the layer + CPU-testable today and mechanically extractable later. + """ + + FORBIDDEN_ROOTS = {"jax", "jaxlib", "flax", "torch", "jax_aiter"} + SELF_PACKAGE = "maxtext.inference.kv_common" + + def _top_level_imports(self, path): + """Root packages imported by `path`, excluding intra-package imports. + + An absolute import of kv_common's own modules is fine; anything else under + `maxtext.` is not, and is reported under its full name so the message is + actionable. + """ + tree = ast.parse(path.read_text()) + roots = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name.startswith(self.SELF_PACKAGE): + continue + roots.add(alias.name.split(".")[0]) + elif isinstance(node, ast.ImportFrom): + if node.level or not node.module: # relative import stays in-package + continue + if node.module.startswith(self.SELF_PACKAGE): + continue + if node.module.startswith("maxtext"): + roots.add(node.module) # full name: this is a rule violation + continue + roots.add(node.module.split(".")[0]) + return roots + + def test_modules_import_only_stdlib_and_numpy(self): + for path in sorted(_KV_COMMON_DIR.glob("*.py")): + roots = self._top_level_imports(path) + forbidden = roots & self.FORBIDDEN_ROOTS + self.assertEqual( + forbidden, + set(), + f"{path.name} imports {sorted(forbidden)}, which the kv_common import " + f"rule forbids", + ) + unexpected = roots - _ALLOWED_IMPORTS + self.assertEqual( + unexpected, + set(), + f"{path.name} imports {sorted(unexpected)}; extend _ALLOWED_IMPORTS " + f"only if the addition is genuinely stdlib or numpy", + ) + + def test_package_init_only_reaches_into_kv_common(self): + """The package __init__ may re-export, but must not reach outside itself.""" + init = _KV_COMMON_DIR / "__init__.py" + tree = ast.parse(init.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module: + self.assertTrue( + node.module.startswith("maxtext.inference.kv_common"), + f"__init__ imports from {node.module}, outside kv_common", + ) + + +if __name__ == "__main__": + unittest.main() From 971db32b6814e81553a00a60c85c126a5d647d11 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Thu, 20 Aug 2026 23:05:06 +0000 Subject: [PATCH 02/21] feat(inference): add a vendor-neutral gpu_paged attention path (M3) Adds attention='gpu_paged', which attends over an external paged KV pool instead of the dense per-layer cache. The pool is a pair of NHD arrays carried as the layer's kv_cache; one step writes the new K/V into it and then attends, so prefill and decode read exactly the pages the append wrote. Nearly all of it lives in layers/gpu_paged_attention.py, leaving one method and one elif in attentions.py. Only a single leaf function knows which kernel provider is in use, so flashinfer on NVIDIA slots in beside aiter on ROCm without a second path through MaxText. Metadata is duck-typed rather than imported. A neutral KvPageTableV1 is used directly; a vLLM-shaped object with block_tables/seq_lens/query_start_loc is converted, packing the row-padded 2D table into the contiguous 1D page list the kernels want, in jnp and shape-static so it survives jit. Importing either producer would defeat the neutrality this exists for, and tpu_inference cannot be installed on an AMD box at all. Two guards. The dense KVCache_0 is not allocated on this path, which would otherwise waste gigabytes beside the pool. And scan_layers is refused outright rather than silently overridden, because stacking the per-layer caches into the scan carry copies the whole pool every step and shows up as a slowdown rather than an error. The parity test compares against MaxText's own dense path and agrees to within one ulp. It earned that bound during development by catching a double-scaling bug: MaxText folds the depth scaling into the query projection's initializer, so the query arriving at a serve path is already scaled and the kernel must be passed 1.0, exactly as forward_serve_vllm does. It is marked gpu_only deliberately -- conftest auto-marks unmarked tests cpu_only and skips those on accelerators, so an unmarked test would report green while never running. 14 tests pass. --- src/maxtext/configs/pyconfig_deprecated.py | 1 + src/maxtext/configs/types.py | 28 ++- src/maxtext/layers/attentions.py | 65 +++++- src/maxtext/layers/gpu_paged_attention.py | 241 +++++++++++++++++++++ tests/unit/attention_test.py | 136 ++++++++++++ tests/unit/gpu_paged_attention_test.py | 201 +++++++++++++++++ 6 files changed, 670 insertions(+), 2 deletions(-) create mode 100644 src/maxtext/layers/gpu_paged_attention.py create mode 100644 tests/unit/gpu_paged_attention_test.py diff --git a/src/maxtext/configs/pyconfig_deprecated.py b/src/maxtext/configs/pyconfig_deprecated.py index c30cc0101b..a3b569d669 100644 --- a/src/maxtext/configs/pyconfig_deprecated.py +++ b/src/maxtext/configs/pyconfig_deprecated.py @@ -107,6 +107,7 @@ def validate_attention_kernel(s: str) -> None: "cudnn_flash_jax", "vllm_rpa", "vllm_batched_rpa", + "gpu_paged", ) if s not in valid_attention_kernels: # currently supported attention raise ValueError("Invalid attention kernel was passed. Valid options ", valid_attention_kernels) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index d7d56469d1..2ce99ad6fb 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -598,7 +598,17 @@ class Attention(BaseModel): attention: str = Field( "autoselected", - description="The attention algorithm to use (dot_product, flash, cudnn_flash_te, vllm_rpa, vllm_batched_rpa, etc).", + description=( + "The attention algorithm to use (dot_product, flash, cudnn_flash_te, vllm_rpa, vllm_batched_rpa, " + "gpu_paged, etc)." + ), + ) + paged_attention_backend: Literal["auto", "aiter", "flashinfer"] = Field( + "auto", + description=( + "Leaf kernel provider for attention='gpu_paged'. Everything above the kernel call is " + "vendor-neutral, so 'auto' picks aiter on ROCm and flashinfer on CUDA." + ), ) attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed", "block_diffusion"] = Field( "global", description="The variant of attention to use." @@ -3158,6 +3168,22 @@ def _load_mesh_config_from_yaml(rule_value: str) -> dict: with open(custom_mesh_path, "r", encoding="utf-8") as f: return yaml.safe_load(f) or {} + @model_validator(mode="after") + def validate_gpu_paged_attention(self) -> "MaxTextConfig": + """Reject `gpu_paged` combinations that would be silently pathological. + + `scan_layers` stacks the per-layer KV caches with `jnp.stack(kv_caches)` so + they can ride the scan carry. For a paged pool that copies the entire pool + on every step, which shows up as a mysterious slowdown rather than an error, + so refuse it outright instead of quietly overriding it. + """ + if self.attention == "gpu_paged" and self.scan_layers: + raise ValueError( + "attention='gpu_paged' requires scan_layers=False. Scanning stacks the per-layer " + "KV caches into the scan carry, which copies the whole paged pool every step." + ) + return self + @model_validator(mode="after") def set_derived_and_validate_values(self) -> "MaxTextConfig": """ diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index a2f4b48afd..a1c0937acd 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -440,7 +440,7 @@ def __init__( self.init_kv_caches(inputs_kv_shape=inputs_kv_shape) if self.model_mode != MODEL_MODE_TRAIN and base_kv_cache - and config.attention not in ("vllm_rpa", "vllm_batched_rpa") + and config.attention not in ("vllm_rpa", "vllm_batched_rpa", "gpu_paged") else None ) @@ -1180,6 +1180,61 @@ def forward_serve_vllm( ) return output, kv_cache + def forward_serve_gpu_paged( + self, + query: Array, + key: Array, + value: Array, + kv_pools: list[Array] | None = None, + metadata: Any = None, + ) -> tuple[Array, list[Array]]: + """Forward function for paged serving on GPU. + + `kv_pools` is `[k_pool, v_pool]`, each NHD + `[num_pages, tokens_per_page, num_kv_heads, head_dim]`. The step writes this + token's K/V into the pool and then attends over it, so prefill and decode + read exactly the pages append wrote. + + Everything vendor-specific lives in `gpu_paged_attention`; this method only + flattens to the ragged layout and hands over. + """ + # pylint: disable=import-outside-toplevel + from maxtext.layers import gpu_paged_attention + + query = query.reshape(-1, query.shape[2], query.shape[3]).astype(self.dtype) + key = key.reshape(-1, key.shape[2], key.shape[3]).astype(self.dtype) + value = value.reshape(-1, value.shape[2], value.shape[3]).astype(self.dtype) + + if not kv_pools or metadata is None: + # Dry run: model initialization and JIT tracing reach here with nothing to + # attend over. Mirrors forward_serve_vllm; without it, init breaks rather + # than inference, which is a confusing way to fail. + return query, [] + + k_pool, v_pool = kv_pools[0], kv_pools[1] + plan = gpu_paged_attention.build_plan( + metadata, + tokens_per_page=k_pool.shape[1], + total_tokens=query.shape[0], + max_seqlen_k=self.max_target_length, + ) + return gpu_paged_attention.paged_attention_step( + query, + key, + value, + k_pool, + v_pool, + plan, + backend=getattr(self.config, "paged_attention_backend", "auto"), + # 1.0, not 1/sqrt(head_dim). MaxText folds the depth scaling into the + # query projection's initializer (see `depth_scaling` in the kernel init) + # and applies `query_pre_attn_scalar` above, so the query arriving here is + # already scaled. forward_serve_vllm passes 1.0 for the same reason; + # letting the kernel apply its own default would scale twice. + scale=1.0, + causal=True, + ) + def __call__( self, inputs_q: Array, @@ -1334,6 +1389,14 @@ def __call__( out = attn_out.reshape(batch, seq_len, num_heads, head_dim) kv_cache = updated_kv + elif self.config.attention == "gpu_paged" and model_mode != MODEL_MODE_TRAIN: + batch, seq_len, num_heads, head_dim = query.shape + attn_out, updated_kv = self.forward_serve_gpu_paged( + query, key, value, kv_pools=kv_cache, metadata=attention_metadata + ) + out = attn_out.reshape(batch, seq_len, num_heads, head_dim) + kv_cache = updated_kv + else: cached_values = [None, None] if model_mode != MODEL_MODE_TRAIN: diff --git a/src/maxtext/layers/gpu_paged_attention.py b/src/maxtext/layers/gpu_paged_attention.py new file mode 100644 index 0000000000..0bf8af57e6 --- /dev/null +++ b/src/maxtext/layers/gpu_paged_attention.py @@ -0,0 +1,241 @@ +"""Paged attention over a GPU KV pool, for `attention: "gpu_paged"`. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Everything here is vendor-neutral except `_call_backend`, which is the single +leaf that knows which kernel provider is in use. That is the whole point of the +split: FlashInfer on NVIDIA slots in beside aiter on ROCm without a second code +path through MaxText. + +The KV pool is a pair of NHD arrays, `[num_pages, tokens_per_page, num_kv_heads, +head_dim]`, carried as the layer's `kv_cache`. One step writes the new K/V into +the pool and then attends over it, so both halves see the same pages and nothing +is repacked when a request moves between prefill and decode. + +Two metadata shapes are accepted, because the serving harness that produces them +differs by platform and neither package can be assumed present: + + * a neutral `KvPageTableV1` (MaxText's own vocabulary, `inference/kv_common/`), + recognised by its `indptr()` / `flat_page_indices()` / `last_page_lens()` + methods. Its fields are concrete host values, so the conversion is numpy. + * a vLLM-shaped metadata object with `block_tables` / `seq_lens` / + `query_start_loc`, as `tpu_inference.AttentionMetadata` supplies. Those are + traced arrays, so that conversion has to be `jnp` and shape-static. + +Both are duck-typed. Importing either producer here would defeat the neutrality +the split exists for. +""" + +from __future__ import annotations + +import dataclasses +from typing import Any + +import jax +import jax.numpy as jnp + + +@dataclasses.dataclass(frozen=True) +class PagedPlan: + """One step's page bookkeeping, as the device arrays the kernels take.""" + + slot_mapping: jax.Array # [total_tokens] where each new token is written + kv_indptr: jax.Array # [num_seqs + 1] prefix sum over page counts + kv_page_indices: jax.Array # [>= total_pages] page ids, packed in request order + kv_last_page_lens: jax.Array # [num_seqs] occupancy of each final page + cu_seqlens_q: jax.Array # [num_seqs + 1] prefix sum over query lengths + max_seqlen_q: int + max_seqlen_k: int + is_decode: bool + + +def is_neutral_page_table(metadata: Any) -> bool: + """True for MaxText's `KvPageTableV1`, without importing it.""" + return all( + callable(getattr(metadata, name, None)) + for name in ("indptr", "flat_page_indices", "last_page_lens", "slot_mapping") + ) + + +def is_vllm_metadata(metadata: Any) -> bool: + """True for a vLLM-shaped metadata object, without importing tpu_inference.""" + return all(hasattr(metadata, name) for name in ("block_tables", "seq_lens", "query_start_loc")) + + +def plan_from_neutral(page_table: Any, tokens_per_page: int, padding_page_id: int = 0) -> PagedPlan: + """Convert a `KvPageTableV1`. Its members are host values, so this is numpy.""" + import numpy as np # pylint: disable=import-outside-toplevel + + page_table.validate(tokens_per_page) + + query_lens = np.asarray(page_table.query_lens, dtype=np.int32) + seq_lens = np.asarray(page_table.seq_lens, dtype=np.int32) + cu_seqlens_q = np.zeros((query_lens.size + 1,), dtype=np.int32) + if query_lens.size: + np.cumsum(query_lens, out=cu_seqlens_q[1:]) + + return PagedPlan( + slot_mapping=jnp.asarray(page_table.slot_mapping(tokens_per_page, padding_page_id), jnp.int32), + kv_indptr=jnp.asarray(page_table.indptr(), jnp.int32), + kv_page_indices=jnp.asarray(page_table.flat_page_indices(), jnp.int32), + kv_last_page_lens=jnp.asarray(page_table.last_page_lens(tokens_per_page), jnp.int32), + cu_seqlens_q=jnp.asarray(cu_seqlens_q, jnp.int32), + max_seqlen_q=int(query_lens.max()) if query_lens.size else 0, + max_seqlen_k=int(seq_lens.max()) if seq_lens.size else 0, + is_decode=bool(query_lens.size and np.all(query_lens == 1)), + ) + + +def plan_from_vllm(metadata: Any, tokens_per_page: int, total_tokens: int, max_seqlen_k: int) -> PagedPlan: + """Convert vLLM-shaped metadata. These are traced arrays, so this is `jnp`. + + The 2D `block_tables` is `[num_seqs, max_blocks_per_seq]` and rows are padded, + while the kernels want the page ids packed contiguously with `kv_indptr` + pointing at each request's run. Packing is a scatter: every valid `(seq, block)` + lands at `kv_indptr[seq] + block`, and invalid entries are sent to a scratch + slot past the end so the shape stays static under `jit`. + """ + block_tables = jnp.asarray(metadata.block_tables, jnp.int32) + seq_lens = jnp.asarray(metadata.seq_lens, jnp.int32) + query_start_loc = jnp.asarray(metadata.query_start_loc, jnp.int32) + + if block_tables.ndim != 2: + raise ValueError(f"block_tables must be [num_seqs, max_blocks_per_seq], got shape {block_tables.shape}") + num_seqs, max_blocks = block_tables.shape + + pages_per_seq = (seq_lens + tokens_per_page - 1) // tokens_per_page + kv_indptr = jnp.concatenate([jnp.zeros((1,), jnp.int32), jnp.cumsum(pages_per_seq).astype(jnp.int32)]) + + block = jnp.arange(max_blocks, dtype=jnp.int32)[None, :] + valid = block < pages_per_seq[:, None] + dest = jnp.where(valid, kv_indptr[:num_seqs, None] + block, num_seqs * max_blocks) + packed = jnp.zeros((num_seqs * max_blocks + 1,), jnp.int32).at[dest.reshape(-1)].set(block_tables.reshape(-1)) + kv_page_indices = packed[:-1] + + # An empty sequence has no final page; an overstated length is how a kernel + # reads bytes left behind by a recycled page's previous occupant. + last_page_lens = jnp.where(seq_lens > 0, ((seq_lens - 1) % tokens_per_page) + 1, 0).astype(jnp.int32) + + # Each new token's absolute position within its own sequence, hence its slot. + query_lens = jnp.diff(query_start_loc) + token = jnp.arange(total_tokens, dtype=jnp.int32) + seq_id = jnp.searchsorted(query_start_loc[1:], token, side="right") + within = token - query_start_loc[seq_id] + position = seq_lens[seq_id] - query_lens[seq_id] + within + page = block_tables[seq_id, position // tokens_per_page] + slot_mapping = (page * tokens_per_page + position % tokens_per_page).astype(jnp.int32) + + return PagedPlan( + slot_mapping=slot_mapping, + kv_indptr=kv_indptr, + kv_page_indices=kv_page_indices, + kv_last_page_lens=last_page_lens, + cu_seqlens_q=query_start_loc, + # Query lengths are traced, so the caller supplies the static bounds the + # kernels bake into their configuration. + max_seqlen_q=total_tokens, + max_seqlen_k=max_seqlen_k, + is_decode=(total_tokens == int(num_seqs)), + ) + + +def build_plan(metadata: Any, tokens_per_page: int, total_tokens: int, max_seqlen_k: int) -> PagedPlan: + """Dispatch on the metadata's shape rather than on its type.""" + if isinstance(metadata, PagedPlan): + return metadata + if is_neutral_page_table(metadata): + return plan_from_neutral(metadata, tokens_per_page) + if is_vllm_metadata(metadata): + return plan_from_vllm(metadata, tokens_per_page, total_tokens, max_seqlen_k) + raise TypeError( + "gpu_paged attention metadata must be either a neutral KvPageTableV1 " + "(indptr/flat_page_indices/last_page_lens/slot_mapping) or a vLLM-shaped object " + f"(block_tables/seq_lens/query_start_loc); got {type(metadata).__name__}." + ) + + +def resolve_backend(configured: str) -> str: + """`auto` follows the platform; anything else is taken literally.""" + if configured != "auto": + return configured + try: + platform = jax.devices()[0].platform + backend = jax.devices()[0].client.platform_version + except Exception: # pylint: disable=broad-exception-caught + return "aiter" + if platform == "gpu" and "cuda" in backend.lower(): + return "flashinfer" + return "aiter" + + +def _call_backend(backend, query, k_pool, v_pool, plan, scale, causal): + """The only vendor-specific code in this file.""" + if backend == "aiter": + # Imported lazily: MaxText must not require jax-aiter to be installed for + # any other attention mode. + # pylint: disable=import-outside-toplevel + from jax_aiter.ops.append_kv import append_kv + from jax_aiter.ops.paged_attention import paged_attention + from jax_aiter.ops.paged_prefill import paged_prefill + + return append_kv, paged_attention, paged_prefill + + if backend == "flashinfer": + raise NotImplementedError( + "paged_attention_backend='flashinfer' is not wired up yet; only 'aiter' is implemented. " + "The MaxText side is backend-agnostic, so adding it is a change to this function alone." + ) + + raise ValueError(f"unknown paged_attention_backend {backend!r}; expected 'auto', 'aiter' or 'flashinfer'") + + +def paged_attention_step(query, key, value, k_pool, v_pool, plan, *, backend="auto", scale=None, causal=True): + """Write this step's K/V into the pool, then attend over it. + + Returns `(output, [k_pool, v_pool])`. The pools are donated and returned + aliased, so callers must rebind rather than keep the originals. + """ + append_kv, paged_attention, paged_prefill = _call_backend( + resolve_backend(backend), query, k_pool, v_pool, plan, scale, causal + ) + + k_pool, v_pool = append_kv(key, value, plan.slot_mapping, k_pool, v_pool) + + if plan.is_decode: + out = paged_attention( + query, + k_pool, + v_pool, + plan.kv_indptr, + plan.kv_page_indices, + plan.kv_last_page_lens, + max_seq_len=plan.max_seqlen_k, + scale=scale, + ) + else: + out = paged_prefill( + query, + k_pool, + v_pool, + plan.cu_seqlens_q, + plan.kv_indptr, + plan.kv_page_indices, + plan.kv_last_page_lens, + max_seqlen_q=plan.max_seqlen_q, + max_seqlen_k=plan.max_seqlen_k, + scale=scale, + causal=causal, + ) + return out, [k_pool, v_pool] diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index e16d0de4aa..6cafc9a4a3 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2544,6 +2544,142 @@ def test_sliding_window_attention(self): ) ) + def _gpu_paged_config(self, seq_len, **overrides): + arguments = self.config_arguments.copy() + arguments["max_target_length"] = seq_len + arguments["max_prefill_predict_length"] = seq_len + arguments["scan_layers"] = False + arguments.update(overrides) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **arguments) + + def _build_attention(self, config, seq_len, model_mode): + dummy = jnp.ones((self.global_batch_size, seq_len, config.base_emb_dim)) + return Attention( + config=config, + num_query_heads=config.num_query_heads, + num_kv_heads=config.num_kv_heads, + head_dim=config.head_dim, + max_target_length=seq_len, + max_prefill_predict_length=seq_len, + inputs_q_shape=dummy.shape, + inputs_kv_shape=dummy.shape, + mesh=self.mesh, + attention_kernel="dot_product", + dtype=config.dtype, + dropout_rate=config.dropout_rate, + model_mode=model_mode, + rngs=self.nnx_rng, + ) + + @pytest.mark.gpu_only + def test_gpu_paged_matches_the_dense_path(self): + """`gpu_paged` prefill must agree with dense attention over the same tokens. + + Both sides see identical weights and identical inputs, and a full prefill of + a whole sequence is exactly causal attention over that sequence, so the two + are comparable directly. The dense side runs in TRAIN mode specifically so + neither module allocates a dense KV cache, which keeps the two parameter + trees identical and lets `nnx.update` copy weights across. + """ + try: + import jax_aiter # pylint: disable=import-outside-toplevel,unused-import + from jax_aiter.ffi.registry import standalone_symbol_available # pylint: disable=import-outside-toplevel + except ImportError: + self.skipTest("jax-aiter is not importable; set PYTHONPATH to the jax-aiter checkout") + for symbol in ("AppendKvJA", "PagedPrefillJA"): + if not standalone_symbol_available(symbol): + self.skipTest(f"{symbol} is not built; run 'make -f Makefile.kv ja_kv' in jax-aiter") + + from maxtext.inference.kv_common import KvPageTableV1 # pylint: disable=import-outside-toplevel + + seq_len, tokens_per_page = 128, 16 + dense_cfg = self._gpu_paged_config(seq_len, attention="dot_product") + paged_cfg = self._gpu_paged_config(seq_len, attention="gpu_paged") + + batch = self.global_batch_size + lnx = jax.random.normal(self.rng, (batch, seq_len, dense_cfg.base_emb_dim), dtype=dense_cfg.dtype) + # One segment per request, positions in order: the plain causal case, which + # is what a full prefill of the sequence computes. + segment_ids = jnp.ones((batch, seq_len), dtype=jnp.int32) + positions = jnp.broadcast_to(jnp.arange(seq_len, dtype=jnp.int32), (batch, seq_len)) + + dense = self._build_attention(dense_cfg, seq_len, MODEL_MODE_TRAIN) + dense_out, _ = dense( + lnx, + lnx, + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + paged = self._build_attention(paged_cfg, seq_len, MODEL_MODE_PREFILL) + nnx.update(paged, nnx.state(dense)) + self.assertIsNone(paged.KVCache_0, "gpu_paged must not allocate the dense KV cache") + + pages_per_request = seq_len // tokens_per_page + # Page 0 is the padding sentinel and is never handed out. + page_ids = [ + [1 + r * pages_per_request + p for p in range(pages_per_request)] for r in range(batch) + ] + page_table = KvPageTableV1( + page_ids=page_ids, + seq_lens=np.full((batch,), seq_len, np.int32), + query_lens=np.full((batch,), seq_len, np.int32), + write_positions=np.tile(np.arange(seq_len, dtype=np.int32), batch), + request_order=np.arange(batch, dtype=np.int32), + ) + pool_shape = (1 + batch * pages_per_request, tokens_per_page, paged_cfg.num_kv_heads, paged_cfg.head_dim) + pools = [jnp.zeros(pool_shape, paged_cfg.dtype), jnp.zeros(pool_shape, paged_cfg.dtype)] + + paged_out, updated = paged( + lnx, + lnx, + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_PREFILL, + kv_cache=pools, + attention_metadata=page_table, + ) + + self.assertEqual(len(updated), 2, "the step must return both pools") + self.assertEqual(paged_out.shape, dense_out.shape) + + got = np.asarray(paged_out, np.float32) + want = np.asarray(dense_out, np.float32) + # One ulp of the output dtype at the scale of the tensor, plus the usual + # relative term. A purely element-wise relative metric is meaningless here: + # attention output passes through zero, where one ulp reads as a 100% miss. + # + # One ulp is what the two paths actually achieve, so the bound is set to it + # rather than to a comfortable multiple of it. That is deliberate — the + # paged path is supposed to be arithmetically the same attention over the + # same numbers, so anything worse than the dtype's own resolution is a real + # difference and should fail. The bound has already earned this: it caught a + # double-scaling bug where passing the kernel's default 1/sqrt(head_dim) on + # top of MaxText's already-scaled query produced an error fifty times larger. + eps = float(jnp.finfo(paged_cfg.dtype).eps) + atol = eps * max(float(np.abs(want).max()), 1e-6) + self.assertTrue( + np.all(np.abs(got - want) <= atol + eps * np.abs(want)), + f"gpu_paged disagrees with the dense path: max abs error " + f"{float(np.abs(got - want).max()):.5f}, atol {atol:.5f}", + ) + + @pytest.mark.gpu_only + def test_gpu_paged_dry_run_passes_through(self): + """Model init and JIT tracing reach the branch with nothing to attend over.""" + seq_len = 128 + cfg = self._gpu_paged_config(seq_len, attention="gpu_paged") + paged = self._build_attention(cfg, seq_len, MODEL_MODE_PREFILL) + + query = jnp.ones((self.global_batch_size, seq_len, cfg.num_query_heads, cfg.head_dim), cfg.dtype) + out, cache = paged.forward_serve_gpu_paged(query, query, query, kv_pools=None, metadata=None) + + self.assertEqual(cache, []) + self.assertEqual(out.shape, (self.global_batch_size * seq_len, cfg.num_query_heads, cfg.head_dim)) + @pytest.mark.skip(reason="Requires `vllm-tpu` package which is not yet a MaxText dependency.") @pytest.mark.tpu_only @mock.patch("tpu_inference.layers.common.attention_interface.sharded_ragged_paged_attention", create=True) diff --git a/tests/unit/gpu_paged_attention_test.py b/tests/unit/gpu_paged_attention_test.py new file mode 100644 index 0000000000..34b6c0935f --- /dev/null +++ b/tests/unit/gpu_paged_attention_test.py @@ -0,0 +1,201 @@ +"""Tests for the vendor-neutral half of `attention: "gpu_paged"`. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +These cover the metadata conversion and the config guards, none of which needs a +kernel. The end-to-end numerical check against the dense path lives in +`attention_test.py`, which already has the mesh and config fixtures. + +Everything is marked `gpu_only`. That is not because these need a GPU -- they do +not -- but because `tests/conftest.py` auto-marks unmarked tests `cpu_only` and +then skips `cpu_only` on any accelerator testbed. An unmarked test here would +report success on the machines we actually care about without ever running. +""" + +import sys +import unittest + +from absl.testing import parameterized +import numpy as np +import pytest + +import jax.numpy as jnp + +from maxtext.configs import pyconfig +from maxtext.inference.kv_common import KvPageTableV1 +from maxtext.layers import gpu_paged_attention as gpa +from tests.utils.test_helpers import get_test_config_path + +TOKENS_PER_PAGE = 16 + + +class _VllmMetadata: + """The three fields the vLLM-shaped path reads, without tpu_inference.""" + + def __init__(self, block_tables, seq_lens, query_start_loc): + self.block_tables = block_tables + self.seq_lens = seq_lens + self.query_start_loc = query_start_loc + + +@pytest.mark.gpu_only +class MetadataConversionTest(parameterized.TestCase): + """Both accepted metadata shapes must produce the same page bookkeeping.""" + + def test_neutral_page_table_round_trips(self): + table = KvPageTableV1( + page_ids=[[1, 2], [3]], + seq_lens=np.asarray([20, 5], np.int32), + query_lens=np.asarray([20, 5], np.int32), + write_positions=np.concatenate([np.arange(20), np.arange(5)]).astype(np.int32), + request_order=np.asarray([0, 1], np.int32), + ) + plan = gpa.plan_from_neutral(table, TOKENS_PER_PAGE) + + np.testing.assert_array_equal(np.asarray(plan.kv_indptr), [0, 2, 3]) + np.testing.assert_array_equal(np.asarray(plan.kv_page_indices), [1, 2, 3]) + np.testing.assert_array_equal(np.asarray(plan.kv_last_page_lens), [4, 5]) + np.testing.assert_array_equal(np.asarray(plan.cu_seqlens_q), [0, 20, 25]) + self.assertEqual(plan.max_seqlen_q, 20) + self.assertEqual(plan.max_seqlen_k, 20) + self.assertFalse(plan.is_decode) + + def test_vllm_metadata_packs_the_padded_block_table(self): + """The 2D table is row-padded; the kernels want the ids packed contiguously.""" + # Two requests of 20 and 5 tokens, holding 2 and 1 pages. Row 1 is padded. + md = _VllmMetadata( + block_tables=jnp.asarray([[1, 2], [3, 0]], jnp.int32), + seq_lens=jnp.asarray([20, 5], jnp.int32), + query_start_loc=jnp.asarray([0, 20, 25], jnp.int32), + ) + plan = gpa.plan_from_vllm(md, TOKENS_PER_PAGE, total_tokens=25, max_seqlen_k=20) + + np.testing.assert_array_equal(np.asarray(plan.kv_indptr), [0, 2, 3]) + # Only the first three entries are addressed by kv_indptr; the tail is slack. + np.testing.assert_array_equal(np.asarray(plan.kv_page_indices)[:3], [1, 2, 3]) + np.testing.assert_array_equal(np.asarray(plan.kv_last_page_lens), [4, 5]) + + def test_both_shapes_agree(self): + """The two conversions describe the same batch, so they must coincide.""" + table = KvPageTableV1( + page_ids=[[5, 9], [7]], + seq_lens=np.asarray([18, 12], np.int32), + query_lens=np.asarray([18, 12], np.int32), + write_positions=np.concatenate([np.arange(18), np.arange(12)]).astype(np.int32), + request_order=np.asarray([0, 1], np.int32), + ) + md = _VllmMetadata( + block_tables=jnp.asarray([[5, 9], [7, 0]], jnp.int32), + seq_lens=jnp.asarray([18, 12], jnp.int32), + query_start_loc=jnp.asarray([0, 18, 30], jnp.int32), + ) + + a = gpa.plan_from_neutral(table, TOKENS_PER_PAGE) + b = gpa.plan_from_vllm(md, TOKENS_PER_PAGE, total_tokens=30, max_seqlen_k=18) + + n_pages = int(np.asarray(a.kv_indptr)[-1]) + np.testing.assert_array_equal(np.asarray(a.kv_indptr), np.asarray(b.kv_indptr)) + np.testing.assert_array_equal( + np.asarray(a.kv_page_indices)[:n_pages], np.asarray(b.kv_page_indices)[:n_pages] + ) + np.testing.assert_array_equal(np.asarray(a.kv_last_page_lens), np.asarray(b.kv_last_page_lens)) + np.testing.assert_array_equal(np.asarray(a.slot_mapping), np.asarray(b.slot_mapping)) + + def test_slot_mapping_places_an_appended_token_after_its_context(self): + """Decode: the new token lands at the first free offset, not at position 0.""" + md = _VllmMetadata( + block_tables=jnp.asarray([[1, 2], [3, 0]], jnp.int32), + seq_lens=jnp.asarray([21, 6], jnp.int32), # one token longer than the prefill above + query_start_loc=jnp.asarray([0, 1, 2], jnp.int32), + ) + plan = gpa.plan_from_vllm(md, TOKENS_PER_PAGE, total_tokens=2, max_seqlen_k=21) + + # request 0: position 20 -> page_ids[20 // 16] = 2, offset 4 -> 2*16 + 4 + # request 1: position 5 -> page_ids[5 // 16] = 3, offset 5 -> 3*16 + 5 + np.testing.assert_array_equal(np.asarray(plan.slot_mapping), [2 * 16 + 4, 3 * 16 + 5]) + self.assertTrue(plan.is_decode) + + def test_unrecognised_metadata_is_rejected_by_shape(self): + with self.assertRaisesRegex(TypeError, "neutral KvPageTableV1|vLLM-shaped"): + gpa.build_plan(object(), TOKENS_PER_PAGE, total_tokens=1, max_seqlen_k=1) + + def test_dispatch_recognises_each_shape(self): + table = KvPageTableV1( + page_ids=[[1]], + seq_lens=np.asarray([4], np.int32), + query_lens=np.asarray([4], np.int32), + write_positions=np.arange(4, dtype=np.int32), + request_order=np.asarray([0], np.int32), + ) + md = _VllmMetadata( + block_tables=jnp.asarray([[1]], jnp.int32), + seq_lens=jnp.asarray([4], jnp.int32), + query_start_loc=jnp.asarray([0, 4], jnp.int32), + ) + self.assertTrue(gpa.is_neutral_page_table(table)) + self.assertFalse(gpa.is_vllm_metadata(table)) + self.assertTrue(gpa.is_vllm_metadata(md)) + self.assertFalse(gpa.is_neutral_page_table(md)) + + +@pytest.mark.gpu_only +class BackendSelectionTest(parameterized.TestCase): + + def test_explicit_backend_is_taken_literally(self): + self.assertEqual(gpa.resolve_backend("aiter"), "aiter") + self.assertEqual(gpa.resolve_backend("flashinfer"), "flashinfer") + + def test_flashinfer_fails_with_a_pointer_rather_than_silently(self): + with self.assertRaisesRegex(NotImplementedError, "flashinfer"): + gpa.paged_attention_step(None, None, None, None, None, None, backend="flashinfer") + + def test_unknown_backend_is_rejected(self): + with self.assertRaisesRegex(ValueError, "unknown paged_attention_backend"): + gpa.paged_attention_step(None, None, None, None, None, None, backend="nope") + + +@pytest.mark.gpu_only +class ConfigGuardTest(parameterized.TestCase): + """`gpu_paged` must be accepted, and its one pathological combination refused.""" + + base_arguments = { + "per_device_batch_size": 1.0, + "run_name": "test_gpu_paged_config", + "enable_checkpointing": False, + "max_target_length": 128, + } + + def _config(self, **overrides): + arguments = dict(self.base_arguments) + arguments.update(overrides) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **arguments) + + def test_gpu_paged_is_an_accepted_attention_value(self): + cfg = self._config(attention="gpu_paged", scan_layers=False) + self.assertEqual(cfg.attention, "gpu_paged") + self.assertEqual(cfg.paged_attention_backend, "auto") + + def test_scan_layers_is_refused(self): + """Scanning stacks the caches, which would copy the whole pool every step.""" + with self.assertRaisesRegex(ValueError, "scan_layers=False"): + self._config(attention="gpu_paged", scan_layers=True) + + def test_other_attention_kernels_are_unaffected_by_the_guard(self): + cfg = self._config(attention="dot_product", scan_layers=True) + self.assertEqual(cfg.attention, "dot_product") + + +if __name__ == "__main__": + unittest.main() From b592e0686edc6e57c7bcc2540798e809ddd08153 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Wed, 26 Aug 2026 01:11:52 +0000 Subject: [PATCH 03/21] feat(inference): add the paged KV control plane, MaxEngine paged path and serving benchmark (M4, M4.5) Adds maxtext/inference/kv_control/ and kv_execution/: the host page allocator, request-to-page map, per-step metadata builder, shape bucketing and continuous-batching driver, plus MaxEngine sibling entry points that run a model on the pool and a serving harness that measures it. kv_control/ is host-only by construction -- stdlib, numpy and kv_common, nothing else. The static AST check that enforced that for kv_common now covers both layers and the direction that only matters once kv_execution exists: neither lower layer may import it, or jax is back in the control plane's dependency graph. That keeps the control plane testable with no accelerator and makes extraction a directory move rather than an archaeology exercise. Reimplemented from sglang-jax's allocator as a reference design rather than ported, with three deliberate divergences. The allocator deals in pages, not token indices. KvPageTableV1 already carries explicit write positions, so the reference's three-part extend fill falls out of position // tokens_per_page and never needs computing; what remains is the page count. The page map is then tokens_per_page times smaller than ReqToTokenPool, 2 MB against 33 MB at 256 requests and 32k context. A double free is diagnosed rather than deduplicated. An allocation bitmap separates the same page appearing twice in one call, which is just token-index deduplication and is fine, from freeing a page that is not currently allocated, which is a use-after-free. Stale request handles are caught the same way, by an epoch on each page-map row. And recycled pages are tracked. A freed page holds the previous occupant's KV until something overwrites it, and neither sglang-jax nor vLLM zeroes them, so this is net-new rather than inherited. A page is dirty from the moment it is freed until a caller confirms it was overwritten, and build_page_table refuses to describe a dirty page. That fixes one order per step -- reserve, scrub, confirm, build, run -- and makes a missed scrub an exception instead of one request reading another's KV and producing plausible tokens. Confirmation is a separate call because it is the single point where the guarantee can be broken. Bucketing pads batch, tokens and the gather table to power-of-two ladders so a churning batch traces a fixed set of shapes; 24 mixed-length requests present 21+ raw shapes and compile 7. Two refinements on the design: the gather table is derived from batch and length rather than given its own ladder, which removes a dimension from the cross product instead of adding one, and max_seqlen_k needs a ladder that was not anticipated, because the kernels take it as a static configuration value. Decode running out of pages preempts the newest request by recomputation. Absent from the design and required by it: stable under churn cannot hold if the loop can deadlock, and with a full pool and every live request needing a page there is otherwise no way forward. MaxEngine gets a parallel surface, not a rewrite: init_paged_runtime, prefill_paged, generate_paged. init_decode_state, _insert_jit and generate are specific to the dense two-region cache, and rebuilding them around pages is a much larger change than it reads as. release(handle) becomes the real API with release_pages(slot) reduced to a shim, because a slot encodes the dense model's assumption that a request owns one fixed reservation. Wiring a whole model through it exposed two bugs in the M3 code that a layer-level test cannot reach. Transformer.__call__ returned kv_caches only for the vLLM modes, so on the gpu_paged path the aliased pool handles were dropped and the caller was left holding a deleted array. And _nnx_run_model, the default pure_nnx path, had no way to carry a pool at all. Parity is teacher-forced rather than a trajectory comparison. Logits are computed in bfloat16, so the top-two gap is quantised and exact ties occur -- not at reproducible places, since bf16 quantises an accumulation whose order XLA may vary between processes. At a tied step argmax is decided by tie-breaking, so two correct implementations diverge and every later token follows the coin flip. Replaying the paged path's own tokens through cacheless forward passes and asserting each was an argmax is tie-tolerant, never diverges, and is the stronger claim. It holds across a 53-token context spanning four pages. The benchmark reports paged at 8.5x dense throughput and 20x better p50 TTFT at an equal KV budget, with 1.14x page fragmentation and no leaks. Getting there needed two corrections worth recording. Counting compiled shapes certified a run that was three-quarters compile time, because padding a variable-length array with jnp compiles once per length and no shape bucket changes; per-call arrays are now built in numpy, and reportability is decided by repeating the workload and comparing passes, which cannot miss a category. And warmup must sweep the sequence-length ladder as well as batch and tokens, since a context growing past a rung presents a new program at an already-compiled batch width. 147 CPU tests and 28 GPU tests, the CPU ones runnable with no accelerator. Co-authored-by: Cursor --- benchmarks/api_server/maxtext_generator.py | 6 +- benchmarks/paged_kv/run_serving_benchmark.py | 263 ++++++ src/maxtext/common/gcloud_stub.py | 104 ++- src/maxtext/configs/types.py | 56 +- src/maxtext/inference/kv_common/page_table.py | 6 +- src/maxtext/inference/kv_control/__init__.py | 71 ++ src/maxtext/inference/kv_control/allocator.py | 310 ++++++++ .../inference/kv_control/control_plane.py | 215 +++++ .../inference/kv_control/logical_block.py | 157 ++++ src/maxtext/inference/kv_control/metadata.py | 104 +++ src/maxtext/inference/kv_control/page_map.py | 229 ++++++ src/maxtext/inference/kv_control/protocols.py | 95 +++ src/maxtext/inference/kv_control/request.py | 101 +++ .../inference/kv_execution/__init__.py | 57 ++ .../inference/kv_execution/benchmark.py | 538 +++++++++++++ .../inference/kv_execution/bucketing.py | 201 +++++ src/maxtext/inference/kv_execution/driver.py | 381 +++++++++ .../inference/kv_execution/engine_adapter.py | 186 +++++ .../inference/kv_execution/layout_builder.py | 74 ++ .../inference/kv_execution/pool_factory.py | 104 +++ .../inference/kv_execution/pool_ops.py | 110 +++ .../inference/kv_execution/step_view.py | 163 ++++ src/maxtext/inference/maxengine/maxengine.py | 443 ++++++++++- src/maxtext/models/models.py | 20 + tests/unit/attention_test.py | 3 + tests/unit/gpu_paged_attention_test.py | 28 +- tests/unit/gpu_paged_decode_parity_test.py | 422 ++++++++++ tests/unit/kv_common_test.py | 133 +--- tests/unit/kv_control_test.py | 746 ++++++++++++++++++ tests/unit/kv_execution_test.py | 526 ++++++++++++ tests/unit/kv_import_rule_test.py | 204 +++++ tests/unit/kv_paged_runtime_test.py | 397 ++++++++++ 32 files changed, 6297 insertions(+), 156 deletions(-) create mode 100644 benchmarks/paged_kv/run_serving_benchmark.py create mode 100644 src/maxtext/inference/kv_control/__init__.py create mode 100644 src/maxtext/inference/kv_control/allocator.py create mode 100644 src/maxtext/inference/kv_control/control_plane.py create mode 100644 src/maxtext/inference/kv_control/logical_block.py create mode 100644 src/maxtext/inference/kv_control/metadata.py create mode 100644 src/maxtext/inference/kv_control/page_map.py create mode 100644 src/maxtext/inference/kv_control/protocols.py create mode 100644 src/maxtext/inference/kv_control/request.py create mode 100644 src/maxtext/inference/kv_execution/__init__.py create mode 100644 src/maxtext/inference/kv_execution/benchmark.py create mode 100644 src/maxtext/inference/kv_execution/bucketing.py create mode 100644 src/maxtext/inference/kv_execution/driver.py create mode 100644 src/maxtext/inference/kv_execution/engine_adapter.py create mode 100644 src/maxtext/inference/kv_execution/layout_builder.py create mode 100644 src/maxtext/inference/kv_execution/pool_factory.py create mode 100644 src/maxtext/inference/kv_execution/pool_ops.py create mode 100644 src/maxtext/inference/kv_execution/step_view.py create mode 100644 tests/unit/gpu_paged_decode_parity_test.py create mode 100644 tests/unit/kv_control_test.py create mode 100644 tests/unit/kv_execution_test.py create mode 100644 tests/unit/kv_import_rule_test.py create mode 100644 tests/unit/kv_paged_runtime_test.py diff --git a/benchmarks/api_server/maxtext_generator.py b/benchmarks/api_server/maxtext_generator.py index 3b5060b476..7e5389230e 100644 --- a/benchmarks/api_server/maxtext_generator.py +++ b/benchmarks/api_server/maxtext_generator.py @@ -446,7 +446,11 @@ def _run_generation_loop(self, streams, decode_state, rng, max_tokens, stop, tem stream.finished = True if is_eos or stop_sequence_found: stream.finish_reason = "stop" - if getattr(self.config, "attention", "") == "paged": + # "paged" is the vestigial value that never allocated a pool; "gpu_paged" + # is the one that does. Both are listed rather than the guard being + # replaced, because release_pages is a no-op without a paged runtime and + # dropping "paged" here would be a behaviour change unrelated to this work. + if getattr(self.config, "attention", "") in ("paged", "gpu_paged"): self.engine.release_pages(slot=slot_idx) return decode_state diff --git a/benchmarks/paged_kv/run_serving_benchmark.py b/benchmarks/paged_kv/run_serving_benchmark.py new file mode 100644 index 0000000000..fe80f242dd --- /dev/null +++ b/benchmarks/paged_kv/run_serving_benchmark.py @@ -0,0 +1,263 @@ +"""Run the paged-versus-dense serving benchmark on one GPU. + +The Section 6.2 A/B: same model, same precision, same request trace, one engine +on the dense two-region cache and one on the page pool. This is the comparison +that proves or refutes the thesis in Section 1.1, and it is deliberately a single +device -- sharding is a separate milestone and would only add a variable. + +Usage, from the MaxText root: + + DECOUPLE_GCLOUD=TRUE JA_ROOT_DIR=/path/to/jax-aiter \\ + PYTHONPATH=src:/path/to/jax-aiter \\ + python3 benchmarks/paged_kv/run_serving_benchmark.py --mode both + +`--mode paged` needs neither JetStream nor a checkpoint. `--mode dense` needs +JetStream for real, because the dense `_prefill_jit` returns its `ResultTokens` +from inside `jit` and the decoupled stub is not a pytree; the paged path builds +its own outside `jit` and is unaffected. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import argparse +import copy +import sys + +import jax +from flax import nnx +from flax.linen import partitioning as nn_partitioning + +from maxtext.common.common_types import MODEL_MODE_PREFILL +from maxtext.configs import pyconfig +from maxtext.inference.kv_execution import benchmark +from maxtext.utils import maxtext_utils, model_creation_utils + +# Small enough to run in a couple of minutes, large enough that the KV footprint +# is the binding constraint rather than a rounding error. head_dim 128 with equal +# query and KV head counts keeps gqa_ratio at 1, inside the prebuilt kernel set. +BASE_CONFIG = { + "base_emb_dim": 512, + "base_mlp_dim": 1024, + "base_num_query_heads": 4, + "base_num_kv_heads": 4, + "base_num_decoder_layers": 4, + "head_dim": 128, + "vocab_size": 256, + "per_device_batch_size": 1.0, + "scan_layers": False, + "sparse_matmul": False, + # The paged kernels take bfloat16 or float16 only, so both arms use bfloat16. + # Comparing a bf16 pool against an fp32 dense cache would measure the dtype. + "dtype": "bfloat16", + "weight_dtype": "float32", + "decode_sampling_strategy": "greedy", + "enable_checkpointing": False, + "skip_jax_distributed_system": True, + "pure_nnx": True, +} + + +def build_config(mode: str, args) -> object: + """Config for one arm of the A/B. + + The two arms are given *equal KV memory*, which is what makes the comparison + fair and is the whole point of the exercise. The dense side commits + `slots x max_target_length`; the paged side is handed the same number of tokens + as a pool. Any concurrency difference is then attributable to fungibility + rather than to one arm having been given more memory. + """ + overrides = dict(BASE_CONFIG) + overrides["max_target_length"] = args.max_context + overrides["max_prefill_predict_length"] = args.max_prompt + # Model scale is the decisive variable for the throughput comparison, not a + # detail. The paged win runs through batch size amortising the per-step weight + # read, so it can only appear once weights dominate a step. At toy scale the + # extra per-step work is all that is left to measure. + overrides["base_num_decoder_layers"] = args.layers + overrides["base_emb_dim"] = args.emb_dim + overrides["base_mlp_dim"] = args.emb_dim * 2 + dense_kv_tokens = args.max_batch * args.max_context + + if mode == "dense": + overrides["attention"] = "dot_product" + overrides["per_device_batch_size"] = float(args.max_batch) + else: + overrides["attention"] = "gpu_paged" + overrides["paged_page_size"] = args.page_size + overrides["paged_num_blocks"] = dense_kv_tokens // args.page_size + overrides["paged_max_context_len"] = args.max_context + overrides["per_device_batch_size"] = float(args.max_batch) + + overrides["run_name"] = f"paged_bench_{mode}" + return pyconfig.initialize([sys.argv[0], args.config_path], **overrides) + + +def build_params(cfg, devices): + """Random weights on one device. No checkpoint, and no network.""" + mesh = jax.sharding.Mesh( + maxtext_utils.create_device_mesh(config=cfg, devices=devices), cfg.mesh_axes + ) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + model = model_creation_utils.create_model( + cfg, mesh, model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(params=0, dropout=0) + ) + _, params_state, _ = nnx.split(model, nnx.Param, ...) + return params_state + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mode", choices=("paged", "dense", "both"), default="paged") + parser.add_argument("--requests", type=int, default=24) + parser.add_argument("--mean-prompt", type=int, default=48) + parser.add_argument("--mean-output", type=int, default=32) + parser.add_argument("--max-batch", type=int, default=8, help="dense slot count, and the KV budget both arms get") + parser.add_argument( + "--paged-max-batch", + type=int, + default=0, + help="paged concurrency cap; 0 derives one high enough that the pool binds instead", + ) + parser.add_argument("--max-context", type=int, default=256) + parser.add_argument("--max-prompt", type=int, default=128) + parser.add_argument("--page-size", type=int, default=16) + parser.add_argument("--layers", type=int, default=4, help="decoder layers; scale this to find the crossover") + parser.add_argument("--emb-dim", type=int, default=512, help="model width; mlp is twice this") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--repeats", + type=int, + default=2, + help="passes over the trace; the last is reported and the spread is the reportability check", + ) + parser.add_argument("--json-out", default=None) + parser.add_argument( + "--config-path", + default="src/maxtext/configs/base.yml", + help="base.yml to layer overrides onto", + ) + args = parser.parse_args() + + # pylint: disable=import-outside-toplevel + from maxtext.inference.maxengine import maxengine + + devices = jax.devices()[:1] + print(f"device: {devices[0].device_kind} ({len(jax.devices())} visible, using 1)") + + results = {} + arms = ("paged", "dense") if args.mode == "both" else (args.mode,) + for arm in arms: + cfg = build_config(arm, args) + params_state = build_params(cfg, devices) + engine = maxengine.MaxEngine(cfg, devices) + params = engine.load_params(params=params_state) + + # A factory rather than one list: each repeat needs requests with no + # timestamps on them, and each arm needs an independent copy. + def fresh_trace(): + return copy.deepcopy( + benchmark.synthetic_trace(args.requests, args.mean_prompt, args.mean_output, seed=args.seed) + ) + + if arm == "paged": + # The batch cap must not bind, or the measurement reports the flag rather + # than the pool. The dense arm's concurrency is structurally its slot + # count; the paged arm's should be whatever the identical KV budget can + # actually hold, which is the entire claim being tested. + # Capped, and the cap is not cosmetic. The derived value grows with the KV + # budget, and a batch bucket above 32 combined with the upper + # sequence-length rungs falls outside the AOT-prebuilt pa_ragged set, so + # every such shape triggers an aiter JIT build. An unbounded derivation + # took over fifty minutes in warmup and produced nothing. Raise it + # deliberately with --paged-max-batch once the prebuild covers the shapes. + derived = max(args.max_batch, (args.max_batch * args.max_context) // max(args.mean_prompt, 1)) + paged_batch = args.paged_max_batch or min(derived, 32) + if not args.paged_max_batch and derived > paged_batch: + print( + f"note: capping paged concurrency at {paged_batch} (pool would allow ~{derived}); " + f"larger batch buckets need an aiter prebuild, so pass --paged-max-batch to override" + ) + engine.init_paged_runtime(max_requests=paged_batch, max_batched_tokens=args.max_prompt) + warmed = benchmark.warmup_paged( + engine, + params, + max_prompt=args.max_prompt, + max_batch=paged_batch, + # The longest context the trace can reach. Warming past it would + # compile shapes the run never presents; stopping short leaves + # compilation inside the measured window. + target_context=min( + args.max_context, int(args.mean_prompt * 1.6) + int(args.mean_output * 1.6) + 2 + ), + ) + summary = benchmark.run_repeated( + engine, + params, + fresh_trace, + max_batch=paged_batch, + warmed_shapes=warmed, + repeats=args.repeats, + ) + summary["kv_tokens_committed"] = summary.get("pool_capacity_tokens") + summary["batch_cap"] = paged_batch + else: + state = engine.init_decode_state() + benchmark.warmup_dense(engine, params, state, prompt_len=min(8, args.max_prompt)) + durations = [] + for _ in range(max(args.repeats, 1)): + summary = benchmark.run_dense(engine, params, fresh_trace(), max_batch=args.max_batch) + durations.append(summary["duration_s"]) + summary["repeat_durations_s"] = durations + summary["stability_ratio"] = benchmark.stability_ratio(durations) + summary["latency_is_reportable"] = benchmark.is_stable(durations) + summary["kv_tokens_committed"] = args.max_batch * args.max_context + summary["batch_cap"] = args.max_batch + + summary["config"] = { + "attention": cfg.attention, + "max_context": args.max_context, + "max_batch": args.max_batch, + "requests": args.requests, + "mean_prompt": args.mean_prompt, + "mean_output": args.mean_output, + } + results[arm] = summary + print(benchmark.report(arm, summary)) + print(f" KV tokens committed: {summary['kv_tokens_committed']}") + + if "paged" in results and "dense" in results: + paged, dense = results["paged"], results["dense"] + print("\n=== paged vs dense, equal KV memory ===") + ratio = paged["output_throughput_tok_per_s"] / max(dense["output_throughput_tok_per_s"], 1e-9) + print(f" output throughput x{ratio:.2f}") + print( + f" TTFT p50 {dense['ttft']['p50_ms']:.2f} -> {paged['ttft']['p50_ms']:.2f} ms" + f" ITL p50 {dense['itl']['p50_ms']:.2f} -> {paged['itl']['p50_ms']:.2f} ms" + ) + occ = paged.get("occupancy", {}) + if occ: + print( + f" concurrency: dense fixed at {dense['config']['max_batch']} slots," + f" paged reached {occ['max_concurrency']} on the same KV budget" + ) + + if args.json_out: + benchmark.write_json(args.json_out, results) + print(f"\nwrote {args.json_out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/maxtext/common/gcloud_stub.py b/src/maxtext/common/gcloud_stub.py index c87ba4123c..f01239085c 100644 --- a/src/maxtext/common/gcloud_stub.py +++ b/src/maxtext/common/gcloud_stub.py @@ -74,36 +74,86 @@ def _import_or_stub( # ---------------- JetStream ----------------- -def _jetstream_stubs(): - """Return lightweight stubs for JetStream modules.""" +class _StubEngine: # minimal base class stub + """Stub Engine accepting any initialization signature.""" - class Engine: # minimal base class stub - """Stub Engine accepting any initialization signature.""" + def __init__(self, *a, **k): # pylint: disable=unused-argument + pass - def __init__(self, *a, **k): # pylint: disable=unused-argument - pass - class ResultTokens: - """Container for result token arrays used by JetStream.""" - - def __init__( - self, - *args, - data=None, - tokens_idx=None, - valid_idx=None, - length_idx=None, - log_prob=None, - samples_per_slot: int | None = None, - **kwargs, - ): - del args, kwargs # unused - self.data = data - self.tokens_idx = tokens_idx - self.valid_idx = valid_idx - self.length_idx = length_idx - self.log_prob = log_prob - self.samples_per_slot = samples_per_slot +class _StubResultTokens: + """Container for result token arrays used by JetStream. + + Defined at module scope rather than inside `_jetstream_stubs` so there is + exactly one type to register as a pytree, however many times the stubs are + requested. Two classes with the same name would be two unrelated types. + """ + + def __init__( + self, + *args, + data=None, + tokens_idx=None, + valid_idx=None, + length_idx=None, + log_prob=None, + samples_per_slot: int | None = None, + **kwargs, + ): + del args, kwargs # unused + self.data = data + self.tokens_idx = tokens_idx + self.valid_idx = valid_idx + self.length_idx = length_idx + self.log_prob = log_prob + self.samples_per_slot = samples_per_slot + + +_stub_result_tokens_registered = False + + +def _register_stub_result_tokens(): + """Make the stub a pytree, because the real `ResultTokens` is one. + + `MaxEngine._prefill_jit` and `_generate_jit` return a `ResultTokens` from + *inside* `jit`, so the type has to be a registered pytree or the call fails + with "not a valid JAX type". Without this the stubs are enough to import + `maxengine` but not to run its dense prefill, which makes decoupled mode look + supported while the main serving path is unusable. + + Registration is lazy and idempotent: this module is imported early and should + not pull in jax by itself, and only the stub type is ever registered, so an + environment with real JetStream installed is unaffected. + """ + global _stub_result_tokens_registered # pylint: disable=global-statement + if _stub_result_tokens_registered: + return + try: + import jax # pylint: disable=import-outside-toplevel + except ImportError: + return + jax.tree_util.register_pytree_node( + _StubResultTokens, + # The arrays are children; the index tuples and sample count are static + # metadata, which is what lets them survive tracing unchanged. + lambda t: ((t.data, t.log_prob), (t.tokens_idx, t.valid_idx, t.length_idx, t.samples_per_slot)), + lambda aux, children: _StubResultTokens( + data=children[0], + log_prob=children[1], + tokens_idx=aux[0], + valid_idx=aux[1], + length_idx=aux[2], + samples_per_slot=aux[3], + ), + ) + _stub_result_tokens_registered = True + + +def _jetstream_stubs(): + """Return lightweight stubs for JetStream modules.""" + _register_stub_result_tokens() + Engine = _StubEngine # pylint: disable=invalid-name + ResultTokens = _StubResultTokens # pylint: disable=invalid-name # Tokenizer placeholders (unused in decoupled tests due to runtime guard). class TokenizerParameters: # pragma: no cover - placeholder diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 2ce99ad6fb..fdfbff77c3 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -610,6 +610,37 @@ class Attention(BaseModel): "vendor-neutral, so 'auto' picks aiter on ROCm and flashinfer on CUDA." ), ) + paged_page_size: PositiveInt = Field( + 16, + description=( + "Tokens per KV page for attention='gpu_paged' ('block_size' in vLLM vocabulary). Smaller " + "wastes less on partial final pages and costs more page-table indirection." + ), + ) + paged_num_blocks: NonNegativeInt = Field( + 0, + description=( + "Usable KV pages in the paged pool, which sets serving capacity: total tokens held is " + "paged_num_blocks * paged_page_size across all requests at once. One further page is " + "allocated and reserved as a padding target, so this is the count actually available. " + "Required when attention='gpu_paged'." + ), + ) + paged_max_context_len: NonNegativeInt = Field( + 0, + description=( + "Longest context a single paged request may reach. Bounds the per-request page-table row " + "and the sequence-length bucket ladder. Defaults to max_target_length when left at 0." + ), + ) + paged_poison_freed_pages: bool = Field( + False, + description=( + "Debug aid for attention='gpu_paged': fill a page with a recognisable sentinel when it is " + "freed, so a missed scrub produces an obviously wrong value rather than a plausible one. " + "Costs a pool write per freed page, so it is off by default." + ), + ) attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed", "block_diffusion"] = Field( "global", description="The variant of attention to use." ) @@ -3177,11 +3208,34 @@ def validate_gpu_paged_attention(self) -> "MaxTextConfig": on every step, which shows up as a mysterious slowdown rather than an error, so refuse it outright instead of quietly overriding it. """ - if self.attention == "gpu_paged" and self.scan_layers: + if self.attention != "gpu_paged": + return self + + if self.scan_layers: raise ValueError( "attention='gpu_paged' requires scan_layers=False. Scanning stacks the per-layer " "KV caches into the scan carry, which copies the whole paged pool every step." ) + + if self.paged_num_blocks < 1: + raise ValueError( + "attention='gpu_paged' requires paged_num_blocks to be set. It is the pool's capacity, so " + "there is no safe default: too small silently caps concurrency, and too large fails at " + f"allocation. Pool tokens = paged_num_blocks * paged_page_size (currently {self.paged_page_size})." + ) + + # Derived rather than required, since max_target_length already states the + # longest sequence the deployment intends to serve. + if self.paged_max_context_len == 0: + self.paged_max_context_len = self.max_target_length + + request_pages = -(-self.paged_max_context_len // self.paged_page_size) + if request_pages > self.paged_num_blocks: + raise ValueError( + f"a single paged_max_context_len={self.paged_max_context_len} request needs {request_pages} " + f"pages but paged_num_blocks is {self.paged_num_blocks}, so no request could ever finish. " + f"Raise paged_num_blocks or lower paged_max_context_len." + ) return self @model_validator(mode="after") diff --git a/src/maxtext/inference/kv_common/page_table.py b/src/maxtext/inference/kv_common/page_table.py index f030ef8117..9dd38d5438 100644 --- a/src/maxtext/inference/kv_common/page_table.py +++ b/src/maxtext/inference/kv_common/page_table.py @@ -42,7 +42,11 @@ class KvPageTableV1: for decode; the uncached suffix length for prefill. write_positions: int32 [num_tokens] absolute token index within its own sequence for each new token, flattened in request order. - request_order: int32 [num_reqs] the order requests appear in the batch. + request_order: int32 [num_reqs] identity of each batch row, as the + producer's own index for the request occupying it. Batch position i + describes request request_order[i]. Kernels do not read it; it is + what lets a caller map a row of output back to a request without + keeping a parallel list in step with the batch. """ page_ids: list[list[int]] = dataclasses.field(default_factory=list) diff --git a/src/maxtext/inference/kv_control/__init__.py b/src/maxtext/inference/kv_control/__init__.py new file mode 100644 index 0000000000..202dd0af79 --- /dev/null +++ b/src/maxtext/inference/kv_control/__init__.py @@ -0,0 +1,71 @@ +"""Semantic control plane for the paged KV runtime: who owns which pages. + +Host logic only, and deliberately so. Admission, page allocation, request-to-page +mapping and per-step metadata are all data-dependent irregular work over small +integer arrays -- the kind of thing that cannot live inside a traced computation +and does not want to. Keeping it here means the whole control plane is unit +testable on a machine with no accelerator. + +Import rule, enforced by CI: this package may import the standard library, +``numpy``, and ``maxtext.inference.kv_common``. It must never import ``jax``, +``jax_aiter``, ``maxtext.layers``, or ``maxtext.models``. Vendor kernel ABIs are +reached only from ``kv_execution``, one layer up. Two things follow, and both are +worth the cost of policing the rule: the layer stays CPU-testable in ordinary CI, +and if a second consumer ever wants it, extraction is a directory move rather +than an archaeology exercise. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from maxtext.inference.kv_control.allocator import DoubleFreeError, PagedBlockAllocator +from maxtext.inference.kv_control.control_plane import DirtyPageError, NativeKvControlPlane +from maxtext.inference.kv_control.logical_block import ( + LogicalBlock, + PageState, + PageStateError, + decode_needs_new_page, + last_page_occupancy, + new_pages_for_extend, + pages_for_tokens, + token_slot, +) +from maxtext.inference.kv_control.metadata import build_decode_table, build_page_table +from maxtext.inference.kv_control.page_map import PageCapacityError, PageMap, StaleRequestHandleError +from maxtext.inference.kv_control.protocols import KvControlPlane +from maxtext.inference.kv_control.request import RequestDescriptor, RequestHandle, RequestState + +__all__ = [ + "DirtyPageError", + "DoubleFreeError", + "KvControlPlane", + "LogicalBlock", + "NativeKvControlPlane", + "PageCapacityError", + "PageMap", + "PageState", + "PageStateError", + "PagedBlockAllocator", + "RequestDescriptor", + "RequestHandle", + "RequestState", + "StaleRequestHandleError", + "build_decode_table", + "build_page_table", + "decode_needs_new_page", + "last_page_occupancy", + "new_pages_for_extend", + "pages_for_tokens", + "token_slot", +] diff --git a/src/maxtext/inference/kv_control/allocator.py b/src/maxtext/inference/kv_control/allocator.py new file mode 100644 index 0000000000..a5940ad160 --- /dev/null +++ b/src/maxtext/inference/kv_control/allocator.py @@ -0,0 +1,310 @@ +"""Page-granular free list for the KV pool. + +Two design choices carried over from the reference implementation, and one +deliberate divergence. + +Carried over: the two-tier free list. `free()` does not return pages to the +allocation list but to a staging area, and the two are merged and re-sorted only +when an allocation would otherwise fail. The sort then happens once, at the +moment it actually buys something, and allocation pops from the front of the +sorted list so a request's pages stay as clustered as the pool's history allows. + +The staging area is a list of arrays rather than one array grown by +concatenation. Concatenating on each free would make the free path cost +proportional to everything already staged, so a workload that frees steadily +while the front list is still long -- a large pool serving short requests, which +is not an exotic case -- pays quadratically in the pool size between merges. That +is the cost the two tiers exist to avoid, so it would be an unfortunate place to +reintroduce it. + +Also carried over: one page id is reserved and never allocated. It is a landing +zone -- padded gather entries default to it and read harmlessly -- which is why +the pool must be zero-initialised rather than merely allocated. It buys nothing +against stale KV in a recycled *real* page; that is a separate obligation +discharged by writing a page's full readable extent before marking it readable. + +The divergence: the reference allocates in token indices, because its +request-to-token table is token-granular. Ours is page-granular and carries +explicit write positions, so the reference's three-part extend fill -- continue +the open page, then whole pages, then a partial page -- falls out of +`position // tokens_per_page` rather than needing to be computed. What remains +is the page *count*, which is `new_pages_for_extend`. `free_token_slots` exists +for callers that only have slots to hand. + +Two things the reference does not do, and this does. + +It diagnoses a double free. The reference set-differences the free lists, which +makes a double free idempotent rather than reported, so a page freed while +another request still holds it is invisible. Here an allocation bitmap separates +the two cases -- the same page appearing twice in one call is just token-index +deduplication and is fine, while freeing a page that is not currently allocated +is a use-after-free and raises. + +And it tracks which pages are dirty. A freed page still holds the previous +occupant's KV until something overwrites it, so handing it to a new request +without scrubbing it means the new request's readable extent covers bytes it +does not own. The reference has no notion of this at all. A page is dirty from +the moment it is freed until a caller confirms it has been scrubbed; since the +pool is zero-initialised, a page that has never been allocated starts clean, so +a fresh pool costs no scrubbing at all. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np + +from maxtext.inference.kv_common import KvStorageLayoutV1 + + +class DoubleFreeError(RuntimeError): + """A page was freed that the allocator does not consider allocated.""" + + +class PagedBlockAllocator: + """Hands out pool pages and takes them back. + + Holds no refcounts. Sharing a page between requests is prefix-index policy, + and keeping it out of here is what lets the free-list mechanics be reasoned + about -- and tested -- on their own. + """ + + def __init__( + self, + num_pages: int, + tokens_per_page: int, + padding_page_id: int = 0, + debug_mode: bool = False, + ): + if num_pages < 1: + raise ValueError(f"num_pages must be at least 1, got {num_pages}") + if tokens_per_page < 1: + raise ValueError(f"tokens_per_page must be at least 1, got {tokens_per_page}") + if not 0 <= padding_page_id < num_pages: + raise ValueError(f"padding_page_id {padding_page_id} is outside the pool's {num_pages} pages") + + self.num_pages = int(num_pages) + self.tokens_per_page = int(tokens_per_page) + self.padding_page_id = int(padding_page_id) + self.debug_mode = bool(debug_mode) + + all_pages = np.arange(self.num_pages, dtype=np.int32) + self._allocatable = all_pages[all_pages != self.padding_page_id] + self._free_pages = self._allocatable.copy() + self._released: list[np.ndarray] = [] + self._num_released = 0 + # int64 so a long-running server cannot wrap an epoch and make a stale + # reference look current again. + self._epochs = np.zeros((self.num_pages,), dtype=np.int64) + self._allocated = np.zeros((self.num_pages,), dtype=bool) + # Clean at startup because the pool is zero-initialised. That is a real + # dependency, not an assumption: `pool_factory` must allocate zeros, and the + # reserved padding page relies on the same thing. + self._dirty = np.zeros((self.num_pages,), dtype=bool) + + @classmethod + def from_layout(cls, layout: KvStorageLayoutV1, debug_mode: bool = False) -> "PagedBlockAllocator": + """Build from pool geometry, which is where these numbers are decided.""" + return cls( + num_pages=layout.num_pages, + tokens_per_page=layout.tokens_per_page, + padding_page_id=layout.padding_page_id, + debug_mode=debug_mode, + ) + + @property + def capacity_pages(self) -> int: + """Allocatable pages, excluding the reserved padding page.""" + return int(self._allocatable.size) + + @property + def num_free_pages(self) -> int: + """Pages available without a merge. Use `available_pages` for the real figure.""" + return int(self._free_pages.size) + + @property + def available_pages(self) -> int: + """Pages an allocation could reach, counting the staging area.""" + return int(self._free_pages.size) + self._num_released + + @property + def available_tokens(self) -> int: + return self.available_pages * self.tokens_per_page + + @property + def num_allocated_pages(self) -> int: + return int(np.count_nonzero(self._allocated)) + + def merge_released(self) -> None: + """Fold the staging area back into the allocation list, sorted. + + Public because a driver may want to pay this cost at a quiet moment rather + than in the middle of a step that is about to fail. + """ + if self._released: + self._free_pages = np.sort(np.concatenate([self._free_pages, *self._released])) + self._released = [] + self._num_released = 0 + + def alloc(self, num_pages: int) -> np.ndarray | None: + """Take `num_pages` pages, or return None if the pool cannot supply them. + + None rather than an exception: exhaustion is a scheduling outcome the driver + must branch on every step, not an error condition. + """ + if num_pages < 0: + raise ValueError(f"num_pages must be non-negative, got {num_pages}") + if num_pages == 0: + return np.empty((0,), dtype=np.int32) + + if num_pages > self._free_pages.size: + self.merge_released() + if num_pages > self._free_pages.size: + return None + + pages = self._free_pages[:num_pages].copy() + self._free_pages = self._free_pages[num_pages:] + self._epochs[pages] += 1 + self._allocated[pages] = True + if self.debug_mode: + self._check_invariants() + return pages + + def free(self, page_ids: Sequence[int] | np.ndarray) -> np.ndarray: + """Return pages to the staging list. + + Returns the deduplicated pages actually released, which is what a caller + needs in order to poison or zero them before they can be handed out again. + """ + pages = np.unique(np.asarray(page_ids, dtype=np.int32)) + if pages.size == 0: + return pages + + if pages[0] < 0 or pages[-1] >= self.num_pages: + raise ValueError(f"page ids {pages[0]}..{pages[-1]} fall outside the pool's {self.num_pages} pages") + if np.any(pages == self.padding_page_id): + raise ValueError( + f"page {self.padding_page_id} is the reserved padding page and is never allocated, " + f"so it cannot be freed" + ) + + unowned = pages[~self._allocated[pages]] + if unowned.size: + raise DoubleFreeError( + f"pages {unowned.tolist()} are not currently allocated. Either they were freed twice, " + f"or a stale handle is still naming pages that were reclaimed." + ) + + self._allocated[pages] = False + self._dirty[pages] = True + self._released.append(pages) + self._num_released += int(pages.size) + if self.debug_mode: + self._check_invariants() + return pages + + def free_token_slots(self, token_slots: Sequence[int] | np.ndarray) -> np.ndarray: + """Free the pages covering `token_slots`. + + Negative slots are dropped rather than rejected: a padded row's slot is -1 + by construction, so a caller passing a `slot_mapping` straight through is + behaving correctly, not making a mistake. + """ + slots = np.asarray(token_slots, dtype=np.int64) + slots = slots[slots >= 0] + if slots.size == 0: + return np.empty((0,), dtype=np.int32) + return self.free(np.unique(slots // self.tokens_per_page).astype(np.int32)) + + def dirty_among(self, page_ids: Sequence[int] | np.ndarray) -> np.ndarray: + """The subset of `page_ids` still holding a previous occupant's KV. + + The caller must overwrite these before any kernel is allowed to read them, + and say so via `mark_scrubbed`. Returning the subset rather than a flag is + what lets a step scrub only the pages it actually recycled: on a fresh pool + that is none of them. + """ + pages = np.asarray(page_ids, dtype=np.int32).reshape(-1) + if pages.size == 0: + return pages + return pages[self._dirty[pages]] + + def is_dirty(self, page_id: int) -> bool: + return bool(self._dirty[page_id]) + + @property + def num_dirty_pages(self) -> int: + return int(np.count_nonzero(self._dirty)) + + def mark_scrubbed(self, page_ids: Sequence[int] | np.ndarray) -> None: + """Record that `page_ids` have been overwritten and are safe to read. + + Only the caller can know this, because the pool is device memory and this + layer never touches a device. Calling it without having done the write is + the one way to defeat the guarantee, which is why it is a separate step + rather than a side effect of allocation. + """ + pages = np.asarray(page_ids, dtype=np.int32).reshape(-1) + if pages.size == 0: + return + if pages.min() < 0 or pages.max() >= self.num_pages: + raise ValueError(f"page ids {pages.min()}..{pages.max()} fall outside the pool's {self.num_pages} pages") + self._dirty[pages] = False + + def epoch_of(self, page_id: int) -> int: + """Current epoch of `page_id`, incremented on each allocation.""" + return int(self._epochs[page_id]) + + def is_allocated(self, page_id: int) -> bool: + return bool(self._allocated[page_id]) + + def holds(self, page_id: int, epoch: int) -> bool: + """Whether a reference taken at `epoch` still names a live allocation. + + False both for a page that has been freed and for one already handed to a + later request, which are the two ways a stale reference goes wrong. + """ + return bool(self._allocated[page_id]) and int(self._epochs[page_id]) == epoch + + def clear(self) -> None: + """Return every page to the free list and invalidate all outstanding references. + + Epochs advance rather than resetting, so a reference taken before the clear + cannot be mistaken for a current one afterwards. Pages that had been handed + out stay dirty: dropping the free list does not overwrite anything. + """ + self._free_pages = self._allocatable.copy() + self._released = [] + self._num_released = 0 + self._dirty |= self._allocated + self._allocated[:] = False + self._epochs += 1 + + def _check_invariants(self) -> None: + """Assert the free lists and the allocation bitmap still agree.""" + free_size = int(self._free_pages.size) + combined = np.concatenate([self._free_pages, *self._released]) + if np.unique(combined).size != free_size + self._num_released: + raise AssertionError("a page appears more than once across the allocation and staging lists") + if combined.size and np.any(self._allocated[combined]): + raise AssertionError("a page is marked allocated while sitting in a free list") + if self.num_allocated_pages + free_size + self._num_released != self.capacity_pages: + raise AssertionError( + f"page accounting lost pages: {self.num_allocated_pages} allocated + {free_size} free + " + f"{self._num_released} staged != {self.capacity_pages} allocatable" + ) diff --git a/src/maxtext/inference/kv_control/control_plane.py b/src/maxtext/inference/kv_control/control_plane.py new file mode 100644 index 0000000000..e0ed2350b7 --- /dev/null +++ b/src/maxtext/inference/kv_control/control_plane.py @@ -0,0 +1,215 @@ +"""The authoritative control plane when MaxText owns the pages. + +Assembles the allocator, the page map, and the metadata builder into the one +object a driver holds. Small on purpose: everything interesting lives in the +parts, and what this adds is two properties none of them can provide alone. + +**A batch reservation either succeeds for every request or changes nothing.** +That matters more than it looks. A reservation that half-succeeds leaves some +requests advanced and some not, and there is no page table describing that +state, so the caller's only recovery would be to unwind bookkeeping it does not +own. Allocating the batch's pages in a single request from the free list, after +checking every request can record them, removes the case entirely. + +**No page table naming a dirty page can be built.** A recycled page holds its +previous occupant's KV until something overwrites it, so a step must scrub what +it recycled before any kernel reads it. That gives a fixed order per step -- +reserve, scrub, confirm, build the table, run -- and `build_page_table` refuses +outright if the confirm has not happened. Making it an error rather than a +convention is the point: the failure it prevents is one request reading +another's KV, which produces plausible tokens and no diagnostic. + +The scrub covers the whole recycled page rather than just the tail beyond the +new writes. Exact last-page lengths already stop a well-behaved kernel reading +past valid data, so the redundancy is deliberate -- it costs one page write per +page-boundary crossing and removes any dependence on every kernel, present and +future, honouring that bound. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np + +from maxtext.inference.kv_common import KvPageTableV1, KvStorageLayoutV1 +from maxtext.inference.kv_control import metadata +from maxtext.inference.kv_control.allocator import PagedBlockAllocator +from maxtext.inference.kv_control.logical_block import new_pages_for_extend, pages_for_tokens +from maxtext.inference.kv_control.page_map import PageCapacityError, PageMap +from maxtext.inference.kv_control.request import RequestDescriptor, RequestHandle + + +class DirtyPageError(RuntimeError): + """A step would have read a page still holding another request's KV.""" + + +class NativeKvControlPlane: + """Page accounting owned by MaxText, satisfying `KvControlPlane`.""" + + def __init__( + self, + layout: KvStorageLayoutV1, + max_requests: int, + max_context_len: int, + debug_mode: bool = False, + ): + if max_context_len < 1: + raise ValueError(f"max_context_len must be at least 1, got {max_context_len}") + + self._layout = layout + self.max_context_len = int(max_context_len) + self.allocator = PagedBlockAllocator.from_layout(layout, debug_mode=debug_mode) + self.page_map = PageMap.from_layout(layout, max_requests=max_requests, max_context_len=max_context_len) + self._pending_scrub: list[np.ndarray] = [] + + # Caught here rather than as a mid-serving allocation failure, which would + # look like ordinary backpressure and never resolve. + pages_needed = pages_for_tokens(max_context_len, layout.tokens_per_page) + if pages_needed > self.allocator.capacity_pages: + raise ValueError( + f"a single {max_context_len}-token request needs {pages_needed} pages but the pool has only " + f"{self.allocator.capacity_pages} allocatable; raise the pool size or lower max_context_len" + ) + + @property + def layout(self) -> KvStorageLayoutV1: + return self._layout + + @property + def available_tokens(self) -> int: + return self.allocator.available_tokens + + @property + def num_live(self) -> int: + return self.page_map.num_live + + def admit(self, descriptor: RequestDescriptor) -> RequestHandle | None: + if descriptor.max_total_len > self.max_context_len: + raise ValueError( + f"request {descriptor.request_id!r} declares a maximum length of {descriptor.max_total_len} " + f"tokens, past the configured max_context_len of {self.max_context_len}" + ) + return self.page_map.admit(descriptor) + + def reserve( + self, + handles: Sequence[RequestHandle], + num_new_tokens: Sequence[int] | np.ndarray, + ) -> bool: + """Reserve pages for `num_new_tokens` per request, all or nothing.""" + counts = np.asarray(num_new_tokens, dtype=np.int64).reshape(-1) + if counts.size != len(handles): + raise ValueError(f"got {len(handles)} handles but {counts.size} token counts") + if counts.size == 0: + return True + if np.any(counts < 0): + raise ValueError(f"token counts must be non-negative, got {counts.tolist()}") + + tokens_per_page = self._layout.tokens_per_page + # Both loops run before anything is mutated: the first would otherwise leave + # a request advanced with no pages recorded for it if a later request in the + # same batch turned out not to fit its row. + per_request_pages = [] + for handle, count in zip(handles, counts): + prefix_len = self.page_map.seq_len(handle) + new_len = prefix_len + int(count) + if pages_for_tokens(new_len, tokens_per_page) > self.page_map.max_pages_per_request: + raise PageCapacityError( + f"request {handle.request_id!r} would reach {new_len} tokens, past the " + f"{self.max_context_len} it was admitted under" + ) + per_request_pages.append(new_pages_for_extend(prefix_len, new_len, tokens_per_page)) + + pages = self.allocator.alloc(int(sum(per_request_pages))) + if pages is None: + return False + + recycled = self.allocator.dirty_among(pages) + if recycled.size: + self._pending_scrub.append(recycled) + + offset = 0 + for handle, count, needed in zip(handles, counts, per_request_pages): + self.page_map.append_pages(handle, pages[offset : offset + needed]) + offset += needed + self.page_map.advance(handle, int(count)) + return True + + def pending_scrub(self) -> np.ndarray: + """Pages reserved this step that must be overwritten before any read. + + Reading this does not discharge the obligation; `confirm_scrubbed` does. + Idempotent, so a caller may inspect it without committing to anything. + """ + if not self._pending_scrub: + return np.empty((0,), dtype=np.int32) + return np.unique(np.concatenate(self._pending_scrub)) + + def confirm_scrubbed(self, page_ids: Sequence[int] | np.ndarray) -> None: + """Record that `page_ids` have been overwritten on the device. + + Call this only after the write has actually been issued. It is the single + point at which the guarantee can be broken, which is why it is explicit + rather than folded into reservation. + """ + self.allocator.mark_scrubbed(page_ids) + scrubbed = set(np.asarray(page_ids, dtype=np.int32).reshape(-1).tolist()) + remaining = [ + pages[~np.isin(pages, list(scrubbed))] for pages in self._pending_scrub + ] + self._pending_scrub = [pages for pages in remaining if pages.size] + + def reserve_decode(self, handles: Sequence[RequestHandle]) -> bool: + return self.reserve(handles, np.ones((len(handles),), dtype=np.int32)) + + def release(self, handle: RequestHandle) -> np.ndarray: + """Reclaim the request's pages and report them. + + The pages are returned so a caller can zero or poison them before they can + be handed out again. Nothing here does that -- the pool is device memory and + this layer never touches a device -- which is why the pages have to leave + through the return value rather than being quietly recycled. + """ + pages = self.page_map.release(handle) + if pages.size: + self.allocator.free(pages) + return pages + + def build_page_table( + self, + handles: Sequence[RequestHandle], + query_lens: Sequence[int] | np.ndarray, + ) -> KvPageTableV1: + """Describe this step, refusing to describe a page that is still dirty. + + The check is over the table's own page list, so it covers exactly the extent + a kernel is about to read -- no more, since pages trimmed as beyond the + current length are not readable this step, and no less. + """ + table = metadata.build_page_table(self.page_map, handles, query_lens) + unscrubbed = self.allocator.dirty_among(table.flat_page_indices()) + if unscrubbed.size: + raise DirtyPageError( + f"pages {unscrubbed[:8].tolist()} still hold a previous request's KV and would be readable " + f"by this step. Scrub the pages from pending_scrub() and call confirm_scrubbed() before " + f"building the table." + ) + return table + + def build_decode_table(self, handles: Sequence[RequestHandle]) -> KvPageTableV1: + return self.build_page_table(handles, np.ones((len(handles),), dtype=np.int32)) diff --git a/src/maxtext/inference/kv_control/logical_block.py b/src/maxtext/inference/kv_control/logical_block.py new file mode 100644 index 0000000000..d576dee479 --- /dev/null +++ b/src/maxtext/inference/kv_control/logical_block.py @@ -0,0 +1,157 @@ +"""Block identity, page states, and the page arithmetic everything else shares. + +The arithmetic here is small enough to look not worth centralising, and that is +exactly why it is centralised: `ceil(seq / page) - ceil(prefix / page)` inlined +in four places is four places to get an off-by-one wrong, and the symptom of +getting it wrong is a kernel reading a page the request does not own. + +One page state matters more than the others. A page is readable only once its +full readable extent has been written, because a recycled page still holds the +previous occupant's KV until something overwrites it. `WRITING` versus `READY` +is where that obligation is expressed, rather than being implicit in whether an +append happened to have run yet. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +import enum + + +class PageState(enum.Enum): + """Whether a page may be allocated, written, or read. + + `WRITING` covers both a fresh allocation and a partial last page being + extended, since in both cases some of the page's readable extent is not yet + valid for this request. A prefill-decode destination page awaiting a transfer + is the same condition with a different filler, and will reuse this state + rather than adding a parallel one. + """ + + FREE = "free" + WRITING = "writing" + READY = "ready" + + +_LEGAL_TRANSITIONS = { + PageState.FREE: frozenset({PageState.WRITING}), + # WRITING -> FREE is a request cancelled or preempted mid-prefill. + PageState.WRITING: frozenset({PageState.READY, PageState.FREE}), + # READY -> WRITING is a partial last page being extended by the next step. + PageState.READY: frozenset({PageState.WRITING, PageState.FREE}), +} + + +class PageStateError(RuntimeError): + """An illegal page state transition, which is always a control-plane bug.""" + + +def check_transition(current: PageState, target: PageState) -> None: + """Raise unless `current -> target` is legal. + + Self-transitions are allowed: re-marking a page WRITING as more tokens arrive + is the normal extend path, not a mistake. + """ + if current is target: + return + if target not in _LEGAL_TRANSITIONS[current]: + raise PageStateError(f"illegal page state transition {current.value} -> {target.value}") + + +@dataclasses.dataclass +class LogicalBlock: + """One page's worth of a request's context. + + Attributes: + page_id: the physical page backing this block. + epoch: the allocator's epoch for `page_id` at the moment it was allocated. + Comparing it against the allocator's current epoch is what turns a + use-after-free into an error instead of a plausible read of someone + else's KV. + num_tokens: tokens of this request written into the page so far. Never the + page's capacity unless the page is genuinely full; an over-stated value is + precisely how a kernel reads a recycled page's tail. + state: see `PageState`. + ref_count: holders of this block. Stays at 1 for the whole of M4, because + only one request can own a page until prefix sharing exists; the prefix + index is what raises it, and it lives outside the allocator on purpose so + that sharing policy and free-list mechanics stay separable. + """ + + page_id: int + epoch: int + num_tokens: int = 0 + state: PageState = PageState.WRITING + ref_count: int = 1 + + def set_state(self, target: PageState) -> None: + check_transition(self.state, target) + self.state = target + + @property + def is_readable(self) -> bool: + return self.state is PageState.READY + + +def pages_for_tokens(num_tokens: int, tokens_per_page: int) -> int: + """Pages needed to hold `num_tokens` contiguous tokens.""" + if num_tokens < 0: + raise ValueError(f"num_tokens must be non-negative, got {num_tokens}") + return -(-num_tokens // tokens_per_page) + + +def new_pages_for_extend(prefix_len: int, seq_len: int, tokens_per_page: int) -> int: + """Pages a request must acquire to grow from `prefix_len` to `seq_len`. + + The already-open last page is continued in place and costs nothing, which is + why this is a difference of two ceilings rather than a ceiling of the + difference. Those disagree exactly when the extend starts mid-page, which is + the common case. + """ + if seq_len < prefix_len: + raise ValueError(f"seq_len {seq_len} is shorter than prefix_len {prefix_len}") + return pages_for_tokens(seq_len, tokens_per_page) - pages_for_tokens(prefix_len, tokens_per_page) + + +def decode_needs_new_page(seq_len_after: int, tokens_per_page: int) -> bool: + """Whether appending the token that brings the context to `seq_len_after` crosses a page boundary.""" + if seq_len_after <= 0: + raise ValueError(f"seq_len_after must be positive, got {seq_len_after}") + return (seq_len_after - 1) % tokens_per_page == 0 + + +def last_page_occupancy(seq_len: int, tokens_per_page: int) -> int: + """Tokens occupying the final page of a `seq_len`-token context. + + A full final page reports the page size, not zero. + """ + if seq_len < 0: + raise ValueError(f"seq_len must be non-negative, got {seq_len}") + if seq_len == 0: + return 0 + remainder = seq_len % tokens_per_page + return remainder if remainder else tokens_per_page + + +def token_slot(page_id: int, position: int, tokens_per_page: int) -> int: + """Absolute pool slot of the token at sequence `position` held in `page_id`. + + Slots are `page_id * tokens_per_page + offset`, and downstream code depends on + that contiguity: it is what lets an append kernel scatter with a single index + array and what lets a page transfer address a page as one byte range. + """ + return page_id * tokens_per_page + position % tokens_per_page diff --git a/src/maxtext/inference/kv_control/metadata.py b/src/maxtext/inference/kv_control/metadata.py new file mode 100644 index 0000000000..0d69b24f84 --- /dev/null +++ b/src/maxtext/inference/kv_control/metadata.py @@ -0,0 +1,104 @@ +"""Turn live page bookkeeping into a `KvPageTableV1` for one step. + +The output is the neutral vocabulary, not a vendor shape. A backend converts it +into whatever flat arrays its kernels take, and that conversion is the only +place a vendor contract appears. + +The one subtlety worth stating plainly, because getting it wrong is silent: +a request's page list is trimmed to exactly the pages its current length needs. +A page table is read together with `kv_last_page_lens`, and that occupancy is +applied to the *last* page in the list. Hand over one page more than the length +requires and the occupancy lands on the wrong page, so the kernel reads a full +page of whatever the previous occupant left behind and then one token of real +data. The page table type permits over-supply -- its own validation only checks +that enough pages are present -- so the trim has to happen here. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np + +from maxtext.inference.kv_common import KvPageTableV1 +from maxtext.inference.kv_control.logical_block import pages_for_tokens +from maxtext.inference.kv_control.page_map import PageMap +from maxtext.inference.kv_control.request import RequestHandle + + +def build_page_table( + page_map: PageMap, + handles: Sequence[RequestHandle], + query_lens: Sequence[int] | np.ndarray, +) -> KvPageTableV1: + """Build one step's page table. + + Args: + page_map: the live bookkeeping. Sequence lengths are read from it and must + already include this step's tokens, since the table describes the state + the kernels will see rather than the state before the step. + handles: the requests in this batch, in the order the batch presents them. + query_lens: new tokens each request contributes this step. All ones for + decode; the uncached suffix length for prefill. + + Returns: + A validated `KvPageTableV1`. + """ + query_lens = np.asarray(query_lens, dtype=np.int32).reshape(-1) + if query_lens.size != len(handles): + raise ValueError(f"got {len(handles)} handles but {query_lens.size} query lengths") + + tokens_per_page = page_map.tokens_per_page + num_requests = len(handles) + page_ids: list[list[int]] = [] + seq_lens = np.zeros((num_requests,), dtype=np.int32) + request_order = np.zeros((num_requests,), dtype=np.int32) + positions: list[np.ndarray] = [] + + for i, handle in enumerate(handles): + seq_len = page_map.seq_len(handle) + query_len = int(query_lens[i]) + if query_len < 0: + raise ValueError(f"request {handle.request_id!r} has negative query length {query_len}") + if query_len > seq_len: + raise ValueError( + f"request {handle.request_id!r} contributes {query_len} tokens but its recorded length is " + f"{seq_len}; advance the length as pages are reserved, not afterwards" + ) + needed = pages_for_tokens(seq_len, tokens_per_page) + page_ids.append(page_map.pages(handle)[:needed].tolist()) + seq_lens[i] = seq_len + request_order[i] = handle.row + positions.append(np.arange(seq_len - query_len, seq_len, dtype=np.int32)) + + write_positions = ( + np.concatenate(positions).astype(np.int32) if positions else np.zeros((0,), dtype=np.int32) + ) + table = KvPageTableV1( + page_ids=page_ids, + seq_lens=seq_lens, + query_lens=query_lens, + write_positions=write_positions, + request_order=request_order, + ) + table.validate(tokens_per_page) + return table + + +def build_decode_table(page_map: PageMap, handles: Sequence[RequestHandle]) -> KvPageTableV1: + """One token per request, which is what makes decode a fixed-shape step.""" + return build_page_table(page_map, handles, np.ones((len(handles),), dtype=np.int32)) diff --git a/src/maxtext/inference/kv_control/page_map.py b/src/maxtext/inference/kv_control/page_map.py new file mode 100644 index 0000000000..29496f9c1a --- /dev/null +++ b/src/maxtext/inference/kv_control/page_map.py @@ -0,0 +1,229 @@ +"""Which pages each live request holds. + +A dense `[max_requests, max_pages_per_request]` int32 table plus a free list of +rows, following the reference design, with two changes. + +The reference's table is token-granular, `[max_requests, max_context_len]`, and +recovers a request's pages by reading its row and filtering zeros. This one is +page-granular, so it is `tokens_per_page` times smaller -- 256 requests at 32k +context is 33 MB there and 2 MB here at 16 tokens per page -- and it tracks each +row's length explicitly instead of filtering a sentinel. Explicit lengths also +mean the table stays correct if the reserved page id is ever something other +than zero, which sentinel filtering silently would not. + +The second change is the epoch. Rows are recycled, so a caller holding a handle +to a released request would otherwise read whichever request now occupies the +row. Every lookup checks the handle's epoch against the row's, which converts +that from a wrong answer into an exception. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np + +from maxtext.inference.kv_common import KvStorageLayoutV1 +from maxtext.inference.kv_control.logical_block import pages_for_tokens +from maxtext.inference.kv_control.request import RequestDescriptor, RequestHandle, RequestState + + +class StaleRequestHandleError(RuntimeError): + """A handle named a request that has been released, and its row reused.""" + + +class PageCapacityError(RuntimeError): + """A request needs more pages than its row can record.""" + + +class PageMap: + """Per-request page bookkeeping, in dense arrays. + + Owns sequence lengths as well as page ids, because the two are inseparable: a + page list without the length is not enough to say how much of the final page a + kernel may read, and that bound is the whole defence against reading a + recycled page's tail. + """ + + def __init__( + self, + max_requests: int, + max_pages_per_request: int, + tokens_per_page: int, + padding_page_id: int = 0, + ): + if max_requests < 1: + raise ValueError(f"max_requests must be at least 1, got {max_requests}") + if max_pages_per_request < 1: + raise ValueError(f"max_pages_per_request must be at least 1, got {max_pages_per_request}") + if tokens_per_page < 1: + raise ValueError(f"tokens_per_page must be at least 1, got {tokens_per_page}") + + self.max_requests = int(max_requests) + self.max_pages_per_request = int(max_pages_per_request) + self.tokens_per_page = int(tokens_per_page) + self.padding_page_id = int(padding_page_id) + + self._pages = np.full((self.max_requests, self.max_pages_per_request), self.padding_page_id, dtype=np.int32) + self._num_pages = np.zeros((self.max_requests,), dtype=np.int32) + self._seq_lens = np.zeros((self.max_requests,), dtype=np.int32) + self._epochs = np.zeros((self.max_requests,), dtype=np.int64) + self._states: list[RequestState | None] = [None] * self.max_requests + self._descriptors: list[RequestDescriptor | None] = [None] * self.max_requests + self._free_rows: list[int] = list(range(self.max_requests)) + + @classmethod + def from_layout(cls, layout: KvStorageLayoutV1, max_requests: int, max_context_len: int) -> "PageMap": + """Size the table from pool geometry and the longest context to be served.""" + return cls( + max_requests=max_requests, + max_pages_per_request=pages_for_tokens(max_context_len, layout.tokens_per_page), + tokens_per_page=layout.tokens_per_page, + padding_page_id=layout.padding_page_id, + ) + + @property + def num_live(self) -> int: + return self.max_requests - len(self._free_rows) + + @property + def available_rows(self) -> int: + return len(self._free_rows) + + def admit(self, descriptor: RequestDescriptor) -> RequestHandle | None: + """Take a row for `descriptor`, or return None if every row is occupied. + + None rather than an exception, for the same reason page exhaustion returns + None: a full batch is a scheduling outcome, not a fault. + """ + if not self._free_rows: + return None + row = self._free_rows.pop(0) + self._num_pages[row] = 0 + self._seq_lens[row] = 0 + self._pages[row, :] = self.padding_page_id + self._states[row] = RequestState.WAITING + self._descriptors[row] = descriptor + return RequestHandle(request_id=descriptor.request_id, row=row, epoch=int(self._epochs[row])) + + def release(self, handle: RequestHandle) -> np.ndarray: + """Give up the row and report the pages it held. + + The pages are returned rather than freed here: this class does not own the + free list, and separating "no longer referenced" from "available again" is + what leaves room to zero or poison a page in between. + """ + row = self._row(handle) + held = self._pages[row, : self._num_pages[row]].copy() + self._num_pages[row] = 0 + self._seq_lens[row] = 0 + self._pages[row, :] = self.padding_page_id + self._states[row] = None + self._descriptors[row] = None + self._epochs[row] += 1 + self._free_rows.append(row) + return held + + def pages(self, handle: RequestHandle) -> np.ndarray: + """The request's pages in sequence order, as a copy.""" + row = self._row(handle) + return self._pages[row, : self._num_pages[row]].copy() + + def num_pages(self, handle: RequestHandle) -> int: + return int(self._num_pages[self._row(handle)]) + + def append_pages(self, handle: RequestHandle, page_ids: Sequence[int] | np.ndarray) -> None: + """Extend the request's page list, in sequence order. + + Order is load-bearing: sequence position `p` is held by page slot + `p // tokens_per_page`, so appending out of order silently misdirects every + subsequent read and write. + """ + row = self._row(handle) + pages = np.asarray(page_ids, dtype=np.int32).reshape(-1) + if pages.size == 0: + return + start = int(self._num_pages[row]) + end = start + pages.size + if end > self.max_pages_per_request: + raise PageCapacityError( + f"request {handle.request_id!r} would hold {end} pages but a row records at most " + f"{self.max_pages_per_request}; raise max_context_len or lower the request's length cap" + ) + self._pages[row, start:end] = pages + self._num_pages[row] = end + + def seq_len(self, handle: RequestHandle) -> int: + return int(self._seq_lens[self._row(handle)]) + + def advance(self, handle: RequestHandle, num_tokens: int) -> int: + """Grow the recorded context by `num_tokens` and return the new length. + + Refuses to advance past what the held pages can address, because that is the + precise condition under which a kernel would read beyond the request's own + data. + """ + if num_tokens < 0: + raise ValueError(f"num_tokens must be non-negative, got {num_tokens}") + row = self._row(handle) + new_len = int(self._seq_lens[row]) + num_tokens + addressable = int(self._num_pages[row]) * self.tokens_per_page + if new_len > addressable: + raise PageCapacityError( + f"request {handle.request_id!r} would reach length {new_len} but its " + f"{int(self._num_pages[row])} pages address only {addressable} tokens: reserve pages " + f"before advancing the length, never after" + ) + self._seq_lens[row] = new_len + return new_len + + def state(self, handle: RequestHandle) -> RequestState: + return self._states[self._row(handle)] + + def set_state(self, handle: RequestHandle, state: RequestState) -> None: + self._states[self._row(handle)] = state + + def descriptor(self, handle: RequestHandle) -> RequestDescriptor: + return self._descriptors[self._row(handle)] + + def live_handles(self) -> list[RequestHandle]: + """Handles for every occupied row, in row order.""" + occupied = set(range(self.max_requests)) - set(self._free_rows) + return [ + RequestHandle( + request_id=self._descriptors[row].request_id, + row=row, + epoch=int(self._epochs[row]), + ) + for row in sorted(occupied) + ] + + def _row(self, handle: RequestHandle) -> int: + """Validate `handle` and return its row.""" + row = handle.row + if not 0 <= row < self.max_requests: + raise StaleRequestHandleError(f"handle row {row} is outside the table's {self.max_requests} rows") + if self._descriptors[row] is None: + raise StaleRequestHandleError( + f"request {handle.request_id!r} has been released; its row {row} is free" + ) + if int(self._epochs[row]) != handle.epoch: + raise StaleRequestHandleError( + f"request {handle.request_id!r} holds epoch {handle.epoch} but row {row} is at epoch " + f"{int(self._epochs[row])}: the row has been reused by another request" + ) + return row diff --git a/src/maxtext/inference/kv_control/protocols.py b/src/maxtext/inference/kv_control/protocols.py new file mode 100644 index 0000000000..d39f62054e --- /dev/null +++ b/src/maxtext/inference/kv_control/protocols.py @@ -0,0 +1,95 @@ +"""The control-plane contract a driver programs against. + +This exists because there will be more than one implementation, and only one may +be authoritative at a time. MaxText's own allocator and page map are one; a +frontend that owns its own page accounting -- vLLM's cache manager, SGLang's +allocator and radix cache -- is another, reached through an adapter. Running two +page owners for the same request means two free lists, two refcounts, and two +disagreeing views of which pages a request holds, so which one is live has to be +a decision asserted at startup rather than something inferred from whichever +object a caller happened to be handed. + +Runtime-checkable so that assertion can be a one-line check. It verifies method +presence only, which is enough to catch an adapter that has drifted from the +contract and not enough to be mistaken for type checking. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Protocol, Sequence, runtime_checkable + +import numpy as np + +from maxtext.inference.kv_common import KvPageTableV1, KvStorageLayoutV1 +from maxtext.inference.kv_control.request import RequestDescriptor, RequestHandle + + +@runtime_checkable +class KvControlPlane(Protocol): + """Admission, page reservation, release, and per-step metadata.""" + + @property + def layout(self) -> KvStorageLayoutV1: + """Pool geometry this control plane is accounting for.""" + + @property + def available_tokens(self) -> int: + """Tokens the pool could still accept. The number backpressure is decided on.""" + + def admit(self, descriptor: RequestDescriptor) -> RequestHandle | None: + """Start tracking a request, or return None if there is no room to track it.""" + + def reserve( + self, + handles: Sequence[RequestHandle], + num_new_tokens: Sequence[int] | np.ndarray, + ) -> bool: + """Reserve pages for the batch's new tokens, all of them or none. + + All-or-nothing because a partial reservation leaves the batch in a state the + caller has no way to describe: some requests advanced, some not, and no + page table that covers both. + """ + + def reserve_decode(self, handles: Sequence[RequestHandle]) -> bool: + """Reserve one token per request.""" + + def pending_scrub(self) -> np.ndarray: + """Pages reserved this step that must be overwritten before any kernel reads them. + + Part of the contract rather than an implementation detail: a frontend that + owns its own page accounting still has to answer this, or the recycled-page + guarantee is only as good as whichever allocator happens to be live. + """ + + def confirm_scrubbed(self, page_ids: Sequence[int] | np.ndarray) -> None: + """Record that `page_ids` have been overwritten on the device.""" + + def release(self, handle: RequestHandle) -> np.ndarray: + """Give up everything the request holds and report the pages reclaimed. + + Request-based, not slot-based. A slot is the dense cache's unit -- one fixed + reservation per request -- whereas a paged request owns a set of pages that + changed size on every step, so the handle is the only durable name for it. + """ + + def build_page_table( + self, + handles: Sequence[RequestHandle], + query_lens: Sequence[int] | np.ndarray, + ) -> KvPageTableV1: + """Describe this step in the neutral vocabulary.""" diff --git a/src/maxtext/inference/kv_control/request.py b/src/maxtext/inference/kv_control/request.py new file mode 100644 index 0000000000..7a7d268030 --- /dev/null +++ b/src/maxtext/inference/kv_control/request.py @@ -0,0 +1,101 @@ +"""Request identity for the paged control plane. + +A dense KV cache identifies a request by the slot it occupies, because one slot +is one fixed-size reservation held for the request's whole life. A paged request +owns a varying set of pages instead, so its identity has to survive that set +growing and has to remain distinguishable from a later request that reuses the +same bookkeeping row. That is the whole reason `RequestHandle` carries an epoch: +without one, a caller holding a released handle reads whichever request now +occupies the row, silently and with plausible-looking results. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +import enum + + +class RequestState(enum.Enum): + """Lifecycle of an admitted request, as the control plane sees it. + + Deliberately coarser than a scheduler's own state machine. These are only the + distinctions that change what the control plane may do with the request's + pages, which is why there is no queued/running/preempted split here. + """ + + WAITING = "waiting" + PREFILL = "prefill" + DECODE = "decode" + FINISHED = "finished" + + +@dataclasses.dataclass(frozen=True) +class RequestHandle: + """A reference to one live request's page bookkeeping. + + Frozen, so it is hashable and can key driver-side state, and so a caller + cannot mutate the row or epoch it was handed. + + Attributes: + request_id: caller-facing identity. Opaque here: the control plane never + interprets it, and two handles with the same id but different epochs are + different requests. + row: dense index into the page map's tables. What makes lookup an array + index rather than a dict probe. + epoch: bumped every time `row` is handed out again. A handle whose epoch no + longer matches the row's is naming a request that has already been + released, and the page map refuses it rather than answering about whoever + holds the row now. + """ + + request_id: str + row: int + epoch: int + + +@dataclasses.dataclass(frozen=True) +class RequestDescriptor: + """What the control plane needs in order to size a request's page demand. + + Lengths only. Prompt token ids are deliberately absent: nothing in M4 reads + them, and admitting them now would force a representation choice -- hashable + tuple, numpy array, shared buffer -- on behalf of prefix sharing, which is the + consumer that will actually have an opinion. Copying every prompt into a tuple + purely to keep this dataclass frozen is a real per-request cost for no present + benefit. + + Attributes: + request_id: caller-facing identity, carried through to the handle. + prompt_len: tokens in the prompt, and so the length of the prefill. + max_new_tokens: upper bound on generated tokens. Only a bound: a request + that stops early simply releases its pages sooner. + """ + + request_id: str + prompt_len: int + max_new_tokens: int + + def __post_init__(self): + if self.prompt_len < 0: + raise ValueError(f"prompt_len must be non-negative, got {self.prompt_len}") + if self.max_new_tokens < 0: + raise ValueError(f"max_new_tokens must be non-negative, got {self.max_new_tokens}") + + @property + def max_total_len(self) -> int: + """Longest context this request can reach, and so its worst-case page count.""" + return self.prompt_len + self.max_new_tokens diff --git a/src/maxtext/inference/kv_execution/__init__.py b/src/maxtext/inference/kv_execution/__init__.py new file mode 100644 index 0000000000..ef1ea896a1 --- /dev/null +++ b/src/maxtext/inference/kv_execution/__init__.py @@ -0,0 +1,57 @@ +"""MaxText-specific execution layer for the paged KV runtime. + +Where the neutral vocabulary and the semantic control plane meet JAX, MaxText's +config, and a vendor's kernels. Unlike `kv_common` and `kv_control`, this package +is *not* extractable and is not meant to be: it exists precisely to hold the +couplings the other two refuse. + +So the import rule inverts here. This package may import `jax`, MaxText config +and layers, and a vendor backend. What it must not do is let any of that leak +downwards -- `kv_control` calling into `kv_execution` would put jax back in the +control plane's dependency graph and undo the split. `kv_import_rule_test.py` +checks that direction statically. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from maxtext.inference.kv_execution.bucketing import ( + StepShape, + batch_ladder, + bucket_up, + token_ladder, +) +from maxtext.inference.kv_execution.driver import PagedDriver, PagedRequest, StepOutcome +from maxtext.inference.kv_execution.layout_builder import build_storage_layout +from maxtext.inference.kv_execution.pool_factory import PagedKvPool, allocate_pool +from maxtext.inference.kv_execution.pool_ops import POISON_SENTINEL, poison_pages, scrub_pages +from maxtext.inference.kv_execution.step_view import StepView, build_step_view + +__all__ = [ + "POISON_SENTINEL", + "PagedDriver", + "PagedKvPool", + "PagedRequest", + "StepOutcome", + "StepShape", + "StepView", + "allocate_pool", + "batch_ladder", + "bucket_up", + "build_step_view", + "build_storage_layout", + "poison_pages", + "scrub_pages", + "token_ladder", +] diff --git a/src/maxtext/inference/kv_execution/benchmark.py b/src/maxtext/inference/kv_execution/benchmark.py new file mode 100644 index 0000000000..7ba0e564c9 --- /dev/null +++ b/src/maxtext/inference/kv_execution/benchmark.py @@ -0,0 +1,538 @@ +"""Serving benchmark for the paged KV runtime, and for the dense path beside it. + +Section 6.1 of the design notes that the MaxText native path has no serving +harness and that this is the one piece that has to be written rather than reused. +This is it. The metric names are chosen to match what `vllm/benchmarks/ +benchmark_serving.py` and `sgl_jax/bench_serving.py` report, so results from the +three land in one table without translation. + +**Half of what this measures is not performance.** Alongside TTFT and ITL it +records pool occupancy against live tokens, concurrency sustained at a fixed pool +size, and peak device memory across the run. Those are capacity and correctness +properties, and they fail silently: a leaked page, a declined donation or an +unbounded compile count all produce plausible tokens and merely worse numbers, so +a harness reporting only latency would pass with any of them present. + +**One methodological warning, learned the hard way here.** Counting compiled +shapes is not enough to know whether a latency figure is real. An earlier version +of this file reported zero unwarmed shapes for a run that was three-quarters +compilation, because padding a variable-length array with `jnp` compiles once per +length and no shape *bucket* changed. Build per-call arrays in numpy, and trust +`run_repeated`: passes that agree with each other cannot both contain a one-off +cost, whereas a compilation counter only ever catches the cases someone +remembered to count. + +**On the memory numbers specifically, to head off a misreading.** The pool is one +fixed allocation made at startup, so paged peak memory is *flat* — it does not +track load, and a run where it grew would indicate a leak rather than a success. +What paging buys is that every byte of that allocation is fungible between +requests, which shows up as concurrency at a fixed budget rather than as a +smaller footprint. `occupancy_*` is the series that tracks live tokens. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +import json +import time +from typing import Any, Sequence + +import numpy as np + +import jax +import jax.numpy as jnp + + +@dataclasses.dataclass +class Request: + """One synthetic request, and the timestamps the run collects for it.""" + + request_id: str + prompt_len: int + max_new_tokens: int + arrival: float = 0.0 + first_token_at: float | None = None + token_times: list[float] = dataclasses.field(default_factory=list) + generated: list[int] = dataclasses.field(default_factory=list) + handle: Any = None + + @property + def ttft(self) -> float | None: + """Arrival to first token. None while the request has not produced one.""" + return None if self.first_token_at is None else self.first_token_at - self.arrival + + def inter_token_latencies(self) -> list[float]: + """Gaps between successive tokens *after* the first. + + The first token's cost is TTFT and belongs to prefill; folding it in here + would flatter ITL on short outputs and inflate it on long ones. + """ + return [b - a for a, b in zip(self.token_times, self.token_times[1:])] + + +def synthetic_trace( + num_requests: int, + mean_prompt: int, + mean_output: int, + *, + seed: int = 0, + length_spread: float = 0.6, +) -> list[Request]: + """A batch of requests with varied lengths, all arriving at once. + + Lengths are what this plan's claims are about, so they vary; arrival times are + not, so every request is available from the start. That measures saturated + throughput rather than a rate-driven trace, which is the right first + measurement: a rate low enough to keep the pool empty would hide exactly the + capacity behaviour under test. + """ + rng = np.random.default_rng(seed) + + def spread(mean: int) -> np.ndarray: + low = max(1, int(mean * (1.0 - length_spread))) + high = max(low + 1, int(mean * (1.0 + length_spread))) + return rng.integers(low, high, size=num_requests) + + prompts, outputs = spread(mean_prompt), spread(mean_output) + return [ + Request(request_id=f"r{i}", prompt_len=int(prompts[i]), max_new_tokens=int(outputs[i])) + for i in range(num_requests) + ] + + +def warmup_paged(engine, params, *, max_prompt: int, max_batch: int, target_context: int) -> set: + """Compile every shape the measured run can present, then discard the results. + + Not optional, and not a detail. XLA compiles per shape, so an uncompiled bucket + pays seconds on its first use and those seconds land in the TTFT and ITL + percentiles. The first measured run here reported a p50 TTFT of 25 seconds and + an ITL p99 of 10 seconds, all of it compilation. Bucketing is what makes + pre-compilation possible at all -- an unbucketed implementation has no finite + shape set to enumerate. + + **Three ladders, not two, and the third is the one that gets missed.** A shape + is `(batch bucket, token bucket, sequence-length bucket)`. Warming the first + two still left multi-second outliers, because the sequence-length rung is a + *static kernel argument*: a request whose context grows past a rung presents a + new program even at an already-compiled batch width. So this sweeps context + length as well, by admitting a full batch and decoding it forward to + `target_context`, touching every batch width at every rung along the way. + + Returns the set of shapes compiled, so a caller can diff it against what the + measured run actually used and see whether the sweep was complete. + """ + runtime = engine.paged_runtime + planner = runtime.planner + batch_rungs = [b for b in planner.batch_rungs if b <= max_batch] or [1] + + # Phase 1: every prefill shape. One prompt per sequence-length rung a prompt + # can reach, since a fresh request's context is exactly its prompt length. + for rung in planner.seqlen_rungs: + prompt_len = min(rung, max_prompt) + handle, _ = engine.prefill_paged( + params=params, + padded_tokens=jnp.ones((prompt_len,), jnp.int32), + true_length=prompt_len, + request_id=f"warmup-prefill-{rung}", + max_new_tokens=1, + ) + if handle is not None: + engine.release(handle) + runtime.control_plane.allocator.merge_released() + + # Phase 2: every decode shape. Admit the widest batch the pool allows, then + # walk it forward so each batch width is exercised at each length rung. + handles, tokens = [], [] + for index in range(max_batch): + handle, result = engine.prefill_paged( + params=params, + padded_tokens=jnp.ones((1,), jnp.int32), + true_length=1, + request_id=f"warmup-decode-{index}", + max_new_tokens=target_context, + ) + if handle is None: + break + handles.append(handle) + tokens.append(int(result.data[0, 0])) + + page_map = runtime.control_plane.page_map + while handles and page_map.seq_len(handles[0]) < target_context: + for width in batch_rungs: + if width > len(handles): + continue + result, ok = engine.generate_paged( + params, handles[:width], next_tokens=jnp.asarray(tokens[:width], jnp.int32) + ) + if not ok: + handles, tokens = handles[:-1], tokens[:-1] + break + fresh = np.asarray(result.data[:, 0]).reshape(-1) + for row in range(width): + tokens[row] = int(fresh[row]) + + for handle in handles: + engine.release(handle) + runtime.control_plane.allocator.merge_released() + return set(runtime.observed_shapes) + + +def warmup_dense(engine, params, decode_state, *, prompt_len: int): + """Compile the dense prefill and generate shapes. Returns the decode state.""" + padded = jnp.asarray( + [1] * prompt_len + [0] * (engine.config.max_prefill_predict_length - prompt_len), jnp.int32 + ) + prefix, _ = engine.prefill(params=params, padded_tokens=padded, true_length=prompt_len) + decode_state = engine.insert(prefix, decode_state, slot=0) + decode_state, _ = engine.generate(params, decode_state) + return decode_state + + +def _peak_bytes() -> int | None: + """Peak device bytes in use, or None where the backend does not report it.""" + try: + stats = jax.local_devices()[0].memory_stats() or {} + except Exception: # pylint: disable=broad-exception-caught + return None + return stats.get("peak_bytes_in_use") + + +def _percentiles(values: Sequence[float]) -> dict[str, float]: + """p50 and p99 in milliseconds, reported together because the tail is the point. + + Paging and prefix sharing show up in tail behaviour; a mean can hide a stall + entirely, which is why Section 6 asks for both. + """ + if not values: + return {"p50_ms": float("nan"), "p99_ms": float("nan"), "mean_ms": float("nan")} + arr = np.asarray(values, dtype=np.float64) * 1e3 + return { + "p50_ms": float(np.percentile(arr, 50)), + "p99_ms": float(np.percentile(arr, 99)), + "mean_ms": float(arr.mean()), + } + + +def run_paged( + engine, + params, + requests: Sequence[Request], + *, + max_batch: int, + warmed_shapes: set | None = None, +) -> dict[str, Any]: + """Serve `requests` on the paged path, admitting greedily up to `max_batch`. + + Prefill one request at a time, then advance every live request by one token. + Prefill is preferred while there is room, which favours TTFT at some cost to + the ITL of requests already running -- the same deliberate policy the standalone + driver uses. + """ + runtime = engine.paged_runtime + allocator = runtime.control_plane.allocator + page_size = runtime.control_plane.layout.tokens_per_page + + waiting, live, done = list(requests), [], [] + occupancy: list[tuple[float, int, int, int]] = [] + steps = 0 + before = set(runtime.observed_shapes) + start = time.perf_counter() + for request in waiting: + request.arrival = start + + while waiting or live: + while waiting and len(live) < max_batch: + candidate = waiting[0] + padded = jnp.asarray( + candidate.generated[:0] + [1] * candidate.prompt_len, dtype=jnp.int32 + ) + handle, result = engine.prefill_paged( + params=params, + padded_tokens=padded, + true_length=candidate.prompt_len, + request_id=candidate.request_id, + max_new_tokens=candidate.max_new_tokens, + ) + if handle is None: + break # backpressure: the pool cannot take another request yet + waiting.pop(0) + candidate.handle = handle + candidate.generated.append(int(result.data[0, 0])) + now = time.perf_counter() + candidate.first_token_at = now + candidate.token_times.append(now) + live.append(candidate) + steps += 1 + + if not live: + if waiting: + raise RuntimeError( + f"{len(waiting)} requests are waiting but nothing is live and the pool is empty; " + f"the pool is too small for even one request in this trace" + ) + break + + result, ok = engine.generate_paged( + params, + [r.handle for r in live], + next_tokens=jnp.asarray([r.generated[-1] for r in live], jnp.int32), + ) + if not ok: + # Pool exhausted mid-decode. Give the newest request's pages back and let + # it be replayed, which is the same recompute-preemption the driver uses. + victim = live.pop() + engine.release(victim.handle) + victim.handle, victim.generated = None, [] + waiting.insert(0, victim) + continue + + steps += 1 + now = time.perf_counter() + tokens = np.asarray(result.data[:, 0]).reshape(-1) + finished = [] + for row, request in enumerate(live): + request.generated.append(int(tokens[row])) + request.token_times.append(now) + if len(request.generated) >= request.max_new_tokens: + finished.append(request) + + live_tokens = sum(r.prompt_len + len(r.generated) for r in live) + occupancy.append((now - start, allocator.num_allocated_pages * page_size, live_tokens, len(live))) + + for request in finished: + engine.release(request.handle) + request.handle = None + live.remove(request) + done.append(request) + + elapsed = time.perf_counter() - start + summary = _summarise(done, occupancy, elapsed, steps, allocator, page_size) + + # Shape accounting, which is necessary but *not sufficient*: it sees bucketed + # `StepShape`s and is blind to eager host-side ops whose values enter a jaxpr as + # literals. An earlier version of this harness reported zero unwarmed shapes + # for a run that was three-quarters compilation, because padding a + # variable-length array with `jnp` compiles per length and no shape bucket + # changed. The empirical check in `run_repeated` is what actually settles it. + used = set(runtime.observed_shapes) + fresh = used - before if warmed_shapes is None else used - set(warmed_shapes) + summary["distinct_shapes_total"] = len(used) + summary["shapes_compiled_during_measurement"] = len(fresh) + summary["all_shapes_prewarmed"] = not fresh + if fresh: + summary["unwarmed_shapes"] = [dataclasses.asdict(s) for s in sorted(fresh, key=str)] + return summary + + +def run_repeated(engine, params, trace_factory, *, max_batch: int, warmed_shapes=None, repeats: int = 2): + """Run the same trace several times and report the last, with the spread. + + This is the reportability check that cannot be fooled. Counting compilations + requires knowing every place one can happen, and the previous attempt at that + missed a whole class; running the identical workload twice does not need to know + anything -- if the first pass was paying for compilation, the second is + materially faster, and the ratio says so. A run is reportable when successive + passes agree. + """ + passes = [] + for _ in range(max(repeats, 1)): + summary = run_paged( + engine, params, trace_factory(), max_batch=max_batch, warmed_shapes=warmed_shapes + ) + passes.append(summary) + + final = passes[-1] + durations = [p["duration_s"] for p in passes] + final["repeat_durations_s"] = durations + final["latency_is_reportable"] = is_stable(durations) and final["all_shapes_prewarmed"] + final["stability_ratio"] = stability_ratio(durations) + return final + + +def stability_ratio(durations: Sequence[float]) -> float: + """Spread across the passes *after* the first, as max over min. + + The first pass is excluded on purpose. Warmup cannot pre-compile everything a + trace touches -- the first request through a code path still pays for whatever + the sweep did not reach -- so comparing the first pass to the last always looks + alarming and says nothing. What matters is whether the passes that follow agree + with each other, because two passes that agree cannot both contain a + one-off cost. + """ + tail = list(durations[1:]) or list(durations) + low, high = min(tail), max(tail) + return high / low if low else float("inf") + + +def is_stable(durations: Sequence[float], tolerance: float = 1.15) -> bool: + """True when repeated passes agree closely enough to report a latency figure.""" + if len(durations) < 2: + return False + return stability_ratio(durations) < tolerance + + +def run_dense(engine, params, requests: Sequence[Request], *, max_batch: int) -> dict[str, Any]: + """Serve `requests` on the dense two-region cache, for the A/B. + + Deliberately the ordinary MaxText path -- `prefill`, `init_decode_state`, + `insert`, `generate` -- because the comparison is against what exists today, + not against an idealised dense implementation. Its batch is a fixed set of + slots, and every slot advances in lockstep whether or not it holds a request, + which is the behaviour under examination rather than an inefficiency to correct. + """ + waiting, done = list(requests), [] + start = time.perf_counter() + for request in waiting: + request.arrival = start + + decode_state = engine.init_decode_state() + slots: dict[int, Request] = {} + steps = 0 + + while waiting or slots: + while waiting and len(slots) < max_batch: + free = next(s for s in range(max_batch) if s not in slots) + candidate = waiting.pop(0) + padded = jnp.asarray( + [1] * candidate.prompt_len + + [0] * (engine.config.max_prefill_predict_length - candidate.prompt_len), + dtype=jnp.int32, + ) + prefix, result = engine.prefill( + params=params, padded_tokens=padded, true_length=candidate.prompt_len + ) + decode_state = engine.insert(prefix, decode_state, slot=free) + candidate.generated.append(int(result.data[0, 0])) + now = time.perf_counter() + candidate.first_token_at = now + candidate.token_times.append(now) + slots[free] = candidate + steps += 1 + + if not slots: + break + + decode_state, result = engine.generate(params, decode_state) + steps += 1 + now = time.perf_counter() + for slot, request in list(slots.items()): + request.generated.append(int(result.data[slot, 0])) + request.token_times.append(now) + if len(request.generated) >= request.max_new_tokens: + del slots[slot] + done.append(request) + + elapsed = time.perf_counter() - start + return _summarise(done, [], elapsed, steps, None, 0) + + +def _summarise( + done: Sequence[Request], + occupancy: Sequence[tuple[float, int, int, int]], + elapsed: float, + steps: int, + allocator: Any, + page_size: int, +) -> dict[str, Any]: + """Collapse a run into the reported schema.""" + output_tokens = sum(len(r.generated) for r in done) + prompt_tokens = sum(r.prompt_len for r in done) + itls = [gap for r in done for gap in r.inter_token_latencies()] + + summary: dict[str, Any] = { + "completed_requests": len(done), + "prompt_tokens": prompt_tokens, + "output_tokens": output_tokens, + "duration_s": elapsed, + "output_throughput_tok_per_s": output_tokens / elapsed if elapsed else 0.0, + "request_throughput_per_s": len(done) / elapsed if elapsed else 0.0, + "engine_steps": steps, + "ttft": _percentiles([r.ttft for r in done if r.ttft is not None]), + "itl": _percentiles(itls), + "peak_device_bytes": _peak_bytes(), + } + + if occupancy: + times, pool_tokens, live_tokens, concurrency = (np.asarray(c) for c in zip(*occupancy)) + summary["occupancy"] = { + # The claim: pages held track tokens held. A ratio pinned above 1 that + # never comes down means pages are not being reclaimed. + "max_pool_tokens_held": int(pool_tokens.max()), + "max_live_tokens": int(live_tokens.max()), + "mean_overhead_ratio": float((pool_tokens / np.maximum(live_tokens, 1)).mean()), + "max_concurrency": int(concurrency.max()), + "mean_concurrency": float(concurrency.mean()), + "samples": len(occupancy), + "series_time_s": times.tolist(), + "series_pool_tokens": pool_tokens.tolist(), + "series_live_tokens": live_tokens.tolist(), + } + if allocator is not None: + summary["pages_leaked"] = int(allocator.num_allocated_pages) + summary["pool_capacity_tokens"] = allocator.capacity_pages * page_size + return summary + + +def report(name: str, summary: dict[str, Any]) -> str: + """A short human-readable form. The JSON is the artifact; this is for reading.""" + lines = [f"=== {name} ==="] + lines.append( + f" requests {summary['completed_requests']:4d}" + f" output_tok {summary['output_tokens']:6d}" + f" {summary['duration_s']:.2f}s" + f" {summary['output_throughput_tok_per_s']:8.1f} tok/s" + f" steps {summary['engine_steps']}" + ) + lines.append( + f" TTFT p50 {summary['ttft']['p50_ms']:8.2f} ms p99 {summary['ttft']['p99_ms']:8.2f} ms" + ) + lines.append( + f" ITL p50 {summary['itl']['p50_ms']:8.2f} ms p99 {summary['itl']['p99_ms']:8.2f} ms" + ) + if summary.get("peak_device_bytes"): + lines.append(f" peak device memory {summary['peak_device_bytes'] / 2**30:.3f} GiB") + occ = summary.get("occupancy") + if occ: + lines.append( + f" pool tokens held max {occ['max_pool_tokens_held']}" + f" vs live tokens max {occ['max_live_tokens']}" + f" (page overhead x{occ['mean_overhead_ratio']:.3f})" + ) + lines.append( + f" concurrency max {occ['max_concurrency']} mean {occ['mean_concurrency']:.2f}" + f" of pool capacity {summary.get('pool_capacity_tokens')} tokens" + ) + if "pages_leaked" in summary: + lines.append(f" pages leaked {summary['pages_leaked']}") + if "distinct_shapes_total" in summary: + lines.append( + f" shapes {summary['distinct_shapes_total']} total," + f" {summary['shapes_compiled_during_measurement']} unwarmed" + ) + if "repeat_durations_s" in summary: + spread = " ".join(f"{d:.3f}" for d in summary["repeat_durations_s"]) + verdict = "reportable" if summary.get("latency_is_reportable") else "NOT reportable" + lines.append( + f" repeats {spread} s (post-first spread x{summary.get('stability_ratio', float('nan')):.3f})" + f" -> latency {verdict}" + ) + return "\n".join(lines) + + +def write_json(path: str, payload: dict[str, Any]) -> None: + with open(path, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2) diff --git a/src/maxtext/inference/kv_execution/bucketing.py b/src/maxtext/inference/kv_execution/bucketing.py new file mode 100644 index 0000000000..333cc377c9 --- /dev/null +++ b/src/maxtext/inference/kv_execution/bucketing.py @@ -0,0 +1,201 @@ +"""Power-of-two shape ladders, so a churning batch traces a fixed set of shapes. + +This is the JAX-specific obligation of the whole design. Every array crossing +into a step is data-dependent in *size*: how many requests are live, how many +tokens they contribute, how many pages they hold. Left alone, a mixed-length +workload under churn presents a new shape almost every step and recompiles +forever, which does not merely cost time -- it makes the steady-state latency the +milestone is supposed to demonstrate unmeasurable. + +Two shape families, because the two phases vary along different axes: + + * **Decode** contributes exactly one token per request, so the token count is + not free -- it *is* the batch bucket. Only the batch size varies. + * **Extend** varies in both request count and total tokens, and independently. + Bucketing both would multiply out, so the batch axis is pinned to its largest + bucket and only the token count varies. + +**One refinement on the plan's three ladders.** The gather table gets no ladder of +its own; it is derived as `num_requests * ceil(max_seqlen_k / tokens_per_page)`, +clamped to the pool. That is a correct upper bound -- no request can hold more +pages than its own capped length needs, and no batch more than the pool has -- +and deriving it removes an entire dimension from the compile cross-product rather +than adding one. What does need a ladder, and is easy to miss, is `max_seqlen_k`: +the kernels take it as a static configuration value, so passing the true maximum +would retrace on nearly every step no matter how well the array shapes were +bucketed. + +Deliberately free of `jax`. Which shapes exist is host arithmetic, and keeping it +that way means the ladders can be reasoned about and tested without a device. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +from typing import Sequence + +# The floor on the token ladder, matching what MaxText's existing prefill +# bucketing uses (`2**i for i in range(6, ...)`). Below 64 tokens a separate +# trace buys nothing. +MIN_TOKEN_BUCKET = 1 << 6 + + +def bucket_up(value: int, ladder: Sequence[int]) -> int: + """Smallest ladder entry at least `value`. + + Raises rather than clamping when `value` exceeds the ladder: silently + bucketing 5000 tokens down to a 4096-token shape would truncate the batch, and + a truncated batch loses tokens instead of running slowly. + """ + for step in ladder: + if value <= step: + return step + raise ValueError(f"{value} exceeds the largest bucket {ladder[-1]}; the ladder is mis-sized for this workload") + + +def _powers_of_two(low: int, high: int) -> tuple[int, ...]: + """Powers of two from `low` up to the first one at or above `high`.""" + rungs, step = [], low + while step < high: + rungs.append(step) + step *= 2 + rungs.append(step) + return tuple(rungs) + + +def batch_ladder(max_batch: int) -> tuple[int, ...]: + """1, 2, 4, ... up to `max_batch`.""" + if max_batch < 1: + raise ValueError(f"max_batch must be at least 1, got {max_batch}") + return _powers_of_two(1, max_batch) + + +def token_ladder(max_tokens: int) -> tuple[int, ...]: + """64, 128, ... up to `max_tokens`, or a single rung if that is below 64.""" + if max_tokens < 1: + raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") + if max_tokens <= MIN_TOKEN_BUCKET: + return (MIN_TOKEN_BUCKET,) + return _powers_of_two(MIN_TOKEN_BUCKET, max_tokens) + + +def seqlen_ladder(tokens_per_page: int, max_context_len: int) -> tuple[int, ...]: + """Page-aligned powers of two up to `max_context_len`. + + Starts at the page size because a shorter context still occupies one whole + page, so a finer rung would describe a shape that cannot occur. + """ + if max_context_len < 1: + raise ValueError(f"max_context_len must be at least 1, got {max_context_len}") + return _powers_of_two(tokens_per_page, max(max_context_len, tokens_per_page)) + + +@dataclasses.dataclass(frozen=True) +class StepShape: + """The padded shape family one step is traced for. + + Hashable, so a driver can count distinct shapes -- which is the direct measure + of whether bucketing is working, and what the milestone's exit criterion is + stated in terms of. + """ + + num_requests: int + num_tokens: int + num_pages: int + max_seqlen_k: int + is_decode: bool + + +class StepShapePlanner: + """Owns the ladders and maps a live batch onto a `StepShape`.""" + + def __init__( + self, + tokens_per_page: int, + max_batch: int, + max_context_len: int, + pool_pages: int, + max_batched_tokens: int | None = None, + ): + """ + Args: + max_batched_tokens: token budget for one extend step, and so the top of + the token ladder. Distinct from `max_context_len`, which bounds a single + request: an extend step batches several requests, so its total can + exceed the longest one. Sizing the ladder from `max_context_len` instead + makes a perfectly legal batch unbucketable. Defaults to + `max_context_len`, which admits one full-length request per step. + """ + self.tokens_per_page = int(tokens_per_page) + self.pool_pages = int(pool_pages) + self.max_batched_tokens = int(max_batched_tokens or max_context_len) + self.batch_rungs = batch_ladder(max_batch) + self.token_rungs = token_ladder(self.max_batched_tokens) + self.seqlen_rungs = seqlen_ladder(tokens_per_page, max_context_len) + + @property + def max_batch_bucket(self) -> int: + return self.batch_rungs[-1] + + def _pages_for(self, num_requests: int, max_seqlen_k: int) -> int: + """Upper bound on pages the batch can reference, clamped to the pool.""" + per_request = -(-max_seqlen_k // self.tokens_per_page) + return min(num_requests * per_request, self.pool_pages) + + def decode_shape(self, num_requests: int, max_seq_len: int) -> StepShape: + """One token per request, so only the batch axis varies.""" + requests = bucket_up(max(num_requests, 1), self.batch_rungs) + seqlen = bucket_up(max(max_seq_len, 1), self.seqlen_rungs) + return StepShape( + num_requests=requests, + num_tokens=requests, + num_pages=self._pages_for(requests, seqlen), + max_seqlen_k=seqlen, + is_decode=True, + ) + + def extend_shape(self, num_tokens: int, max_seq_len: int, num_requests: int | None = None) -> StepShape: + """Batch pinned to its largest bucket; only the token count varies. + + Args: + num_requests: pass a count to bucket the batch axis instead of pinning it. + Pinning exists to stop the request count and the token count multiplying + out into a cross-product of shapes, so it buys nothing for a caller whose + request count is fixed -- `MaxEngine.prefill_paged` always prefills one + prompt — and there it only pads the per-request arrays to the maximum + batch for no reason. + """ + requests = self.max_batch_bucket if num_requests is None else bucket_up(max(num_requests, 1), self.batch_rungs) + tokens = bucket_up(max(num_tokens, 1), self.token_rungs) + seqlen = bucket_up(max(max_seq_len, 1), self.seqlen_rungs) + return StepShape( + num_requests=requests, + num_tokens=tokens, + num_pages=self._pages_for(requests, seqlen), + max_seqlen_k=seqlen, + is_decode=False, + ) + + def max_distinct_shapes(self) -> int: + """Upper bound on shapes this planner can ever produce. + + Worth being able to state, because "bounded" is only meaningful with a + number attached. Decode contributes batch times seqlen rungs, extend + contributes token times seqlen rungs. + """ + seqlens = len(self.seqlen_rungs) + return len(self.batch_rungs) * seqlens + len(self.token_rungs) * seqlens diff --git a/src/maxtext/inference/kv_execution/driver.py b/src/maxtext/inference/kv_execution/driver.py new file mode 100644 index 0000000000..adfd63b560 --- /dev/null +++ b/src/maxtext/inference/kv_execution/driver.py @@ -0,0 +1,381 @@ +"""Page-based continuous batching, beside the slot-based loop rather than inside it. + +`InferenceWorker._run_continuous_batching` is built on a static set of integer +slots: one slot is one fixed reservation, held for a request's whole life, handed +out by `empty_decode_slots.pop()` and returned by the detokenisation thread. That +is the right shape for the dense two-region cache and the wrong shape here, where +a request owns a set of pages that changes size on every step and where admitting +one depends on how much *pool* is free rather than on whether a slot is spare. +Generalising the existing loop to cover both would make it unreviewable and put +the working dense path at risk, so this is a sibling. + +**Execution is injected.** The driver owns admission, reservation, scrubbing, +shape selection and release; it does not own the forward pass. A caller supplies +a step function taking the padded `StepView` and the pool and returning one token +per active request. That keeps the scheduling logic testable without a model, and +it is the seam a real engine adapter plugs into. + +Two policies are worth naming because they are choices, not consequences: + + * **Prefill is preferred over decode when both are possible.** It favours time + to first token at the cost of inter-token latency for requests already + running. A production scheduler would make this configurable; the milestone + only needs it to be deliberate. + * **Decode running out of pages preempts the newest request by recomputation.** + Its pages are released and it returns to the queue to be prefilled again. + This loses work, but it is the simplest policy that cannot deadlock, and a + loop that can deadlock under churn fails the stability criterion outright. + Preempting the newest rather than the oldest keeps the requests closest to + finishing. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +from typing import Callable, Iterable, Sequence + +import numpy as np + +from maxtext.inference.kv_control import ( + NativeKvControlPlane, + RequestDescriptor, + RequestHandle, + RequestState, + pages_for_tokens, +) +from maxtext.inference.kv_execution.bucketing import StepShape, StepShapePlanner +from maxtext.inference.kv_execution.pool_factory import PagedKvPool +from maxtext.inference.kv_execution.pool_ops import poison_pages, scrub_pages +from maxtext.inference.kv_execution.step_view import StepView, build_step_view + + +@dataclasses.dataclass(eq=False) +class PagedRequest: + """A request as the driver tracks it, across preemption and requeueing. + + Compared by identity, not by field values. The driver moves these between the + waiting, live and done lists by `remove`, and two requests that happen to + agree on every field -- same length, same tokens generated so far, which is + entirely possible -- would otherwise be interchangeable to `list.remove` and + the wrong one would be dropped. + """ + + request_id: str + prompt_len: int + max_new_tokens: int + handle: RequestHandle | None = None + generated: list[int] = dataclasses.field(default_factory=list) + state: RequestState = RequestState.WAITING + finish_reason: str | None = None + preemptions: int = 0 + + @property + def is_finished(self) -> bool: + return self.finish_reason is not None + + def descriptor(self) -> RequestDescriptor: + """The remaining work, not the original request. + + After a preemption the prompt has to be recomputed but the tokens already + generated are part of the context, so the prompt to replay is longer than + the one submitted. Describing the original would under-reserve. + """ + return RequestDescriptor( + request_id=self.request_id, + prompt_len=self.prompt_len + len(self.generated), + max_new_tokens=max(self.max_new_tokens - len(self.generated), 0), + ) + + +@dataclasses.dataclass(frozen=True) +class StepOutcome: + """What one executed step did, for a caller that wants to observe the loop.""" + + shape: StepShape + is_decode: bool + num_requests: int + num_tokens: int + preempted: int = 0 + + +# Takes the padded view and the pool; returns one token per *active* request and +# the rebound pool arrays, which alias the ones passed in. +StepFn = Callable[[StepView, PagedKvPool], np.ndarray] + + +class PagedDriver: + """Admits, schedules, reserves for and releases paged requests.""" + + def __init__( + self, + control_plane: NativeKvControlPlane, + pool: PagedKvPool, + step_fn: StepFn, + *, + max_batch: int, + max_batched_tokens: int | None = None, + eos_ids: Iterable[int] = (), + poison_on_free: bool = False, + ): + layout = control_plane.layout + self.plane = control_plane + self.pool = pool + self.step_fn = step_fn + self.max_batch = int(max_batch) + self.eos_ids = frozenset(int(t) for t in eos_ids) + self.poison_on_free = bool(poison_on_free) + self.planner = StepShapePlanner( + tokens_per_page=layout.tokens_per_page, + max_batch=max_batch, + max_context_len=control_plane.max_context_len, + pool_pages=layout.num_pages, + max_batched_tokens=max_batched_tokens, + ) + self._waiting: list[PagedRequest] = [] + self._live: list[PagedRequest] = [] + self._done: list[PagedRequest] = [] + self.observed_shapes: set[StepShape] = set() + + @property + def num_distinct_shapes(self) -> int: + """Distinct traced shapes so far. The direct measure of whether bucketing works.""" + return len(self.observed_shapes) + + @property + def num_waiting(self) -> int: + return len(self._waiting) + + @property + def num_live(self) -> int: + return len(self._live) + + def submit(self, requests: Sequence[PagedRequest]) -> None: + self._waiting.extend(requests) + + def run(self, max_steps: int = 100_000) -> list[PagedRequest]: + """Drive every submitted request to completion. + + `max_steps` is a guard, not a schedule: hitting it means the loop failed to + make progress, and raising says so rather than returning half an answer. + """ + steps = 0 + while self._waiting or self._live: + if steps >= max_steps: + raise RuntimeError( + f"the driver made no progress in {max_steps} steps with {len(self._waiting)} waiting and " + f"{len(self._live)} live requests" + ) + if self.step() is None: + break + steps += 1 + return list(self._done) + + def step(self) -> StepOutcome | None: + """Run one step. Prefill if anything can be admitted, else decode. + + Returns None when there is nothing left to do. + """ + admitted = self._admit() + if admitted: + return self._run_phase(admitted, [r.descriptor().prompt_len for r in admitted], is_decode=False) + if self._live: + return self._run_phase(self._live, [1] * len(self._live), is_decode=True) + return None + + def _admit(self) -> list[PagedRequest]: + """Take as many waiting requests as rows, pages and the token bucket allow. + + Budgeted against `available_pages` before anything is admitted, because + reservation is all-or-nothing: admitting a request the pool cannot back would + fail the whole batch rather than just that request. + """ + if not self._waiting: + return [] + + tokens_per_page = self.plane.layout.tokens_per_page + token_budget = self.planner.token_rungs[-1] + page_budget = self.plane.allocator.available_pages + room = self.max_batch - len(self._live) + + admitted: list[PagedRequest] = [] + tokens = 0 + for request in list(self._waiting): + if len(admitted) >= room or self.plane.page_map.available_rows == 0: + break + prompt_len = request.descriptor().prompt_len + needed = pages_for_tokens(prompt_len, tokens_per_page) + if needed > page_budget or tokens + prompt_len > token_budget: + break + handle = self.plane.admit(request.descriptor()) + if handle is None: + break + request.handle = handle + request.state = RequestState.PREFILL + page_budget -= needed + tokens += prompt_len + admitted.append(request) + self._waiting.remove(request) + return admitted + + def _run_phase( + self, + requests: Sequence[PagedRequest], + query_lens: Sequence[int], + *, + is_decode: bool, + ) -> StepOutcome: + """Reserve, scrub, build, execute, and retire -- in that order. + + The order is not incidental. Scrubbing has to happen after reservation, + since that is what decides which pages were recycled, and before the table is + built, because the control plane refuses to describe a page it still + considers dirty. + """ + batch = list(requests) + lens = list(query_lens) + preempted = 0 + + while batch and not self.plane.reserve([r.handle for r in batch], lens): + if not is_decode: + # A prefill batch was budgeted against the free pool, so a failure here + # means the budget was wrong rather than that the pool is under pressure. + raise RuntimeError( + f"reservation failed for an admitted prefill batch of {len(batch)} requests; " + f"{self.plane.available_tokens} tokens available" + ) + self._preempt_newest(batch) + preempted += 1 + lens = [1] * len(batch) + + if not batch: + return StepOutcome( + shape=StepShape(0, 0, 0, 0, is_decode), + is_decode=is_decode, + num_requests=0, + num_tokens=0, + preempted=preempted, + ) + + self._scrub_recycled_pages() + + handles = [r.handle for r in batch] + table = self.plane.build_page_table(handles, lens) + max_seq_len = int(table.seq_lens.max()) if table.num_requests else 1 + shape = ( + self.planner.decode_shape(len(batch), max_seq_len) + if is_decode + else self.planner.extend_shape(sum(lens), max_seq_len) + ) + self.observed_shapes.add(shape) + + view = build_step_view( + table, + shape, + tokens_per_page=self.plane.layout.tokens_per_page, + padding_page_id=self.plane.layout.padding_page_id, + ) + next_tokens = np.asarray(self.step_fn(view, self.pool)).reshape(-1) + if next_tokens.size < len(batch): + raise ValueError(f"the step function returned {next_tokens.size} tokens for {len(batch)} requests") + + self._retire(batch, next_tokens[: len(batch)], is_decode=is_decode) + return StepOutcome( + shape=shape, + is_decode=is_decode, + num_requests=len(batch), + num_tokens=sum(lens), + preempted=preempted, + ) + + def _scrub_recycled_pages(self) -> None: + """Zero every page this step recycled, then say so. + + Both halves matter. Without the write a later read sees the previous + occupant's KV; without the confirmation the control plane refuses to build + the table, which is the mechanism that stops the write being forgotten. + """ + pending = self.plane.pending_scrub() + if not pending.size: + return + for layer in range(self.pool.num_layers): + k, v = scrub_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], pending) + self.pool.replace_layer(layer, k, v) + self.plane.confirm_scrubbed(pending) + + def _preempt_newest(self, batch: list[PagedRequest]) -> None: + """Return the newest request's pages to the pool and requeue it. + + Its generated tokens are kept, so the replay is a longer prompt rather than + lost output. What is lost is the compute that built its KV, which is the + price of not deadlocking. + """ + if not batch: + return + victim = batch.pop() + self._release_pages(victim) + victim.state = RequestState.WAITING + victim.preemptions += 1 + if victim in self._live: + self._live.remove(victim) + self._waiting.insert(0, victim) + + def _release_pages(self, request: PagedRequest) -> None: + """Give up a request's pages, poisoning them first when asked. + + Poison is applied *before* the pages leave the request, and does not count as + a scrub: the control plane keeps them dirty, so the next occupant still has + to zero them. That ordering is what makes the sentinel a detector of a missed + scrub rather than a substitute for one. + """ + if request.handle is None: + return + pages = self.plane.page_map.pages(request.handle) + if self.poison_on_free and pages.size: + for layer in range(self.pool.num_layers): + k, v = poison_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], pages) + self.pool.replace_layer(layer, k, v) + self.plane.release(request.handle) + request.handle = None + + def _retire(self, batch: Sequence[PagedRequest], tokens: np.ndarray, *, is_decode: bool) -> None: + """Record this step's tokens and release whatever finished. + + Releasing is deferred to a second pass because the stop conditions read the + request's recorded length, which only exists while it still holds a handle. + """ + finished: list[PagedRequest] = [] + for request, token in zip(batch, tokens.tolist()): + request.generated.append(int(token)) + request.state = RequestState.DECODE + if int(token) in self.eos_ids: + request.finish_reason = "stop" + elif len(request.generated) >= request.max_new_tokens: + request.finish_reason = "length" + elif self.plane.page_map.seq_len(request.handle) >= self.plane.max_context_len: + request.finish_reason = "context" + + if request.is_finished: + finished.append(request) + elif not is_decode: + self._live.append(request) + + for request in finished: + request.state = RequestState.FINISHED + self._release_pages(request) + if request in self._live: + self._live.remove(request) + self._done.append(request) diff --git a/src/maxtext/inference/kv_execution/engine_adapter.py b/src/maxtext/inference/kv_execution/engine_adapter.py new file mode 100644 index 0000000000..df77cb7e14 --- /dev/null +++ b/src/maxtext/inference/kv_execution/engine_adapter.py @@ -0,0 +1,186 @@ +"""Wiring between `MaxEngine` and the paged control plane. + +`MaxEngine.release_pages(slot)` was a no-op that printed a warning, with a +docstring referring to a `PageManager` that exists nowhere in the file, and three +call sites that all pass a fixed slot integer and ignore the result. It is +tempting to just implement it. That would be a mistake, because the signature +encodes the dense cache's assumption that a request *is* a slot -- one fixed +reservation for its whole life -- and building the new runtime around that would +bake in the abstraction the paging work exists to replace. + +So `release(handle)` is the real API and `release_pages(slot)` becomes a shim over +it. The shim needs a slot-to-handle map, which this adapter keeps; that map is the +entire cost of keeping the three legacy call sites working, and it is confined to +one object that a paged deployment can eventually stop constructing. + +`MaxEngine.prefill` already accepts a `request_id` it never forwards, which is the +natural anchor for a handle and is why this adapter keys on request id as well as +slot. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np + +from maxtext.inference.kv_control import NativeKvControlPlane, RequestDescriptor, RequestHandle +from maxtext.inference.kv_execution.bucketing import StepShape, StepShapePlanner +from maxtext.inference.kv_execution.pool_factory import PagedKvPool +from maxtext.inference.kv_execution.pool_ops import poison_pages, scrub_pages +from maxtext.inference.kv_execution.step_view import StepView, build_step_view + + +class PagedRuntime: + """The paged control plane and pool, as one attachable object. + + `MaxEngine` holds at most one of these and delegates to it. Keeping the two + together matters because releasing pages and poisoning them are the same + operation from a caller's point of view, and poisoning needs the pool. + + It also owns `prepare_step`, which is the fixed per-step order — reserve, + scrub, confirm, build, pad — expressed once. `MaxEngine` should not be + reimplementing that sequence, and neither should a second driver: getting the + order wrong is how a kernel ends up reading a page it does not own. + """ + + def __init__( + self, + control_plane: NativeKvControlPlane, + pool: PagedKvPool, + planner: StepShapePlanner | None = None, + poison_on_free: bool = False, + ): + self.control_plane = control_plane + self.pool = pool + self.planner = planner + self.poison_on_free = bool(poison_on_free) + # Every bucketed shape this runtime has ever asked for. One compiled program + # per entry, so this is the compile count -- the Section 6.4 measurement, and + # the only way to tell a warmup that covered the shape space from one that + # merely ran for a while. + self.observed_shapes: set[StepShape] = set() + self._by_slot: dict[int, RequestHandle] = {} + self._by_request_id: dict[str, RequestHandle] = {} + + def admit(self, request_id: str, prompt_len: int, max_new_tokens: int) -> RequestHandle | None: + """Start tracking a request, returning None when there is no room.""" + return self.control_plane.admit( + RequestDescriptor(request_id=request_id, prompt_len=prompt_len, max_new_tokens=max_new_tokens) + ) + + def prepare_step( + self, + handles: Sequence[RequestHandle], + query_lens: Sequence[int], + *, + is_decode: bool, + num_requests: int | None = None, + ) -> StepView | None: + """Reserve pages for this step and return the padded device arrays. + + Returns None if the pool cannot back the step, which is backpressure rather + than an error. The scrub sits between reservation and table construction + because reservation is what decides which pages were recycled, and the + control plane refuses to describe a page it still considers dirty. + """ + if self.planner is None: + raise ValueError("this runtime has no StepShapePlanner, so it cannot decide a bucketed shape") + if not self.control_plane.reserve(handles, query_lens): + return None + + self.scrub_recycled() + + table = self.control_plane.build_page_table(handles, query_lens) + max_seq_len = int(table.seq_lens.max()) if table.num_requests else 1 + shape = ( + self.planner.decode_shape(len(handles), max_seq_len) + if is_decode + else self.planner.extend_shape(int(sum(query_lens)), max_seq_len, num_requests=num_requests) + ) + self.observed_shapes.add(shape) + layout = self.control_plane.layout + return build_step_view( + table, shape, tokens_per_page=layout.tokens_per_page, padding_page_id=layout.padding_page_id + ) + + def scrub_recycled(self) -> np.ndarray: + """Zero every page this step recycled, then record that it was done. + + Both halves are necessary: without the write a later read sees the previous + occupant's KV, and without the confirmation the control plane refuses to + build the table -- which is the mechanism that stops the write being skipped. + """ + pending = self.control_plane.pending_scrub() + if not pending.size: + return pending + for layer in range(self.pool.num_layers): + k, v = scrub_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], pending) + self.pool.replace_layer(layer, k, v) + self.control_plane.confirm_scrubbed(pending) + return pending + + def track(self, handle: RequestHandle, slot: int | None = None) -> None: + """Record a handle so it can be found again by request id, or by slot. + + The slot argument exists only for the legacy API. A caller that has a handle + should keep it and pass it to `release` directly. + """ + self._by_request_id[handle.request_id] = handle + if slot is not None: + self._by_slot[int(slot)] = handle + + def handle_for_slot(self, slot: int) -> RequestHandle | None: + return self._by_slot.get(int(slot)) + + def handle_for_request(self, request_id: str) -> RequestHandle | None: + return self._by_request_id.get(str(request_id)) + + def release(self, handle: RequestHandle) -> np.ndarray: + """Reclaim everything the request holds. The canonical API. + + Idempotent at this level: an untracked handle is a no-op rather than an + error, because the three legacy call sites fire on sequence termination and + do not coordinate with each other. The control plane underneath is *not* + idempotent, and that is the right split -- a double release through a handle + the adapter still knows about is a bug worth reporting, while a release of + something already forgotten is just a duplicate notification. + """ + known = self._by_request_id.get(handle.request_id) + if known is None or known != handle: + return np.empty((0,), dtype=np.int32) + if self.poison_on_free: + pages = self.control_plane.page_map.pages(handle) + for layer in range(self.pool.num_layers): + k, v = poison_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], pages) + self.pool.replace_layer(layer, k, v) + pages = self.control_plane.release(handle) + self._forget(handle) + return pages + + def release_slot(self, slot: int) -> np.ndarray: + """Shim for `release_pages(slot)`.""" + handle = self._by_slot.get(int(slot)) + if handle is None: + return np.empty((0,), dtype=np.int32) + return self.release(handle) + + def _forget(self, handle: RequestHandle) -> None: + self._by_request_id.pop(handle.request_id, None) + for slot, tracked in list(self._by_slot.items()): + if tracked == handle: + del self._by_slot[slot] diff --git a/src/maxtext/inference/kv_execution/layout_builder.py b/src/maxtext/inference/kv_execution/layout_builder.py new file mode 100644 index 0000000000..7ea74331e6 --- /dev/null +++ b/src/maxtext/inference/kv_execution/layout_builder.py @@ -0,0 +1,74 @@ +"""MaxText config plus mesh into a `KvStorageLayoutV1`. + +One function, and its only interesting job is deciding `kv_head_shards` from the +mesh rather than from a config field, because a mismatch between the two is +exactly the kind of thing that produces a pool a factor of TP too small and a +wrong-answer failure much later. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Any + +from maxtext.inference.kv_common import KvStorageLayoutV1 + +# MaxText spells the KV-head tensor-parallel axes several ways depending on the +# model; these are the ones that shard `num_kv_heads`. +_KV_HEAD_MESH_AXES = ("tensor", "tensor_transpose", "tensor_sequence") + + +def kv_head_shards(mesh: Any | None) -> int: + """Product of the mesh axes that shard KV heads, or 1 with no mesh.""" + if mesh is None: + return 1 + shape = dict(getattr(mesh, "shape", {}) or {}) + shards = 1 + for axis in _KV_HEAD_MESH_AXES: + shards *= int(shape.get(axis, 1)) + return max(shards, 1) + + +def build_storage_layout(config: Any, mesh: Any | None = None) -> KvStorageLayoutV1: + """Derive pool geometry from a MaxText config. + + `num_pages` includes the reserved padding page, so the pool holds + `paged_num_blocks` usable pages plus one. Sizing the pool as exactly + `paged_num_blocks` and then reserving one out of it would silently cost a page + of capacity relative to what the config asked for. + """ + num_kv_heads = int(getattr(config, "num_kv_heads", 0) or 0) + head_dim = int(getattr(config, "head_dim", 0) or 0) + num_layers = int(getattr(config, "num_decoder_layers", 0) or getattr(config, "base_num_decoder_layers", 0) or 0) + tokens_per_page = int(getattr(config, "paged_page_size", 16)) + usable_pages = int(getattr(config, "paged_num_blocks", 0) or 0) + if usable_pages < 1: + raise ValueError( + f"paged_num_blocks must be at least 1 for attention='gpu_paged', got {usable_pages}" + ) + + dtype = str(getattr(config, "dtype", "bfloat16")) + + return KvStorageLayoutV1( + tokens_per_page=tokens_per_page, + num_pages=usable_pages + 1, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + dtype=dtype, + kv_head_shards=kv_head_shards(mesh), + padding_page_id=0, + ) diff --git a/src/maxtext/inference/kv_execution/pool_factory.py b/src/maxtext/inference/kv_execution/pool_factory.py new file mode 100644 index 0000000000..af57b32ad6 --- /dev/null +++ b/src/maxtext/inference/kv_execution/pool_factory.py @@ -0,0 +1,104 @@ +"""Allocate the paged KV pool. + +The pool is a pair of NHD arrays per layer, `[num_pages, tokens_per_page, +heads_per_shard, head_dim]`, matching what the M3 attention path already reads so +nothing is repacked between append, prefill and decode. + +**Zeros, not `empty`.** Two separate guarantees depend on it. The reserved +padding page is a landing zone that padded gather rows read from, and it must +read as zeros rather than as whatever was in that memory. And the allocator +treats a never-allocated page as clean, so it issues no scrub for it -- which is +only sound if the page really does start zeroed. `jnp.zeros` under `jit` is +materialised by the compiler rather than transferred, so this costs a kernel +launch, not a host copy. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +from typing import Any + +import jax +import jax.numpy as jnp + +from maxtext.inference.kv_common import KvStorageLayoutV1 + + +@dataclasses.dataclass +class PagedKvPool: + """The per-layer K and V arrays, plus the geometry that describes them. + + Mutable because the arrays are rebound on every step: the append op donates + them and returns aliased results, so holding the originals after a step is how + a caller ends up reading a stale buffer. + """ + + layout: KvStorageLayoutV1 + k_pages: list[jax.Array] + v_pages: list[jax.Array] + + @property + def num_layers(self) -> int: + return len(self.k_pages) + + @property + def page_shape(self) -> tuple[int, ...]: + layout = self.layout + return (layout.num_pages, layout.tokens_per_page, layout.heads_per_shard(), layout.head_dim) + + def bytes_per_shard(self) -> int: + return self.layout.pool_bytes_per_shard() + + def replace_layer(self, layer: int, k: jax.Array, v: jax.Array) -> None: + """Rebind one layer's pages after an aliased write.""" + self.k_pages[layer] = k + self.v_pages[layer] = v + + +def allocate_pool( + layout: KvStorageLayoutV1, + sharding: Any | None = None, + dtype: Any | None = None, +) -> PagedKvPool: + """Allocate a zero-initialised pool for every layer. + + Args: + layout: pool geometry. `heads_per_shard()` is already the per-device head + count, including the replication case where TP exceeds the KV head count, + so the shape below is the local shard rather than the global pool. + sharding: optional sharding for each array. Passed through to + `jax.device_put`; a `NamedSharding` carrying a `memory_kind` is how the + pool would later be placed in a collective memory space. + dtype: overrides the layout's dtype. Only useful for tests that want a + dtype numpy can print. + + Returns: + A `PagedKvPool` whose arrays are all zeros. + """ + shape = (layout.num_pages, layout.tokens_per_page, layout.heads_per_shard(), layout.head_dim) + resolved = jnp.dtype(dtype) if dtype is not None else jnp.dtype(layout.dtype) + + k_pages, v_pages = [], [] + for _ in range(layout.num_layers): + k = jnp.zeros(shape, resolved) + v = jnp.zeros(shape, resolved) + if sharding is not None: + k = jax.device_put(k, sharding) + v = jax.device_put(v, sharding) + k_pages.append(k) + v_pages.append(v) + return PagedKvPool(layout=layout, k_pages=k_pages, v_pages=v_pages) diff --git a/src/maxtext/inference/kv_execution/pool_ops.py b/src/maxtext/inference/kv_execution/pool_ops.py new file mode 100644 index 0000000000..409eaa1d7e --- /dev/null +++ b/src/maxtext/inference/kv_execution/pool_ops.py @@ -0,0 +1,110 @@ +"""Device-side page hygiene: scrub a recycled page, or poison a freed one. + +The control plane can say which pages hold another request's KV, but it cannot +do anything about it -- the pool is device memory and `kv_control` never touches +a device. These two functions are where that obligation is actually discharged. + +`scrub_pages` zeroes. That is what makes a recycled page safe to hand on. +`poison_pages` fills with a recognisable sentinel instead, which is strictly a +debugging aid: it makes an unscrubbed read *loud* rather than plausible. Poison +is not a scrub, and the control plane deliberately keeps a poisoned page marked +dirty so it cannot be read until it has really been zeroed. + +**The page-count axis is bucketed, and it has to be.** These run inside a serving +loop, so a fresh trace for every distinct number of recycled pages would defeat +the whole of Step 5. The index array is padded up to a power of two with repeats +of a page that is already being written, so the padding is idempotent rather than +needing a mask. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import functools +from typing import Sequence + +import numpy as np + +import jax +import jax.numpy as jnp + +# A power of two, and that is the whole reason for this value. A sentinel is only +# useful if a test can say "this byte is the sentinel" exactly, and every dtype a +# KV pool might use -- bfloat16, float16, both fp8 variants -- carries far too +# little mantissa to store a round decimal like -8e4 without rounding it. Powers +# of two are exact in all of them, and -256 is inside fp8 e4m3's range of 448 +# where -65536 would overflow float16 outright. +# +# Otherwise: large, finite, negative, and nothing a real activation produces. Not +# a NaN, which would propagate through any downstream reduction and destroy the +# evidence of where the read actually happened. +POISON_SENTINEL = -256.0 + + +def _pad_page_indices(page_ids: Sequence[int] | np.ndarray) -> np.ndarray: + """Pad to a power of two by repeating the first page. + + Repetition rather than a sentinel plus a mask: both fills are idempotent, so + writing a page twice is free and needs no branch in the kernel. + """ + pages = np.asarray(page_ids, dtype=np.int32).reshape(-1) + if pages.size == 0: + return pages + padded_size = 1 << (int(pages.size) - 1).bit_length() + if padded_size == pages.size: + return pages + return np.concatenate([pages, np.full((padded_size - pages.size,), pages[0], dtype=np.int32)]) + + +@functools.partial(jax.jit, donate_argnums=(0, 1)) +def _fill(k_pages: jax.Array, v_pages: jax.Array, page_ids: jax.Array, value: jax.Array): + """Set every element of the named pages to `value`, in place.""" + fill_shape = (page_ids.shape[0],) + k_pages.shape[1:] + block = jnp.full(fill_shape, value, dtype=k_pages.dtype) + return k_pages.at[page_ids].set(block), v_pages.at[page_ids].set(block) + + +def scrub_pages( + k_pages: jax.Array, + v_pages: jax.Array, + page_ids: Sequence[int] | np.ndarray, +) -> tuple[jax.Array, jax.Array]: + """Zero the named pages. Returns the rebound arrays, which alias the inputs. + + A no-op for an empty list, which is the common case on a pool that has not + wrapped around yet. + """ + pages = _pad_page_indices(page_ids) + if pages.size == 0: + return k_pages, v_pages + return _fill(k_pages, v_pages, jnp.asarray(pages), jnp.zeros((), k_pages.dtype)) + + +def poison_pages( + k_pages: jax.Array, + v_pages: jax.Array, + page_ids: Sequence[int] | np.ndarray, + value: float = POISON_SENTINEL, +) -> tuple[jax.Array, jax.Array]: + """Fill the named pages with a sentinel, for debugging only. + + Does not discharge a scrub: a poisoned page is still dirty, and the control + plane will still refuse to build a page table naming it. + """ + pages = _pad_page_indices(page_ids) + if pages.size == 0: + return k_pages, v_pages + return _fill(k_pages, v_pages, jnp.asarray(pages), jnp.asarray(value, k_pages.dtype)) diff --git a/src/maxtext/inference/kv_execution/step_view.py b/src/maxtext/inference/kv_execution/step_view.py new file mode 100644 index 0000000000..12ebd55ed6 --- /dev/null +++ b/src/maxtext/inference/kv_execution/step_view.py @@ -0,0 +1,163 @@ +"""A `KvPageTableV1` as the padded device arrays a step actually consumes. + +The padding convention is the whole content of this module, and it is chosen so +that a padded row is *inert* rather than merely ignored, because "ignored" +depends on every kernel agreeing to ignore it: + + * `slot_mapping` pads with -1, which the append kernel drops. Not with a real + slot, which would write a padded row's garbage into a live page. + * `kv_indptr` and `cu_seqlens_q` repeat their final value, so every padded + request has a zero-length page range and a zero-length query. A kernel + iterating the full bucketed batch does no work for them without needing a + mask. + * `kv_page_indices` pads with the reserved page, and `seq_lens` and + `kv_last_page_lens` pad with zero. The reserved page reads as zeros, so even + a kernel that addressed a padded entry despite the flat indptr would read + zeros rather than another request's KV. + +Every one of those is a second line of defence behind the flat indptr. That is +deliberate: padding bugs are silent, and the failure mode is a wrong token rather +than a crash. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +from typing import Any + +import numpy as np + +import jax +import jax.numpy as jnp + +from maxtext.inference.kv_common import KvPageTableV1 +from maxtext.inference.kv_execution.bucketing import StepShape + + +@dataclasses.dataclass(frozen=True) +class StepView: + """One step's page bookkeeping, padded to a bucketed shape. + + The active counts are carried alongside because they are what a caller needs + to slice a result back down; they are host ints, not traced, so using one in a + shape would defeat the bucketing. + """ + + slot_mapping: jax.Array # int32 [num_tokens] + kv_indptr: jax.Array # int32 [num_requests + 1] + kv_page_indices: jax.Array # int32 [num_pages] + kv_last_page_lens: jax.Array # int32 [num_requests] + cu_seqlens_q: jax.Array # int32 [num_requests + 1] + seq_lens: jax.Array # int32 [num_requests] + shape: StepShape + num_active_requests: int + num_active_tokens: int + max_seqlen_q: int + + def to_paged_plan(self) -> Any: + """Adapt to the M3 attention path's `PagedPlan`. + + Imported lazily so the control and execution layers stay usable without + pulling in the attention module, which reaches a vendor backend. + """ + from maxtext.layers.gpu_paged_attention import PagedPlan # pylint: disable=import-outside-toplevel + + return PagedPlan( + slot_mapping=self.slot_mapping, + kv_indptr=self.kv_indptr, + kv_page_indices=self.kv_page_indices, + kv_last_page_lens=self.kv_last_page_lens, + cu_seqlens_q=self.cu_seqlens_q, + max_seqlen_q=self.max_seqlen_q, + max_seqlen_k=self.shape.max_seqlen_k, + is_decode=self.shape.is_decode, + ) + + +def _pad_to(values: np.ndarray, size: int, fill: int, name: str) -> np.ndarray: + if values.size > size: + raise ValueError(f"{name} has {values.size} entries, past the bucketed {size}") + if values.size == size: + return values + return np.concatenate([values, np.full((size - values.size,), fill, dtype=np.int32)]) + + +def _pad_cumulative(values: np.ndarray, size: int, name: str) -> np.ndarray: + """Extend a prefix-sum array by repeating its last entry. + + Repetition is what makes a padded request zero-length rather than + out-of-range, which is the difference between a kernel skipping it and a + kernel reading whatever sits at index zero. + """ + if values.size > size: + raise ValueError(f"{name} has {values.size} entries, past the bucketed {size}") + if values.size == size: + return values + tail = int(values[-1]) if values.size else 0 + return np.concatenate([values, np.full((size - values.size,), tail, dtype=np.int32)]) + + +def build_step_view( + table: KvPageTableV1, + shape: StepShape, + tokens_per_page: int, + padding_page_id: int = 0, +) -> StepView: + """Pad `table` out to `shape` and move it to the device. + + Validates the table first: an inconsistent table padded into a static shape is + considerably harder to diagnose than one rejected on the spot. + """ + table.validate(tokens_per_page) + + num_requests = shape.num_requests + query_lens = np.asarray(table.query_lens, dtype=np.int32) + seq_lens = np.asarray(table.seq_lens, dtype=np.int32) + + cu_seqlens_q = np.zeros((query_lens.size + 1,), dtype=np.int32) + if query_lens.size: + np.cumsum(query_lens, out=cu_seqlens_q[1:]) + + active_tokens = int(query_lens.sum()) + max_seq_len = int(seq_lens.max()) if seq_lens.size else 0 + if max_seq_len > shape.max_seqlen_k: + raise ValueError( + f"a request is {max_seq_len} tokens long but the step shape is configured for " + f"{shape.max_seqlen_k}; the sequence-length ladder is mis-sized" + ) + + return StepView( + slot_mapping=jnp.asarray( + _pad_to(table.slot_mapping(tokens_per_page, padding_page_id), shape.num_tokens, -1, "slot_mapping"), + jnp.int32, + ), + kv_indptr=jnp.asarray(_pad_cumulative(table.indptr(), num_requests + 1, "kv_indptr"), jnp.int32), + kv_page_indices=jnp.asarray( + _pad_to(table.flat_page_indices(), shape.num_pages, padding_page_id, "kv_page_indices"), jnp.int32 + ), + kv_last_page_lens=jnp.asarray( + _pad_to(table.last_page_lens(tokens_per_page), num_requests, 0, "kv_last_page_lens"), jnp.int32 + ), + cu_seqlens_q=jnp.asarray(_pad_cumulative(cu_seqlens_q, num_requests + 1, "cu_seqlens_q"), jnp.int32), + seq_lens=jnp.asarray(_pad_to(seq_lens, num_requests, 0, "seq_lens"), jnp.int32), + shape=shape, + num_active_requests=table.num_requests, + num_active_tokens=active_tokens, + # Static, from the bucket rather than from the data: a traced maximum would + # be a fresh kernel configuration on almost every step. + max_seqlen_q=1 if shape.is_decode else shape.num_tokens, + ) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 62c5fe0345..bfb4e5d3a6 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -25,6 +25,7 @@ from jax.sharding import PartitionSpec as P import jax import jax.numpy as jnp +import numpy as np if jax.__version_info__ >= (0, 6, 3): from jax.experimental.layout import Layout as DLL # type: ignore @@ -181,6 +182,10 @@ def __init__(self, config: Any, devices: Any | None = None): self.rng = None self._compiled_initialize_fn = None self._compiled_init_cache_fn = None + # Set by whatever builds the `attention="gpu_paged"` path. Left None here so + # that `release`/`release_pages` degrade to a log line on every other + # attention mode rather than requiring a paged runtime to exist. + self.paged_runtime = None def print_stats(self, label: str): max_utils.print_mem_stats(label) @@ -231,8 +236,16 @@ def _nnx_run_model( encoder_video_masks=None, encoder_video_grid_thw=None, encoder_audios=None, + kv_caches=None, + attention_metadata=None, ): - """NNX equivalent of `model.apply(..., mutable=["cache"])`. Returns (logits, new_cache_dict).""" + """NNX equivalent of `model.apply(..., mutable=["cache"])`. + + Returns `(logits, new_cache_dict)` normally, and + `(logits, new_cache_dict, kv_caches)` when a paged pool is threaded through. + The third value is not optional bookkeeping: the pool is donated and comes + back aliased, so a caller that does not rebind it is holding a deleted array. + """ cache_state = self._nnx_cache_state_template(mode=model_mode) nnx.replace_by_pure_dict(cache_state, cache_dict) # Merge with the graphdef built for this mode. Layers that captured their @@ -244,7 +257,7 @@ def _nnx_run_model( model = nnx.merge( # pyrefly: ignore[no-matching-overload] graphdef, params, cache_state, self._nnx_rest_state, copy=True ) # pyrefly: ignore[no-matching-overload] - logits = model( + outputs = model( decoder_input_tokens, decoder_positions, decoder_segment_ids=decoder_segment_ids, @@ -259,9 +272,14 @@ def _nnx_run_model( previous_chunk=previous_chunk, true_length=true_length, slot=slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, ) new_cache = nnx.to_pure_dict(nnx.state(model, nnx.Cache)) - return logits, new_cache + if kv_caches is not None: + logits, updated_pools = outputs + return logits, new_cache, updated_pools + return outputs, new_cache def generate_aot( self, params: Params, decode_state: DecodeState, rng: PRNGKeyType | None = None @@ -1863,9 +1881,424 @@ def copy(path, partial_cache, full_cache, annotations): "token_logp": inserted_token_logp, } + def release(self, request_handle: Any): + """Reclaim the pages a request holds. The canonical release API. + + Request-based rather than slot-based on purpose. A slot is the dense cache's + unit -- one fixed reservation held for a request's whole life -- whereas a + paged request owns a set of pages that changed size on every step, so the + handle is its only durable name. + + Returns the page ids reclaimed, or an empty array when no paged runtime is + attached. + """ + if self.paged_runtime is None: + max_logging.log("release: paged attention is not configured, so there are no pages to release.") + return np.empty((0,), dtype=np.int32) + return self.paged_runtime.release(request_handle) + def release_pages(self, slot: int): - """Releases pages associated with a specific slot (page group) via the PageManager.""" - print(f"Warning: release_pages called for slot {slot} but paged attention is not configured.") + """Compatibility shim over `release`, for callers that only have a slot. + + Kept because three existing call sites fire on sequence termination with a + fixed slot integer. Implementing the paged runtime around this signature + instead would have baked the dense cache's one-slot-per-request assumption + into the new code. + """ + if self.paged_runtime is None: + max_logging.log(f"release_pages: slot {slot} ignored because paged attention is not configured.") + return np.empty((0,), dtype=np.int32) + return self.paged_runtime.release_slot(slot) + + # --------------------------------------------------------------------------- + # Paged serving: a sibling entry path, not a rewrite of the dense one. + # + # `init_decode_state`, `_insert_jit` and `generate` are specific to the dense + # two-region cache -- one fixed slot per request, every slot advancing in + # lockstep, and a cache copied into a batch row at insert time. None of that + # survives contact with a page pool, where a request owns a varying set of + # pages and each one is at its own position. Rebuilding those three around + # pages would be a much larger change than it reads as, and it would put the + # working dense path at risk, so this is a parallel surface instead: + # `init_paged_runtime`, `prefill_paged`, `generate_paged`. Nothing above + # touches them and they touch nothing above. + # --------------------------------------------------------------------------- + + def init_paged_runtime( + self, + max_requests: int | None = None, + max_batched_tokens: int | None = None, + ): + """Allocate the paged pool and control plane. The sibling of `init_decode_state`. + + Args: + max_requests: concurrent requests to size the page-map rows for. Defaults + to the dense path's batch width, so a paged deployment is no narrower + than the dense one it replaces. + max_batched_tokens: token budget for one prefill step, and the top of the + token bucket ladder. Defaults to the per-request context limit. + + Returns: + The `PagedRuntime`, also stored on `self.paged_runtime`. + """ + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_control import NativeKvControlPlane + from maxtext.inference.kv_execution.bucketing import StepShapePlanner + from maxtext.inference.kv_execution.engine_adapter import PagedRuntime + from maxtext.inference.kv_execution.layout_builder import build_storage_layout + from maxtext.inference.kv_execution.pool_factory import allocate_pool + + if self.config.attention != "gpu_paged": + raise ValueError( + f"init_paged_runtime requires attention='gpu_paged', but this engine is configured for " + f"'{self.config.attention}'." + ) + + layout = build_storage_layout(self.config, self._mesh) + if layout.dtype not in ("bfloat16", "float16"): + raise ValueError( + f"the paged attention kernels accept bfloat16 and float16 only, but dtype is " + f"'{layout.dtype}'. Set dtype=bfloat16 for attention='gpu_paged'." + ) + + max_requests = int(max_requests or self.max_concurrent_decodes) + control_plane = NativeKvControlPlane( + layout=layout, + max_requests=max_requests, + max_context_len=self.config.paged_max_context_len, + ) + planner = StepShapePlanner( + tokens_per_page=layout.tokens_per_page, + max_batch=max_requests, + max_context_len=self.config.paged_max_context_len, + pool_pages=layout.num_pages, + max_batched_tokens=max_batched_tokens, + ) + self.paged_runtime = PagedRuntime( + control_plane=control_plane, + pool=allocate_pool(layout, sharding=self.kv_pool_sharding()), + planner=planner, + poison_on_free=bool(getattr(self.config, "paged_poison_freed_pages", False)), + ) + return self.paged_runtime + + def kv_pool_sharding(self): + """Sharding for one layer's pool array, or None while unsharded. + + Pages are the leading axis and are never sharded -- a page belongs to one + request and splitting it would put half a token's KV on another device. The + KV-head axis is the one tensor parallelism divides, which is M6's work; until + then the pool is replicated and this returns None. + """ + return None + + def _paged_pool_arrays(self): + """The pool as the nested list `kv_caches` expects: one `[k, v]` per layer.""" + pool = self.paged_runtime.pool + return [[pool.k_pages[i], pool.v_pages[i]] for i in range(pool.num_layers)] + + def _rebind_paged_pool(self, updated): + """Store the aliased pool handles returned by a step. + + Not optional. The pool is donated into the step, so the arrays passed in are + invalid afterwards and only these handles refer to live memory. + """ + pool = self.paged_runtime.pool + for layer, pair in enumerate(updated): + pool.replace_layer(layer, pair[0], pair[1]) + + @functools.partial( + jax.jit, + static_argnums=(0,), + static_argnames=("model_mode", "max_seqlen_k", "is_decode", "algorithm", "topk", "nucleus_topp"), + donate_argnames=("kv_caches",), + ) + def _paged_forward_jit( + self, + *, + params: Params, + kv_caches, + tokens: jax.Array, + positions: jax.Array, + segment_ids: jax.Array | None, + slot_mapping: jax.Array, + kv_indptr: jax.Array, + kv_page_indices: jax.Array, + kv_last_page_lens: jax.Array, + cu_seqlens_q: jax.Array, + sample_at: jax.Array, + rng: PRNGKeyType, + model_mode: str, + max_seqlen_k: int, + is_decode: bool, + algorithm: str | None = None, + topk: int | None = None, + nucleus_topp: float | None = None, + temperature: float | None = None, + ): + """One paged step: write this step's K/V into the pool, attend, and sample. + + The page bookkeeping arrives as plain int32 arrays rather than as a + `PagedPlan`, because the plan is not a pytree and could not cross a `jit` + boundary; it is reassembled here where the arrays are already tracers. + + `sample_at` is the per-row index whose logits to sample, traced rather than + static so a new prompt length does not force a new trace. + """ + # pylint: disable=import-outside-toplevel + from maxtext.layers.gpu_paged_attention import PagedPlan + + plan = PagedPlan( + slot_mapping=slot_mapping, + kv_indptr=kv_indptr, + kv_page_indices=kv_page_indices, + kv_last_page_lens=kv_last_page_lens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1 if is_decode else int(tokens.shape[1]), + max_seqlen_k=max_seqlen_k, + is_decode=is_decode, + ) + + rng, run_rng = jax.random.split(rng) + with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + if self.config.pure_nnx: + logits, _, updated_pools = self._nnx_run_model( + params=params, + cache_dict=self._nnx_init_cache_dict(mode=model_mode), + decoder_input_tokens=tokens, + decoder_positions=positions, + decoder_segment_ids=segment_ids, + enable_dropout=False, + model_mode=model_mode, + kv_caches=kv_caches, + attention_metadata=plan, + ) + else: + (logits, updated_pools), _ = self.model.apply( + params, + tokens, + positions, + decoder_segment_ids=segment_ids, + enable_dropout=False, + model_mode=model_mode, + rngs={"params": run_rng}, + mutable=["cache"], + kv_caches=kv_caches, + attention_metadata=plan, + ) + + # One row per request, each sampled at its own index: for prefill that is the + # last prompt token, for decode the single new token. Doing it with a gather + # rather than a slice is what lets a decode batch hold requests at different + # positions, which is the whole point of paging. + rows = jnp.arange(logits.shape[0]) + selected = logits[rows, sample_at][:, None, :] + selected = jax.lax.with_sharding_constraint(selected, self.replicated_sharding) + + rng, sample_rng = jax.random.split(rng) + sampled = inference_utils.sampling( + selected, + sample_rng, + algorithm if algorithm is not None else self.config.decode_sampling_strategy, + topk=topk if topk is not None else self.config.decode_sampling_top_k, + nucleus_topp=nucleus_topp if nucleus_topp is not None else self.config.decode_sampling_nucleus_p, + temperature=temperature if temperature is not None else self.config.decode_sampling_temperature, + ) + return sampled, selected, updated_pools + + def _paged_result_tokens(self, sampled: jax.Array, selected: jax.Array, step: int): + """Wrap sampled tokens in the same `ResultTokens` shape the dense path returns.""" + all_valid = jnp.ones(sampled.shape, dtype=jnp.int8) + lengths = jnp.full(sampled.shape, step, dtype=jnp.int32) + if self.config.return_log_prob: + token_logp = inference_utils.log_prob_of_chosen_token(selected, sampled) + else: + token_logp = jnp.zeros(sampled.shape, dtype=jnp.float32) + return engine_api.ResultTokens( + data=jnp.concatenate((sampled, all_valid, lengths), axis=1), + tokens_idx=(0, 1), + valid_idx=(1, 2), + length_idx=(2, 3), + log_prob=token_logp, + samples_per_slot=1, + ) + + def prefill_paged( + self, + *, + params: Params, + padded_tokens: jax.Array, + true_length: int, + request_id: str | None = None, + max_new_tokens: int | None = None, + slot: int | None = None, + rng: PRNGKeyType | None = None, + algorithm: str | None = None, + topk: int | None = None, + nucleus_topp: float | None = None, + temperature: float | None = None, + ): + """Prefill one prompt into the page pool and sample its first token. + + The sibling of `prefill`. It returns a `RequestHandle` rather than a prefix + dict, because there is no cache to carry: the K/V is already in the pool and + the handle is what names the pages holding it. + + Returns: + `(handle, result_tokens)`, or `(None, None)` if the pool could not admit + the request -- backpressure, not an error. + """ + if self.paged_runtime is None: + raise ValueError("call init_paged_runtime() before prefill_paged()") + + runtime = self.paged_runtime + prompt_len = int(true_length) + budget = self.config.paged_max_context_len - prompt_len + handle = runtime.admit( + request_id=request_id if request_id is not None else f"req-{uuid.uuid4()}", + prompt_len=prompt_len, + max_new_tokens=int(max_new_tokens) if max_new_tokens is not None else max(budget, 0), + ) + if handle is None: + return None, None + + view = runtime.prepare_step([handle], [prompt_len], is_decode=False, num_requests=1) + if view is None: + runtime.control_plane.release(handle) + return None, None + + # The flattened query the attention layer sees is batch times sequence, so + # the padded prompt has to be exactly as long as the token bucket. + # + # Built in numpy, deliberately. The obvious `jnp` spelling of these three + # lines compiles a fresh program per *prompt length*: the input to a pad is a + # different shape each time, and `arange(n) < prompt_len` bakes `prompt_len` + # into the jaxpr as a literal. A serving trace with two dozen distinct prompt + # lengths then pays two dozen compilations that no shape-bucket accounting + # can see, which is exactly how an earlier measurement here ended up + # three-quarters compile time while reporting zero unwarmed shapes. numpy has + # no such cache to miss, and these arrays are a few hundred bytes. + width = view.shape.num_tokens + prompt = np.asarray(padded_tokens, dtype=np.int32).reshape(-1)[:width] + token_row = np.zeros((1, width), np.int32) + token_row[0, : prompt.size] = prompt + index = np.arange(width, dtype=np.int32) + tokens = jnp.asarray(token_row) + positions = jnp.asarray(index[None, :]) + segment_ids = jnp.asarray( + np.where(index < prompt_len, DECODING_ACTIVE_SEQUENCE_INDICATOR, 0).astype(np.int32)[None, :] + ) + + if rng is None: + if self.rng is None: + self.rng = jax.random.PRNGKey(0) + self.rng, rng = jax.random.split(self.rng) + + sampled, selected, updated = self._paged_forward_jit( + params=params, + kv_caches=self._paged_pool_arrays(), + tokens=tokens, + positions=positions, + segment_ids=segment_ids, + slot_mapping=view.slot_mapping, + kv_indptr=view.kv_indptr, + kv_page_indices=view.kv_page_indices, + kv_last_page_lens=view.kv_last_page_lens, + cu_seqlens_q=view.cu_seqlens_q, + sample_at=jnp.asarray([prompt_len - 1], jnp.int32), + rng=rng, + model_mode=MODEL_MODE_PREFILL, + max_seqlen_k=view.shape.max_seqlen_k, + is_decode=False, + algorithm=algorithm, + topk=topk, + nucleus_topp=nucleus_topp, + temperature=temperature, + ) + self._rebind_paged_pool(updated) + runtime.track(handle, slot=slot) + return handle, self._paged_result_tokens(sampled, selected, step=0) + + def generate_paged( + self, + params: Params, + handles, + next_tokens: jax.Array, + rng: PRNGKeyType | None = None, + algorithm: str | None = None, + topk: int | None = None, + nucleus_topp: float | None = None, + temperature: float | None = None, + ): + """Advance every live request by one token. The sibling of `generate`. + + Unlike the dense path there is no lockstep: each request contributes one + token at its own position, which is read from the control plane rather than + from a shared counter. + + Returns: + `(result_tokens, ok)`. `ok` is False when the pool could not back the step, + which the caller resolves by releasing or preempting something. + """ + if self.paged_runtime is None: + raise ValueError("call init_paged_runtime() before generate_paged()") + + runtime = self.paged_runtime + handles = list(handles) + if not handles: + raise ValueError("generate_paged needs at least one live request") + + # Read before reserving: this is each request's position for the token about + # to be written, and reserving advances the recorded length past it. + positions = [runtime.control_plane.page_map.seq_len(h) for h in handles] + + view = runtime.prepare_step(handles, [1] * len(handles), is_decode=True) + if view is None: + return None, False + + # numpy for the same reason as `prefill_paged`: padding a variable-length + # array with `jnp` compiles once per distinct live-request count, and the + # live count changes constantly under continuous batching. + width = view.shape.num_requests + token_col = np.zeros((width, 1), np.int32) + supplied = np.asarray(next_tokens, dtype=np.int32).reshape(-1)[:width] + token_col[: supplied.size, 0] = supplied + position_col = np.zeros((width, 1), np.int32) + recorded = np.asarray(positions, dtype=np.int32).reshape(-1)[:width] + position_col[: recorded.size, 0] = recorded + tokens = jnp.asarray(token_col) + pos = jnp.asarray(position_col) + + if rng is None: + if self.rng is None: + self.rng = jax.random.PRNGKey(0) + self.rng, rng = jax.random.split(self.rng) + + sampled, selected, updated = self._paged_forward_jit( + params=params, + kv_caches=self._paged_pool_arrays(), + tokens=tokens, + positions=pos, + # None, not zeros: the model refuses segment ids in autoregressive mode, + # where every token is by definition in the active sequence. + segment_ids=None, + slot_mapping=view.slot_mapping, + kv_indptr=view.kv_indptr, + kv_page_indices=view.kv_page_indices, + kv_last_page_lens=view.kv_last_page_lens, + cu_seqlens_q=view.cu_seqlens_q, + sample_at=jnp.zeros((width,), jnp.int32), + rng=rng, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + max_seqlen_k=view.shape.max_seqlen_k, + is_decode=True, + algorithm=algorithm, + topk=topk, + nucleus_topp=nucleus_topp, + temperature=temperature, + ) + self._rebind_paged_pool(updated) + return self._paged_result_tokens(sampled, selected, step=1), True def get_prefix_destination_sharding(self) -> Any: return { diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index a70ad780a3..5e941d49bf 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -257,6 +257,16 @@ def __call__( # In vLLM, logits are computed separately after updating the KV cache. return hidden_state, kv_caches + if self.config.attention == "gpu_paged" and kv_caches is not None: + # The paged pool is donated into the step and comes back aliased, so the + # updated handles have to leave the model. Dropping them invalidates the + # caller's arrays and the next step reads a deleted buffer. Unlike the vLLM + # path this still returns logits, because MaxText's own serving loop + # samples inside the same jitted step rather than in a separate sampler. + # Gated on `kv_caches` so model init and the dry-run path, which thread no + # pool, keep the plain single-value return every other caller expects. + return logits, kv_caches + return logits @@ -616,4 +626,14 @@ def __call__( # In vLLM, logits are computed separately after updating the KV cache. return hidden_state, kv_caches + if self.config.attention == "gpu_paged" and kv_caches is not None: + # The paged pool is donated into the step and comes back aliased, so the + # updated handles have to leave the model. Dropping them invalidates the + # caller's arrays and the next step reads a deleted buffer. Unlike the vLLM + # path this still returns logits, because MaxText's own serving loop + # samples inside the same jitted step rather than in a separate sampler. + # Gated on `kv_caches` so model init and the dry-run path, which thread no + # pool, keep the plain single-value return every other caller expects. + return logits, kv_caches + return logits diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 6cafc9a4a3..46880c5896 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2549,6 +2549,9 @@ def _gpu_paged_config(self, seq_len, **overrides): arguments["max_target_length"] = seq_len arguments["max_prefill_predict_length"] = seq_len arguments["scan_layers"] = False + # Pool capacity has no safe default, so `gpu_paged` requires it. Enough here + # for one request of this length per batch row. + arguments["paged_num_blocks"] = 8 * (seq_len // 16 + 1) arguments.update(overrides) return pyconfig.initialize([sys.argv[0], get_test_config_path()], **arguments) diff --git a/tests/unit/gpu_paged_attention_test.py b/tests/unit/gpu_paged_attention_test.py index 34b6c0935f..ae565aecd9 100644 --- a/tests/unit/gpu_paged_attention_test.py +++ b/tests/unit/gpu_paged_attention_test.py @@ -182,19 +182,41 @@ def _config(self, **overrides): arguments.update(overrides) return pyconfig.initialize([sys.argv[0], get_test_config_path()], **arguments) + def _paged(self, **overrides): + arguments = {"attention": "gpu_paged", "scan_layers": False, "paged_num_blocks": 256} + arguments.update(overrides) + return self._config(**arguments) + def test_gpu_paged_is_an_accepted_attention_value(self): - cfg = self._config(attention="gpu_paged", scan_layers=False) + cfg = self._paged() self.assertEqual(cfg.attention, "gpu_paged") self.assertEqual(cfg.paged_attention_backend, "auto") + self.assertEqual(cfg.paged_page_size, TOKENS_PER_PAGE) def test_scan_layers_is_refused(self): """Scanning stacks the caches, which would copy the whole pool every step.""" with self.assertRaisesRegex(ValueError, "scan_layers=False"): - self._config(attention="gpu_paged", scan_layers=True) + self._paged(scan_layers=True) + + def test_a_pool_size_is_required(self): + """Capacity has no safe default: too small caps concurrency silently.""" + with self.assertRaisesRegex(ValueError, "paged_num_blocks"): + self._config(attention="gpu_paged", scan_layers=False) + + def test_the_context_length_defaults_to_the_target_length(self): + """max_target_length already states the longest sequence to be served.""" + cfg = self._paged() + self.assertEqual(cfg.paged_max_context_len, cfg.max_target_length) + + def test_a_pool_too_small_for_one_request_is_refused(self): + """Otherwise this surfaces as backpressure that never clears.""" + with self.assertRaisesRegex(ValueError, "no request could ever finish"): + self._paged(paged_num_blocks=2, paged_max_context_len=1024) - def test_other_attention_kernels_are_unaffected_by_the_guard(self): + def test_other_attention_kernels_are_unaffected_by_the_guards(self): cfg = self._config(attention="dot_product", scan_layers=True) self.assertEqual(cfg.attention, "dot_product") + self.assertEqual(cfg.paged_num_blocks, 0) if __name__ == "__main__": diff --git a/tests/unit/gpu_paged_decode_parity_test.py b/tests/unit/gpu_paged_decode_parity_test.py new file mode 100644 index 0000000000..521b69c4d1 --- /dev/null +++ b/tests/unit/gpu_paged_decode_parity_test.py @@ -0,0 +1,422 @@ +"""Model-level greedy decode: `gpu_paged` must match the dense path token for token. + +This is M3's actual exit criterion, and it is a stronger claim than the +attention-layer parity test in `attention_test.py`. That one compares one layer +on one forward pass with weights copied across. This one runs a whole model +through `MaxEngine` twice, for a prompt and several generated tokens, and asks +whether the *sequence* agrees. The difference matters because a paged decode +gets its context from pages written on earlier steps, so an error in slot +arithmetic, page ordering or last-page occupancy shows up only after the context +crosses a page boundary — which a single forward pass never does. + +It is also the wiring a benchmark needs. Both paths go through `MaxEngine`: +dense via `prefill`/`init_decode_state`/`insert`/`generate`, paged via the +sibling `init_paged_runtime`/`prefill_paged`/`generate_paged`. Nothing here +reaches around the engine into the model, so what passes here is what a serving +harness would drive. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import sys +import unittest + +from absl.testing import parameterized +import numpy as np +import pytest + +import jax +import jax.numpy as jnp +from flax import nnx +from flax.linen import partitioning as nn_partitioning + +from maxtext.common.common_types import MODEL_MODE_PREFILL, MODEL_MODE_TRAIN +from maxtext.configs import pyconfig +from maxtext.utils import maxtext_utils, model_creation_utils + +try: + from maxtext.inference.maxengine import maxengine +except ModuleNotFoundError as _exc: # pragma: no cover - environment dependent + # `maxengine` reaches JetStream for its engine base class and tokenizer types. + # Naming both remedies matters: installing google-jetstream pulls TensorFlow and + # two dozen other packages, which is a heavy intervention in a container built + # around a self-built ROCm jaxlib, whereas DECOUPLE_GCLOUD=TRUE substitutes + # stubs that carry everything this test reads. + pytest.skip( + f"MaxEngine needs JetStream ({_exc}). Either install google-jetstream under a constraints file " + f"pinning jax and jaxlib, or run with DECOUPLE_GCLOUD=TRUE to use the built-in stubs.", + allow_module_level=True, + ) + +from tests.utils.test_helpers import get_test_config_path # pylint: disable=wrong-import-position + +PROMPT = [3, 17, 42, 5, 9, 21, 33, 2] +# 8 prompt tokens plus 45 generated is 53, so the context spans four pages and +# crosses three boundaries. That is the point: a decode that stayed inside its +# first page would exercise no page ordering and no gather, and would pass +# whatever the page arithmetic did. +STEPS = 45 +PAGE = 16 + +# How close two logits may be and still count as tied. One bfloat16 ulp at unit +# magnitude, which is the resolution the reference itself is computed at. +# +# This exists because comparing two greedy *trajectories* token for token does +# not work here, and the reason took a while to see. Logits are computed in +# bfloat16, so the gap between the top two candidates is quantised, and exact +# ties are common — with this model they occur every few dozen steps and *not at +# reproducible places*, because bf16 quantises an accumulation whose order XLA is +# free to vary between processes. At a tied step argmax is decided by +# tie-breaking, two arithmetically correct implementations can pick differently, +# and every later token is downstream of that coin flip. A trajectory comparison +# is therefore flaky by construction: it looked at various times like a page +# ordering bug, a numerics drift and a clean pass, from the same code. +# +# So the comparison below is teacher-forced instead. It replays the paged path's +# own tokens through a full forward pass and asks, at every step, whether the +# token the paged path chose was *an* argmax. That is tie-tolerant, never +# diverges, and is a stronger statement than trajectory equality. +TIE_TOLERANCE = 2.0**-8 + +_COMMON = { + "base_emb_dim": 512, + "base_mlp_dim": 512, + "base_num_query_heads": 4, + "base_num_kv_heads": 4, + "base_num_decoder_layers": 2, + # head_dim 128 with equal query and KV head counts puts gqa_ratio at 1, which + # is inside the prebuilt pa_ragged configuration set. head_dim 128 is also the + # only size the ASM path accepts. + "head_dim": 128, + # 256 rather than 64. A small vocabulary over randomly initialised weights + # makes exact logit ties likely and makes the model collapse to emitting one + # token, both of which defeat a token-for-token comparison. See + # MIN_LOGIT_MARGIN. + "vocab_size": 256, + "max_prefill_predict_length": 64, + "max_target_length": 128, + "per_device_batch_size": 1, + "scan_layers": False, + "sparse_matmul": False, + # bfloat16, not float32: the paged kernels accept bfloat16 and float16 only, + # and comparing a bf16 paged pool against an fp32 dense cache would measure + # the dtype rather than the paging. + "dtype": "bfloat16", + "weight_dtype": "float32", + "matmul_precision": "highest", + "decode_sampling_strategy": "greedy", + "enable_checkpointing": False, + "skip_jax_distributed_system": True, + "pure_nnx": True, +} + +_DENSE = {"attention": "dot_product"} +_PAGED = { + "attention": "gpu_paged", + "paged_page_size": PAGE, + # 1024 tokens of pool, comfortably more than one request needs, so nothing + # here depends on recycling. The recycling path has its own tests. + "paged_num_blocks": 64, +} + + +def _require_kernels(): + """Skip unless jax-aiter is importable and its KV shims are built.""" + try: + from jax_aiter.ffi.registry import standalone_symbol_available # pylint: disable=import-outside-toplevel + except ImportError as exc: + raise unittest.SkipTest("jax-aiter is not importable; set PYTHONPATH to the jax-aiter checkout") from exc + for symbol in ("AppendKvJA", "PagedAttentionJA", "PagedPrefillJA"): + if not standalone_symbol_available(symbol): + raise unittest.SkipTest(f"{symbol} is not built; run 'make -f Makefile.kv ja_kv' and set JA_ROOT_DIR") + + +def _config(**overrides): + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **(_COMMON | overrides)) + + +def _devices(): + """One device, explicitly, however many the host exposes. + + Parity is a claim about arithmetic, and sharding is a separate milestone, so a + single device is the right scope. It is also the only scope that runs here: an + 8-way sharded model forward pass puts a collective in the program, and RCCL + clique initialisation aborts in this container with a rocprofiler + double-registration fatal once torch is loaded (transformers pulls it in). That + is unrelated to paging, but it would present as a core dump in this test, so + the device list is pinned rather than inherited. + """ + return jax.devices()[:1] + + +def _mesh(cfg): + return jax.sharding.Mesh( + maxtext_utils.create_device_mesh(config=cfg, devices=_devices()), cfg.mesh_axes + ) + + +@pytest.mark.gpu_only +class GpuPagedDecodeParityTest(parameterized.TestCase): + """A paged decode must reproduce the dense decode exactly.""" + + def setUp(self): + super().setUp() + _require_kernels() + + def _build_params(self, cfg): + """One set of random weights, to be loaded into both engines. + + Built once and shared rather than built twice from the same seed: sharing is + what makes the comparison about paging instead of about whether two + initialisations happened to agree. + """ + mesh = _mesh(cfg) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + model = model_creation_utils.create_model( + cfg, mesh, model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(params=0, dropout=0) + ) + _, params_state, _ = nnx.split(model, nnx.Param, ...) + return params_state + + def _verify_tokens_are_argmax(self, cfg, params_state, generated): + """Replay `generated` through full forward passes and check each was an argmax. + + Teacher-forced, so the reference follows the sequence under test rather than + running its own trajectory. That removes the failure mode a trajectory + comparison cannot avoid: one tied step no longer invalidates everything after + it, because the reference is re-anchored on the actual prefix every step. + + The reference has no KV cache at all, so it cannot share a bug with the thing + under test. Returns the per-step margin between the chosen token and the best + alternative, for reporting. + """ + mesh = _mesh(cfg) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + model = model_creation_utils.create_model( + cfg, mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=0) + ) + nnx.update(model, params_state) + + pad = cfg.max_target_length + # Train mode shards the batch axis over fsdp, which absorbs every device, so + # a batch of one cannot be laid out at all. Every row is the same prompt and + # only row 0 is read. + batch = cfg.micro_batch_size_to_train_on + margins = [] + for step, chosen in enumerate(generated): + context = list(PROMPT) + list(generated[:step]) + row = context + [0] * (pad - len(context)) + ids = jnp.tile(jnp.asarray([row], dtype=jnp.int32), (batch, 1)) + positions = jnp.tile(jnp.asarray([list(range(pad))], dtype=jnp.int32), (batch, 1)) + mask = [1] * len(context) + [0] * (pad - len(context)) + segment_ids = jnp.tile(jnp.asarray([mask], dtype=jnp.int32), (batch, 1)) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + logits = model( + ids, positions, decoder_segment_ids=segment_ids, enable_dropout=False, model_mode=MODEL_MODE_TRAIN + ) + scores = np.asarray(logits[0, len(context) - 1].astype(jnp.float32)) + best = float(scores.max()) + self.assertGreaterEqual( + float(scores[chosen]), + best - TIE_TOLERANCE, + f"at generated token {step} (context length {len(context)}, page boundary every {PAGE}) the " + f"paged path chose {chosen} scoring {float(scores[chosen]):.5f}, but the best was " + f"{int(scores.argmax())} scoring {best:.5f} — a gap of {best - float(scores[chosen]):.5f}, far " + f"beyond the {TIE_TOLERANCE:.5f} a bfloat16 tie can explain", + ) + margins.append(best - float(np.partition(scores, -2)[-2])) + return margins + + def _dense_rollout(self, cfg, params_state, steps=STEPS): + """Greedy decode on the dense two-region cache, through MaxEngine.""" + engine = maxengine.MaxEngine(cfg, _devices()) + params = engine.load_params(params=params_state) + padded = jnp.asarray( + PROMPT + [0] * (cfg.max_prefill_predict_length - len(PROMPT)), dtype=jnp.int32 + ) + prefix, first = engine.prefill(params=params, padded_tokens=padded, true_length=len(PROMPT)) + generated = [int(first.data[0, 0])] + + decode_state = engine.init_decode_state() + decode_state = engine.insert(prefix, decode_state, slot=0) + for _ in range(steps): + decode_state, result = engine.generate(params, decode_state) + generated.append(int(result.data[0, 0])) + return generated + + def _paged_rollout(self, cfg, params_state, steps=STEPS): + """Greedy decode on the page pool, through the sibling entry points.""" + engine = maxengine.MaxEngine(cfg, _devices()) + params = engine.load_params(params=params_state) + runtime = engine.init_paged_runtime() + + padded = jnp.asarray( + PROMPT + [0] * (cfg.max_prefill_predict_length - len(PROMPT)), dtype=jnp.int32 + ) + handle, first = engine.prefill_paged( + params=params, padded_tokens=padded, true_length=len(PROMPT), request_id="parity" + ) + self.assertIsNotNone(handle, "the pool refused a single request, so it is mis-sized for this test") + generated = [int(first.data[0, 0])] + + for _ in range(steps): + result, ok = engine.generate_paged(params, [handle], next_tokens=jnp.asarray([generated[-1]], jnp.int32)) + self.assertTrue(ok, "the pool ran out of pages during a single-request decode") + generated.append(int(result.data[0, 0])) + + # Housekeeping is part of the contract, so assert it rather than assume it. + pages_held = runtime.control_plane.page_map.num_pages(handle) + self.assertGreaterEqual( + pages_held, 3, f"the context only reached {pages_held} pages, so page ordering was barely exercised" + ) + engine.release(handle) + self.assertEqual(runtime.control_plane.allocator.num_allocated_pages, 0) + return generated + + def _assert_same_tokens(self, reference, actual, label): + """Compare sequences, reporting *where* they diverge. + + The index is the diagnosis. Token 0 comes from prefill, so a difference there + points at the projection or the prefill kernel; a difference first appearing + near a multiple of the page size points at page arithmetic; a drift that + starts late and grows points at numerics. + """ + self.assertEqual(len(reference), len(actual)) + first_diff = next((i for i, (a, b) in enumerate(zip(reference, actual)) if a != b), None) + self.assertIsNone( + first_diff, + f"paged decode diverged from {label} at generated token {first_diff} " + f"(context length {len(PROMPT) + (first_diff or 0)}, page boundary every {PAGE}):" + f"\n {label} = {reference}\n paged = {actual}", + ) + + def test_every_paged_token_is_an_argmax_of_a_full_forward_pass(self): + """M3's exit criterion, in the form that is actually decidable. + + The strong statement: at every step, over a 53-token context spanning four + pages, the token the paged path produced was the argmax of a cacheless + forward pass over exactly the prefix the paged path had built. A page + ordering error, a slot-arithmetic error or a stale last-page length would all + put a wrong token somewhere in that sequence. + """ + params_state = self._build_params(_config(**_DENSE)) + paged = self._paged_rollout(_config(**_PAGED), params_state) + margins = self._verify_tokens_are_argmax(_config(**_DENSE), params_state, paged) + + self.assertEqual(len(paged), STEPS + 1) + # Guard against the check passing because everything was tied. A run where + # most steps are decidable is a run where the assertion above meant something. + decidable = sum(1 for m in margins if m > TIE_TOLERANCE) + self.assertGreater( + decidable, + len(margins) // 2, + f"only {decidable} of {len(margins)} steps had a decidable argmax, so this proved little", + ) + + def test_paged_decode_matches_the_dense_engine_path_token_for_token(self): + """The same comparison against the dense cached path, engine to engine. + + The dense `_prefill_jit` returns its `ResultTokens` from inside `jit`, so the + type has to be a registered pytree. The paged siblings build theirs outside + `jit` and are unaffected, which is why the forward-pass comparison above runs + everywhere. + + The skip tests that capability directly rather than asking whether JetStream + is stubbed. Those are different questions: the stub is registered as a pytree + precisely so this path works without the real package, and a + provenance-based check would skip a comparison that is perfectly able to run. + """ + probe = maxengine.engine_api.ResultTokens( + data=jnp.zeros((1, 3), jnp.int32), + tokens_idx=(0, 1), + valid_idx=(1, 2), + length_idx=(2, 3), + log_prob=None, + samples_per_slot=1, + ) + if not jax.tree_util.tree_leaves(probe): + self.skipTest( + "engine_api.ResultTokens is not a registered pytree, so the dense prefill path cannot " + "return one from inside jit; install google-jetstream or register the stub" + ) + if jax.device_count() != len(_devices()): + # `pyconfig` derives the batch from `jax.device_count()`, so it sizes the + # dense cache for every visible device while the engine mesh is pinned to + # one. Prefill then runs at batch 1 against a batch-N cache and asserts + # deep inside the attention op. The paged path is unaffected because its + # pool geometry comes from the layout rather than from the device count. + self.skipTest( + f"the dense path needs the process to see exactly the devices the engine uses; " + f"{jax.device_count()} are visible and the engine is pinned to {len(_devices())}. " + f"Re-run with HIP_VISIBLE_DEVICES=0 (or CUDA_VISIBLE_DEVICES=0)." + ) + params_state = self._build_params(_config(**_DENSE)) + dense = self._dense_rollout(_config(**_DENSE), params_state) + # Verified against the forward pass rather than against each other, because a + # single tied step would otherwise make the two trajectories diverge for + # reasons that have nothing to do with paging. + self._verify_tokens_are_argmax(_config(**_DENSE), params_state, dense) + paged = self._paged_rollout(_config(**_PAGED), params_state) + self._verify_tokens_are_argmax(_config(**_DENSE), params_state, paged) + + def test_the_paged_path_allocates_no_dense_cache(self): + """A dense cache alongside the pool would waste gigabytes at real sizes.""" + cfg = _config(**_PAGED) + engine = maxengine.MaxEngine(cfg, _devices()) + params = engine.load_params(params=self._build_params(cfg)) + runtime = engine.init_paged_runtime() + + expected_bytes = runtime.control_plane.layout.pool_bytes_per_shard() + actual = sum( + int(np.prod(a.shape)) * a.dtype.itemsize + for layer in range(runtime.pool.num_layers) + for a in (runtime.pool.k_pages[layer], runtime.pool.v_pages[layer]) + ) + self.assertEqual(actual, expected_bytes) + del params + + def test_the_pool_is_mutated_in_place_rather_than_replaced(self): + """The M0 aliasing invariant, observed at the top of the stack. + + A pool that is copied rather than aliased still produces correct tokens, so + nothing else in this file would notice. Checking that the arrays change + identity exactly once per step -- because donation hands back a new handle to + the same buffer -- is not the invariant either. What is observable here and + worth pinning is that the pool does not grow: a replacement allocation per + step would show up as a rising live-buffer count. + """ + cfg = _config(**_PAGED) + engine = maxengine.MaxEngine(cfg, _devices()) + params = engine.load_params(params=self._build_params(cfg)) + runtime = engine.init_paged_runtime() + + padded = jnp.asarray(PROMPT + [0] * (cfg.max_prefill_predict_length - len(PROMPT)), dtype=jnp.int32) + handle, first = engine.prefill_paged( + params=params, padded_tokens=padded, true_length=len(PROMPT), request_id="alias" + ) + shapes_before = [(a.shape, a.dtype) for a in runtime.pool.k_pages] + token = int(first.data[0, 0]) + for _ in range(8): + result, ok = engine.generate_paged(params, [handle], next_tokens=jnp.asarray([token], jnp.int32)) + self.assertTrue(ok) + token = int(result.data[0, 0]) + + self.assertEqual([(a.shape, a.dtype) for a in runtime.pool.k_pages], shapes_before) + for array in runtime.pool.k_pages + runtime.pool.v_pages: + self.assertFalse(array.is_deleted(), "a donated pool array was never rebound") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/kv_common_test.py b/tests/unit/kv_common_test.py index 6553d45333..d6ad7e8d6d 100644 --- a/tests/unit/kv_common_test.py +++ b/tests/unit/kv_common_test.py @@ -20,67 +20,31 @@ limitations under the License. """ -import ast -import importlib.util -import pathlib -import sys import unittest import numpy as np -# Loaded straight from their files rather than through `maxtext.inference.kv_common` -# on purpose. The neutral layer is meant to run with no accelerator and no -# framework, but importing it by package path executes `maxtext/__init__.py`, -# which pulls in the full config stack. Loading the modules directly is what -# actually demonstrates the property, and it keeps these tests runnable in plain -# CI. `_ImportRuleTest` below pins the rule statically. -_KV_COMMON_DIR = ( - pathlib.Path(__file__).resolve().parents[2] - / "src" - / "maxtext" - / "inference" - / "kv_common" -) - - -def _load(module_name: str): - path = _KV_COMMON_DIR / f"{module_name}.py" - spec = importlib.util.spec_from_file_location(f"_kv_common_{module_name}", path) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -KvStorageLayoutV1 = _load("storage_layout").KvStorageLayoutV1 -KvPageTableV1 = _load("page_table").KvPageTableV1 - -_ALLOWED_IMPORTS = { - "numpy", - "dataclasses", - "annotations", # from __future__ - "__future__", - "ast", - "typing", - "math", - "enum", - "collections", - "hashlib", -} +from maxtext.inference.kv_common import KvPageTableV1, KvStorageLayoutV1 + +# Imported by package path, which is cheap: `maxtext/__init__.py` resolves its +# heavy exports through a lazy `__getattr__`, so reaching a submodule loads +# neither jax nor the config stack. `kv_import_rule_test.py` is what holds that +# true, statically over the whole layer and once against a fresh interpreter. class KvStorageLayoutTest(unittest.TestCase): """Pool geometry and sizing.""" def _layout(self, **kw): - base = dict( - tokens_per_page=16, - num_pages=1024, - num_layers=32, - num_kv_heads=8, - head_dim=128, - dtype="bfloat16", - ) + """A realistic pool geometry, overridable field by field.""" + base = { + "tokens_per_page": 16, + "num_pages": 1024, + "num_layers": 32, + "num_kv_heads": 8, + "head_dim": 128, + "dtype": "bfloat16", + } base.update(kw) return KvStorageLayoutV1(**base) @@ -238,72 +202,5 @@ def test_slot_mapping_rejects_position_beyond_held_pages(self): table.slot_mapping(tokens_per_page=16) -class ImportRuleTest(unittest.TestCase): - """The neutral layer must stay free of framework and vendor dependencies. - - Checked statically, so it holds without importing anything. A convention no - one can verify erodes, and this one is load-bearing: it is what makes the layer - CPU-testable today and mechanically extractable later. - """ - - FORBIDDEN_ROOTS = {"jax", "jaxlib", "flax", "torch", "jax_aiter"} - SELF_PACKAGE = "maxtext.inference.kv_common" - - def _top_level_imports(self, path): - """Root packages imported by `path`, excluding intra-package imports. - - An absolute import of kv_common's own modules is fine; anything else under - `maxtext.` is not, and is reported under its full name so the message is - actionable. - """ - tree = ast.parse(path.read_text()) - roots = set() - for node in ast.walk(tree): - if isinstance(node, ast.Import): - for alias in node.names: - if alias.name.startswith(self.SELF_PACKAGE): - continue - roots.add(alias.name.split(".")[0]) - elif isinstance(node, ast.ImportFrom): - if node.level or not node.module: # relative import stays in-package - continue - if node.module.startswith(self.SELF_PACKAGE): - continue - if node.module.startswith("maxtext"): - roots.add(node.module) # full name: this is a rule violation - continue - roots.add(node.module.split(".")[0]) - return roots - - def test_modules_import_only_stdlib_and_numpy(self): - for path in sorted(_KV_COMMON_DIR.glob("*.py")): - roots = self._top_level_imports(path) - forbidden = roots & self.FORBIDDEN_ROOTS - self.assertEqual( - forbidden, - set(), - f"{path.name} imports {sorted(forbidden)}, which the kv_common import " - f"rule forbids", - ) - unexpected = roots - _ALLOWED_IMPORTS - self.assertEqual( - unexpected, - set(), - f"{path.name} imports {sorted(unexpected)}; extend _ALLOWED_IMPORTS " - f"only if the addition is genuinely stdlib or numpy", - ) - - def test_package_init_only_reaches_into_kv_common(self): - """The package __init__ may re-export, but must not reach outside itself.""" - init = _KV_COMMON_DIR / "__init__.py" - tree = ast.parse(init.read_text()) - for node in ast.walk(tree): - if isinstance(node, ast.ImportFrom) and node.module: - self.assertTrue( - node.module.startswith("maxtext.inference.kv_common"), - f"__init__ imports from {node.module}, outside kv_common", - ) - - if __name__ == "__main__": unittest.main() diff --git a/tests/unit/kv_control_test.py b/tests/unit/kv_control_test.py new file mode 100644 index 0000000000..82178b1fcc --- /dev/null +++ b/tests/unit/kv_control_test.py @@ -0,0 +1,746 @@ +"""CPU-only tests for the paged KV control plane. + +All of this is host logic over small integer arrays, so all of it is testable +with no accelerator. That is the point of the layering, and these tests are what +make the claim more than an assertion. + +A trap worth knowing about before trusting a green run here: `tests/conftest.py` +auto-marks any test without a hardware marker as `cpu_only`, and `cpu_only` tests +are skipped whenever an accelerator is visible. On a GPU box this whole file +therefore *skips* rather than passes, silently and quickly. Run it with +`JAX_PLATFORMS=cpu`, and read the count, not the colour. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest + +import numpy as np + +from maxtext.inference.kv_common import KvStorageLayoutV1 +from maxtext.inference.kv_control import ( + DirtyPageError, + DoubleFreeError, + KvControlPlane, + NativeKvControlPlane, + PageCapacityError, + PageMap, + PagedBlockAllocator, + PageState, + PageStateError, + RequestDescriptor, + RequestHandle, + RequestState, + StaleRequestHandleError, + build_decode_table, + build_page_table, + decode_needs_new_page, + last_page_occupancy, + new_pages_for_extend, + pages_for_tokens, + token_slot, +) +from maxtext.inference.kv_control.logical_block import LogicalBlock, check_transition + +PAGE = 16 + + +def _layout(**kw) -> KvStorageLayoutV1: + """A small but realistic pool geometry, overridable field by field.""" + base = { + "tokens_per_page": PAGE, + "num_pages": 64, + "num_layers": 4, + "num_kv_heads": 8, + "head_dim": 128, + "dtype": "bfloat16", + } + base.update(kw) + return KvStorageLayoutV1(**base) + + +class PageArithmeticTest(unittest.TestCase): + """The ceilings and offsets every other module leans on.""" + + def test_pages_for_tokens(self): + self.assertEqual(pages_for_tokens(0, PAGE), 0) + self.assertEqual(pages_for_tokens(1, PAGE), 1) + self.assertEqual(pages_for_tokens(PAGE, PAGE), 1) + self.assertEqual(pages_for_tokens(PAGE + 1, PAGE), 2) + + def test_extend_continues_the_open_page_rather_than_paying_for_it(self): + """A difference of ceilings, not a ceiling of the difference. + + Growing 8 -> 25 tokens needs one new page: the first page already holds + positions 0-15, so eight of the seventeen new tokens land in it for free. + Costing the extend as ceil(17/16) would claim two and over-allocate on every + mid-page extend, which is the common case. + """ + self.assertEqual(new_pages_for_extend(8, 25, PAGE), 1) + self.assertNotEqual(new_pages_for_extend(8, 25, PAGE), pages_for_tokens(25 - 8, PAGE)) + + def test_extend_from_empty_is_a_plain_ceiling(self): + self.assertEqual(new_pages_for_extend(0, 33, PAGE), 3) + + def test_extend_within_the_open_page_is_free(self): + self.assertEqual(new_pages_for_extend(1, PAGE, PAGE), 0) + + def test_extend_backwards_is_rejected(self): + with self.assertRaises(ValueError): + new_pages_for_extend(20, 10, PAGE) + + def test_decode_takes_a_page_only_when_it_crosses_a_boundary(self): + self.assertTrue(decode_needs_new_page(1, PAGE)) # the very first token + self.assertFalse(decode_needs_new_page(PAGE, PAGE)) # fills the first page + self.assertTrue(decode_needs_new_page(PAGE + 1, PAGE)) # opens the second + self.assertFalse(decode_needs_new_page(PAGE + 2, PAGE)) + + def test_last_page_occupancy_reports_a_full_page_as_full(self): + self.assertEqual(last_page_occupancy(0, PAGE), 0) + self.assertEqual(last_page_occupancy(PAGE, PAGE), PAGE) + self.assertEqual(last_page_occupancy(PAGE + 1, PAGE), 1) + self.assertEqual(last_page_occupancy(33, PAGE), 1) + + def test_token_slot_is_page_major(self): + self.assertEqual(token_slot(3, 32, PAGE), 48) + self.assertEqual(token_slot(5, 19, PAGE), 83) + + +class PageStateTest(unittest.TestCase): + """The readable-extent gate, expressed as a transition.""" + + def test_a_fresh_page_is_not_readable_until_written(self): + with self.assertRaises(PageStateError): + check_transition(PageState.FREE, PageState.READY) + + def test_written_then_ready_is_the_normal_path(self): + block = LogicalBlock(page_id=7, epoch=1) + self.assertFalse(block.is_readable) + block.set_state(PageState.READY) + self.assertTrue(block.is_readable) + + def test_a_ready_page_may_be_reopened_for_more_tokens(self): + block = LogicalBlock(page_id=7, epoch=1, state=PageState.READY) + block.set_state(PageState.WRITING) + self.assertFalse(block.is_readable) + + def test_restating_the_current_state_is_not_an_error(self): + block = LogicalBlock(page_id=7, epoch=1) + block.set_state(PageState.WRITING) + self.assertIs(block.state, PageState.WRITING) + + +class AllocatorTest(unittest.TestCase): + """Free-list mechanics, including the parts the reference gets away without.""" + + def _allocator(self, num_pages=8, **kw): + return PagedBlockAllocator(num_pages=num_pages, tokens_per_page=PAGE, debug_mode=True, **kw) + + def test_the_padding_page_is_never_handed_out(self): + alloc = self._allocator(num_pages=8) + self.assertEqual(alloc.capacity_pages, 7) + pages = alloc.alloc(7) + self.assertIsNotNone(pages) + self.assertNotIn(0, pages.tolist()) + self.assertEqual(pages.tolist(), [1, 2, 3, 4, 5, 6, 7]) + + def test_a_non_zero_reserved_page_is_also_respected(self): + alloc = self._allocator(num_pages=8, padding_page_id=3) + self.assertEqual(alloc.alloc(7).tolist(), [0, 1, 2, 4, 5, 6, 7]) + + def test_allocation_pops_the_front_of_a_sorted_list(self): + alloc = self._allocator() + self.assertEqual(alloc.alloc(3).tolist(), [1, 2, 3]) + self.assertEqual(alloc.alloc(2).tolist(), [4, 5]) + + def test_exhaustion_returns_none_rather_than_raising(self): + alloc = self._allocator() + self.assertIsNotNone(alloc.alloc(7)) + self.assertIsNone(alloc.alloc(1)) + + def test_freeing_stages_the_pages_instead_of_sorting_immediately(self): + """The cheap half of the two-tier list: no sort on the free path.""" + alloc = self._allocator() + alloc.alloc(7) + self.assertEqual(alloc.num_free_pages, 0) + alloc.free([5, 2]) + self.assertEqual(alloc.num_free_pages, 0, "freed pages must not reach the allocation list yet") + self.assertEqual(alloc.available_pages, 2) + + def test_the_merge_happens_when_allocation_would_otherwise_fail(self): + alloc = self._allocator() + alloc.alloc(7) + alloc.free([5, 2]) + self.assertEqual(alloc.alloc(2).tolist(), [2, 5], "the merge must also restore sorted order") + self.assertEqual(alloc.num_free_pages, 0) + + def test_an_allocation_that_fits_the_front_list_skips_the_merge(self): + alloc = self._allocator() + alloc.alloc(3) # pages 1-3, leaving 4-7 in the front list + alloc.free([2]) + self.assertEqual(alloc.alloc(1).tolist(), [4], "page 2 was staged, so it must not be reused yet") + self.assertEqual(alloc.num_free_pages, 3, "pages 5-7 remain in the front list") + self.assertEqual(alloc.available_pages, 4, "page 2 is still staged, and still counts as available") + + def test_many_separate_frees_merge_into_one_sorted_list(self): + """The staging path with many entries, not the single-free case above. + + Freed one call at a time and in descending order, so a merge that failed to + consider every staged entry, or to re-sort, would show up here. + """ + alloc = self._allocator(num_pages=32) + pages = alloc.alloc(31).tolist() + for page in reversed(pages): + alloc.free([page]) + self.assertEqual(alloc.num_free_pages, 0) + self.assertEqual(alloc.available_pages, 31) + self.assertEqual(alloc.alloc(31).tolist(), sorted(pages)) + + def test_available_tokens_counts_staged_pages(self): + alloc = self._allocator() + alloc.alloc(7) + alloc.free([1, 2]) + self.assertEqual(alloc.available_tokens, 2 * PAGE) + + def test_a_double_free_is_diagnosed_not_absorbed(self): + """The reference deduplicates this silently; a page freed twice is a bug.""" + alloc = self._allocator() + pages = alloc.alloc(3) + alloc.free(pages) + with self.assertRaises(DoubleFreeError): + alloc.free(pages) + + def test_duplicates_within_one_free_call_are_fine(self): + """They are what converting token slots to page ids produces.""" + alloc = self._allocator() + alloc.alloc(3) + self.assertEqual(alloc.free([2, 2, 3]).tolist(), [2, 3]) + + def test_freeing_the_reserved_page_is_rejected(self): + alloc = self._allocator() + with self.assertRaises(ValueError): + alloc.free([0]) + + def test_freeing_outside_the_pool_is_rejected(self): + alloc = self._allocator() + with self.assertRaises(ValueError): + alloc.free([99]) + + def test_free_token_slots_drops_the_padding_sentinel(self): + """A slot_mapping carries -1 for padded rows, which is correct, not an error.""" + alloc = self._allocator() + alloc.alloc(4) + freed = alloc.free_token_slots([-1, 2 * PAGE, 2 * PAGE + 5, 3 * PAGE, -1]) + self.assertEqual(freed.tolist(), [2, 3]) + + def test_free_token_slots_on_an_all_padding_step_frees_nothing(self): + alloc = self._allocator() + self.assertEqual(alloc.free_token_slots([-1, -1]).size, 0) + + def test_the_epoch_advances_when_a_page_is_reused(self): + alloc = self._allocator() + page = int(alloc.alloc(1)[0]) + first = alloc.epoch_of(page) + self.assertTrue(alloc.holds(page, first)) + + alloc.free([page]) + self.assertFalse(alloc.holds(page, first), "a freed page must not satisfy a live reference") + + alloc.merge_released() + self.assertEqual(int(alloc.alloc(1)[0]), page) + self.assertEqual(alloc.epoch_of(page), first + 1) + self.assertFalse(alloc.holds(page, first), "the previous owner's reference must not match") + + def test_clear_invalidates_every_outstanding_reference(self): + alloc = self._allocator() + page = int(alloc.alloc(1)[0]) + epoch = alloc.epoch_of(page) + alloc.clear() + self.assertEqual(alloc.available_pages, alloc.capacity_pages) + self.assertFalse(alloc.holds(page, epoch)) + + def test_churn_returns_every_page(self): + """The M4 exit criterion in miniature: no leak under repeated turnover.""" + alloc = self._allocator(num_pages=32) + rng = np.random.default_rng(0) + for _ in range(500): + want = int(rng.integers(1, 8)) + pages = alloc.alloc(want) + if pages is None: + continue + alloc.free(pages) + self.assertEqual(alloc.num_allocated_pages, 0) + self.assertEqual(alloc.available_pages, alloc.capacity_pages) + + def test_from_layout_takes_the_pool_geometry(self): + alloc = PagedBlockAllocator.from_layout(_layout(num_pages=64)) + self.assertEqual(alloc.capacity_pages, 63) + self.assertEqual(alloc.tokens_per_page, PAGE) + + def test_a_never_allocated_page_is_clean(self): + """The pool is zero-initialised, so a fresh pool costs no scrubbing at all.""" + alloc = self._allocator() + self.assertEqual(alloc.num_dirty_pages, 0) + self.assertEqual(alloc.dirty_among(alloc.alloc(3)).size, 0) + + def test_freeing_makes_a_page_dirty(self): + alloc = self._allocator() + pages = alloc.alloc(3) + alloc.free(pages) + self.assertEqual(alloc.num_dirty_pages, 3) + for page in pages.tolist(): + self.assertTrue(alloc.is_dirty(page)) + + def test_a_recycled_page_is_reported_dirty_on_reallocation(self): + """The point of the tracking: this page holds someone else's KV.""" + alloc = self._allocator() + first = alloc.alloc(2) + alloc.free(first) + alloc.merge_released() + second = alloc.alloc(2) + np.testing.assert_array_equal(np.sort(alloc.dirty_among(second)), np.sort(first)) + + def test_marking_scrubbed_clears_the_obligation(self): + alloc = self._allocator() + pages = alloc.alloc(2) + alloc.free(pages) + alloc.mark_scrubbed(pages) + self.assertEqual(alloc.num_dirty_pages, 0) + alloc.merge_released() + self.assertEqual(alloc.dirty_among(alloc.alloc(2)).size, 0) + + def test_clear_leaves_handed_out_pages_dirty(self): + """Dropping the free list overwrites nothing, so it cannot clean anything.""" + alloc = self._allocator() + pages = alloc.alloc(3) + alloc.clear() + for page in pages.tolist(): + self.assertTrue(alloc.is_dirty(page)) + + +class PageMapTest(unittest.TestCase): + """Rows, epochs, and the length bound that keeps reads inside a request.""" + + def _map(self, max_requests=2, max_pages_per_request=4): + return PageMap( + max_requests=max_requests, + max_pages_per_request=max_pages_per_request, + tokens_per_page=PAGE, + ) + + def _descriptor(self, request_id="r0", prompt_len=16, max_new_tokens=16): + return RequestDescriptor(request_id=request_id, prompt_len=prompt_len, max_new_tokens=max_new_tokens) + + def test_admit_then_release_returns_the_row(self): + page_map = self._map() + handle = page_map.admit(self._descriptor()) + self.assertEqual(page_map.num_live, 1) + self.assertIs(page_map.state(handle), RequestState.WAITING) + page_map.release(handle) + self.assertEqual(page_map.num_live, 0) + self.assertEqual(page_map.available_rows, 2) + + def test_admit_returns_none_when_every_row_is_taken(self): + page_map = self._map(max_requests=1) + self.assertIsNotNone(page_map.admit(self._descriptor("a"))) + self.assertIsNone(page_map.admit(self._descriptor("b"))) + + def test_release_reports_the_pages_that_were_held(self): + page_map = self._map() + handle = page_map.admit(self._descriptor()) + page_map.append_pages(handle, [4, 9]) + np.testing.assert_array_equal(page_map.release(handle), np.array([4, 9], dtype=np.int32)) + + def test_a_released_handle_is_refused(self): + page_map = self._map() + handle = page_map.admit(self._descriptor()) + page_map.release(handle) + with self.assertRaises(StaleRequestHandleError): + page_map.pages(handle) + + def test_a_handle_to_a_reused_row_is_refused(self): + """The case an index alone cannot catch, and the reason for the epoch.""" + page_map = self._map(max_requests=1) + first = page_map.admit(self._descriptor("first")) + page_map.append_pages(first, [3]) + page_map.release(first) + + second = page_map.admit(self._descriptor("second")) + self.assertEqual(second.row, first.row) + self.assertNotEqual(second.epoch, first.epoch) + page_map.append_pages(second, [7]) + with self.assertRaises(StaleRequestHandleError): + page_map.pages(first) + np.testing.assert_array_equal(page_map.pages(second), np.array([7], dtype=np.int32)) + + def test_a_forged_handle_is_refused(self): + page_map = self._map() + page_map.admit(self._descriptor()) + with self.assertRaises(StaleRequestHandleError): + page_map.pages(RequestHandle(request_id="r0", row=0, epoch=99)) + with self.assertRaises(StaleRequestHandleError): + page_map.pages(RequestHandle(request_id="r0", row=17, epoch=0)) + + def test_pages_are_returned_in_append_order(self): + page_map = self._map() + handle = page_map.admit(self._descriptor()) + page_map.append_pages(handle, [9]) + page_map.append_pages(handle, [4, 1]) + np.testing.assert_array_equal(page_map.pages(handle), np.array([9, 4, 1], dtype=np.int32)) + + def test_pages_are_returned_as_a_copy(self): + page_map = self._map() + handle = page_map.admit(self._descriptor()) + page_map.append_pages(handle, [9]) + pages = page_map.pages(handle) + pages[0] = 123 + np.testing.assert_array_equal(page_map.pages(handle), np.array([9], dtype=np.int32)) + + def test_exceeding_a_row_is_reported(self): + page_map = self._map(max_pages_per_request=2) + handle = page_map.admit(self._descriptor()) + with self.assertRaises(PageCapacityError): + page_map.append_pages(handle, [1, 2, 3]) + + def test_advancing_past_the_held_pages_is_refused(self): + """The bound that stops a kernel reading a page the request does not own.""" + page_map = self._map() + handle = page_map.admit(self._descriptor()) + page_map.append_pages(handle, [1]) + self.assertEqual(page_map.advance(handle, PAGE), PAGE) + with self.assertRaises(PageCapacityError): + page_map.advance(handle, 1) + + def test_live_handles_are_current(self): + page_map = self._map() + first = page_map.admit(self._descriptor("a")) + second = page_map.admit(self._descriptor("b")) + self.assertEqual([h.request_id for h in page_map.live_handles()], ["a", "b"]) + page_map.release(first) + self.assertEqual([h.request_id for h in page_map.live_handles()], ["b"]) + self.assertEqual(page_map.live_handles()[0], second) + + def test_from_layout_sizes_rows_from_the_context_length(self): + page_map = PageMap.from_layout(_layout(), max_requests=4, max_context_len=33) + self.assertEqual(page_map.max_pages_per_request, 3) + + +class MetadataTest(unittest.TestCase): + """Building a `KvPageTableV1` from live bookkeeping.""" + + def _map_with(self, specs): + """`specs` is (request_id, pages, seq_len) per request.""" + page_map = PageMap(max_requests=8, max_pages_per_request=8, tokens_per_page=PAGE) + handles = [] + for request_id, pages, seq_len in specs: + handle = page_map.admit( + RequestDescriptor(request_id=request_id, prompt_len=seq_len, max_new_tokens=0) + ) + page_map.append_pages(handle, pages) + page_map.advance(handle, seq_len) + handles.append(handle) + return page_map, handles + + def test_decode_table_matches_a_hand_calculation(self): + page_map, handles = self._map_with([("a", [1, 2, 3], 33), ("b", [4, 5], 20)]) + table = build_decode_table(page_map, handles) + + np.testing.assert_array_equal(table.seq_lens, np.array([33, 20], dtype=np.int32)) + np.testing.assert_array_equal(table.query_lens, np.array([1, 1], dtype=np.int32)) + np.testing.assert_array_equal(table.write_positions, np.array([32, 19], dtype=np.int32)) + np.testing.assert_array_equal(table.indptr(), np.array([0, 3, 5], dtype=np.int32)) + np.testing.assert_array_equal(table.flat_page_indices(), np.array([1, 2, 3, 4, 5], dtype=np.int32)) + np.testing.assert_array_equal(table.last_page_lens(PAGE), np.array([1, 4], dtype=np.int32)) + # page 3 offset 0, and page 5 offset 3 + np.testing.assert_array_equal(table.slot_mapping(PAGE), np.array([48, 83], dtype=np.int32)) + + def test_pages_beyond_the_current_length_are_trimmed(self): + """Over-supply is silent corruption, because occupancy lands on the last page. + + Three pages for a 17-token context would put a last-page length of 1 on + page 3, so a kernel would read all of page 2 -- whatever the page's previous + owner left there -- and one token of real data. + """ + page_map, handles = self._map_with([("a", [1, 2], 17)]) + page_map.append_pages(handles[0], [3]) + self.assertEqual(page_map.num_pages(handles[0]), 3) + + table = build_decode_table(page_map, handles) + self.assertEqual(table.page_ids, [[1, 2]]) + np.testing.assert_array_equal(table.last_page_lens(PAGE), np.array([1], dtype=np.int32)) + + def test_prefill_positions_cover_the_uncached_suffix(self): + page_map, handles = self._map_with([("a", [1, 2], 18)]) + table = build_page_table(page_map, handles, [18]) + np.testing.assert_array_equal(table.write_positions, np.arange(18, dtype=np.int32)) + expected = [1 * PAGE + i for i in range(PAGE)] + [2 * PAGE, 2 * PAGE + 1] + np.testing.assert_array_equal(table.slot_mapping(PAGE), np.array(expected, dtype=np.int32)) + + def test_a_partial_prefill_starts_where_the_prefix_ended(self): + """A chunked or prefix-cached request contributes only its suffix.""" + page_map, handles = self._map_with([("a", [1, 2], 20)]) + table = build_page_table(page_map, handles, [4]) + np.testing.assert_array_equal(table.write_positions, np.array([16, 17, 18, 19], dtype=np.int32)) + + def test_request_order_names_the_row_behind_each_batch_position(self): + page_map, handles = self._map_with([("a", [1], 5), ("b", [2], 5)]) + table = build_page_table(page_map, [handles[1], handles[0]], [1, 1]) + np.testing.assert_array_equal(table.request_order, np.array([1, 0], dtype=np.int32)) + + def test_a_query_longer_than_the_context_is_rejected(self): + page_map, handles = self._map_with([("a", [1], 4)]) + with self.assertRaises(ValueError): + build_page_table(page_map, handles, [5]) + + def test_mismatched_lengths_are_rejected(self): + page_map, handles = self._map_with([("a", [1], 4)]) + with self.assertRaises(ValueError): + build_page_table(page_map, handles, [1, 1]) + + def test_an_empty_batch_builds_an_empty_table(self): + page_map, _ = self._map_with([]) + table = build_page_table(page_map, [], []) + self.assertEqual(table.num_requests, 0) + self.assertEqual(table.num_tokens, 0) + np.testing.assert_array_equal(table.indptr(), np.zeros((1,), dtype=np.int32)) + + +class NativeKvControlPlaneTest(unittest.TestCase): + """Admission, reservation, release, and the all-or-nothing property.""" + + def _plane(self, num_pages=16, max_requests=4, max_context_len=64, **kw): + return NativeKvControlPlane( + layout=_layout(num_pages=num_pages), + max_requests=max_requests, + max_context_len=max_context_len, + debug_mode=True, + **kw, + ) + + def _admit(self, plane, request_id="r0", prompt_len=16, max_new_tokens=16): + return plane.admit( + RequestDescriptor(request_id=request_id, prompt_len=prompt_len, max_new_tokens=max_new_tokens) + ) + + def test_it_satisfies_the_control_plane_protocol(self): + self.assertIsInstance(self._plane(), KvControlPlane) + + def test_a_request_longer_than_the_configured_context_is_refused(self): + plane = self._plane(max_context_len=32) + with self.assertRaises(ValueError): + self._admit(plane, prompt_len=32, max_new_tokens=1) + + def test_a_pool_too_small_for_one_request_fails_at_construction(self): + """Better than surfacing later as backpressure that never clears.""" + with self.assertRaises(ValueError): + self._plane(num_pages=3, max_context_len=64) + + def test_prefill_reserves_exactly_the_pages_the_prompt_needs(self): + plane = self._plane() + handle = self._admit(plane, prompt_len=33) + before = plane.allocator.available_pages + self.assertTrue(plane.reserve([handle], [33])) + self.assertEqual(plane.allocator.available_pages, before - 3) + self.assertEqual(plane.page_map.seq_len(handle), 33) + + def test_decode_takes_a_page_only_on_a_boundary(self): + plane = self._plane() + handle = self._admit(plane, prompt_len=PAGE) + self.assertTrue(plane.reserve([handle], [PAGE])) + self.assertEqual(plane.page_map.num_pages(handle), 1) + + self.assertTrue(plane.reserve_decode([handle])) # position 16 opens page 2 + self.assertEqual(plane.page_map.num_pages(handle), 2) + + self.assertTrue(plane.reserve_decode([handle])) # position 17 continues it + self.assertEqual(plane.page_map.num_pages(handle), 2) + self.assertEqual(plane.page_map.seq_len(handle), PAGE + 2) + + def test_a_batch_reservation_that_cannot_fit_changes_nothing(self): + """A partial reservation would leave a batch no page table can describe.""" + plane = self._plane(num_pages=4, max_requests=3, max_context_len=32) + handles = [self._admit(plane, f"r{i}", prompt_len=PAGE, max_new_tokens=PAGE) for i in range(3)] + self.assertTrue(plane.reserve(handles, [PAGE] * 3)) + self.assertEqual(plane.allocator.available_pages, 0) + + self.assertFalse(plane.reserve_decode(handles), "three boundary crossings cannot fit in zero pages") + for handle in handles: + self.assertEqual(plane.page_map.seq_len(handle), PAGE) + self.assertEqual(plane.page_map.num_pages(handle), 1) + + def test_generating_past_the_admitted_bound_is_reported(self): + plane = self._plane(max_context_len=PAGE) + handle = self._admit(plane, prompt_len=PAGE, max_new_tokens=0) + self.assertTrue(plane.reserve([handle], [PAGE])) + with self.assertRaises(PageCapacityError): + plane.reserve_decode([handle]) + + def test_release_returns_the_pages_to_the_pool(self): + plane = self._plane() + capacity = plane.allocator.capacity_pages + handle = self._admit(plane, prompt_len=33) + plane.reserve([handle], [33]) + reclaimed = plane.release(handle) + self.assertEqual(reclaimed.size, 3) + self.assertEqual(plane.allocator.available_pages, capacity) + self.assertEqual(plane.allocator.num_allocated_pages, 0) + + def test_releasing_twice_is_reported(self): + plane = self._plane() + handle = self._admit(plane) + plane.reserve([handle], [16]) + plane.release(handle) + with self.assertRaises(StaleRequestHandleError): + plane.release(handle) + + def test_a_reserve_of_zero_tokens_is_a_no_op(self): + plane = self._plane() + handle = self._admit(plane) + before = plane.allocator.available_pages + self.assertTrue(plane.reserve([handle], [0])) + self.assertEqual(plane.allocator.available_pages, before) + self.assertEqual(plane.page_map.seq_len(handle), 0) + + def test_an_empty_batch_reserves_successfully(self): + self.assertTrue(self._plane().reserve([], [])) + + def test_a_mixed_length_batch_under_churn_leaks_nothing(self): + """The M4 exit criterion, host side: turnover must return every page. + + Requests of assorted lengths arrive, decode a while, and leave. If either + the free list or the page map dropped a page, the pool would not come back + to its full capacity. + """ + plane = self._plane(num_pages=64, max_requests=4, max_context_len=128) + capacity = plane.allocator.capacity_pages + rng = np.random.default_rng(0) + live: list[RequestHandle] = [] + + for step in range(200): + if len(live) < 4 and rng.random() < 0.5: + prompt_len = int(rng.integers(1, 40)) + handle = self._admit(plane, f"r{step}", prompt_len=prompt_len, max_new_tokens=32) + if handle is not None and plane.reserve([handle], [prompt_len]): + plane.confirm_scrubbed(plane.pending_scrub()) + plane.page_map.set_state(handle, RequestState.DECODE) + live.append(handle) + elif handle is not None: + plane.release(handle) + + if live: + if plane.reserve_decode(live): + # Standing in for the driver's device-side scrub. Without it the next + # line raises, which is the whole point of the gate. + plane.confirm_scrubbed(plane.pending_scrub()) + table = plane.build_decode_table(live) + table.validate(PAGE) + self.assertEqual(table.num_tokens, len(live)) + finished = [h for h in live if rng.random() < 0.25] + for handle in finished: + plane.release(handle) + live = [h for h in live if h != handle] + + for handle in live: + plane.release(handle) + self.assertEqual(plane.num_live, 0) + self.assertEqual(plane.allocator.num_allocated_pages, 0) + self.assertEqual(plane.allocator.available_pages, capacity) + + def test_a_fresh_pool_needs_no_scrubbing(self): + plane = self._plane() + handle = self._admit(plane, prompt_len=33) + self.assertTrue(plane.reserve([handle], [33])) + self.assertEqual(plane.pending_scrub().size, 0) + plane.build_page_table([handle], [33]) # must not raise + + def test_a_recycled_page_must_be_scrubbed_before_it_can_be_described(self): + """The acceptance gate, host side: a dirty page cannot reach a kernel. + + A page table is read together with its last-page lengths, so a page holding + a previous request's KV inside a live extent produces plausible attention + output and no diagnostic. Refusing to build the table is what turns that into + an error at the one point where it is still cheap to notice. + """ + # Exactly one allocatable page, so the second request is forced to take the + # first one's. A larger pool would hand out a fresh page and never recycle. + plane = self._plane(num_pages=2, max_requests=2, max_context_len=PAGE) + first = self._admit(plane, "first", prompt_len=PAGE, max_new_tokens=0) + plane.reserve([first], [PAGE]) + plane.release(first) + + second = self._admit(plane, "second", prompt_len=PAGE, max_new_tokens=0) + self.assertTrue(plane.reserve([second], [PAGE])) + recycled = plane.pending_scrub() + self.assertEqual(recycled.size, 1, "the reused page must be reported as needing a scrub") + + with self.assertRaises(DirtyPageError): + plane.build_page_table([second], [PAGE]) + + plane.confirm_scrubbed(recycled) + self.assertEqual(plane.pending_scrub().size, 0) + plane.build_page_table([second], [PAGE]) # now permitted + + def test_a_partial_confirmation_still_blocks(self): + """Scrubbing some of what was recycled must not clear the whole obligation.""" + plane = self._plane(num_pages=5, max_requests=2, max_context_len=64) + first = self._admit(plane, "first", prompt_len=3 * PAGE, max_new_tokens=0) + plane.reserve([first], [3 * PAGE]) + plane.release(first) + + second = self._admit(plane, "second", prompt_len=3 * PAGE, max_new_tokens=0) + plane.reserve([second], [3 * PAGE]) + recycled = plane.pending_scrub() + self.assertEqual(recycled.size, 3) + + plane.confirm_scrubbed(recycled[:2]) + self.assertEqual(plane.pending_scrub().size, 1) + with self.assertRaises(DirtyPageError): + plane.build_page_table([second], [3 * PAGE]) + + def test_a_dirty_page_trimmed_as_unreadable_does_not_block(self): + """The check covers the readable extent, not every page a request holds.""" + plane = self._plane(num_pages=8, max_requests=2, max_context_len=64) + donor = self._admit(plane, "donor", prompt_len=PAGE, max_new_tokens=0) + plane.reserve([donor], [PAGE]) + dirtied = int(plane.page_map.pages(donor)[0]) + plane.release(donor) + self.assertTrue(plane.allocator.is_dirty(dirtied)) + + reader = self._admit(plane, "reader", prompt_len=PAGE, max_new_tokens=PAGE) + plane.reserve([reader], [PAGE]) + plane.confirm_scrubbed(plane.pending_scrub()) + + # Held, but past the recorded length, so the table trims it away and no + # kernel can reach it this step. + plane.page_map.append_pages(reader, [dirtied]) + table = plane.build_page_table([reader], [PAGE]) + self.assertNotIn(dirtied, table.flat_page_indices().tolist()) + + def test_the_table_a_step_produces_is_internally_consistent(self): + plane = self._plane() + handles = [self._admit(plane, f"r{i}", prompt_len=17 + i, max_new_tokens=4) for i in range(3)] + self.assertTrue(plane.reserve(handles, [17, 18, 19])) + table = plane.build_page_table(handles, [17, 18, 19]) + + table.validate(PAGE) + self.assertEqual(table.num_requests, 3) + self.assertEqual(table.num_tokens, 17 + 18 + 19) + slots = table.slot_mapping(PAGE) + self.assertEqual(np.unique(slots).size, slots.size, "two tokens must never share a pool slot") + self.assertTrue(np.all(slots >= PAGE), "no token may land on the reserved padding page") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/kv_execution_test.py b/tests/unit/kv_execution_test.py new file mode 100644 index 0000000000..1847490a07 --- /dev/null +++ b/tests/unit/kv_execution_test.py @@ -0,0 +1,526 @@ +"""Tests for the paged KV execution layer: bucketing, step views, and the driver. + +These need jax, but only on CPU: the step function is injected, so the scheduling +loop is exercised without a model or a kernel. The real-kernel checks -- the +poisoned-page acceptance gate and the compile-count bound -- live in +`kv_paged_runtime_test.py` and are `gpu_only`. + +Same marking trap as the rest of this work: `tests/conftest.py` auto-marks +unmarked tests `cpu_only` and skips `cpu_only` on any accelerator testbed, so on a +GPU box this file skips rather than passes. Run it with `JAX_PLATFORMS=cpu` and +read the count. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest + +import numpy as np + +from maxtext.inference.kv_common import KvPageTableV1, KvStorageLayoutV1 +from maxtext.inference.kv_control import NativeKvControlPlane, RequestHandle +from maxtext.inference.kv_execution import allocate_pool, PagedDriver, PagedRequest +from maxtext.inference.kv_execution.bucketing import ( + StepShapePlanner, + batch_ladder, + bucket_up, + seqlen_ladder, + token_ladder, +) +from maxtext.inference.kv_execution.engine_adapter import PagedRuntime +from maxtext.inference.kv_execution.layout_builder import build_storage_layout, kv_head_shards +from maxtext.inference.kv_execution.pool_ops import POISON_SENTINEL, poison_pages, scrub_pages +from maxtext.inference.kv_execution.step_view import build_step_view + +PAGE = 16 + + +def _layout(**kw) -> KvStorageLayoutV1: + """A small pool geometry, overridable field by field.""" + base = { + "tokens_per_page": PAGE, + "num_pages": 65, + "num_layers": 2, + "num_kv_heads": 4, + "head_dim": 32, + "dtype": "float32", + } + base.update(kw) + return KvStorageLayoutV1(**base) + + +class _FakeMesh: + def __init__(self, shape): + self.shape = shape + + +class _FakeConfig: + """The handful of attributes `build_storage_layout` reads.""" + + def __init__(self, **kw): + self.num_kv_heads = 4 + self.head_dim = 128 + self.num_decoder_layers = 8 + self.paged_page_size = PAGE + self.paged_num_blocks = 256 + self.dtype = "bfloat16" + for key, value in kw.items(): + setattr(self, key, value) + + +class LadderTest(unittest.TestCase): + """The power-of-two rungs that decide which shapes can exist.""" + + def test_bucket_up_finds_the_smallest_sufficient_rung(self): + ladder = (1, 2, 4, 8) + self.assertEqual(bucket_up(1, ladder), 1) + self.assertEqual(bucket_up(3, ladder), 4) + self.assertEqual(bucket_up(8, ladder), 8) + + def test_exceeding_the_ladder_raises_rather_than_clamping(self): + """Clamping would silently drop the tokens that did not fit.""" + with self.assertRaisesRegex(ValueError, "mis-sized"): + bucket_up(9, (1, 2, 4, 8)) + + def test_batch_ladder_starts_at_one(self): + self.assertEqual(batch_ladder(8), (1, 2, 4, 8)) + self.assertEqual(batch_ladder(5), (1, 2, 4, 8)) + + def test_token_ladder_starts_at_sixty_four(self): + """Matching MaxText's existing prefill bucketing; finer rungs buy nothing.""" + self.assertEqual(token_ladder(256), (64, 128, 256)) + self.assertEqual(token_ladder(10), (64,)) + + def test_seqlen_ladder_starts_at_the_page_size(self): + """A shorter context still occupies a whole page, so a finer rung cannot occur.""" + self.assertEqual(seqlen_ladder(16, 128), (16, 32, 64, 128)) + + +class StepShapeTest(unittest.TestCase): + """Which shapes the two phases produce, and how many there can be.""" + + def _planner(self, max_batch=8, max_context_len=128, pool_pages=65, max_batched_tokens=None): + return StepShapePlanner( + tokens_per_page=PAGE, + max_batch=max_batch, + max_context_len=max_context_len, + pool_pages=pool_pages, + max_batched_tokens=max_batched_tokens, + ) + + def test_decode_ties_the_token_count_to_the_batch_bucket(self): + """One token per request, so the token axis is not independently free.""" + shape = self._planner().decode_shape(num_requests=3, max_seq_len=40) + self.assertEqual(shape.num_requests, 4) + self.assertEqual(shape.num_tokens, 4) + self.assertTrue(shape.is_decode) + + def test_extend_pins_the_batch_and_varies_only_tokens(self): + planner = self._planner(max_batch=8, max_batched_tokens=512) + first = planner.extend_shape(num_tokens=100, max_seq_len=100) + second = planner.extend_shape(num_tokens=200, max_seq_len=100) + self.assertEqual(first.num_requests, 8) + self.assertEqual(second.num_requests, 8) + self.assertEqual(first.num_tokens, 128) + self.assertEqual(second.num_tokens, 256) + + def test_the_token_budget_is_a_batch_budget_not_a_request_length(self): + """An extend step batches requests, so its total exceeds the longest one. + + Sizing the token ladder from max_context_len made a perfectly legal batch -- + four 100-token prompts against a 128-token cap -- impossible to bucket. + """ + planner = self._planner(max_context_len=128, max_batched_tokens=1024) + self.assertEqual(planner.extend_shape(num_tokens=400, max_seq_len=128).num_tokens, 512) + + def test_a_batch_past_the_token_budget_is_refused(self): + planner = self._planner(max_context_len=128, max_batched_tokens=128) + with self.assertRaisesRegex(ValueError, "mis-sized"): + planner.extend_shape(num_tokens=400, max_seq_len=128) + + def test_the_gather_table_is_an_upper_bound_not_a_guess(self): + """Derived from batch and length, so it cannot under-size the page list.""" + planner = self._planner(max_batch=4, max_context_len=64, pool_pages=1024) + shape = planner.decode_shape(num_requests=4, max_seq_len=64) + self.assertEqual(shape.num_pages, 4 * 4) + + def test_the_gather_table_is_clamped_to_the_pool(self): + planner = self._planner(max_batch=256, max_context_len=4096, pool_pages=40) + self.assertEqual(planner.decode_shape(256, 4096).num_pages, 40) + + def test_a_bucketed_shape_is_reached_by_a_range_of_batches(self): + """The property that bounds compilation: many inputs, one shape.""" + planner = self._planner() + shapes = {planner.decode_shape(n, 40) for n in (5, 6, 7, 8)} + self.assertEqual(len(shapes), 1) + + def test_the_shape_count_has_a_stated_bound(self): + planner = self._planner(max_batch=8, max_context_len=128) + # 4 batch rungs + 2 token rungs, each against 4 seqlen rungs + self.assertEqual(planner.max_distinct_shapes(), (4 + 2) * 4) + + +class StepViewTest(unittest.TestCase): + """Padding a page table into a static shape, inertly.""" + + def _table(self): + return KvPageTableV1( + page_ids=[[1, 2, 3], [4, 5]], + seq_lens=np.array([33, 20], dtype=np.int32), + query_lens=np.array([1, 1], dtype=np.int32), + write_positions=np.array([32, 19], dtype=np.int32), + request_order=np.array([0, 1], dtype=np.int32), + ) + + def _view(self, num_requests=4, num_tokens=4, num_pages=8, max_seqlen_k=64): + """The two-request decode table above, padded to a given bucketed shape.""" + from maxtext.inference.kv_execution.bucketing import StepShape # pylint: disable=import-outside-toplevel + + shape = StepShape( + num_requests=num_requests, + num_tokens=num_tokens, + num_pages=num_pages, + max_seqlen_k=max_seqlen_k, + is_decode=True, + ) + return build_step_view(self._table(), shape, tokens_per_page=PAGE) + + def test_the_arrays_have_exactly_the_bucketed_shapes(self): + view = self._view() + self.assertEqual(view.slot_mapping.shape, (4,)) + self.assertEqual(view.kv_indptr.shape, (5,)) + self.assertEqual(view.kv_page_indices.shape, (8,)) + self.assertEqual(view.kv_last_page_lens.shape, (4,)) + self.assertEqual(view.cu_seqlens_q.shape, (5,)) + self.assertEqual(view.num_active_requests, 2) + self.assertEqual(view.num_active_tokens, 2) + + def test_padded_writes_go_to_the_skip_sentinel(self): + """Not to a real slot, which would write garbage into a live page.""" + slots = np.asarray(self._view().slot_mapping) + np.testing.assert_array_equal(slots[:2], [48, 83]) + np.testing.assert_array_equal(slots[2:], [-1, -1]) + + def test_padded_requests_get_zero_length_page_ranges(self): + """A repeated final indptr entry, so a kernel does no work for them.""" + view = self._view() + np.testing.assert_array_equal(np.asarray(view.kv_indptr), [0, 3, 5, 5, 5]) + np.testing.assert_array_equal(np.asarray(view.cu_seqlens_q), [0, 1, 2, 2, 2]) + + def test_padded_gather_entries_point_at_the_reserved_page(self): + """Belt and braces behind the flat indptr: the reserved page reads as zeros.""" + view = self._view() + indices = np.asarray(view.kv_page_indices) + np.testing.assert_array_equal(indices[:5], [1, 2, 3, 4, 5]) + np.testing.assert_array_equal(indices[5:], [0, 0, 0]) + + def test_padded_lengths_are_zero(self): + view = self._view() + np.testing.assert_array_equal(np.asarray(view.kv_last_page_lens), [1, 4, 0, 0]) + np.testing.assert_array_equal(np.asarray(view.seq_lens), [33, 20, 0, 0]) + + def test_a_bucket_too_small_for_the_batch_is_rejected(self): + with self.assertRaisesRegex(ValueError, "kv_page_indices"): + self._view(num_pages=2) + + def test_a_sequence_longer_than_the_configured_bucket_is_rejected(self): + """Silently truncating max_seqlen_k would under-configure the kernel.""" + with self.assertRaisesRegex(ValueError, "ladder is mis-sized"): + self._view(max_seqlen_k=32) + + def test_it_converts_to_the_attention_path_plan(self): + plan = self._view().to_paged_plan() + self.assertTrue(plan.is_decode) + self.assertEqual(plan.max_seqlen_q, 1) + self.assertEqual(plan.max_seqlen_k, 64) + + +class PoolTest(unittest.TestCase): + """Allocation and the two hygiene operations.""" + + def test_the_pool_starts_zeroed(self): + """Load-bearing: the reserved page and every unscrubbed page rely on it.""" + pool = allocate_pool(_layout(num_pages=8, num_layers=3)) + self.assertEqual(pool.num_layers, 3) + self.assertEqual(pool.page_shape, (8, PAGE, 4, 32)) + for layer in range(3): + self.assertTrue(bool((np.asarray(pool.k_pages[layer]) == 0).all())) + self.assertTrue(bool((np.asarray(pool.v_pages[layer]) == 0).all())) + + def test_scrubbing_zeroes_only_the_named_pages(self): + pool = allocate_pool(_layout(num_pages=8, num_layers=1)) + k, v = poison_pages(pool.k_pages[0], pool.v_pages[0], [1, 2, 3]) + k, v = scrub_pages(k, v, [2]) + k, v = np.asarray(k), np.asarray(v) + self.assertTrue(bool((k[2] == 0).all())) + self.assertTrue(bool((k[1] == POISON_SENTINEL).all())) + self.assertTrue(bool((v[3] == POISON_SENTINEL).all())) + + def test_scrubbing_nothing_is_a_no_op(self): + pool = allocate_pool(_layout(num_pages=8, num_layers=1)) + k, v = scrub_pages(pool.k_pages[0], pool.v_pages[0], []) + self.assertTrue(bool((np.asarray(k) == 0).all())) + self.assertTrue(bool((np.asarray(v) == 0).all())) + + def test_an_odd_page_count_is_padded_idempotently(self): + """Padding repeats a page already being written, so no mask is needed.""" + pool = allocate_pool(_layout(num_pages=8, num_layers=1)) + k, v = poison_pages(pool.k_pages[0], pool.v_pages[0], [1, 2, 3, 4, 5]) + k = np.asarray(k) + for page in (1, 2, 3, 4, 5): + self.assertTrue(bool((k[page] == POISON_SENTINEL).all()), f"page {page} not filled") + for page in (0, 6, 7): + self.assertTrue(bool((k[page] == 0).all()), f"page {page} was filled and should not have been") + self.assertTrue(bool((np.asarray(v)[5] == POISON_SENTINEL).all())) + + +class LayoutBuilderTest(unittest.TestCase): + """Config and mesh into pool geometry.""" + + def test_it_reads_the_model_dimensions(self): + layout = build_storage_layout(_FakeConfig()) + self.assertEqual(layout.tokens_per_page, PAGE) + self.assertEqual(layout.num_kv_heads, 4) + self.assertEqual(layout.head_dim, 128) + self.assertEqual(layout.num_layers, 8) + + def test_the_reserved_page_is_added_rather_than_taken_out_of_capacity(self): + """paged_num_blocks is what the operator asked to be usable.""" + layout = build_storage_layout(_FakeConfig(paged_num_blocks=256)) + self.assertEqual(layout.num_pages, 257) + self.assertEqual(layout.max_tokens(), 256 * PAGE) + + def test_shards_come_from_the_mesh_not_from_a_config_field(self): + """Two sources of truth for TP width is how a pool ends up a factor too small.""" + mesh = _FakeMesh({"data": 2, "tensor": 4}) + self.assertEqual(kv_head_shards(mesh), 4) + self.assertEqual(build_storage_layout(_FakeConfig(), mesh).kv_head_shards, 4) + + def test_multiple_tensor_axes_multiply(self): + self.assertEqual(kv_head_shards(_FakeMesh({"tensor": 2, "tensor_sequence": 2})), 4) + + def test_no_mesh_means_unsharded(self): + self.assertEqual(kv_head_shards(None), 1) + + def test_an_unset_pool_size_is_refused(self): + with self.assertRaisesRegex(ValueError, "paged_num_blocks"): + build_storage_layout(_FakeConfig(paged_num_blocks=0)) + + +class _CountingStep: + """A step function that records what it was handed.""" + + def __init__(self, token=7): + self.token = token + self.views = [] + + def __call__(self, view, pool): + del pool + self.views.append(view) + return np.full((view.num_active_requests,), self.token, dtype=np.int32) + + +class PagedDriverTest(unittest.TestCase): + """Admission, scheduling, scrubbing and release, without a model.""" + + def _driver(self, *, num_pages=65, max_requests=4, max_context_len=128, step=None, **kw): + layout = _layout(num_pages=num_pages) + plane = NativeKvControlPlane( + layout=layout, max_requests=max_requests, max_context_len=max_context_len, debug_mode=True + ) + pool = allocate_pool(layout) + return PagedDriver(plane, pool, step or _CountingStep(), max_batch=max_requests, **kw) + + def _requests(self, count, prompt_len=20, max_new_tokens=8): + return [ + PagedRequest(request_id=f"r{i}", prompt_len=prompt_len, max_new_tokens=max_new_tokens) + for i in range(count) + ] + + def test_a_single_request_runs_to_its_length_cap(self): + driver = self._driver() + driver.submit(self._requests(1, prompt_len=20, max_new_tokens=5)) + done = driver.run() + self.assertEqual(len(done), 1) + self.assertEqual(len(done[0].generated), 5) + self.assertEqual(done[0].finish_reason, "length") + + def test_an_eos_token_stops_generation_early(self): + driver = self._driver(step=_CountingStep(token=99), eos_ids=(99,)) + driver.submit(self._requests(1, max_new_tokens=50)) + done = driver.run() + self.assertEqual(len(done[0].generated), 1) + self.assertEqual(done[0].finish_reason, "stop") + + def test_prefill_is_preferred_over_decode(self): + """A deliberate policy, so it is worth pinning.""" + step = _CountingStep() + driver = self._driver(step=step) + driver.submit(self._requests(2)) + driver.step() + self.assertFalse(step.views[0].shape.is_decode) + + def test_a_mixed_length_batch_completes_without_leaking(self): + driver = self._driver(max_requests=4) + lengths = [7, 19, 33, 48, 64, 5] + driver.submit( + [PagedRequest(request_id=f"r{i}", prompt_len=n, max_new_tokens=6) for i, n in enumerate(lengths)] + ) + done = driver.run() + + self.assertEqual(len(done), len(lengths)) + self.assertTrue(all(len(r.generated) == 6 for r in done)) + self.assertEqual(driver.plane.allocator.num_allocated_pages, 0) + self.assertEqual(driver.plane.allocator.available_pages, driver.plane.allocator.capacity_pages) + self.assertEqual(driver.plane.num_live, 0) + + def test_churn_traces_a_bounded_set_of_shapes(self): + """The exit criterion, host side. Without bucketing this grows with the trace.""" + driver = self._driver(max_requests=8, max_context_len=256, num_pages=257) + rng = np.random.default_rng(0) + driver.submit( + [ + PagedRequest( + request_id=f"r{i}", + prompt_len=int(rng.integers(4, 120)), + max_new_tokens=int(rng.integers(1, 30)), + ) + for i in range(60) + ] + ) + driver.run() + self.assertLessEqual(driver.num_distinct_shapes, driver.planner.max_distinct_shapes()) + self.assertLessEqual(driver.num_distinct_shapes, 20, "many more shapes than rungs means bucketing broke") + self.assertEqual(driver.plane.allocator.available_pages, driver.plane.allocator.capacity_pages) + + def test_a_pool_under_pressure_preempts_rather_than_deadlocking(self): + """Recompute preemption: work is lost, but the loop always progresses.""" + # 8 usable pages of 16 tokens, four concurrent requests each wanting 55. + driver = self._driver(num_pages=9, max_requests=4, max_context_len=64) + driver.submit( + [PagedRequest(request_id=f"r{i}", prompt_len=30, max_new_tokens=25) for i in range(5)] + ) + done = driver.run() + + self.assertEqual(len(done), 5) + self.assertTrue(all(len(r.generated) == 25 for r in done)) + self.assertGreater(sum(r.preemptions for r in done), 0, "this pool is too small not to preempt") + self.assertEqual(driver.plane.allocator.num_allocated_pages, 0) + + def test_a_preempted_request_keeps_the_tokens_it_generated(self): + """Preemption replays a longer prompt; it does not discard output.""" + driver = self._driver(num_pages=9, max_requests=4, max_context_len=64) + driver.submit( + [PagedRequest(request_id=f"r{i}", prompt_len=30, max_new_tokens=25) for i in range(5)] + ) + done = driver.run() + preempted = [r for r in done if r.preemptions] + self.assertTrue(preempted) + for request in preempted: + self.assertEqual(len(request.generated), 25) + + def test_recycled_pages_are_scrubbed_before_the_step_reads_them(self): + """The driver's half of the acceptance gate. + + A pool small enough to wrap forces recycling. If the driver failed to scrub, + the control plane would refuse to build the table and this would raise. + """ + driver = self._driver(num_pages=6, max_requests=2, max_context_len=64) + driver.submit(self._requests(8, prompt_len=20, max_new_tokens=3)) + done = driver.run() + self.assertEqual(len(done), 8) + self.assertEqual(driver.plane.pending_scrub().size, 0) + + def test_poisoning_leaves_pages_dirty_so_they_must_still_be_scrubbed(self): + """Poison is a detector, not a substitute for zeroing.""" + driver = self._driver(num_pages=6, max_requests=2, max_context_len=64, poison_on_free=True) + driver.submit(self._requests(2, prompt_len=20, max_new_tokens=2)) + driver.run() + self.assertGreater(driver.plane.allocator.num_dirty_pages, 0) + + def test_run_reports_a_stalled_loop_rather_than_returning_partial_output(self): + driver = self._driver() + driver.submit(self._requests(4, max_new_tokens=100)) + with self.assertRaisesRegex(RuntimeError, "no progress"): + driver.run(max_steps=3) + + def test_an_empty_driver_does_nothing(self): + self.assertIsNone(self._driver().step()) + self.assertEqual(self._driver().run(), []) + + +class PagedRuntimeTest(unittest.TestCase): + """The engine adapter, and the slot shim over the request-based API.""" + + def _runtime(self): + layout = _layout() + plane = NativeKvControlPlane(layout=layout, max_requests=4, max_context_len=128, debug_mode=True) + return PagedRuntime(plane, allocate_pool(layout)), plane + + def _admit(self, plane, request_id="r0", prompt_len=20): + from maxtext.inference.kv_control import RequestDescriptor # pylint: disable=import-outside-toplevel + + handle = plane.admit(RequestDescriptor(request_id=request_id, prompt_len=prompt_len, max_new_tokens=4)) + plane.reserve([handle], [prompt_len]) + return handle + + def test_release_by_handle_reclaims_the_pages(self): + runtime, plane = self._runtime() + handle = self._admit(plane, prompt_len=33) + runtime.track(handle) + self.assertEqual(runtime.release(handle).size, 3) + self.assertEqual(plane.allocator.num_allocated_pages, 0) + + def test_the_slot_shim_reaches_the_same_pages(self): + """What keeps the three legacy release_pages call sites working.""" + runtime, plane = self._runtime() + handle = self._admit(plane, prompt_len=33) + runtime.track(handle, slot=5) + self.assertEqual(runtime.release_slot(5).size, 3) + self.assertEqual(plane.allocator.num_allocated_pages, 0) + + def test_an_unknown_slot_is_a_no_op(self): + """The legacy call sites fire on termination and do not coordinate.""" + runtime, _ = self._runtime() + self.assertEqual(runtime.release_slot(11).size, 0) + + def test_a_second_release_through_the_adapter_is_absorbed(self): + runtime, plane = self._runtime() + handle = self._admit(plane) + runtime.track(handle, slot=0) + runtime.release(handle) + self.assertEqual(runtime.release(handle).size, 0) + self.assertEqual(runtime.release_slot(0).size, 0) + + def test_a_handle_from_another_epoch_is_not_honoured(self): + runtime, plane = self._runtime() + handle = self._admit(plane) + runtime.track(handle) + stale = RequestHandle(request_id=handle.request_id, row=handle.row, epoch=handle.epoch + 1) + self.assertEqual(runtime.release(stale).size, 0) + + def test_handles_are_findable_by_request_id(self): + runtime, plane = self._runtime() + handle = self._admit(plane, request_id="abc") + runtime.track(handle, slot=2) + self.assertEqual(runtime.handle_for_request("abc"), handle) + self.assertEqual(runtime.handle_for_slot(2), handle) + runtime.release(handle) + self.assertIsNone(runtime.handle_for_request("abc")) + self.assertIsNone(runtime.handle_for_slot(2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/kv_import_rule_test.py b/tests/unit/kv_import_rule_test.py new file mode 100644 index 0000000000..9bd65e87d8 --- /dev/null +++ b/tests/unit/kv_import_rule_test.py @@ -0,0 +1,204 @@ +"""The paged KV runtime's layering rule, enforced rather than documented. + +Two packages, two allowances: + + * ``kv_common`` -- the neutral vocabulary -- may import the standard library + and ``numpy``, and nothing else. + * ``kv_control`` -- the semantic control plane -- may also import + ``kv_common``. + +Neither may import ``jax``, ``jax_aiter``, ``maxtext.layers`` or +``maxtext.models``. Vendor kernel ABIs are reached only from ``kv_execution``, +one layer up. + +Both properties this buys are load-bearing, which is why the rule is checked and +not merely written down. The layers stay testable with no accelerator present, +so their logic is covered by ordinary CI rather than by a GPU job. And if a +second consumer ever wants the control plane, extracting it is a directory move +with no vendor dependency to untangle first. + +Checked statically over the AST, so a violation is reported without importing +anything, and one subprocess check confirms the static rule matches what the +interpreter actually loads. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import ast +import dataclasses +import pathlib +import subprocess +import sys +import unittest + +_INFERENCE_DIR = pathlib.Path(__file__).resolve().parents[2] / "src" / "maxtext" / "inference" + +# A package importing any of these has broken the rule outright, and saying so +# by name gives a better failure than "unexpected import". +FORBIDDEN_ROOTS = frozenset({"jax", "jaxlib", "flax", "torch", "jax_aiter", "tensorflow"}) + +# Everything either layer legitimately needs. Extend only for something that is +# genuinely stdlib or numpy. +ALLOWED_ROOTS = frozenset( + { + "__future__", + "collections", + "dataclasses", + "enum", + "hashlib", + "math", + "numpy", + "typing", + } +) + + +@dataclasses.dataclass(frozen=True) +class _Layer: + """One package and the dotted prefixes under `maxtext` it may reach.""" + + package: str + allowed_maxtext: frozenset[str] + + @property + def directory(self) -> pathlib.Path: + return _INFERENCE_DIR / self.package.rsplit(".", 1)[1] + + def permits_maxtext(self, module: str) -> bool: + return any(module == prefix or module.startswith(prefix + ".") for prefix in self.allowed_maxtext) + + +_KV_COMMON = "maxtext.inference.kv_common" +_KV_CONTROL = "maxtext.inference.kv_control" +# Deliberately not a checked layer: `kv_execution` is where jax, MaxText config +# and the vendor backend legitimately live. What matters is that nothing below it +# imports it, which the test below asserts. +_KV_EXECUTION = "maxtext.inference.kv_execution" + +LAYERS = ( + _Layer(package=_KV_COMMON, allowed_maxtext=frozenset({_KV_COMMON})), + _Layer(package=_KV_CONTROL, allowed_maxtext=frozenset({_KV_CONTROL, _KV_COMMON})), +) + + +def _imports(path: pathlib.Path, layer: _Layer) -> set[str]: + """Root packages imported by `path`, with permitted `maxtext` imports elided. + + A `maxtext` import the layer does not permit is reported under its full dotted + name rather than as the root `maxtext`, so the failure message names the actual + violation instead of making the reader go and find it. + """ + roots: set[str] = set() + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.Import): + for alias in node.names: + if layer.permits_maxtext(alias.name): + continue + roots.add(alias.name if alias.name.startswith("maxtext") else alias.name.split(".")[0]) + elif isinstance(node, ast.ImportFrom): + if node.level or not node.module: # a relative import stays inside the package + continue + if layer.permits_maxtext(node.module): + continue + roots.add(node.module if node.module.startswith("maxtext") else node.module.split(".")[0]) + return roots + + +class ImportRuleTest(unittest.TestCase): + """Each layer imports only what its position in the stack allows.""" + + def test_layer_directories_exist(self): + """Guards against the rule passing vacuously if a package is moved.""" + for layer in LAYERS: + self.assertTrue(layer.directory.is_dir(), f"{layer.package} not found at {layer.directory}") + self.assertTrue(sorted(layer.directory.glob("*.py")), f"{layer.package} has no modules to check") + + def test_no_forbidden_imports(self): + for layer in LAYERS: + for path in sorted(layer.directory.glob("*.py")): + forbidden = _imports(path, layer) & FORBIDDEN_ROOTS + self.assertEqual( + forbidden, + set(), + f"{layer.package}/{path.name} imports {sorted(forbidden)}, which the import rule forbids", + ) + + def test_no_unexpected_imports(self): + for layer in LAYERS: + for path in sorted(layer.directory.glob("*.py")): + unexpected = _imports(path, layer) - ALLOWED_ROOTS + self.assertEqual( + unexpected, + set(), + f"{layer.package}/{path.name} imports {sorted(unexpected)}; extend ALLOWED_ROOTS only if " + f"the addition is genuinely stdlib or numpy, and never to admit a maxtext module", + ) + + def test_kv_common_does_not_reach_up_into_kv_control(self): + """The dependency runs one way. A cycle here would defeat extraction.""" + common = LAYERS[0] + for path in sorted(common.directory.glob("*.py")): + self.assertNotIn( + _KV_CONTROL, + _imports(path, common), + f"kv_common/{path.name} imports kv_control, inverting the layering", + ) + + def test_neither_lower_layer_reaches_up_into_kv_execution(self): + """`kv_execution` is where jax and the vendor backend live. + + It exists precisely to hold the couplings the two layers below refuse, so a + single import in this direction would put jax back into the control plane's + dependency graph and undo the whole split. Named separately from the + allow-list check because this is the failure most likely to be introduced by + someone reaching for something convenient. + """ + for layer in LAYERS: + for path in sorted(layer.directory.glob("*.py")): + self.assertNotIn( + _KV_EXECUTION, + _imports(path, layer), + f"{layer.package}/{path.name} imports kv_execution, inverting the layering", + ) + + def test_package_inits_only_re_export_their_own_modules(self): + for layer in LAYERS: + tree = ast.parse((layer.directory / "__init__.py").read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module: + self.assertTrue( + node.module.startswith(layer.package), + f"{layer.package}/__init__.py imports from {node.module}, outside its own package", + ) + + def test_importing_the_control_plane_does_not_load_a_framework(self): + """What the static rule is actually for, confirmed against the interpreter. + + A fresh interpreter, because by the time this test runs pytest's own + collection has already imported jax, so `sys.modules` in-process proves + nothing. + """ + code = ( + "import sys, maxtext.inference.kv_control, maxtext.inference.kv_common;" + "print(sorted({m.split('.')[0] for m in sys.modules} & " + "{'jax', 'jaxlib', 'flax', 'torch', 'tensorflow'}))" + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + self.assertEqual(result.returncode, 0, f"importing the KV layers failed:\n{result.stderr}") + self.assertEqual(result.stdout.strip(), "[]", f"a framework was loaded: {result.stdout.strip()}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/kv_paged_runtime_test.py b/tests/unit/kv_paged_runtime_test.py new file mode 100644 index 0000000000..12a0060b70 --- /dev/null +++ b/tests/unit/kv_paged_runtime_test.py @@ -0,0 +1,397 @@ +"""Acceptance gates for the paged KV runtime, against the real kernels. + +Two properties that cannot be established on the host alone. + +**No recycled bytes are readable by a subsequent request.** A freed page keeps the +previous occupant's KV until something overwrites it, so this checks the whole +chain: the control plane refuses to describe a dirty page, the driver zeroes what +it recycled, and the page a new request receives is genuinely clean. The negative +control matters as much as the positive one -- a test that passes because the +sentinel was never written in the first place proves nothing. + +**Mixed-length churn traces a bounded set of shapes.** Counted by incrementing a +counter inside the traced function, so it measures actual retraces rather than the +driver's own opinion of how many shapes it produced. + +`gpu_only`, and additionally skipped when jax-aiter is absent or its shims are not +built, following the M3 pattern. Note that `PYTHONPATH` must include the jax-aiter +checkout and `JA_ROOT_DIR` must point at it, or the FFI shims will not be found. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest + +from absl.testing import parameterized +import numpy as np +import pytest + +import jax +import jax.numpy as jnp + +from maxtext.inference.kv_common import KvStorageLayoutV1 +from maxtext.inference.kv_control import DirtyPageError, NativeKvControlPlane, RequestDescriptor +from maxtext.inference.kv_control import metadata as kv_metadata +from maxtext.inference.kv_execution import allocate_pool, PagedDriver, PagedRequest +from maxtext.inference.kv_execution.pool_ops import POISON_SENTINEL, poison_pages, scrub_pages +from maxtext.inference.kv_execution.step_view import build_step_view + +PAGE = 16 +# The AOT pa_ragged configurations cover head_size 128, block_size 16 and +# gqa_ratio in {1, 4, 8}. Equal query and KV head counts put the ratio at 1, so +# these shapes are inside the prebuilt set. +HEAD_DIM = 128 +NUM_HEADS = 8 +DONOR_VALUE = 100.0 +READER_VALUE = 3.0 + + +def _require_kernels(): + """Skip unless jax-aiter is importable and its KV shims are built.""" + try: + from jax_aiter.ffi.registry import standalone_symbol_available # pylint: disable=import-outside-toplevel + except ImportError as exc: + raise unittest.SkipTest( + "jax-aiter is not importable; set PYTHONPATH to the jax-aiter checkout" + ) from exc + for symbol in ("AppendKvJA", "PagedAttentionJA", "PagedPrefillJA"): + if not standalone_symbol_available(symbol): + raise unittest.SkipTest( + f"{symbol} is not built; run 'make -f Makefile.kv ja_kv' in jax-aiter and set JA_ROOT_DIR" + ) + + +def _layout(num_pages, num_layers=1) -> KvStorageLayoutV1: + return KvStorageLayoutV1( + tokens_per_page=PAGE, + num_pages=num_pages, + num_layers=num_layers, + num_kv_heads=NUM_HEADS, + head_dim=HEAD_DIM, + dtype="bfloat16", + ) + + +def _append(pool, layer, slot_mapping, value): + """Write a constant K and V at every non-padded slot.""" + from jax_aiter.ops.append_kv import append_kv # pylint: disable=import-outside-toplevel + + count = int(np.asarray(slot_mapping).shape[0]) + block = jnp.full((count, NUM_HEADS, HEAD_DIM), value, jnp.bfloat16) + k, v = append_kv(block, block, jnp.asarray(slot_mapping, jnp.int32), pool.k_pages[layer], pool.v_pages[layer]) + pool.replace_layer(layer, k, v) + + +def _decode_attend(pool, layer, view): + """One decode step over the pool, returning float32 output.""" + from jax_aiter.ops.paged_attention import paged_attention # pylint: disable=import-outside-toplevel + + query = jnp.ones((view.shape.num_tokens, NUM_HEADS, HEAD_DIM), jnp.bfloat16) + out = paged_attention( + query, + pool.k_pages[layer], + pool.v_pages[layer], + view.kv_indptr, + view.kv_page_indices, + view.kv_last_page_lens, + max_seq_len=view.shape.max_seqlen_k, + scale=1.0, + ) + return np.asarray(out.astype(jnp.float32)) + + +@pytest.mark.gpu_only +class PoisonedPageTest(parameterized.TestCase): + """The acceptance gate: a recycled page must carry nothing forward.""" + + def setUp(self): + super().setUp() + _require_kernels() + # Exactly one allocatable page, so the second request is forced to take the + # first one's. Any larger pool would hand out a fresh page and the test would + # pass without ever exercising recycling. + self.layout = _layout(num_pages=2) + self.plane = NativeKvControlPlane( + layout=self.layout, max_requests=2, max_context_len=PAGE, debug_mode=True + ) + self.pool = allocate_pool(self.layout) + + def _admit(self, request_id, prompt_len): + handle = self.plane.admit( + RequestDescriptor(request_id=request_id, prompt_len=prompt_len, max_new_tokens=0) + ) + self.assertIsNotNone(handle, "admission must succeed for this test to mean anything") + self.assertTrue(self.plane.reserve([handle], [prompt_len])) + return handle + + def _shape(self, num_requests, num_tokens, num_pages, is_decode=True): + from maxtext.inference.kv_execution.bucketing import StepShape # pylint: disable=import-outside-toplevel + + return StepShape( + num_requests=num_requests, + num_tokens=num_tokens, + num_pages=num_pages, + max_seqlen_k=PAGE, + is_decode=is_decode, + ) + + def _run_donor(self): + """Fill the single page with a distinctive value, then free it poisoned.""" + donor = self._admit("donor", PAGE) + self.plane.confirm_scrubbed(self.plane.pending_scrub()) + table = self.plane.build_page_table([donor], [PAGE]) + _append(self.pool, 0, table.slot_mapping(PAGE), DONOR_VALUE) + + page = int(self.plane.page_map.pages(donor)[0]) + written = np.asarray(self.pool.k_pages[0].astype(jnp.float32))[page] + self.assertTrue(bool((written == DONOR_VALUE).all()), "the donor did not fill its page") + + k, v = poison_pages(self.pool.k_pages[0], self.pool.v_pages[0], [page]) + self.pool.replace_layer(0, k, v) + self.plane.release(donor) + return page + + def test_the_sentinel_survives_the_pool_dtype_exactly(self): + """Otherwise every comparison against it is approximate and proves little. + + bfloat16 carries eight bits of mantissa, so a round decimal sentinel is + stored rounded and `== POISON_SENTINEL` is false even on a page that was + definitely poisoned -- which reads as a passing scrub test. + """ + stored = np.asarray(jnp.full((1,), POISON_SENTINEL, jnp.bfloat16).astype(jnp.float32)) + self.assertEqual(float(stored[0]), POISON_SENTINEL) + + def test_the_donor_really_did_contaminate_the_page(self): + """Without this the other tests could pass vacuously.""" + page = self._run_donor() + written = np.asarray(self.pool.k_pages[0].astype(jnp.float32))[page] + self.assertTrue(bool((written == POISON_SENTINEL).all())) + self.assertTrue(self.plane.allocator.is_dirty(page)) + + def test_poisoning_does_not_count_as_scrubbing(self): + """A sentinel makes a missed scrub loud; it does not make the page safe.""" + self._run_donor() + self._admit("reader", 1) + self.assertEqual(self.plane.pending_scrub().size, 1) + + def test_the_control_plane_refuses_to_describe_the_dirty_page(self): + """The enforcement point, with the real pool behind it.""" + self._run_donor() + reader = self._admit("reader", 1) + with self.assertRaises(DirtyPageError): + self.plane.build_page_table([reader], [1]) + + def test_after_scrubbing_no_recycled_byte_remains_in_the_page(self): + """The property the milestone asks for, stated over the whole page. + + Checked across the entire page rather than only the readable extent. Exact + last-page lengths already keep a well-behaved kernel inside valid data; this + is what makes an over-read harmless too, and over-reads are the failure the + scrub actually exists for. + """ + page = self._run_donor() + reader = self._admit("reader", 1) + + pending = self.plane.pending_scrub() + k, v = scrub_pages(self.pool.k_pages[0], self.pool.v_pages[0], pending) + self.pool.replace_layer(0, k, v) + self.plane.confirm_scrubbed(pending) + + table = self.plane.build_page_table([reader], [1]) + self.assertEqual(int(self.plane.page_map.pages(reader)[0]), page, "the page must have been recycled") + _append(self.pool, 0, table.slot_mapping(PAGE), READER_VALUE) + + written = np.asarray(self.pool.k_pages[0].astype(jnp.float32))[page] + self.assertTrue(bool((written[0] == READER_VALUE).all()), "the reader's own token must be present") + self.assertFalse(bool((written == POISON_SENTINEL).any()), "the poison sentinel survived the scrub") + self.assertFalse(bool((written == DONOR_VALUE).any()), "the donor's KV survived the scrub") + self.assertTrue(bool((written[1:] == 0.0).all()), "everything the reader did not write must be zero") + + def test_attention_over_a_recycled_page_returns_the_readers_own_value(self): + """A single-token context makes the correct answer exact rather than close. + + Softmax over one key is 1.0 whatever the query, so the output must equal that + token's V exactly. Any contribution from the recycled page would show up as a + mixture. + """ + self._run_donor() + reader = self._admit("reader", 1) + + pending = self.plane.pending_scrub() + k, v = scrub_pages(self.pool.k_pages[0], self.pool.v_pages[0], pending) + self.pool.replace_layer(0, k, v) + self.plane.confirm_scrubbed(pending) + + table = self.plane.build_page_table([reader], [1]) + _append(self.pool, 0, table.slot_mapping(PAGE), READER_VALUE) + + view = build_step_view(table, self._shape(1, 1, 1), tokens_per_page=PAGE) + out = _decode_attend(self.pool, 0, view) + np.testing.assert_array_equal(out, np.full_like(out, READER_VALUE)) + + def test_without_the_scrub_the_contamination_is_still_there(self): + """The negative control: the assertions above have teeth. + + Bypasses the control plane's gate by building the table through the metadata + builder directly, which is the one way to reach a dirty page. The sentinel is + then still in the page, so the check that it is absent after a scrub is + testing something real. + """ + page = self._run_donor() + reader = self._admit("reader", 1) + + table = kv_metadata.build_page_table(self.plane.page_map, [reader], [1]) + _append(self.pool, 0, table.slot_mapping(PAGE), READER_VALUE) + + written = np.asarray(self.pool.k_pages[0].astype(jnp.float32))[page] + self.assertTrue(bool((written[0] == READER_VALUE).all())) + self.assertTrue( + bool((written[1:] == POISON_SENTINEL).all()), + "the unscrubbed page should still hold the previous occupant's bytes", + ) + + +class _TracingStep: + """Runs the real append-and-attend under `jit`, counting actual retraces. + + The counter increments in the traced body, which executes once per trace, so + this measures compilation rather than the driver's own shape bookkeeping. The + two agreeing is the result worth having. + """ + + def __init__(self): + self.traces = 0 + # What the shapes would have been with no bucketing, so a test can show the + # collapse rather than merely assert that a number came out small. + self.raw_shapes: set[tuple[int, int, int, bool]] = set() + + def body(query, key, value, k_pool, v_pool, slot_mapping, kv_indptr, kv_page_indices, + kv_last_page_lens, cu_seqlens_q, *, max_seqlen_q, max_seqlen_k, is_decode): + self.traces += 1 + # pylint: disable=import-outside-toplevel + from maxtext.layers.gpu_paged_attention import PagedPlan, paged_attention_step + + plan = PagedPlan( + slot_mapping=slot_mapping, + kv_indptr=kv_indptr, + kv_page_indices=kv_page_indices, + kv_last_page_lens=kv_last_page_lens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + is_decode=is_decode, + ) + out, pools = paged_attention_step( + query, key, value, k_pool, v_pool, plan, backend="aiter", scale=1.0, causal=True + ) + return out, pools[0], pools[1] + + self._jitted = jax.jit(body, static_argnames=("max_seqlen_q", "max_seqlen_k", "is_decode")) + + def __call__(self, view, pool): + seq_lens = np.asarray(view.seq_lens) + self.raw_shapes.add( + ( + view.num_active_requests, + view.num_active_tokens, + int(seq_lens.max()) if seq_lens.size else 0, + view.shape.is_decode, + ) + ) + tokens = view.shape.num_tokens + block = jnp.ones((tokens, NUM_HEADS, HEAD_DIM), jnp.bfloat16) + _, k, v = self._jitted( + block, + block, + block, + pool.k_pages[0], + pool.v_pages[0], + view.slot_mapping, + view.kv_indptr, + view.kv_page_indices, + view.kv_last_page_lens, + view.cu_seqlens_q, + max_seqlen_q=view.max_seqlen_q, + max_seqlen_k=view.shape.max_seqlen_k, + is_decode=view.shape.is_decode, + ) + pool.replace_layer(0, k, v) + return np.full((view.num_active_requests,), 7, dtype=np.int32) + + +@pytest.mark.gpu_only +class CompileCountTest(parameterized.TestCase): + """Mixed-length churn must not retrace without bound.""" + + def test_churn_traces_no_more_shapes_than_the_ladders_allow(self): + _require_kernels() + layout = _layout(num_pages=129) + plane = NativeKvControlPlane(layout=layout, max_requests=4, max_context_len=64, debug_mode=True) + step = _TracingStep() + driver = PagedDriver(plane, allocate_pool(layout), step, max_batch=4, max_batched_tokens=64) + + rng = np.random.default_rng(0) + driver.submit( + [ + PagedRequest( + request_id=f"r{i}", + prompt_len=int(rng.integers(4, 48)), + max_new_tokens=int(rng.integers(1, 12)), + ) + for i in range(24) + ] + ) + done = driver.run() + + self.assertEqual(len(done), 24) + self.assertGreater(step.traces, 0, "nothing was traced, so nothing was measured") + self.assertEqual( + step.traces, + driver.num_distinct_shapes, + "each bucketed shape must be traced exactly once; a mismatch means something " + "outside the bucketing is varying", + ) + self.assertLessEqual(step.traces, driver.planner.max_distinct_shapes()) + # The bucketing has to be doing the work, not the workload happening to be + # uniform: many raw shapes must be collapsing onto few traced ones. + self.assertGreater( + len(step.raw_shapes), + 3 * step.traces, + f"only {len(step.raw_shapes)} distinct raw shapes for {step.traces} traces, so this workload " + f"would not have retraced much anyway and proves little", + ) + # Every page returned, so the churn was real rather than a single long batch. + self.assertEqual(plane.allocator.num_allocated_pages, 0) + self.assertEqual(plane.allocator.available_pages, plane.allocator.capacity_pages) + + def test_repeating_a_shape_does_not_retrace(self): + """The property bucketing exists for, isolated from the scheduling loop.""" + _require_kernels() + layout = _layout(num_pages=65) + plane = NativeKvControlPlane(layout=layout, max_requests=2, max_context_len=64, debug_mode=True) + step = _TracingStep() + driver = PagedDriver(plane, allocate_pool(layout), step, max_batch=2, max_batched_tokens=64) + + driver.submit([PagedRequest(request_id="a", prompt_len=20, max_new_tokens=10)]) + driver.run() + after_first = step.traces + + driver.submit([PagedRequest(request_id="b", prompt_len=20, max_new_tokens=10)]) + driver.run() + self.assertEqual(step.traces, after_first, "an identical second request retraced") + + +if __name__ == "__main__": + unittest.main() From 90b8572898b46a9b12618a54954221760bd1ffde Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Wed, 26 Aug 2026 14:27:21 +0000 Subject: [PATCH 04/21] feat(inference): share KV pages between requests with a common prompt prefix (M5) Adds kv_control/prefix_index.py and kv_common/namespace.py: a page trie mapping token prefixes to the pages already holding their K/V, so a request whose prompt starts with tokens someone else already computed reads those pages instead of recomputing them. Off by default, behind paged_enable_prefix_cache. On a trace of 24 requests sharing a 512-token prefix, 52.8% of prefill tokens are not recomputed and TTFT p50 goes from 7.74ms to 4.27ms. Read the token count as the result and the latency as an indication: tokens avoided is arithmetic and scale-free, while the 1.81x is sub-linear in the 52.8% of work removed because this model is small enough for fixed per-step cost to still matter. Reimplemented from sglang-jax's radix cache as a reference design rather than ported, with three divergences that page granularity permits. A node is one page, not a variable-length token run. Node splitting disappears entirely, which is most of the reference's complexity, and costs nothing here because only whole pages are ever published; the reference needs splitting because it matches at token granularity within a page. The cache namespace is folded into the hash chain rather than compared beside it. Each node's key is a hash chained from its parent, and the chain starts at the namespace digest instead of a constant, so two configurations do not share a root and a mismatch is structurally unable to hit. The reference compares extra_key and dp_rank as a side check, and a check is a thing that can be forgotten. That namespace covers everything which changes the K/V for identical token ids: weights fingerprint and revision, tokenizer, adapter, tenant, RoPE, KV dtype and quantisation, layout, sharding, prompt embeddings, multimodal inputs. Its digest enumerates its own dataclass fields rather than listing them, and the negative test is generated the same way. A hand-written digest is a second place to remember every field, and the one occasion someone adds a field and forgets is the occasion two incompatible configurations share K/V. And recency is a monotonic counter, not time.monotonic(). A clock ties at its resolution when many nodes are touched in one step, which makes eviction order depend on how the heap broke the tie; a counter cannot tie, so the order is reproducible and a test can assert against it. A match deliberately stops one page short even when the whole prompt is cached, because a request with nothing left to run has no query token to produce a logit from. The failure this has to rule out is arithmetic rather than bookkeeping. After a hit the step runs the prompt's suffix, which sits at absolute positions cached..prompt_len, and RoPE encodes absolute position -- so running that suffix from position zero yields K/V rotated as though it began the sequence. Pages laid out correctly, page table correct, nothing leaked, output wrong, and no host-side test can see it. gpu_paged_prefix_cache_test.py asserts a warm rollout is token-identical to a cold one, which is the only form of the claim that catches it. Two page-lifetime bugs that only exist once pages outlive requests are fixed here and will need re-checking when vLLM owns allocation. Poison-on-free now follows what release actually freed rather than everything the request held, since poisoning a page the index just adopted destroys K/V about to be read as valid. And preemption no longer publishes: it exists to reclaim pages, and the cache retains what it adopts. Symmetrically, admission budgets against free plus evictable pages, because reservation evicts on shortfall -- budgeting against the free list alone stalls a loop that could still progress, which would turn an optimisation into a reason requests stop being served. 175 CPU tests and 4 GPU tests, including a generated negative test per namespace field confirming that varying one alone defeats the match. run_prefix_cache_benchmark.py measures two arms of the same engine on the same trace, warming with a discarded pass over that trace rather than a shape sweep. That is not only simpler: warmup_paged currently fails in this container with an aiter allocation error which reproduces on the pre-M5 default configuration, so it is unrelated to prefix caching but does need diagnosing separately. Co-authored-by: Cursor --- .../paged_kv/run_prefix_cache_benchmark.py | 269 +++++++++++ benchmarks/paged_kv/run_serving_benchmark.py | 48 +- src/maxtext/configs/types.py | 11 + src/maxtext/inference/kv_common/__init__.py | 3 + src/maxtext/inference/kv_common/namespace.py | 118 +++++ src/maxtext/inference/kv_control/__init__.py | 19 +- .../inference/kv_control/control_plane.py | 191 +++++++- .../inference/kv_control/prefix_index.py | 397 +++++++++++++++ .../inference/kv_execution/benchmark.py | 104 +++- src/maxtext/inference/kv_execution/driver.py | 100 +++- .../inference/kv_execution/engine_adapter.py | 41 +- src/maxtext/inference/maxengine/maxengine.py | 39 +- tests/unit/gpu_paged_prefix_cache_test.py | 246 ++++++++++ tests/unit/kv_execution_test.py | 163 ++++++- tests/unit/kv_import_rule_test.py | 1 + tests/unit/kv_prefix_cache_test.py | 454 ++++++++++++++++++ 16 files changed, 2155 insertions(+), 49 deletions(-) create mode 100644 benchmarks/paged_kv/run_prefix_cache_benchmark.py create mode 100644 src/maxtext/inference/kv_common/namespace.py create mode 100644 src/maxtext/inference/kv_control/prefix_index.py create mode 100644 tests/unit/gpu_paged_prefix_cache_test.py create mode 100644 tests/unit/kv_prefix_cache_test.py diff --git a/benchmarks/paged_kv/run_prefix_cache_benchmark.py b/benchmarks/paged_kv/run_prefix_cache_benchmark.py new file mode 100644 index 0000000000..c60a146f22 --- /dev/null +++ b/benchmarks/paged_kv/run_prefix_cache_benchmark.py @@ -0,0 +1,269 @@ +"""M5's measurement: what prefix sharing costs and what it saves. + +Two arms of the *same* paged engine on the *same* trace, differing only in +whether the prefix cache is on. That isolation is the point: a paged-versus-dense +comparison mixes in every other difference between the two paths, whereas this +changes one flag. + +Two things are reported, and the distinction matters. + +**Prefill tokens avoided** is exact, and it is arithmetic rather than a +measurement -- prompt tokens minus tokens actually run. It does not depend on +model scale, kernel quality or what else the machine is doing, so it is the +number to quote when the question is whether sharing works. + +**Time to first token** is the number anyone actually cares about, and it is only +meaningful once prefill compute dominates a step. At toy width the run is launch +bound, the avoided compute is a rounding error against fixed per-step overhead, +and the TTFT delta will understate the saving badly. Scale the model with +`--layers` and `--emb-dim` before believing it. This is the same trap the M4.5 +throughput comparison fell into. + +Warmup is a full discarded pass over the trace rather than a shape sweep. That +compiles exactly the shapes the measured pass will present -- no more, and +crucially no fewer -- and it sidesteps `warmup_paged`, whose enumeration of the +shape space currently trips an aiter failure in this container that has nothing +to do with paging. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import time + +import numpy as np + +import jax +import jax.numpy as jnp +from flax import nnx +from flax.linen import partitioning as nn_partitioning + +from maxtext.common.common_types import MODEL_MODE_PREFILL +from maxtext.configs import pyconfig +from maxtext.inference.kv_common import CacheNamespace +from maxtext.inference.kv_execution import benchmark +from maxtext.utils import maxtext_utils, model_creation_utils + +NAMESPACE = CacheNamespace(model_fingerprint="prefix-bench", tokenizer="synthetic") + +BASE_CONFIG = { + "base_num_query_heads": 8, + "base_num_kv_heads": 8, + "head_dim": 128, + "vocab_size": 32000, + "per_device_batch_size": 1.0, + "scan_layers": False, + "sparse_matmul": False, + "dtype": "bfloat16", + "weight_dtype": "bfloat16", + "decode_sampling_strategy": "greedy", + "enable_checkpointing": False, + "skip_jax_distributed_system": True, + "pure_nnx": True, + "attention": "gpu_paged", +} + + +def build_config(args, prefix_cache: bool): + overrides = dict(BASE_CONFIG) + overrides["base_num_decoder_layers"] = args.layers + overrides["base_emb_dim"] = args.emb_dim + overrides["base_mlp_dim"] = args.emb_dim * 4 + overrides["max_target_length"] = args.max_context + overrides["max_prefill_predict_length"] = args.max_prompt + overrides["paged_page_size"] = args.page_size + overrides["paged_num_blocks"] = args.pool_tokens // args.page_size + overrides["paged_max_context_len"] = args.max_context + overrides["paged_enable_prefix_cache"] = prefix_cache + overrides["run_name"] = f"prefix_bench_{'on' if prefix_cache else 'off'}" + return pyconfig.initialize([sys.argv[0], args.config_path], **overrides) + + +def build_params(cfg, devices): + mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(config=cfg, devices=devices), cfg.mesh_axes) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh: + model = model_creation_utils.create_model( + cfg, mesh, model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(params=0, dropout=0) + ) + _, params_state, _ = nnx.split(model, nnx.Param, ...) + return params_state + + +def serve(engine, params, requests, *, max_batch: int, measure: bool): + """One pass over the trace. Returns per-request TTFT and prefill token counts. + + Requests are admitted greedily and every live request is advanced one token per + step, which is the same policy the driver and the A/B harness use. + """ + runtime = engine.paged_runtime + waiting, live = list(requests), [] + ttfts, prefill_tokens = [], [] + + while waiting or live: + while waiting and len(live) < max_batch: + candidate = waiting[0] + tokens = candidate.token_ids + started = time.perf_counter() + handle, first = engine.prefill_paged( + params=params, + padded_tokens=jnp.asarray(tokens, jnp.int32), + true_length=int(tokens.size), + request_id=candidate.request_id, + max_new_tokens=candidate.max_new_tokens, + prompt_token_ids=tokens, + namespace=NAMESPACE, + ) + if handle is None: + break + token = int(first.data[0, 0]) # blocks, so the timing is of real work + if measure: + ttfts.append(time.perf_counter() - started) + prefill_tokens.append(int(tokens.size) - runtime.cached_tokens(handle)) + waiting.pop(0) + candidate.handle = handle + candidate.generated.append(token) + live.append(candidate) + + if not live: + raise RuntimeError("nothing is live and nothing could be admitted; the pool is too small") + + result, ok = engine.generate_paged( + params, [r.handle for r in live], next_tokens=jnp.asarray([r.generated[-1] for r in live], jnp.int32) + ) + if not ok: + victim = live.pop() + engine.release(victim.handle) + victim.handle, victim.generated = None, [] + waiting.insert(0, victim) + continue + + tokens_out = np.asarray(result.data[:, 0]).reshape(-1) + for row, request in enumerate(live): + request.generated.append(int(tokens_out[row])) + for request in [r for r in live if len(r.generated) >= r.max_new_tokens]: + engine.release(request.handle, request.context_tokens()) + request.handle = None + live.remove(request) + + return ttfts, prefill_tokens + + +def run_arm(args, devices, prefix_cache: bool): + # pylint: disable=import-outside-toplevel + from maxtext.inference.maxengine import maxengine + + cfg = build_config(args, prefix_cache) + engine = maxengine.MaxEngine(cfg, devices) + params = engine.load_params(params=build_params(cfg, devices)) + engine.init_paged_runtime(max_requests=args.max_batch, max_batched_tokens=args.max_prompt) + + def trace(): + return benchmark.shared_prefix_trace( + args.requests, + args.shared_prefix, + args.unique, + args.output, + seed=args.seed, + num_variants=args.variants, + ) + + # Discarded, and it compiles every shape the measured pass will present. + serve(engine, params, trace(), max_batch=args.max_batch, measure=False) + # A cold cache for the measured pass, so the reported saving comes from + # requests sharing with each other rather than with the warmup. + plane = engine.paged_runtime.control_plane + plane.evict_cached(plane.prefix_index.num_cached_pages) + + start = time.perf_counter() + ttfts, prefill_tokens = serve(engine, params, trace(), max_batch=args.max_batch, measure=True) + duration = time.perf_counter() - start + + prompted = sum(int(r.token_ids.size) for r in trace()) + return { + "prefix_cache": prefix_cache, + "duration_s": duration, + "ttft_p50_ms": statistics.median(ttfts) * 1e3, + "ttft_mean_ms": statistics.fmean(ttfts) * 1e3, + "prompt_tokens": prompted, + "prefill_tokens_run": sum(prefill_tokens), + "prefill_tokens_saved": prompted - sum(prefill_tokens), + "prefill_saving_fraction": (prompted - sum(prefill_tokens)) / prompted, + "page_hit_rate": plane.prefix_index.hit_rate, + "pages_retained": plane.prefix_index.num_cached_pages, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--requests", type=int, default=24) + parser.add_argument("--shared-prefix", type=int, default=512, help="tokens of common context per request") + parser.add_argument("--unique", type=int, default=128, help="mean tokens of per-request context") + parser.add_argument("--variants", type=int, default=1, help="distinct shared prefixes to spread across") + parser.add_argument("--output", type=int, default=16) + parser.add_argument("--max-batch", type=int, default=8) + parser.add_argument("--max-context", type=int, default=1024) + parser.add_argument("--max-prompt", type=int, default=1024) + parser.add_argument("--page-size", type=int, default=16) + parser.add_argument("--pool-tokens", type=int, default=16384) + parser.add_argument("--layers", type=int, default=8) + parser.add_argument("--emb-dim", type=int, default=1024) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--json-out", default=None) + parser.add_argument("--config-path", default="src/maxtext/configs/base.yml") + args = parser.parse_args() + + devices = jax.devices()[:1] + print(f"device: {devices[0].device_kind} model: {args.layers}L x {args.emb_dim}d") + print( + f"trace: {args.requests} requests, {args.shared_prefix} shared + ~{args.unique} unique prompt tokens, " + f"{args.variants} prefix variant(s), {args.output} output tokens" + ) + + results = {} + for prefix_cache in (False, True): + arm = "on" if prefix_cache else "off" + summary = run_arm(args, devices, prefix_cache) + results[arm] = summary + print(f"\n=== prefix cache {arm} ===") + print(f" prompt tokens {summary['prompt_tokens']}") + print(f" prefill tokens run {summary['prefill_tokens_run']}") + print(f" prefill saved {summary['prefill_tokens_saved']} ({summary['prefill_saving_fraction']:.1%})") + print(f" TTFT p50 {summary['ttft_p50_ms']:.2f} ms") + print(f" wall clock {summary['duration_s']:.2f} s") + + off, on = results["off"], results["on"] + print("\n=== prefix sharing, same engine and same trace ===") + print( + f" prefill work {off['prefill_tokens_run']} -> {on['prefill_tokens_run']} tokens " + f"({on['prefill_saving_fraction']:.1%} avoided)" + ) + print(f" TTFT p50 {off['ttft_p50_ms']:.2f} -> {on['ttft_p50_ms']:.2f} ms") + print(f" wall clock {off['duration_s']:.2f} -> {on['duration_s']:.2f} s") + print(f" page hit rate {on['page_hit_rate']:.1%}, pages retained {on['pages_retained']}") + + if args.json_out: + with open(args.json_out, "w", encoding="utf-8") as handle: + json.dump(results, handle, indent=2) + print(f"\nwrote {args.json_out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/paged_kv/run_serving_benchmark.py b/benchmarks/paged_kv/run_serving_benchmark.py index fe80f242dd..970f6d96fe 100644 --- a/benchmarks/paged_kv/run_serving_benchmark.py +++ b/benchmarks/paged_kv/run_serving_benchmark.py @@ -99,6 +99,7 @@ def build_config(mode: str, args) -> object: overrides["paged_num_blocks"] = dense_kv_tokens // args.page_size overrides["paged_max_context_len"] = args.max_context overrides["per_device_batch_size"] = float(args.max_batch) + overrides["paged_enable_prefix_cache"] = bool(args.prefix_cache) overrides["run_name"] = f"paged_bench_{mode}" return pyconfig.initialize([sys.argv[0], args.config_path], **overrides) @@ -120,6 +121,26 @@ def build_params(cfg, devices): def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--mode", choices=("paged", "dense", "both"), default="paged") + parser.add_argument( + "--shared-prefix", + type=int, + default=0, + help=( + "tokens of common leading context per request; 0 uses the independent-prompt trace. " + "Non-zero is the workload prefix sharing targets, so pair it with --prefix-cache" + ), + ) + parser.add_argument( + "--prefix-variants", + type=int, + default=1, + help="distinct shared prefixes to spread requests across, as a deployment serving several system prompts", + ) + parser.add_argument( + "--prefix-cache", + action="store_true", + help="share already-computed pages between requests with a common prefix (paged arm only)", + ) parser.add_argument("--requests", type=int, default=24) parser.add_argument("--mean-prompt", type=int, default=48) parser.add_argument("--mean-output", type=int, default=32) @@ -167,9 +188,24 @@ def main() -> int: # A factory rather than one list: each repeat needs requests with no # timestamps on them, and each arm needs an independent copy. def fresh_trace(): - return copy.deepcopy( - benchmark.synthetic_trace(args.requests, args.mean_prompt, args.mean_output, seed=args.seed) - ) + if args.shared_prefix: + # The unique tail is what is left of the mean prompt once the shared + # block is accounted for, so both traces present the same total prompt + # length and the comparison is about sharing rather than about size. + unique = max(args.mean_prompt - args.shared_prefix, args.page_size) + trace = benchmark.shared_prefix_trace( + args.requests, + args.shared_prefix, + unique, + args.mean_output, + seed=args.seed, + num_variants=args.prefix_variants, + ) + else: + trace = benchmark.synthetic_trace( + args.requests, args.mean_prompt, args.mean_output, seed=args.seed + ) + return copy.deepcopy(trace) if arm == "paged": # The batch cap must not bind, or the measurement reports the flag rather @@ -236,6 +272,12 @@ def fresh_trace(): results[arm] = summary print(benchmark.report(arm, summary)) print(f" KV tokens committed: {summary['kv_tokens_committed']}") + if summary.get("prefix_cache_enabled"): + print( + f" prefix cache: {summary['prefill_tokens_saved']} of {summary['prompt_tokens']} prompt tokens " + f"not recomputed ({summary['prefill_saving_fraction']:.1%}), " + f"page hit rate {summary['prefix_cache_page_hit_rate']:.1%}" + ) if "paged" in results and "dense" in results: paged, dense = results["paged"], results["dense"] diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index fdfbff77c3..a11122add0 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -641,6 +641,17 @@ class Attention(BaseModel): "Costs a pool write per freed page, so it is off by default." ), ) + paged_enable_prefix_cache: bool = Field( + False, + description=( + "Share already-computed KV pages between requests whose prompts begin with the same tokens, " + "for attention='gpu_paged'. Skips the prefill of the shared prefix, which is the dominant " + "cost for a workload with a common system prompt or multi-turn conversations. Off by " + "default because a hit is only sound if the caller supplies a cache namespace that " + "distinguishes everything affecting the KV -- weights, adapter, tokenizer, RoPE, quantisation " + "-- and defaulting it on would make that omission silent rather than opt-in." + ), + ) attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed", "block_diffusion"] = Field( "global", description="The variant of attention to use." ) diff --git a/src/maxtext/inference/kv_common/__init__.py b/src/maxtext/inference/kv_common/__init__.py index 08570bd315..5fa881272a 100644 --- a/src/maxtext/inference/kv_common/__init__.py +++ b/src/maxtext/inference/kv_common/__init__.py @@ -25,6 +25,7 @@ limitations under the License. """ +from maxtext.inference.kv_common.namespace import CACHE_NAMESPACE_VERSION, CacheNamespace from maxtext.inference.kv_common.page_table import KV_PAGE_TABLE_VERSION, KvPageTableV1 from maxtext.inference.kv_common.storage_layout import ( KV_STORAGE_LAYOUT_VERSION, @@ -32,8 +33,10 @@ ) __all__ = [ + "CACHE_NAMESPACE_VERSION", "KV_PAGE_TABLE_VERSION", "KV_STORAGE_LAYOUT_VERSION", + "CacheNamespace", "KvPageTableV1", "KvStorageLayoutV1", ] diff --git a/src/maxtext/inference/kv_common/namespace.py b/src/maxtext/inference/kv_common/namespace.py new file mode 100644 index 0000000000..abf8bcffea --- /dev/null +++ b/src/maxtext/inference/kv_common/namespace.py @@ -0,0 +1,118 @@ +"""Cache identity: everything that can change the K/V produced for the same tokens. + +A prefix cache answers "have I already computed the K/V for these token ids?". +That question is only well posed relative to everything *else* that determines the +answer, and the failure mode when the set is incomplete is the worst kind: a +correct-looking cache hit that returns another configuration's K/V. Nothing +crashes, nothing is logged, and the output is merely wrong. + +So the namespace is folded into the block hash chain rather than compared +alongside it. Two requests in different namespaces do not traverse the same +subtree at all, because their very first block hash differs. A mismatch cannot +produce a hit, as opposed to producing one that a later check is relied upon to +catch. + +**The digest iterates the dataclass fields rather than listing them.** That is the +load-bearing detail. A hand-written digest is a second place to remember every +field, and the one time someone adds a field and forgets is the one time two +incompatible configurations collide. Adding a field here puts it in the hash +automatically. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +import hashlib + +CACHE_NAMESPACE_VERSION = 1 + + +@dataclasses.dataclass(frozen=True) +class CacheNamespace: + """Identity of the configuration that produced a cached K/V page. + + Every field is a string so the digest is defined without per-type handling, and + so a caller can put whatever summary is meaningful in it -- a checkpoint hash, a + serialised RoPE config, a digest of image inputs. What matters is that two + configurations which would produce different K/V produce different strings. + + Attributes: + model_fingerprint: identity of the weights themselves, ideally a checkpoint + digest rather than a name. Two finetunes of one base model share a name and + must not share a cache. + model_revision: the version of those weights, for the case where the + fingerprint is a mutable reference. + tokenizer: tokenizer identity. The same text maps to different ids under a + different tokenizer, and the same ids to different text. + adapter: LoRA or other adapter identity. An adapter changes the projections + and therefore the K/V for identical tokens. + tenant: cache domain. Not a correctness field but an isolation one -- two + tenants may be unwilling to share pages even when sharing is sound. + rope: RoPE configuration. Scaling factors and base frequency change the K/V + for a token at a given position, and llama3-style scaling makes this a real + variant rather than a theoretical one. + kv_dtype: storage dtype of the pool. + kv_quantization: quantisation scheme and scales, if any. + layout: page size and physical form. A page cached under one layout is not + readable under another. + sharding: tensor-parallel width and axis mapping. Under replication the same + logical head lives on several devices, and a page cached at one TP width is + not valid at another. + prompt_embeddings: digest of any soft-prompt or prefix-tuning embeddings, + which change the K/V without changing a single token id. + multimodal: digest of image, audio or video inputs interleaved with the text, + for the same reason. + """ + + model_fingerprint: str = "" + model_revision: str = "" + tokenizer: str = "" + adapter: str = "" + tenant: str = "" + rope: str = "" + kv_dtype: str = "" + kv_quantization: str = "" + layout: str = "" + sharding: str = "" + prompt_embeddings: str = "" + multimodal: str = "" + version: int = CACHE_NAMESPACE_VERSION + + def digest(self) -> bytes: + """A stable 32-byte digest over *every* field of this dataclass. + + Field names are hashed alongside their values, so renaming a field or + reordering the declaration changes the digest. That is deliberate: either + change means the stored meaning has moved, and silently keeping old entries + valid across it would be a correctness bug rather than a convenience. + """ + hasher = hashlib.blake2b(digest_size=32) + for field in dataclasses.fields(self): + hasher.update(field.name.encode("utf-8")) + hasher.update(b"\x00") + hasher.update(str(getattr(self, field.name)).encode("utf-8")) + hasher.update(b"\x01") + return hasher.digest() + + def describe(self) -> str: + """The non-empty fields, for a log line that explains a cache miss.""" + parts = [ + f"{f.name}={getattr(self, f.name)}" + for f in dataclasses.fields(self) + if f.name != "version" and getattr(self, f.name) + ] + return ", ".join(parts) if parts else "" diff --git a/src/maxtext/inference/kv_control/__init__.py b/src/maxtext/inference/kv_control/__init__.py index 202dd0af79..6b176eaed5 100644 --- a/src/maxtext/inference/kv_control/__init__.py +++ b/src/maxtext/inference/kv_control/__init__.py @@ -30,7 +30,11 @@ """ from maxtext.inference.kv_control.allocator import DoubleFreeError, PagedBlockAllocator -from maxtext.inference.kv_control.control_plane import DirtyPageError, NativeKvControlPlane +from maxtext.inference.kv_control.control_plane import ( + DirtyPageError, + NativeKvControlPlane, + SharedPageWriteError, +) from maxtext.inference.kv_control.logical_block import ( LogicalBlock, PageState, @@ -43,6 +47,13 @@ ) from maxtext.inference.kv_control.metadata import build_decode_table, build_page_table from maxtext.inference.kv_control.page_map import PageCapacityError, PageMap, StaleRequestHandleError +from maxtext.inference.kv_control.prefix_index import ( + PrefixIndex, + PrefixMatch, + PrefixNode, + PublishResult, + block_hash, +) from maxtext.inference.kv_control.protocols import KvControlPlane from maxtext.inference.kv_control.request import RequestDescriptor, RequestHandle, RequestState @@ -57,10 +68,16 @@ "PageState", "PageStateError", "PagedBlockAllocator", + "PrefixIndex", + "PrefixMatch", + "PrefixNode", + "PublishResult", "RequestDescriptor", "RequestHandle", "RequestState", + "SharedPageWriteError", "StaleRequestHandleError", + "block_hash", "build_decode_table", "build_page_table", "decode_needs_new_page", diff --git a/src/maxtext/inference/kv_control/control_plane.py b/src/maxtext/inference/kv_control/control_plane.py index e0ed2350b7..c38c14ee1c 100644 --- a/src/maxtext/inference/kv_control/control_plane.py +++ b/src/maxtext/inference/kv_control/control_plane.py @@ -25,6 +25,14 @@ page-boundary crossing and removes any dependence on every kernel, present and future, honouring that bound. +**Prefix sharing, when enabled, changes which pages a request owns but not how a +step is described.** A shared page is attached to the request's row like any +other, so it appears in the page table and the kernel reads it. What keeps it +read-only is that the request's query length covers only its uncached suffix, +and write positions are derived from that -- so a shared page cannot enter a +`slot_mapping` without the query length being wrong first. Under `debug_mode` +that implication is checked rather than assumed. + Copyright 2026 Advanced Micro Devices, Inc. Licensed under the Apache License, Version 2.0 (the "License"); @@ -42,22 +50,48 @@ from __future__ import annotations +import dataclasses from typing import Sequence import numpy as np -from maxtext.inference.kv_common import KvPageTableV1, KvStorageLayoutV1 +from maxtext.inference.kv_common import CacheNamespace, KvPageTableV1, KvStorageLayoutV1 from maxtext.inference.kv_control import metadata from maxtext.inference.kv_control.allocator import PagedBlockAllocator from maxtext.inference.kv_control.logical_block import new_pages_for_extend, pages_for_tokens from maxtext.inference.kv_control.page_map import PageCapacityError, PageMap +from maxtext.inference.kv_control.prefix_index import PrefixIndex, PrefixMatch, PrefixNode from maxtext.inference.kv_control.request import RequestDescriptor, RequestHandle +_DEFAULT_NAMESPACE = CacheNamespace() + class DirtyPageError(RuntimeError): """A step would have read a page still holding another request's KV.""" +class SharedPageWriteError(RuntimeError): + """A step would have written into a page another request is reading.""" + + +@dataclasses.dataclass +class _PrefixState: + """What a request borrowed from the prefix index, and under whose identity. + + Held per live request because neither the page map nor the allocator can + distinguish a borrowed page from an owned one, and releasing a request must + free exactly the ones it owns. + + Recorded even when the lookup missed, because the namespace is needed again at + release to publish under. Inferring it then would mean guessing, and a wrong + guess publishes one configuration's K/V where another will find it. + """ + + namespace: CacheNamespace + node: PrefixNode | None = None + shared_pages: np.ndarray = dataclasses.field(default_factory=lambda: np.empty((0,), dtype=np.int32)) + + class NativeKvControlPlane: """Page accounting owned by MaxText, satisfying `KvControlPlane`.""" @@ -67,15 +101,22 @@ def __init__( max_requests: int, max_context_len: int, debug_mode: bool = False, + enable_prefix_cache: bool = False, ): if max_context_len < 1: raise ValueError(f"max_context_len must be at least 1, got {max_context_len}") self._layout = layout self.max_context_len = int(max_context_len) + self.debug_mode = bool(debug_mode) self.allocator = PagedBlockAllocator.from_layout(layout, debug_mode=debug_mode) self.page_map = PageMap.from_layout(layout, max_requests=max_requests, max_context_len=max_context_len) self._pending_scrub: list[np.ndarray] = [] + # Off unless asked for: sharing is only sound when the caller supplies a + # namespace that genuinely distinguishes its configurations, and defaulting + # it on would make that someone else's problem to discover. + self.prefix_index = PrefixIndex(layout.tokens_per_page, enabled=enable_prefix_cache) + self._prefix_state: dict[RequestHandle, _PrefixState] = {} # Caught here rather than as a mid-serving allocation failure, which would # look like ordinary backpressure and never resolve. @@ -106,6 +147,51 @@ def admit(self, descriptor: RequestDescriptor) -> RequestHandle | None: ) return self.page_map.admit(descriptor) + @property + def prefix_cache_enabled(self) -> bool: + return self.prefix_index.enabled + + def attach_prefix( + self, + handle: RequestHandle, + token_ids: Sequence[int] | np.ndarray, + namespace: CacheNamespace = _DEFAULT_NAMESPACE, + ) -> PrefixMatch: + """Lend the request every already-computed page of its prompt. + + The shared pages are attached to the request's row and its length advanced + over them, so the request begins as though it had already prefilled that + much. The caller then prefills only `prompt_len - match.num_tokens` tokens, + and that difference is the entire benefit. + + Must be called before the request holds any pages, since a prefix is by + definition the front of the sequence and the page list is positional. + """ + if self.page_map.num_pages(handle) or self.page_map.seq_len(handle): + raise ValueError( + f"request {handle.request_id!r} already holds " + f"{self.page_map.num_pages(handle)} pages at length {self.page_map.seq_len(handle)}; " + "a prefix can only be attached to a request that has not yet reserved anything" + ) + + match = self.prefix_index.match(token_ids, namespace) + # Recorded on a miss too: this is where the caller states which + # configuration the request belongs to, and release needs it to publish + # under the same one. + self._prefix_state[handle] = _PrefixState(namespace=namespace) + if not match: + return match + + self.prefix_index.acquire(match.node) + self.page_map.append_pages(handle, match.pages) + self.page_map.advance(handle, match.num_tokens) + self._prefix_state[handle] = _PrefixState( + namespace=namespace, + node=match.node, + shared_pages=match.pages, + ) + return match + def reserve( self, handles: Sequence[RequestHandle], @@ -135,9 +221,19 @@ def reserve( ) per_request_pages.append(new_pages_for_extend(prefix_len, new_len, tokens_per_page)) - pages = self.allocator.alloc(int(sum(per_request_pages))) + needed = int(sum(per_request_pages)) + pages = self.allocator.alloc(needed) if pages is None: - return False + # Cached pages are the one reclaimable thing left: they hold K/V nothing + # currently needs, and recomputing them is strictly better than refusing + # to make progress. Evicting only the shortfall keeps the rest of the + # cache for whoever asks next. + reclaimed = self.evict_cached(needed - self.allocator.available_pages) + if reclaimed.size == 0: + return False + pages = self.allocator.alloc(needed) + if pages is None: + return False recycled = self.allocator.dirty_among(pages) if recycled.size: @@ -177,15 +273,67 @@ def confirm_scrubbed(self, page_ids: Sequence[int] | np.ndarray) -> None: def reserve_decode(self, handles: Sequence[RequestHandle]) -> bool: return self.reserve(handles, np.ones((len(handles),), dtype=np.int32)) - def release(self, handle: RequestHandle) -> np.ndarray: - """Reclaim the request's pages and report them. + def evict_cached(self, num_pages: int) -> np.ndarray: + """Drop `num_pages` unreferenced cached pages and return them to the pool. - The pages are returned so a caller can zero or poison them before they can - be handed out again. Nothing here does that -- the pool is device memory and - this layer never touches a device -- which is why the pages have to leave - through the return value rather than being quietly recycled. + The pages come back dirty, as any freed page does, so a later step that + reserves them will scrub them through the ordinary path. """ - pages = self.page_map.release(handle) + reclaimed = self.prefix_index.evict(num_pages) + if reclaimed.size: + self.allocator.free(reclaimed) + return reclaimed + + def release( + self, + handle: RequestHandle, + token_ids: Sequence[int] | np.ndarray | None = None, + num_valid_tokens: int | None = None, + namespace: CacheNamespace | None = None, + ) -> np.ndarray: + """Reclaim the request's pages and report the ones actually freed. + + With prefix sharing on and `token_ids` supplied, the request's full pages are + offered to the index first, and the ones it adopts stay allocated for the + next request with the same prefix. Those are excluded from the return value, + which continues to mean "pages that are now free and may be poisoned" -- so a + caller that poisons what it gets back cannot corrupt a page it just donated. + + Pages borrowed at admission are likewise never freed here: the index owns + them and this request was only reading them. + + Publishing needs to know which configuration produced the K/V. That comes + from the `attach_prefix` call that admitted the request, or from `namespace` + for a caller that never made one. It is never defaulted: publishing under a + namespace the request did not belong to files its K/V where a different + configuration will find and trust it. + """ + state = self._prefix_state.pop(handle, None) + held = self.page_map.release(handle) + + keep = np.zeros((0,), dtype=np.int32) + if state is not None: + if state.node is not None: + self.prefix_index.release(state.node) + keep = state.shared_pages + if namespace is not None and namespace != state.namespace: + raise ValueError( + f"request {handle.request_id!r} was admitted under namespace ({state.namespace.describe()}) " + f"but is being released under a different one ({namespace.describe()})" + ) + namespace = state.namespace + + if token_ids is not None and self.prefix_index.enabled: + if namespace is None: + raise ValueError( + f"cannot publish request {handle.request_id!r} without a namespace. Either admit it via " + "attach_prefix, which records one, or pass namespace= here." + ) + published = self.prefix_index.publish(token_ids, held, namespace, num_valid_tokens) + if published.adopted.size: + keep = np.union1d(keep, published.adopted) + + pages = held[~np.isin(held, keep)] if keep.size else held if pages.size: self.allocator.free(pages) return pages @@ -209,7 +357,30 @@ def build_page_table( f"by this step. Scrub the pages from pending_scrub() and call confirm_scrubbed() before " f"building the table." ) + if self.debug_mode and self._prefix_state: + self._check_no_shared_writes(table) return table + def _check_no_shared_writes(self, table: KvPageTableV1) -> None: + """Confirm this step writes into no page it is only borrowing. + + The property already follows from query lengths covering only the uncached + suffix, so this can never fire against correct arithmetic. It is here + because the failure it would catch -- one request overwriting the cached + prefix that several others are reading -- corrupts those requests silently + and at a distance, and is worth paying a debug-mode set intersection to rule + out while the arithmetic is still new. + """ + borrowed = np.unique(np.concatenate([s.shared_pages for s in self._prefix_state.values()])) + slots = table.slot_mapping(self._layout.tokens_per_page, self.allocator.padding_page_id).reshape(-1) + written = np.unique(slots[slots >= 0] // self._layout.tokens_per_page).astype(np.int32) + clash = np.intersect1d(written, borrowed) + if clash.size: + raise SharedPageWriteError( + f"pages {clash[:8].tolist()} are shared prefix pages that other requests are reading, but " + f"this step's slot_mapping writes into them. A query length is covering tokens that were " + f"served from the prefix cache." + ) + def build_decode_table(self, handles: Sequence[RequestHandle]) -> KvPageTableV1: return self.build_page_table(handles, np.ones((len(handles),), dtype=np.int32)) diff --git a/src/maxtext/inference/kv_control/prefix_index.py b/src/maxtext/inference/kv_control/prefix_index.py new file mode 100644 index 0000000000..0a274296fe --- /dev/null +++ b/src/maxtext/inference/kv_control/prefix_index.py @@ -0,0 +1,397 @@ +"""Sharing pages between requests that begin with the same tokens. + +The K/V for a token depends on that token and everything before it, so two +requests with a common prefix compute bit-identical K/V for it. A shared system +prompt or a multi-turn conversation replays the same leading tokens on every +turn, and prefill is quadratic in prompt length, so not recomputing that prefix +is the largest single win available to a paged runtime. + +The index is a trie of pages. Each node owns one page and is keyed by a hash +chained from its parent, so a node's key identifies not just its own tokens but +the entire path that produced them -- which is exactly the condition under which +its K/V is valid. The chain starts at the namespace digest rather than at a +constant, so two configurations do not share a root and a namespace mismatch is +structurally unable to produce a hit. See `kv_common/namespace.py`. + +Three decisions differ from the sglang-jax radix cache this follows in outline: + +**Nodes are one page, not a variable-length token run.** That removes node +splitting entirely, which is most of the reference's complexity, and costs +nothing here because only whole pages are ever published. + +**Only full pages are published.** A partial page is still being appended to, so +sharing it would mean two requests writing different tokens into the same page. +The tail is left private and recomputed, which is at most `tokens_per_page - 1` +tokens of duplicated prefill. + +**Recency is a counter, not a clock.** `time.monotonic()` ties at clock +resolution when many nodes are touched in one step, which makes eviction order +depend on how the heap happened to break the tie. A counter cannot tie, so +eviction order is reproducible and a test can assert against it. + +Stored tokens are compared on a hash hit. A 256-bit chained hash makes collision +a non-argument, but the comparison is a page of integers against a cached page of +integers, and what it buys is that the one bug class here which produces wrong +output rather than a crash cannot occur at all. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import heapq +from typing import Iterable, Sequence + +import numpy as np + +from maxtext.inference.kv_common import CacheNamespace + +_EMPTY_PAGES = np.empty((0,), dtype=np.int32) + + +def block_hash(parent_hash: bytes, token_ids: Sequence[int]) -> bytes: + """Hash one page's tokens into the chain ending at `parent_hash`. + + Chaining rather than hashing the page alone is what makes a node's key mean + "these tokens, at this position, after exactly this history". Hashing the page + in isolation would let a page match at the wrong depth or after a different + prefix, and its K/V would be wrong in both cases. + """ + hasher = hashlib.blake2b(digest_size=32) + hasher.update(parent_hash) + hasher.update(np.asarray(token_ids, dtype=np.int64).tobytes()) + return hasher.digest() + + +@dataclasses.dataclass(eq=False) +class PrefixNode: + """One cached page, and the path that produced it.""" + + block_hash: bytes + page_id: int + tokens: np.ndarray + depth: int + parent: "PrefixNode | None" + children: dict[bytes, "PrefixNode"] = dataclasses.field(default_factory=dict) + ref_count: int = 0 + last_access: int = 0 + hit_count: int = 0 + + @property + def is_leaf(self) -> bool: + return not self.children + + +@dataclasses.dataclass(frozen=True) +class PrefixMatch: + """What a lookup found. `pages` are readable but must never be written.""" + + pages: np.ndarray + num_tokens: int + node: PrefixNode | None + + @property + def num_pages(self) -> int: + return int(self.pages.size) + + def __bool__(self) -> bool: + return self.pages.size > 0 + + +@dataclasses.dataclass(frozen=True) +class PublishResult: + """The outcome of offering a request's pages to the index. + + `adopted` are now owned by the index and must not be freed by the caller. + `duplicate` held content some other request had already cached; the index kept + what it had, so these are the caller's to free. Both are page ids the caller + passed in, partitioned -- returning them rather than a count is what lets the + caller free exactly the right set without tracking the split itself. + """ + + adopted: np.ndarray + duplicate: np.ndarray + + @property + def num_adopted(self) -> int: + return int(self.adopted.size) + + +class PrefixIndex: + """A page trie mapping token prefixes to the pages already holding their K/V. + + The index owns every page it holds. A page enters through `publish` and leaves + only through `evict` or `reset`, both of which hand it back so the caller can + return it to the allocator. Nothing here allocates or frees; this layer never + touches a device or a free list. + """ + + def __init__(self, tokens_per_page: int, enabled: bool = True): + if tokens_per_page < 1: + raise ValueError(f"tokens_per_page must be at least 1, got {tokens_per_page}") + self.tokens_per_page = int(tokens_per_page) + self.enabled = bool(enabled) + self._root = PrefixNode(block_hash=b"", page_id=-1, tokens=_EMPTY_PAGES, depth=0, parent=None) + self._num_nodes = 0 + self._protected_nodes = 0 + self._clock = 0 + self.num_queries = 0 + self.num_hit_pages = 0 + self.num_queried_pages = 0 + + # -- statistics ----------------------------------------------------------- + + @property + def num_cached_pages(self) -> int: + return self._num_nodes + + @property + def protected_pages(self) -> int: + """Pages a live request depends on. Not evictable at any pressure.""" + return self._protected_nodes + + @property + def evictable_pages(self) -> int: + return self._num_nodes - self._protected_nodes + + @property + def hit_rate(self) -> float: + """Fraction of queried pages served from the cache, across all lookups.""" + if not self.num_queried_pages: + return 0.0 + return self.num_hit_pages / self.num_queried_pages + + def stats(self) -> dict[str, float | int]: + return { + "cached_pages": self.num_cached_pages, + "protected_pages": self.protected_pages, + "evictable_pages": self.evictable_pages, + "queries": self.num_queries, + "queried_pages": self.num_queried_pages, + "hit_pages": self.num_hit_pages, + "hit_rate": self.hit_rate, + } + + # -- lookup --------------------------------------------------------------- + + def match(self, token_ids: Sequence[int], namespace: CacheNamespace) -> PrefixMatch: + """Find the longest cached prefix of `token_ids` under `namespace`. + + Matching stops one page short of the end even on a total match. A request + whose every token is cached would have no query tokens left to run, and the + step would have nothing to compute a next-token logit from; leaving the final + page to be recomputed keeps every request with at least one token of work. + """ + tokens = np.asarray(token_ids, dtype=np.int64).reshape(-1) + full_pages = int(tokens.size) // self.tokens_per_page + # A prompt that is an exact multiple of the page size would otherwise match + # to its own end; hold back the last page in every case. + matchable = max(full_pages - 1, 0) if tokens.size % self.tokens_per_page == 0 else full_pages + + self.num_queries += 1 + self.num_queried_pages += matchable + if not self.enabled or matchable == 0: + return PrefixMatch(pages=_EMPTY_PAGES, num_tokens=0, node=None) + + matched: list[int] = [] + node = self._root + parent_hash = namespace.digest() + for page_index in range(matchable): + start = page_index * self.tokens_per_page + page_tokens = tokens[start : start + self.tokens_per_page] + key = block_hash(parent_hash, page_tokens) + child = node.children.get(key) + if child is None or not np.array_equal(child.tokens, page_tokens): + break + node = child + parent_hash = key + matched.append(child.page_id) + + if not matched: + return PrefixMatch(pages=_EMPTY_PAGES, num_tokens=0, node=None) + + self._clock += 1 + self.num_hit_pages += len(matched) + walk: PrefixNode | None = node + while walk is not None and walk is not self._root: + walk.last_access = self._clock + walk.hit_count += 1 + walk = walk.parent + + return PrefixMatch( + pages=np.asarray(matched, dtype=np.int32), + num_tokens=len(matched) * self.tokens_per_page, + node=node, + ) + + # -- publication ---------------------------------------------------------- + + def publish( + self, + token_ids: Sequence[int], + page_ids: Sequence[int] | np.ndarray, + namespace: CacheNamespace, + num_valid_tokens: int | None = None, + ) -> PublishResult: + """Offer a request's computed pages to the index. + + Only pages fully covered by `num_valid_tokens` are taken: a page holding + tokens the request has not written yet would be published with K/V that does + not exist. `page_ids` is the request's page list in sequence order, and the + caller must have written the K/V for every token it declares valid. + """ + tokens = np.asarray(token_ids, dtype=np.int64).reshape(-1) + pages = np.asarray(page_ids, dtype=np.int32).reshape(-1) + valid = int(tokens.size if num_valid_tokens is None else num_valid_tokens) + if valid > tokens.size: + raise ValueError(f"declared {valid} valid tokens but was given only {tokens.size} token ids") + if valid < 0: + raise ValueError(f"num_valid_tokens must be non-negative, got {valid}") + + publishable = valid // self.tokens_per_page + if publishable > pages.size: + raise ValueError( + f"{valid} valid tokens span {publishable} pages but the request holds only {pages.size}; " + "the page list and the token list disagree about the request's length" + ) + if not self.enabled or publishable == 0: + return PublishResult(adopted=_EMPTY_PAGES, duplicate=_EMPTY_PAGES) + + adopted: list[int] = [] + duplicate: list[int] = [] + node = self._root + parent_hash = namespace.digest() + self._clock += 1 + for page_index in range(publishable): + start = page_index * self.tokens_per_page + page_tokens = tokens[start : start + self.tokens_per_page] + key = block_hash(parent_hash, page_tokens) + child = node.children.get(key) + if child is None: + child = PrefixNode( + block_hash=key, + page_id=int(pages[page_index]), + tokens=page_tokens.copy(), + depth=page_index + 1, + parent=node, + last_access=self._clock, + ) + node.children[key] = child + self._num_nodes += 1 + adopted.append(int(pages[page_index])) + # A new node under a protected parent inherits nothing: the request that + # locked the parent never referenced this page, so it starts evictable. + else: + child.last_access = self._clock + if int(pages[page_index]) != child.page_id: + duplicate.append(int(pages[page_index])) + node = child + parent_hash = key + + return PublishResult( + adopted=np.asarray(adopted, dtype=np.int32) if adopted else _EMPTY_PAGES, + duplicate=np.asarray(duplicate, dtype=np.int32) if duplicate else _EMPTY_PAGES, + ) + + # -- reference counting --------------------------------------------------- + + def acquire(self, node: PrefixNode | None) -> None: + """Protect `node` and its ancestors from eviction while a request reads them. + + Without this a page could be evicted, freed, reallocated and overwritten + while a live request's page table still names it -- which is the same + use-after-free the dirty-page gate exists to prevent, arriving by a different + route. + """ + while node is not None and node is not self._root: + if node.ref_count == 0: + self._protected_nodes += 1 + node.ref_count += 1 + node = node.parent + + def release(self, node: PrefixNode | None) -> None: + """Undo one `acquire`, making the path evictable again once nothing holds it.""" + while node is not None and node is not self._root: + if node.ref_count <= 0: + raise RuntimeError( + f"prefix node at depth {node.depth} (page {node.page_id}) was released more times than " + "it was acquired, so some live request's pages are now evictable" + ) + node.ref_count -= 1 + if node.ref_count == 0: + self._protected_nodes -= 1 + node = node.parent + + # -- eviction ------------------------------------------------------------- + + def evict(self, num_pages: int) -> np.ndarray: + """Drop up to `num_pages` least-recently-used pages and return them. + + Only leaves are eligible, so a page is never dropped while a longer cached + prefix still depends on it. Freeing the returned pages is the caller's job; + the index has already forgotten them by the time this returns. + """ + if num_pages <= 0: + return _EMPTY_PAGES + + leaves = [(n.last_access, id(n), n) for n in self._collect_leaves() if n.ref_count == 0] + heapq.heapify(leaves) + + evicted: list[int] = [] + while leaves and len(evicted) < num_pages: + _, _, node = heapq.heappop(leaves) + # A node reached earlier in this pass may have gained a reference or, more + # commonly, is stale because we pushed its parent after deleting it. + if node.ref_count > 0 or not node.is_leaf or node is self._root: + continue + parent = node.parent + evicted.append(node.page_id) + del parent.children[node.block_hash] + node.parent = None + self._num_nodes -= 1 + if parent is not self._root and parent.is_leaf and parent.ref_count == 0: + heapq.heappush(leaves, (parent.last_access, id(parent), parent)) + + return np.asarray(evicted, dtype=np.int32) if evicted else _EMPTY_PAGES + + def reset(self) -> np.ndarray: + """Forget everything and hand back every page the index held. + + Refuses while any page is referenced: a live request's page table would + still name pages the caller is about to free. + """ + if self._protected_nodes: + raise RuntimeError( + f"{self._protected_nodes} cached pages are still referenced by live requests; " + "release them before resetting the index" + ) + pages = [node.page_id for node in self._walk(self._root) if node is not self._root] + self._root = PrefixNode(block_hash=b"", page_id=-1, tokens=_EMPTY_PAGES, depth=0, parent=None) + self._num_nodes = 0 + return np.asarray(pages, dtype=np.int32) if pages else _EMPTY_PAGES + + # -- internals ------------------------------------------------------------ + + def _walk(self, node: PrefixNode) -> Iterable[PrefixNode]: + stack = [node] + while stack: + current = stack.pop() + yield current + stack.extend(current.children.values()) + + def _collect_leaves(self) -> list[PrefixNode]: + return [n for n in self._walk(self._root) if n is not self._root and n.is_leaf] diff --git a/src/maxtext/inference/kv_execution/benchmark.py b/src/maxtext/inference/kv_execution/benchmark.py index 7ba0e564c9..17cc8f9c31 100644 --- a/src/maxtext/inference/kv_execution/benchmark.py +++ b/src/maxtext/inference/kv_execution/benchmark.py @@ -69,6 +69,16 @@ class Request: token_times: list[float] = dataclasses.field(default_factory=list) generated: list[int] = dataclasses.field(default_factory=list) handle: Any = None + token_ids: np.ndarray | None = None + prefill_tokens: int = 0 + + def context_tokens(self) -> np.ndarray | None: + """Prompt plus generated, which is what the prefix cache is offered.""" + if self.token_ids is None: + return None + if not self.generated: + return self.token_ids + return np.concatenate([self.token_ids, np.asarray(self.generated, dtype=np.int64)]) @property def ttft(self) -> float | None: @@ -108,12 +118,64 @@ def spread(mean: int) -> np.ndarray: return rng.integers(low, high, size=num_requests) prompts, outputs = spread(mean_prompt), spread(mean_output) + # Token ids are distinct per request on purpose. A constant filler would make + # every prompt a prefix of every other, so switching the prefix cache on would + # report a hit rate this workload does not represent -- the baseline has to be a + # workload with nothing to share. return [ - Request(request_id=f"r{i}", prompt_len=int(prompts[i]), max_new_tokens=int(outputs[i])) + Request( + request_id=f"r{i}", + prompt_len=int(prompts[i]), + max_new_tokens=int(outputs[i]), + token_ids=rng.integers(1, 30000, size=int(prompts[i]), dtype=np.int64), + ) for i in range(num_requests) ] +def shared_prefix_trace( + num_requests: int, + shared_prefix: int, + mean_unique: int, + mean_output: int, + *, + seed: int = 0, + length_spread: float = 0.6, + num_variants: int = 1, +) -> list[Request]: + """Requests that begin with a common block, as a served system prompt does. + + This is the workload prefix sharing exists for, and the one where a paged + runtime's advantage over a dense one is largest: the shared block is prefilled + once and read by everything after it. `num_variants` splits the population + across several distinct prefixes, which is the realistic case -- a deployment + serves a handful of system prompts, not one -- and it also keeps the measurement + honest about eviction, since several prefixes compete for the same pool. + """ + rng = np.random.default_rng(seed) + prefixes = [rng.integers(1, 30000, size=shared_prefix, dtype=np.int64) for _ in range(max(num_variants, 1))] + + def spread(mean: int) -> np.ndarray: + low = max(1, int(mean * (1.0 - length_spread))) + high = max(low + 1, int(mean * (1.0 + length_spread))) + return rng.integers(low, high, size=num_requests) + + uniques, outputs = spread(mean_unique), spread(mean_output) + requests = [] + for i in range(num_requests): + tail = rng.integers(1, 30000, size=int(uniques[i]), dtype=np.int64) + tokens = np.concatenate([prefixes[i % len(prefixes)], tail]) + requests.append( + Request( + request_id=f"r{i}", + prompt_len=int(tokens.size), + max_new_tokens=int(outputs[i]), + token_ids=tokens, + ) + ) + return requests + + def warmup_paged(engine, params, *, max_prompt: int, max_batch: int, target_context: int) -> set: """Compile every shape the measured run can present, then discard the results. @@ -257,18 +319,25 @@ def run_paged( while waiting or live: while waiting and len(live) < max_batch: candidate = waiting[0] - padded = jnp.asarray( - candidate.generated[:0] + [1] * candidate.prompt_len, dtype=jnp.int32 + prompt = ( + candidate.token_ids + if candidate.token_ids is not None + else np.ones((candidate.prompt_len,), dtype=np.int64) ) handle, result = engine.prefill_paged( params=params, - padded_tokens=padded, + padded_tokens=jnp.asarray(prompt, dtype=jnp.int32), true_length=candidate.prompt_len, request_id=candidate.request_id, max_new_tokens=candidate.max_new_tokens, + prompt_token_ids=candidate.token_ids, ) if handle is None: break # backpressure: the pool cannot take another request yet + # The tokens this prefill actually ran, which is the prompt minus whatever + # the cache supplied. Recorded per request so the saving can be reported as + # work avoided rather than inferred from a latency difference alone. + candidate.prefill_tokens = candidate.prompt_len - runtime.cached_tokens(handle) waiting.pop(0) candidate.handle = handle candidate.generated.append(int(result.data[0, 0])) @@ -295,6 +364,8 @@ def run_paged( # Pool exhausted mid-decode. Give the newest request's pages back and let # it be replayed, which is the same recompute-preemption the driver uses. victim = live.pop() + # Nothing published: preemption is trying to reclaim pages, and the cache + # would hold onto whatever it adopted. engine.release(victim.handle) victim.handle, victim.generated = None, [] waiting.insert(0, victim) @@ -314,13 +385,14 @@ def run_paged( occupancy.append((now - start, allocator.num_allocated_pages * page_size, live_tokens, len(live))) for request in finished: - engine.release(request.handle) + engine.release(request.handle, request.context_tokens()) request.handle = None live.remove(request) done.append(request) elapsed = time.perf_counter() - start summary = _summarise(done, occupancy, elapsed, steps, allocator, page_size) + summary.update(_summarise_prefix_cache(done, runtime.control_plane)) # Shape accounting, which is necessary but *not sufficient*: it sees bucketed # `StepShape`s and is blind to eager host-side ops whose values enter a jaxpr as @@ -338,6 +410,28 @@ def run_paged( return summary +def _summarise_prefix_cache(requests: Sequence[Request], control_plane) -> dict[str, Any]: + """What sharing avoided, in tokens rather than in seconds. + + Reported separately from latency because it is the direct measurement: a + latency difference between two runs mixes in the pool pressure the retained + pages cause, whereas prompt tokens minus prefilled tokens is exactly the work + that did not happen. + """ + prompted = sum(r.prompt_len for r in requests) + prefilled = sum(r.prefill_tokens for r in requests) + index = control_plane.prefix_index + return { + "prefix_cache_enabled": index.enabled, + "prompt_tokens": prompted, + "prefill_tokens_run": prefilled, + "prefill_tokens_saved": prompted - prefilled, + "prefill_saving_fraction": (prompted - prefilled) / prompted if prompted else 0.0, + "prefix_cache_page_hit_rate": index.hit_rate, + "prefix_cache_pages_retained": index.num_cached_pages, + } + + def run_repeated(engine, params, trace_factory, *, max_batch: int, warmed_shapes=None, repeats: int = 2): """Run the same trace several times and report the last, with the spread. diff --git a/src/maxtext/inference/kv_execution/driver.py b/src/maxtext/inference/kv_execution/driver.py index adfd63b560..17f624fe43 100644 --- a/src/maxtext/inference/kv_execution/driver.py +++ b/src/maxtext/inference/kv_execution/driver.py @@ -50,6 +50,7 @@ import numpy as np +from maxtext.inference.kv_common import CacheNamespace from maxtext.inference.kv_control import ( NativeKvControlPlane, RequestDescriptor, @@ -62,6 +63,8 @@ from maxtext.inference.kv_execution.pool_ops import poison_pages, scrub_pages from maxtext.inference.kv_execution.step_view import StepView, build_step_view +_DEFAULT_NAMESPACE = CacheNamespace() + @dataclasses.dataclass(eq=False) class PagedRequest: @@ -82,6 +85,12 @@ class PagedRequest: state: RequestState = RequestState.WAITING finish_reason: str | None = None preemptions: int = 0 + # Optional, and only consulted when the control plane's prefix cache is on. A + # workload that never repeats a prefix pays nothing for leaving these unset, + # which is why the driver treats them as opt-in rather than required. + prompt_tokens: np.ndarray | None = None + namespace: CacheNamespace = _DEFAULT_NAMESPACE + cached_tokens: int = 0 @property def is_finished(self) -> bool: @@ -100,6 +109,26 @@ def descriptor(self) -> RequestDescriptor: max_new_tokens=max(self.max_new_tokens - len(self.generated), 0), ) + def token_ids(self) -> np.ndarray | None: + """The full context: prompt plus whatever has been generated. + + Generated tokens are included because they are context like any other, so a + preempted request replaying a longer prompt can match its own prefix from + before the preemption. That makes recovery from backpressure cheaper on + exactly the workload where backpressure is most likely. + """ + if self.prompt_tokens is None: + return None + if not self.generated: + return self.prompt_tokens + return np.concatenate( + [np.asarray(self.prompt_tokens, dtype=np.int64), np.asarray(self.generated, dtype=np.int64)] + ) + + def prefill_len(self) -> int: + """Tokens this request must actually compute, after any cache hit.""" + return self.descriptor().prompt_len - self.cached_tokens + @dataclasses.dataclass(frozen=True) class StepOutcome: @@ -191,7 +220,7 @@ def step(self) -> StepOutcome | None: """ admitted = self._admit() if admitted: - return self._run_phase(admitted, [r.descriptor().prompt_len for r in admitted], is_decode=False) + return self._run_phase(admitted, [r.prefill_len() for r in admitted], is_decode=False) if self._live: return self._run_phase(self._live, [1] * len(self._live), is_decode=True) return None @@ -208,29 +237,49 @@ def _admit(self) -> list[PagedRequest]: tokens_per_page = self.plane.layout.tokens_per_page token_budget = self.planner.token_rungs[-1] - page_budget = self.plane.allocator.available_pages room = self.max_batch - len(self._live) admitted: list[PagedRequest] = [] tokens = 0 + committed = 0 for request in list(self._waiting): if len(admitted) >= room or self.plane.page_map.available_rows == 0: break prompt_len = request.descriptor().prompt_len needed = pages_for_tokens(prompt_len, tokens_per_page) - if needed > page_budget or tokens + prompt_len > token_budget: + # Cached pages count towards the budget because reservation evicts on + # shortfall, so they are reclaimable rather than spoken for. Budgeting + # against the free list alone would stall a loop that could still make + # progress by giving up cache entries -- turning an optimisation into a + # reason requests stop being served. Recomputed each time round because + # attaching a prefix protects pages, which takes them out of the figure. + budget = self.plane.allocator.available_pages + self.plane.prefix_index.evictable_pages - committed + if needed > budget or tokens + prompt_len > token_budget: break handle = self.plane.admit(request.descriptor()) if handle is None: break request.handle = handle request.state = RequestState.PREFILL - page_budget -= needed - tokens += prompt_len + # Charged at full price above and refunded here: the hit is only knowable + # once the handle exists, and over-estimating admits fewer requests than it + # might, where under-estimating would fail the whole batch's reservation. + request.cached_tokens = self._attach_prefix(request) + committed += needed - request.cached_tokens // tokens_per_page + tokens += prompt_len - request.cached_tokens admitted.append(request) self._waiting.remove(request) return admitted + def _attach_prefix(self, request: PagedRequest) -> int: + """Lend `request` any cached pages of its context. Returns tokens skipped.""" + if not self.plane.prefix_cache_enabled: + return 0 + tokens = request.token_ids() + if tokens is None: + return 0 + return self.plane.attach_prefix(request.handle, tokens, request.namespace).num_tokens + def _run_phase( self, requests: Sequence[PagedRequest], @@ -322,34 +371,53 @@ def _preempt_newest(self, batch: list[PagedRequest]) -> None: Its generated tokens are kept, so the replay is a longer prompt rather than lost output. What is lost is the compute that built its KV, which is the price of not deadlocking. + + Nothing is published on the way out. Preemption happens precisely because + pages are scarce, and the prefix cache holds onto what it adopts -- so + publishing here would hand back fewer pages than the preemption was trying + to reclaim, and could leave the retry no better off than the attempt that + triggered it. """ if not batch: return victim = batch.pop() - self._release_pages(victim) + self._release_pages(victim, publish=False) victim.state = RequestState.WAITING victim.preemptions += 1 if victim in self._live: self._live.remove(victim) self._waiting.insert(0, victim) - def _release_pages(self, request: PagedRequest) -> None: - """Give up a request's pages, poisoning them first when asked. + def _release_pages(self, request: PagedRequest, publish: bool = True) -> None: + """Give up a request's pages, poisoning the ones that actually came free. + + Poison is applied to what `release` reports rather than to everything the + request held, because with prefix sharing those differ: a page the index + adopted, or one this request only borrowed, stays live for the next reader + and poisoning it would destroy K/V that is about to be trusted. + + The pages are still dirty when poisoned -- `release` marks them so, and the + poison does not count as a scrub -- so the next occupant must still zero + them. That is what makes the sentinel a detector of a missed scrub rather + than a substitute for one. - Poison is applied *before* the pages leave the request, and does not count as - a scrub: the control plane keeps them dirty, so the next occupant still has - to zero them. That ordering is what makes the sentinel a detector of a missed - scrub rather than a substitute for one. + `publish` is what the prefix cache is offered. The recorded sequence length + is the written extent by definition, since a step advances it by exactly the + tokens it is about to write, so it is the right bound to publish under: the + final generated token is in the token list but its K/V will not be computed + until a step that now never runs. """ if request.handle is None: return - pages = self.plane.page_map.pages(request.handle) - if self.poison_on_free and pages.size: + tokens = request.token_ids() if publish else None + valid = self.plane.page_map.seq_len(request.handle) + freed = self.plane.release(request.handle, tokens, num_valid_tokens=valid) + if self.poison_on_free and freed.size: for layer in range(self.pool.num_layers): - k, v = poison_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], pages) + k, v = poison_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], freed) self.pool.replace_layer(layer, k, v) - self.plane.release(request.handle) request.handle = None + request.cached_tokens = 0 def _retire(self, batch: Sequence[PagedRequest], tokens: np.ndarray, *, is_decode: bool) -> None: """Record this step's tokens and release whatever finished. diff --git a/src/maxtext/inference/kv_execution/engine_adapter.py b/src/maxtext/inference/kv_execution/engine_adapter.py index df77cb7e14..88aaa592b9 100644 --- a/src/maxtext/inference/kv_execution/engine_adapter.py +++ b/src/maxtext/inference/kv_execution/engine_adapter.py @@ -76,6 +76,7 @@ def __init__( self.observed_shapes: set[StepShape] = set() self._by_slot: dict[int, RequestHandle] = {} self._by_request_id: dict[str, RequestHandle] = {} + self._cached_tokens: dict[str, int] = {} def admit(self, request_id: str, prompt_len: int, max_new_tokens: int) -> RequestHandle | None: """Start tracking a request, returning None when there is no room.""" @@ -83,6 +84,26 @@ def admit(self, request_id: str, prompt_len: int, max_new_tokens: int) -> Reques RequestDescriptor(request_id=request_id, prompt_len=prompt_len, max_new_tokens=max_new_tokens) ) + def attach_prefix(self, handle: RequestHandle, token_ids, namespace=None) -> int: + """Lend the request any cached pages of its prompt. Returns tokens skipped. + + The caller must then shorten the step's query to the tokens that remain, and + offset their positions by what was skipped -- the suffix sits at absolute + positions `cached..prompt_len`, and RoPE is not translation invariant. + """ + if not self.control_plane.prefix_cache_enabled or token_ids is None: + return 0 + if namespace is None: + cached = self.control_plane.attach_prefix(handle, token_ids).num_tokens + else: + cached = self.control_plane.attach_prefix(handle, token_ids, namespace).num_tokens + self._cached_tokens[handle.request_id] = cached + return cached + + def cached_tokens(self, handle: RequestHandle) -> int: + """Tokens this request did not have to prefill. Zero without a hit.""" + return self._cached_tokens.get(handle.request_id, 0) + def prepare_step( self, handles: Sequence[RequestHandle], @@ -150,7 +171,7 @@ def handle_for_slot(self, slot: int) -> RequestHandle | None: def handle_for_request(self, request_id: str) -> RequestHandle | None: return self._by_request_id.get(str(request_id)) - def release(self, handle: RequestHandle) -> np.ndarray: + def release(self, handle: RequestHandle, token_ids=None) -> np.ndarray: """Reclaim everything the request holds. The canonical API. Idempotent at this level: an untracked handle is a no-op rather than an @@ -159,18 +180,25 @@ def release(self, handle: RequestHandle) -> np.ndarray: idempotent, and that is the right split -- a double release through a handle the adapter still knows about is a bug worth reporting, while a release of something already forgotten is just a duplicate notification. + + `token_ids` is the full context, offered to the prefix cache. The recorded + sequence length bounds what may be published, since that is exactly the + extent whose K/V has been written; the final sampled token is in the list but + its own K/V never got computed. """ known = self._by_request_id.get(handle.request_id) if known is None or known != handle: return np.empty((0,), dtype=np.int32) - if self.poison_on_free: - pages = self.control_plane.page_map.pages(handle) + valid = self.control_plane.page_map.seq_len(handle) + freed = self.control_plane.release(handle, token_ids, num_valid_tokens=valid) + # Poisoning what came free rather than what the request held: with sharing on + # they differ, and a page the cache adopted is about to be read as valid K/V. + if self.poison_on_free and freed.size: for layer in range(self.pool.num_layers): - k, v = poison_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], pages) + k, v = poison_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], freed) self.pool.replace_layer(layer, k, v) - pages = self.control_plane.release(handle) self._forget(handle) - return pages + return freed def release_slot(self, slot: int) -> np.ndarray: """Shim for `release_pages(slot)`.""" @@ -181,6 +209,7 @@ def release_slot(self, slot: int) -> np.ndarray: def _forget(self, handle: RequestHandle) -> None: self._by_request_id.pop(handle.request_id, None) + self._cached_tokens.pop(handle.request_id, None) for slot, tracked in list(self._by_slot.items()): if tracked == handle: del self._by_slot[slot] diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index bfb4e5d3a6..f6addd4c01 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -1881,7 +1881,7 @@ def copy(path, partial_cache, full_cache, annotations): "token_logp": inserted_token_logp, } - def release(self, request_handle: Any): + def release(self, request_handle: Any, token_ids=None): """Reclaim the pages a request holds. The canonical release API. Request-based rather than slot-based on purpose. A slot is the dense cache's @@ -1889,13 +1889,20 @@ def release(self, request_handle: Any): paged request owns a set of pages that changed size on every step, so the handle is its only durable name. + `token_ids` is the request's full context, offered to the prefix cache when + one is enabled. Pages it adopts stay allocated for the next request with the + same prefix, and are therefore *not* among the ids returned -- the return + value continues to mean "free now, safe to overwrite". Omitting the tokens + publishes nothing, which is the right default for a caller that has not + thought about cache identity. + Returns the page ids reclaimed, or an empty array when no paged runtime is attached. """ if self.paged_runtime is None: max_logging.log("release: paged attention is not configured, so there are no pages to release.") return np.empty((0,), dtype=np.int32) - return self.paged_runtime.release(request_handle) + return self.paged_runtime.release(request_handle, token_ids) def release_pages(self, slot: int): """Compatibility shim over `release`, for callers that only have a slot. @@ -1966,6 +1973,7 @@ def init_paged_runtime( layout=layout, max_requests=max_requests, max_context_len=self.config.paged_max_context_len, + enable_prefix_cache=self.config.paged_enable_prefix_cache, ) planner = StepShapePlanner( tokens_per_page=layout.tokens_per_page, @@ -2132,6 +2140,8 @@ def prefill_paged( request_id: str | None = None, max_new_tokens: int | None = None, slot: int | None = None, + prompt_token_ids=None, + namespace=None, rng: PRNGKeyType | None = None, algorithm: str | None = None, topk: int | None = None, @@ -2144,6 +2154,11 @@ def prefill_paged( dict, because there is no cache to carry: the K/V is already in the pool and the handle is what names the pages holding it. + Passing `prompt_token_ids` opts the request into prefix sharing, when the + control plane has it enabled. The prompt's already-computed leading pages are + then lent to the request and only the remainder is run, which is where the + milestone's saving actually comes from. + Returns: `(handle, result_tokens)`, or `(None, None)` if the pool could not admit the request -- backpressure, not an error. @@ -2162,7 +2177,9 @@ def prefill_paged( if handle is None: return None, None - view = runtime.prepare_step([handle], [prompt_len], is_decode=False, num_requests=1) + cached = runtime.attach_prefix(handle, prompt_token_ids, namespace) + query_len = prompt_len - cached + view = runtime.prepare_step([handle], [query_len], is_decode=False, num_requests=1) if view is None: runtime.control_plane.release(handle) return None, None @@ -2178,15 +2195,21 @@ def prefill_paged( # can see, which is exactly how an earlier measurement here ended up # three-quarters compile time while reporting zero unwarmed shapes. numpy has # no such cache to miss, and these arrays are a few hundred bytes. + # + # With a prefix hit the query is the prompt's *suffix*, so the tokens are + # sliced from `cached` and the positions start there too. Absolute positions + # are what RoPE encodes, so restarting them at zero would rotate the suffix + # as though it began the sequence and produce K/V that does not belong after + # the cached prefix -- a wrong answer rather than a slow one. width = view.shape.num_tokens - prompt = np.asarray(padded_tokens, dtype=np.int32).reshape(-1)[:width] + prompt = np.asarray(padded_tokens, dtype=np.int32).reshape(-1)[cached : cached + width] token_row = np.zeros((1, width), np.int32) token_row[0, : prompt.size] = prompt index = np.arange(width, dtype=np.int32) tokens = jnp.asarray(token_row) - positions = jnp.asarray(index[None, :]) + positions = jnp.asarray((index + cached)[None, :]) segment_ids = jnp.asarray( - np.where(index < prompt_len, DECODING_ACTIVE_SEQUENCE_INDICATOR, 0).astype(np.int32)[None, :] + np.where(index < query_len, DECODING_ACTIVE_SEQUENCE_INDICATOR, 0).astype(np.int32)[None, :] ) if rng is None: @@ -2205,7 +2228,9 @@ def prefill_paged( kv_page_indices=view.kv_page_indices, kv_last_page_lens=view.kv_last_page_lens, cu_seqlens_q=view.cu_seqlens_q, - sample_at=jnp.asarray([prompt_len - 1], jnp.int32), + # Indexed within the query window, not the sequence: the logits this step + # produces cover only the tokens it ran. + sample_at=jnp.asarray([query_len - 1], jnp.int32), rng=rng, model_mode=MODEL_MODE_PREFILL, max_seqlen_k=view.shape.max_seqlen_k, diff --git a/tests/unit/gpu_paged_prefix_cache_test.py b/tests/unit/gpu_paged_prefix_cache_test.py new file mode 100644 index 0000000000..ebae929bb0 --- /dev/null +++ b/tests/unit/gpu_paged_prefix_cache_test.py @@ -0,0 +1,246 @@ +"""M5's exit criterion: a prefix cache hit must change the cost, not the answer. + +The host-side tests in `kv_prefix_cache_test.py` prove that the right pages are +shared, that a namespace mismatch cannot hit, and that nothing leaks. None of +them can catch the failure that matters most here, because it is arithmetic +inside the model rather than bookkeeping around it. + +**Position offset.** After a hit the step runs the prompt's *suffix*, and those +tokens sit at absolute positions `cached..prompt_len`. RoPE encodes absolute +position, so feeding the suffix at positions starting from zero produces K/V +rotated as though the suffix began the sequence. The pages would be laid out +correctly, the page table would be correct, nothing would leak, and the output +would be wrong. That is invisible to every host-side test and is precisely what +this file checks. + +The claim is a comparison against the same engine with nothing cached: a warm +rollout must produce the identical token sequence to a cold one. Identical, not +close -- the two differ only in how much of the prompt was recomputed, so any +divergence at all is a bug rather than numerics. That also makes the test immune +to the tie-breaking flakiness the parity test has to work around, since both +sides here run the same kernels on the same pool. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import sys +import unittest + +from absl.testing import parameterized +import numpy as np +import pytest + +import jax +import jax.numpy as jnp +from flax import nnx +from flax.linen import partitioning as nn_partitioning + +from maxtext.common.common_types import MODEL_MODE_PREFILL +from maxtext.configs import pyconfig +from maxtext.inference.kv_common import CacheNamespace +from maxtext.utils import maxtext_utils, model_creation_utils + +try: + from maxtext.inference.maxengine import maxengine +except ModuleNotFoundError as _exc: # pragma: no cover - environment dependent + pytest.skip( + f"MaxEngine needs JetStream ({_exc}). Install google-jetstream under a constraints file pinning " + f"jax and jaxlib, or run with DECOUPLE_GCLOUD=TRUE to use the built-in stubs.", + allow_module_level=True, + ) + +from tests.utils.test_helpers import get_test_config_path # pylint: disable=wrong-import-position + +PAGE = 16 +STEPS = 12 + +# 40 tokens is two and a half pages, so a hit shares whole pages and leaves a +# partial one to recompute. A prompt of one page or less would make the +# publication trivial and the suffix empty, and the position offset -- the thing +# most worth testing -- would never be exercised. +PROMPT = [3, 17, 42, 5, 9, 21, 33, 2, 11, 27] * 4 + +_COMMON = { + "base_emb_dim": 512, + "base_mlp_dim": 512, + "base_num_query_heads": 4, + "base_num_kv_heads": 4, + "base_num_decoder_layers": 2, + "head_dim": 128, + "vocab_size": 256, + "max_prefill_predict_length": 64, + "max_target_length": 128, + "per_device_batch_size": 1, + "scan_layers": False, + "sparse_matmul": False, + "dtype": "bfloat16", + "weight_dtype": "float32", + "matmul_precision": "highest", + "decode_sampling_strategy": "greedy", + "enable_checkpointing": False, + "skip_jax_distributed_system": True, + "pure_nnx": True, +} + +_PAGED = { + "attention": "gpu_paged", + "paged_page_size": PAGE, + "paged_num_blocks": 64, + "paged_enable_prefix_cache": True, +} + +NAMESPACE = CacheNamespace(model_fingerprint="test-random-weights", tokenizer="test") + + +def _require_kernels(): + """Skip unless jax-aiter is importable and its KV shims are built.""" + try: + from jax_aiter.ffi.registry import standalone_symbol_available # pylint: disable=import-outside-toplevel + except ImportError as exc: + raise unittest.SkipTest("jax-aiter is not importable; set PYTHONPATH to the jax-aiter checkout") from exc + for symbol in ("AppendKvJA", "PagedAttentionJA", "PagedPrefillJA"): + if not standalone_symbol_available(symbol): + raise unittest.SkipTest(f"{symbol} is not built; run 'make -f Makefile.kv ja_kv' and set JA_ROOT_DIR") + + +def _config(**overrides): + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **(_COMMON | overrides)) + + +def _devices(): + """One device, for the same container reason as the parity test.""" + return jax.devices()[:1] + + +def _mesh(cfg): + return jax.sharding.Mesh(maxtext_utils.create_device_mesh(config=cfg, devices=_devices()), cfg.mesh_axes) + + +@pytest.mark.gpu_only +class GpuPagedPrefixCacheTest(parameterized.TestCase): + """Sharing a prefix must not change a single token.""" + + def setUp(self): + super().setUp() + _require_kernels() + self.cfg = _config(**_PAGED) + mesh = _mesh(self.cfg) + with nn_partitioning.axis_rules(self.cfg.logical_axis_rules), mesh: + model = model_creation_utils.create_model( + self.cfg, mesh, model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(params=0, dropout=0) + ) + _, self.params_state, _ = nnx.split(model, nnx.Param, ...) + + def _engine(self): + engine = maxengine.MaxEngine(self.cfg, _devices()) + params = engine.load_params(params=self.params_state) + runtime = engine.init_paged_runtime() + return engine, params, runtime + + def _rollout(self, engine, params, request_id, prompt=None, namespace=NAMESPACE, steps=STEPS): + """One greedy rollout, offering its context to the cache on the way out.""" + tokens = np.asarray(prompt if prompt is not None else PROMPT, dtype=np.int64) + padded = jnp.asarray( + list(tokens) + [0] * (self.cfg.max_prefill_predict_length - tokens.size), dtype=jnp.int32 + ) + handle, first = engine.prefill_paged( + params=params, + padded_tokens=padded, + true_length=int(tokens.size), + request_id=request_id, + prompt_token_ids=tokens, + namespace=namespace, + ) + self.assertIsNotNone(handle, "the pool refused a single request, so it is mis-sized for this test") + cached = engine.paged_runtime.cached_tokens(handle) + + generated = [int(first.data[0, 0])] + for _ in range(steps): + result, ok = engine.generate_paged( + params, [handle], next_tokens=jnp.asarray([generated[-1]], jnp.int32) + ) + self.assertTrue(ok, "the pool ran out of pages during a single-request decode") + generated.append(int(result.data[0, 0])) + + context = np.concatenate([tokens, np.asarray(generated, dtype=np.int64)]) + engine.release(handle, context) + return generated, cached + + def test_a_warm_rollout_produces_the_same_tokens_as_a_cold_one(self): + """The exit criterion. A wrong position offset fails here and nowhere else. + + Both rollouts run in one engine so they share weights, kernels and pool + exactly. The only difference is that the second one finds its prefix already + computed, so an identical token sequence is the whole claim. + """ + engine, params, runtime = self._engine() + + cold, cached_cold = self._rollout(engine, params, "cold") + self.assertEqual(cached_cold, 0, "nothing was cached yet, so this rollout must have paid in full") + self.assertGreater(runtime.control_plane.prefix_index.num_cached_pages, 0, "nothing was published") + + warm, cached_warm = self._rollout(engine, params, "warm") + self.assertGreater(cached_warm, 0, "the repeated prompt did not hit the cache, so this proved nothing") + self.assertEqual( + cold, + warm, + f"a cache hit changed the output. {cached_warm} of {len(PROMPT)} prompt tokens were served from " + f"the cache, and the suffix must be run at absolute positions {cached_warm}.. for RoPE to agree:" + f"\n cold = {cold}\n warm = {warm}", + ) + + def test_a_namespace_mismatch_recomputes_and_still_agrees(self): + """Two guarantees at once: the miss is real, and the miss path is unchanged.""" + engine, params, _ = self._engine() + cold, _ = self._rollout(engine, params, "cold") + + other = CacheNamespace(model_fingerprint="test-random-weights", tokenizer="different") + miss, cached = self._rollout(engine, params, "other-namespace", namespace=other) + self.assertEqual(cached, 0, "a different tokenizer must not share pages") + self.assertEqual(cold, miss) + + def test_a_shared_prefix_with_a_different_tail_agrees_with_a_cold_run(self): + """The realistic case: a common system prompt and a divergent question. + + Stronger than the repeated-prompt test, because here the cache supplies part + of the context and the step computes a genuine suffix, so the boundary + between borrowed and freshly written pages falls mid-request. + """ + engine, params, _ = self._engine() + tail = [101, 202, 303, 44, 55, 66, 77, 88] + + baseline, cached_baseline = self._rollout(engine, params, "baseline", prompt=PROMPT + tail) + self.assertEqual(cached_baseline, 0) + + self._rollout(engine, params, "publisher") + shared, cached_shared = self._rollout(engine, params, "sharer", prompt=PROMPT + tail) + + self.assertGreater(cached_shared, 0, "the common prefix was not shared") + self.assertLess(cached_shared, len(PROMPT) + len(tail), "the divergent tail must still be computed") + self.assertEqual(baseline, shared, "sharing part of the context changed the output") + + def test_the_pool_is_accounted_for_once_the_cache_is_dropped(self): + engine, params, runtime = self._engine() + self._rollout(engine, params, "r0") + self._rollout(engine, params, "r1") + + plane = runtime.control_plane + self.assertEqual(plane.allocator.num_allocated_pages, plane.prefix_index.num_cached_pages) + plane.evict_cached(plane.prefix_index.num_cached_pages) + self.assertEqual(plane.allocator.num_allocated_pages, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/kv_execution_test.py b/tests/unit/kv_execution_test.py index 1847490a07..d90ec342f8 100644 --- a/tests/unit/kv_execution_test.py +++ b/tests/unit/kv_execution_test.py @@ -29,7 +29,7 @@ import numpy as np -from maxtext.inference.kv_common import KvPageTableV1, KvStorageLayoutV1 +from maxtext.inference.kv_common import CacheNamespace, KvPageTableV1, KvStorageLayoutV1 from maxtext.inference.kv_control import NativeKvControlPlane, RequestHandle from maxtext.inference.kv_execution import allocate_pool, PagedDriver, PagedRequest from maxtext.inference.kv_execution.bucketing import ( @@ -461,6 +461,167 @@ def test_an_empty_driver_does_nothing(self): self.assertEqual(self._driver().run(), []) +class DriverPrefixCacheTest(unittest.TestCase): + """The driver's use of the prefix index: fewer prefill tokens, no lost pages. + + The index itself is covered in `kv_prefix_cache_test.py`. What is at stake here + is the wiring -- that the saving reaches the step's query length, and that the + two page-lifetime hazards sharing introduces are handled. + """ + + def _driver(self, *, num_pages=65, max_requests=4, max_context_len=128, step=None, **kw): + layout = _layout(num_pages=num_pages) + plane = NativeKvControlPlane( + layout=layout, + max_requests=max_requests, + max_context_len=max_context_len, + debug_mode=True, + enable_prefix_cache=True, + ) + return PagedDriver(plane, allocate_pool(layout), step or _CountingStep(), max_batch=max_requests, **kw) + + def _request(self, request_id, prompt, max_new_tokens=2): + return PagedRequest( + request_id=request_id, + prompt_len=len(prompt), + max_new_tokens=max_new_tokens, + prompt_tokens=np.asarray(prompt, dtype=np.int64), + ) + + def test_a_repeated_prompt_prefills_fewer_tokens(self): + """The milestone's whole point, measured as tokens the step actually computes. + + Not as the bucketed shape: the token ladder floors at 64, so a saving smaller + than that is real but invisible there. The prompt here is long enough that + both figures move, and both are checked -- the token count because it is the + work avoided, the bucket because it is what the work costs. + """ + step = _CountingStep() + driver = self._driver(step=step, max_requests=1, max_context_len=512) + prompt = list(range(256)) + + driver.submit([self._request("first", prompt)]) + first = driver.step() + driver.run() + first_bucket = step.views[0].shape.num_tokens + + step.views.clear() + driver.submit([self._request("second", prompt)]) + second = driver.step() + + self.assertEqual(first.num_tokens, 256) + self.assertEqual(second.num_tokens, 16, "only the held-back final page should be recomputed") + self.assertLess(step.views[0].shape.num_tokens, first_bucket) + + def test_a_diverging_prompt_shares_only_what_it_has_in_common(self): + driver = self._driver(max_requests=1, max_context_len=512) + shared = list(range(128)) + + driver.submit([self._request("first", shared + list(range(900, 964)))]) + driver.run() + + driver.submit([self._request("second", shared + list(range(500, 564)))]) + outcome = driver.step() + self.assertEqual( + outcome.num_tokens, 64, "exactly the divergent tail should be computed, and the shared 128 skipped" + ) + + def test_a_request_without_tokens_is_unaffected(self): + """Supplying prompt ids is opt-in, so the cache must tolerate their absence.""" + driver = self._driver() + driver.submit([PagedRequest(request_id=f"r{i}", prompt_len=20, max_new_tokens=4) for i in range(3)]) + done = driver.run() + self.assertEqual(len(done), 3) + self.assertEqual(driver.plane.prefix_index.num_cached_pages, 0) + self.assertEqual(driver.plane.allocator.num_allocated_pages, 0) + + def test_only_pages_with_computed_kv_are_published(self): + """The final generated token has no K/V: its step never ran. + + Publishing it would cache a page whose tail is whatever the scrub left, and + the next request to match that prefix would attend over zeros as though they + were real keys. + """ + driver = self._driver(max_requests=1) + prompt = list(range(16)) + driver.submit([self._request("r0", prompt, max_new_tokens=5)]) + done = driver.run() + + written = len(prompt) + len(done[0].generated) - 1 + self.assertLessEqual( + driver.plane.prefix_index.num_cached_pages * PAGE, + written, + "published more tokens than had their K/V computed", + ) + + def test_nothing_leaks_once_the_cache_is_dropped(self): + driver = self._driver(max_requests=2) + prompts = [list(range(i * 100, i * 100 + 48)) for i in range(6)] + driver.submit([self._request(f"r{i}", p, max_new_tokens=3) for i, p in enumerate(prompts)]) + driver.run() + + driver.plane.evict_cached(driver.plane.prefix_index.num_cached_pages) + self.assertEqual(driver.plane.prefix_index.num_cached_pages, 0) + self.assertEqual(driver.plane.allocator.num_allocated_pages, 0) + self.assertEqual(driver.plane.allocator.available_pages, driver.plane.allocator.capacity_pages) + + def test_the_cache_does_not_stop_a_tight_pool_making_progress(self): + """Retained pages must be reclaimable, or sharing turns into a deadlock.""" + driver = self._driver(num_pages=9, max_requests=4, max_context_len=64) + driver.submit( + [self._request(f"r{i}", list(range(i * 100, i * 100 + 30)), max_new_tokens=25) for i in range(5)] + ) + done = driver.run() + self.assertEqual(len(done), 5) + self.assertTrue(all(len(r.generated) == 25 for r in done)) + + def test_poison_on_free_does_not_destroy_a_retained_page(self): + """Poison must follow what was freed, not what the request held.""" + driver = self._driver(max_requests=1, poison_on_free=True) + prompt = list(range(64)) + driver.submit([self._request("first", prompt)]) + driver.run() + + cached = driver.plane.prefix_index.match(prompt, CacheNamespace()).pages + self.assertGreater(cached.size, 0) + k = np.asarray(driver.pool.k_pages[0])[cached] + self.assertFalse(np.any(k == POISON_SENTINEL), "poisoned a page the prefix cache had adopted") + + def test_a_preempted_request_does_not_retain_its_pages(self): + """Preemption exists to free pages; publishing them would work against it. + + Checked at the first preemption rather than at the end of the run, because by + then a completed request has published legitimately and the two sources of + cached pages are no longer distinguishable. + """ + driver = self._driver(num_pages=9, max_requests=4, max_context_len=64) + requests = [ + self._request(f"r{i}", list(range(i * 100, i * 100 + 30)), max_new_tokens=25) for i in range(5) + ] + driver.submit(requests) + while not any(r.preemptions for r in requests): + if driver.step() is None: + self.fail("this pool is too small not to preempt") + + self.assertFalse(any(r.is_finished for r in requests), "no request should have completed this early") + self.assertEqual( + driver.plane.prefix_index.num_cached_pages, + 0, + "a preempted request published its pages, so the preemption reclaimed fewer than it should", + ) + + def test_the_only_pages_left_allocated_are_the_ones_the_cache_holds(self): + """The leak invariant, restated for a runtime that deliberately retains pages.""" + driver = self._driver(num_pages=9, max_requests=4, max_context_len=64) + driver.submit( + [self._request(f"r{i}", list(range(i * 100, i * 100 + 30)), max_new_tokens=25) for i in range(5)] + ) + driver.run() + self.assertEqual( + driver.plane.allocator.num_allocated_pages, driver.plane.prefix_index.num_cached_pages + ) + + class PagedRuntimeTest(unittest.TestCase): """The engine adapter, and the slot shim over the request-based API.""" diff --git a/tests/unit/kv_import_rule_test.py b/tests/unit/kv_import_rule_test.py index 9bd65e87d8..00e792f8c5 100644 --- a/tests/unit/kv_import_rule_test.py +++ b/tests/unit/kv_import_rule_test.py @@ -58,6 +58,7 @@ "dataclasses", "enum", "hashlib", + "heapq", "math", "numpy", "typing", diff --git a/tests/unit/kv_prefix_cache_test.py b/tests/unit/kv_prefix_cache_test.py new file mode 100644 index 0000000000..6119d9c05b --- /dev/null +++ b/tests/unit/kv_prefix_cache_test.py @@ -0,0 +1,454 @@ +"""Prefix sharing: the index, the namespace, and their use by the control plane. + +Host-only and CPU-only, like the rest of `kv_control`. Nothing here allocates a +device array, so the whole file runs in under a second and can be a pre-commit +check rather than a nightly one. + +The tests that matter most are the ones asserting a *negative*: that a namespace +mismatch cannot hit, that a shared page is never written, and that a page a live +request is reading cannot be evicted. Those are the failures which produce +plausible tokens instead of a crash, so they are the ones worth spending test +surface on. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import dataclasses +import unittest + +import numpy as np + +from maxtext.inference.kv_common import CacheNamespace, KvStorageLayoutV1 +from maxtext.inference.kv_control import ( + NativeKvControlPlane, + PrefixIndex, + RequestDescriptor, + SharedPageWriteError, + block_hash, +) + +PAGE = 4 +NS = CacheNamespace(model_fingerprint="sha256:abc", tokenizer="llama3") + + +def _layout(**kw) -> KvStorageLayoutV1: + base = { + "tokens_per_page": PAGE, + "num_pages": 64, + "num_layers": 4, + "num_kv_heads": 8, + "head_dim": 128, + "dtype": "bfloat16", + } + base.update(kw) + return KvStorageLayoutV1(**base) + + +def _tokens(n, start=0): + return list(range(start, start + n)) + + +class CacheNamespaceTest(unittest.TestCase): + """The identity that decides whether a hit is sound.""" + + def test_the_digest_covers_every_field_without_being_told_them(self): + """The property that survives someone adding a field and forgetting this test. + + Enumerating the dataclass rather than hard-coding names is what makes the + digest exhaustive by construction. This asserts the mechanism, so a field + added later is covered whether or not anyone writes a test for it. + """ + baseline = CacheNamespace() + for field in dataclasses.fields(baseline): + value = 99 if field.name == "version" else f"changed-{field.name}" + varied = dataclasses.replace(baseline, **{field.name: value}) + self.assertNotEqual( + baseline.digest(), + varied.digest(), + f"changing {field.name} alone left the digest unchanged, so two configurations that " + f"differ only in {field.name} would share cached K/V", + ) + + def test_the_digest_is_stable_across_instances(self): + self.assertEqual(CacheNamespace(tenant="a").digest(), CacheNamespace(tenant="a").digest()) + + def test_field_names_are_hashed_so_a_value_cannot_migrate_between_fields(self): + """Two fields holding each other's values must not digest alike.""" + swapped = CacheNamespace(adapter="x", tenant="y") + other = CacheNamespace(adapter="y", tenant="x") + self.assertNotEqual(swapped.digest(), other.digest()) + + def test_describe_names_the_fields_that_are_set(self): + text = CacheNamespace(tenant="acme", adapter="lora-7").describe() + self.assertIn("tenant=acme", text) + self.assertIn("adapter=lora-7", text) + self.assertNotIn("tokenizer", text) + + +class BlockHashTest(unittest.TestCase): + """The chain that makes a page's key mean "after exactly this history".""" + + def test_the_same_page_after_a_different_prefix_hashes_differently(self): + page = [7, 8, 9, 10] + self.assertNotEqual(block_hash(b"parent-a", page), block_hash(b"parent-b", page)) + + def test_the_chain_is_order_sensitive(self): + self.assertNotEqual( + block_hash(block_hash(b"", [1, 2]), [3, 4]), + block_hash(block_hash(b"", [3, 4]), [1, 2]), + ) + + +class PrefixIndexTest(unittest.TestCase): + """Matching, publication, refcounts and eviction, without a control plane.""" + + def _index(self, enabled=True): + return PrefixIndex(tokens_per_page=PAGE, enabled=enabled) + + def test_a_cold_index_matches_nothing(self): + self.assertFalse(self._index().match(_tokens(16), NS)) + + def test_publish_then_match_returns_the_same_pages(self): + index = self._index() + tokens = _tokens(16) + result = index.publish(tokens, [10, 11, 12, 13], NS) + self.assertEqual(result.adopted.tolist(), [10, 11, 12, 13]) + + match = index.match(tokens, NS) + # The last page is held back so the request has something left to compute. + self.assertEqual(match.pages.tolist(), [10, 11, 12]) + self.assertEqual(match.num_tokens, 12) + + def test_a_fully_cached_prompt_still_has_a_page_left_to_compute(self): + """Otherwise the step has no query tokens and nothing to predict from. + + The held-back page is the difference between "this prompt is cached" and + "this prompt needs no work", and only the first of those is true. + """ + index = self._index() + tokens = _tokens(16) + index.publish(tokens, [10, 11, 12, 13], NS) + match = index.match(tokens, NS) + self.assertEqual(match.num_tokens, 12) + self.assertLess(match.num_tokens, len(tokens)) + + def test_a_partial_page_is_never_published(self): + """Sharing a page still being appended to would give two writers one page.""" + index = self._index() + result = index.publish(_tokens(10), [10, 11, 12], NS) + self.assertEqual(result.adopted.tolist(), [10, 11]) + + def test_only_computed_tokens_are_published(self): + index = self._index() + result = index.publish(_tokens(16), [10, 11, 12, 13], NS, num_valid_tokens=9) + self.assertEqual(result.adopted.tolist(), [10, 11]) + + def test_a_longer_prompt_matches_the_shared_prefix_and_stops(self): + index = self._index() + index.publish(_tokens(16), [10, 11, 12, 13], NS) + match = index.match(_tokens(12) + [900, 901, 902, 903] + _tokens(4), NS) + self.assertEqual(match.pages.tolist(), [10, 11, 12]) + + def test_a_divergent_prompt_matches_only_up_to_the_divergence(self): + index = self._index() + index.publish(_tokens(16), [10, 11, 12, 13], NS) + diverged = _tokens(8) + [500, 501, 502, 503] + _tokens(4, 12) + self.assertEqual(index.match(diverged, NS).pages.tolist(), [10, 11]) + + def test_a_second_request_donating_the_same_content_keeps_the_original(self): + index = self._index() + index.publish(_tokens(16), [10, 11, 12, 13], NS) + again = index.publish(_tokens(16), [20, 21, 22, 23], NS) + self.assertEqual(again.adopted.size, 0) + self.assertEqual(again.duplicate.tolist(), [20, 21, 22, 23]) + self.assertEqual(index.match(_tokens(16), NS).pages.tolist(), [10, 11, 12]) + + def test_a_branch_shares_its_common_prefix(self): + index = self._index() + index.publish(_tokens(8) + _tokens(8, 100), [10, 11, 12, 13], NS) + index.publish(_tokens(8) + _tokens(8, 200), [10, 11, 30, 31], NS) + self.assertEqual(index.num_cached_pages, 6) + + def test_the_index_is_inert_when_disabled(self): + index = self._index(enabled=False) + self.assertEqual(index.publish(_tokens(16), [10, 11, 12, 13], NS).adopted.size, 0) + self.assertFalse(index.match(_tokens(16), NS)) + + def test_publishing_more_tokens_than_pages_is_refused(self): + with self.assertRaises(ValueError): + self._index().publish(_tokens(16), [10, 11], NS) + + +class NamespaceIsolationTest(unittest.TestCase): + """Varying one namespace field alone must defeat the match. + + The plan asks for a negative test per field, and generating them from the + dataclass rather than writing twelve near-identical methods means a field + added later is covered on the day it is added. + """ + + def test_every_field_in_isolation_defeats_a_hit(self): + tokens = _tokens(16) + for field in dataclasses.fields(CacheNamespace): + with self.subTest(field=field.name): + index = PrefixIndex(tokens_per_page=PAGE) + index.publish(tokens, [10, 11, 12, 13], NS) + self.assertTrue(index.match(tokens, NS), "the control case should hit") + + value = 99 if field.name == "version" else f"other-{field.name}" + varied = dataclasses.replace(NS, **{field.name: value}) + self.assertFalse( + index.match(tokens, varied), + f"a request differing only in {field.name} was served another configuration's K/V", + ) + + def test_two_namespaces_coexist_without_evicting_each_other(self): + index = PrefixIndex(tokens_per_page=PAGE) + other = dataclasses.replace(NS, tenant="second") + index.publish(_tokens(16), [10, 11, 12, 13], NS) + index.publish(_tokens(16), [20, 21, 22, 23], other) + self.assertEqual(index.num_cached_pages, 8) + self.assertEqual(index.match(_tokens(16), NS).pages.tolist(), [10, 11, 12]) + self.assertEqual(index.match(_tokens(16), other).pages.tolist(), [20, 21, 22]) + + +class RefCountAndEvictionTest(unittest.TestCase): + """What may be dropped, and what may not be dropped at any pressure.""" + + def _loaded(self): + """Two unrelated three-page sequences: pages 10-12 and 20-22.""" + index = PrefixIndex(tokens_per_page=PAGE) + index.publish(_tokens(12), [10, 11, 12], NS) + index.publish(_tokens(12, 100), [20, 21, 22], NS) + return index + + def test_eviction_follows_recency_across_unrelated_sequences(self): + """Using one sequence should cost the other its pages. + + Tail pages go first in both orderings, because a tail is never part of a + match -- it is held back by design, so it really is the least useful page in + the index rather than merely the coldest by accident. + """ + index = self._loaded() + index.match(_tokens(12), NS) + self.assertEqual(index.evict(2).tolist(), [12, 22]) + + def test_pressure_drains_the_untouched_sequence_in_full_first(self): + """A page orphaned by eviction keeps its own recency, not its child's. + + So the cold sequence is reclaimed end to end -- deepest page first, since a + parent only becomes eligible once its children are gone -- before the + recently matched one gives up anything. + """ + index = self._loaded() + index.match(_tokens(12, 100), NS) + self.assertEqual(index.evict(3).tolist(), [12, 11, 10]) + self.assertEqual(index.match(_tokens(12, 100), NS).num_pages, 2) + + def test_eviction_never_orphans_a_prefix(self): + """A parent is only reachable for eviction once its children are gone.""" + index = PrefixIndex(tokens_per_page=PAGE) + index.publish(_tokens(12), [10, 11, 12], NS) + self.assertEqual(index.evict(2).tolist(), [12, 11]) + self.assertEqual(index.num_cached_pages, 1) + + def test_a_referenced_path_survives_maximum_pressure(self): + index = self._loaded() + index.acquire(index.match(_tokens(12), NS).node) + self.assertEqual(index.protected_pages, 2) + self.assertEqual(sorted(index.evict(100).tolist()), [12, 20, 21, 22]) + self.assertEqual(index.num_cached_pages, 2) + + def test_releasing_makes_the_path_evictable_again(self): + index = self._loaded() + node = index.match(_tokens(12), NS).node + index.acquire(node) + index.release(node) + self.assertEqual(index.protected_pages, 0) + self.assertEqual(index.evict(100).size, 6) + + def test_over_release_is_an_error_rather_than_a_silent_unprotect(self): + index = self._loaded() + node = index.match(_tokens(12), NS).node + index.acquire(node) + index.release(node) + with self.assertRaises(RuntimeError): + index.release(node) + + def test_two_readers_of_one_prefix_both_have_to_leave(self): + index = self._loaded() + node = index.match(_tokens(12), NS).node + index.acquire(node) + index.acquire(node) + index.release(node) + self.assertEqual(index.protected_pages, 2) + self.assertEqual(index.num_cached_pages - index.evict(100).size, 2) + + def test_reset_refuses_while_a_request_still_reads(self): + index = self._loaded() + index.acquire(index.match(_tokens(12), NS).node) + with self.assertRaises(RuntimeError): + index.reset() + + def test_reset_hands_back_every_page(self): + index = self._loaded() + self.assertEqual(sorted(index.reset().tolist()), [10, 11, 12, 20, 21, 22]) + self.assertEqual(index.num_cached_pages, 0) + + +class ControlPlaneWithPrefixCacheTest(unittest.TestCase): + """The end-to-end property: a repeated prompt does less work.""" + + def _plane(self, num_pages=32, max_context_len=64, **kw): + return NativeKvControlPlane( + layout=_layout(num_pages=num_pages), + max_requests=4, + max_context_len=max_context_len, + debug_mode=True, + enable_prefix_cache=True, + **kw, + ) + + def _run(self, plane, request_id, tokens, namespace=NS): + """Admit, attach any cached prefix, prefill the rest, release. + + The whole point of the milestone in six lines: the step's query length is the + prompt minus whatever the cache supplied, and everything downstream follows + from that one subtraction. + """ + handle = plane.admit( + RequestDescriptor(request_id=request_id, prompt_len=len(tokens), max_new_tokens=0) + ) + match = plane.attach_prefix(handle, tokens, namespace) + to_prefill = len(tokens) - match.num_tokens + self.assertTrue(plane.reserve([handle], [to_prefill])) + plane.confirm_scrubbed(plane.pending_scrub()) + table = plane.build_page_table([handle], [to_prefill]) + plane.release(handle, tokens) + return to_prefill, table + + def test_a_repeated_prompt_skips_the_prefill_it_already_paid_for(self): + plane = self._plane() + tokens = _tokens(16) + first, _ = self._run(plane, "r0", tokens) + second, _ = self._run(plane, "r1", tokens) + self.assertEqual(first, 16) + self.assertEqual(second, 4, "the cached 12 tokens should not have been prefilled again") + + def test_a_shared_page_never_appears_in_a_slot_mapping(self): + """The M5 exit criterion, checked against the array a kernel actually writes.""" + plane = self._plane() + tokens = _tokens(16) + self._run(plane, "r0", tokens) + + handle = plane.admit(RequestDescriptor(request_id="r1", prompt_len=16, max_new_tokens=0)) + match = plane.attach_prefix(handle, tokens, NS) + self.assertEqual(match.num_pages, 3) + self.assertTrue(plane.reserve([handle], [16 - match.num_tokens])) + plane.confirm_scrubbed(plane.pending_scrub()) + table = plane.build_page_table([handle], [16 - match.num_tokens]) + + slots = table.slot_mapping(PAGE) + written = set((slots[slots >= 0] // PAGE).tolist()) + self.assertTrue( + written.isdisjoint(set(match.pages.tolist())), + f"step wrote into shared pages {written & set(match.pages.tolist())}", + ) + self.assertEqual(set(match.pages.tolist()) - set(table.flat_page_indices().tolist()), set()) + + def test_the_debug_gate_catches_a_query_length_that_would_overwrite_a_prefix(self): + """A wrong query length is the only way a shared page can be written.""" + plane = self._plane() + tokens = _tokens(16) + self._run(plane, "r0", tokens) + + handle = plane.admit(RequestDescriptor(request_id="r1", prompt_len=16, max_new_tokens=0)) + plane.attach_prefix(handle, tokens, NS) + self.assertTrue(plane.reserve([handle], [4])) + plane.confirm_scrubbed(plane.pending_scrub()) + with self.assertRaises(SharedPageWriteError): + # Claiming the whole prompt as this step's query, as a caller would if it + # ignored the match, walks write positions back over the shared pages. + plane.build_page_table([handle], [16]) + + def test_a_different_namespace_pays_full_price(self): + plane = self._plane() + tokens = _tokens(16) + self._run(plane, "r0", tokens) + other, _ = self._run(plane, "r1", tokens, namespace=dataclasses.replace(NS, adapter="lora-2")) + self.assertEqual(other, 16) + + def test_cached_pages_stay_allocated_after_the_request_leaves(self): + plane = self._plane() + before = plane.allocator.available_pages + self._run(plane, "r0", _tokens(16)) + self.assertEqual(plane.prefix_index.num_cached_pages, 4) + self.assertEqual(plane.allocator.available_pages, before - 4) + + def test_pressure_evicts_the_cache_rather_than_refusing_to_serve(self): + """A cache must never be the reason a request cannot run.""" + plane = self._plane(num_pages=10, max_context_len=32) + self._run(plane, "r0", _tokens(16)) + self.assertEqual(plane.prefix_index.num_cached_pages, 4) + + handle = plane.admit(RequestDescriptor(request_id="big", prompt_len=32, max_new_tokens=0)) + self.assertTrue(plane.reserve([handle], [32])) + self.assertLess(plane.prefix_index.num_cached_pages, 4) + + def test_a_borrowed_page_is_not_freed_when_the_borrower_leaves(self): + plane = self._plane() + tokens = _tokens(16) + self._run(plane, "r0", tokens) + cached = plane.prefix_index.match(tokens, NS).pages + + handle = plane.admit(RequestDescriptor(request_id="r1", prompt_len=16, max_new_tokens=0)) + plane.attach_prefix(handle, tokens, NS) + self.assertTrue(plane.reserve([handle], [4])) + freed = plane.release(handle, tokens) + self.assertTrue( + set(freed.tolist()).isdisjoint(set(cached.tolist())), + "released a page the index still owns and other requests will read", + ) + + def test_no_pages_leak_across_repeated_churn(self): + plane = self._plane() + before = plane.allocator.available_pages + for i in range(20): + self._run(plane, f"r{i}", _tokens(16, start=i * 100)) + plane.evict_cached(plane.prefix_index.num_cached_pages) + self.assertEqual(plane.allocator.available_pages, before) + + def test_attaching_a_prefix_to_a_started_request_is_refused(self): + plane = self._plane() + tokens = _tokens(16) + self._run(plane, "r0", tokens) + handle = plane.admit(RequestDescriptor(request_id="r1", prompt_len=16, max_new_tokens=0)) + self.assertTrue(plane.reserve([handle], [8])) + with self.assertRaises(ValueError): + plane.attach_prefix(handle, tokens, NS) + + def test_the_cache_is_off_unless_asked_for(self): + plane = NativeKvControlPlane( + layout=_layout(), max_requests=2, max_context_len=64, debug_mode=True + ) + self.assertFalse(plane.prefix_cache_enabled) + handle = plane.admit(RequestDescriptor(request_id="r0", prompt_len=16, max_new_tokens=0)) + self.assertTrue(plane.reserve([handle], [16])) + freed = plane.release(handle, _tokens(16)) + self.assertEqual(freed.size, 4, "nothing should have been retained by a disabled cache") + + +if __name__ == "__main__": + unittest.main() From 2c00f9f23658a1be773f49e7d8581957692546f4 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Wed, 26 Aug 2026 14:45:56 +0000 Subject: [PATCH 05/21] fix(inference): correct three prefix-cache accounting errors in the benchmark harness Found by running the shared-prefix path through run_serving_benchmark.py for the first time, where it reported 86 leaked pages and a 95.2% prefill saving. Both figures were wrong. pages_leaked counted the pages the prefix cache is deliberately holding. It reported a leak on a run that leaked nothing, and an alarm that fires on every run with sharing enabled is one nobody reads -- which would hide the genuine leak the metric exists to catch. It now subtracts what the index retains and reports that separately. run_repeated now clears the cache between passes. Otherwise a later pass finds the earlier pass's pages waiting for it, which inflates the saving past what the trace's own requests share with each other and makes the passes incomparable, defeating the stability check that is the entire reason for repeating. The 95.2% was almost all cross-pass reuse; the same trace measures 0% once each pass starts cold, which is correct, because that harness admits every request at once and nothing can reuse pages nobody has released yet. The run also went from "NOT reportable" to reportable at a 1.014 spread, which is the change paying for itself. And the occupancy ratio now reads as a sharing dividend when it falls below 1. Pages held can be fewer than the tokens requests collectively address, because several of them are reading the same pages; calling that "page overhead" reports the benefit as a cost. --prefix-cache now says in its help that this harness is the wrong shape to measure sharing and points at run_prefix_cache_benchmark.py, which staggers admission through a batch cap below the request count -- that staggering is what creates anything to share, and is why the dedicated script exists. No change to the headline result: 52.8% of prefill tokens avoided, unchanged across runs because it is arithmetic rather than a measurement. Co-authored-by: Cursor --- benchmarks/paged_kv/run_serving_benchmark.py | 7 +++- .../inference/kv_execution/benchmark.py | 37 ++++++++++++++++--- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/benchmarks/paged_kv/run_serving_benchmark.py b/benchmarks/paged_kv/run_serving_benchmark.py index 970f6d96fe..74c6ab0d13 100644 --- a/benchmarks/paged_kv/run_serving_benchmark.py +++ b/benchmarks/paged_kv/run_serving_benchmark.py @@ -139,7 +139,12 @@ def main() -> int: parser.add_argument( "--prefix-cache", action="store_true", - help="share already-computed pages between requests with a common prefix (paged arm only)", + help=( + "share already-computed pages between requests with a common prefix (paged arm only). " + "Expect little or no saving here: this harness admits the whole trace at once, and a " + "request can only reuse pages another request has already finished with. Use " + "run_prefix_cache_benchmark.py to measure sharing" + ), ) parser.add_argument("--requests", type=int, default=24) parser.add_argument("--mean-prompt", type=int, default=48) diff --git a/src/maxtext/inference/kv_execution/benchmark.py b/src/maxtext/inference/kv_execution/benchmark.py index 17cc8f9c31..e7c5a50089 100644 --- a/src/maxtext/inference/kv_execution/benchmark.py +++ b/src/maxtext/inference/kv_execution/benchmark.py @@ -391,7 +391,10 @@ def run_paged( done.append(request) elapsed = time.perf_counter() - start - summary = _summarise(done, occupancy, elapsed, steps, allocator, page_size) + index = runtime.control_plane.prefix_index + summary = _summarise( + done, occupancy, elapsed, steps, allocator, page_size, retained_pages=index.num_cached_pages + ) summary.update(_summarise_prefix_cache(done, runtime.control_plane)) # Shape accounting, which is necessary but *not sufficient*: it sees bucketed @@ -442,8 +445,15 @@ def run_repeated(engine, params, trace_factory, *, max_batch: int, warmed_shapes materially faster, and the ratio says so. A run is reportable when successive passes agree. """ + plane = engine.paged_runtime.control_plane passes = [] for _ in range(max(repeats, 1)): + # Each pass starts with a cold prefix cache. Otherwise a later pass finds the + # previous pass's pages waiting for it, which both inflates the reported + # saving past what this trace's own requests share with each other and makes + # the passes incomparable -- defeating the stability check, which is the + # whole reason for repeating. + plane.evict_cached(plane.prefix_index.num_cached_pages) summary = run_paged( engine, params, trace_factory(), max_batch=max_batch, warmed_shapes=warmed_shapes ) @@ -541,8 +551,16 @@ def _summarise( steps: int, allocator: Any, page_size: int, + retained_pages: int = 0, ) -> dict[str, Any]: - """Collapse a run into the reported schema.""" + """Collapse a run into the reported schema. + + `retained_pages` is what the prefix cache is deliberately still holding once + every request has finished. Those pages are allocated on purpose, so counting + them as leaked would report a leak on every run with sharing enabled -- and an + alarm that always fires is one nobody reads, which would hide the real leak + this metric exists to catch. + """ output_tokens = sum(len(r.generated) for r in done) prompt_tokens = sum(r.prompt_len for r in done) itls = [gap for r in done for gap in r.inter_token_latencies()] @@ -576,7 +594,8 @@ def _summarise( "series_live_tokens": live_tokens.tolist(), } if allocator is not None: - summary["pages_leaked"] = int(allocator.num_allocated_pages) + summary["pages_retained_by_cache"] = int(retained_pages) + summary["pages_leaked"] = int(allocator.num_allocated_pages) - int(retained_pages) summary["pool_capacity_tokens"] = allocator.capacity_pages * page_size return summary @@ -601,17 +620,23 @@ def report(name: str, summary: dict[str, Any]) -> str: lines.append(f" peak device memory {summary['peak_device_bytes'] / 2**30:.3f} GiB") occ = summary.get("occupancy") if occ: + ratio = occ["mean_overhead_ratio"] + # Below 1 means the pool is holding fewer tokens than the requests + # collectively address, which only happens when several of them are reading + # the same pages. Calling that "overhead" would report the benefit as a cost. + label = f"page overhead x{ratio:.3f}" if ratio >= 1 else f"sharing dividend x{1 / ratio:.3f}" lines.append( f" pool tokens held max {occ['max_pool_tokens_held']}" - f" vs live tokens max {occ['max_live_tokens']}" - f" (page overhead x{occ['mean_overhead_ratio']:.3f})" + f" vs live tokens max {occ['max_live_tokens']} ({label})" ) lines.append( f" concurrency max {occ['max_concurrency']} mean {occ['mean_concurrency']:.2f}" f" of pool capacity {summary.get('pool_capacity_tokens')} tokens" ) if "pages_leaked" in summary: - lines.append(f" pages leaked {summary['pages_leaked']}") + retained = summary.get("pages_retained_by_cache", 0) + suffix = f" (plus {retained} deliberately retained by the prefix cache)" if retained else "" + lines.append(f" pages leaked {summary['pages_leaked']}{suffix}") if "distinct_shapes_total" in summary: lines.append( f" shapes {summary['distinct_shapes_total']} total," From 65ceaa5f7589b983d5283cf434f4675f4373541b Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Wed, 26 Aug 2026 18:24:38 +0000 Subject: [PATCH 06/21] fix(inference): correct heads-per-shard at and above the TP boundary heads_per_shard() returned the full KV head count whenever kv_head_shards >= num_kv_heads, which is wrong in both regimes it covers. At the boundary -- kv_head_shards == num_kv_heads, which is the common case of TP=8 on a model with 8 KV heads -- the heads partition exactly one per shard. Returning 8 sized every shard's pool a factor of TP too large, and contradicted replication_factor(), which correctly reported no replication. A TP=8 smoke test allocated (513, 16, 8, 128) per device where (513, 16, 1, 128) was called for. Above the boundary it was wrong in the same direction. Each rank computes a subset of the query heads and therefore needs exactly the one KV head those map to, not every head, so the shard holds one head and it is the number of copies that grows. Counting both the heads and the copies double-counted the footprint by the replication factor. Both regimes are now the same expression with a floor of one, and the test for the replication case asserts the two ways of counting total_pool_bytes agree, since pool sizing is exactly what a disagreement would corrupt. Co-authored-by: Cursor --- .../inference/kv_common/storage_layout.py | 20 +++++++++----- tests/unit/kv_common_test.py | 26 +++++++++++++++++-- 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/src/maxtext/inference/kv_common/storage_layout.py b/src/maxtext/inference/kv_common/storage_layout.py index 6142245fa4..d8499094f6 100644 --- a/src/maxtext/inference/kv_common/storage_layout.py +++ b/src/maxtext/inference/kv_common/storage_layout.py @@ -114,14 +114,20 @@ def heads_per_shard(self) -> int: """KV heads held by one shard. Two regimes, and conflating them mis-sizes the pool. When the head count - divides the shard count the heads are partitioned. When the shard count - exceeds the head count -- GQA at high TP, and MQA always -- each head is - *replicated*, so a shard still holds at least one whole head and the - total footprint is multiplied rather than divided. + divides the shard count the heads are partitioned, so a shard holds + ``num_kv_heads // kv_head_shards`` of them. When the shard count exceeds + the head count -- GQA at high TP, and MQA always -- there is nothing left + to divide, so each shard holds exactly one head and it is the *number of + copies* that grows. `replication_factor` reports that growth, and + `total_pool_bytes` multiplies by it. + + The floor of one is what makes both regimes the same expression, and the + boundary case is the one worth stating: at ``kv_head_shards == + num_kv_heads`` every shard holds a single head and nothing is replicated. """ - if self.kv_head_shards >= self.num_kv_heads: - return max(self.num_kv_heads, 1) if self.num_kv_heads else 0 - return self.num_kv_heads // self.kv_head_shards + if not self.num_kv_heads: + return 0 + return max(self.num_kv_heads // self.kv_head_shards, 1) def replication_factor(self) -> int: """How many shards hold a copy of the same KV head.""" diff --git a/tests/unit/kv_common_test.py b/tests/unit/kv_common_test.py index d6ad7e8d6d..1cd6701f13 100644 --- a/tests/unit/kv_common_test.py +++ b/tests/unit/kv_common_test.py @@ -69,15 +69,37 @@ def test_clean_partition_when_heads_divide_shards(self): self.assertEqual(layout.heads_per_shard(), 2) self.assertEqual(layout.replication_factor(), 1) + def test_one_head_per_shard_at_the_boundary(self): + """`kv_head_shards == num_kv_heads` partitions exactly, it does not replicate. + + The case an earlier version got wrong, and it is the common one: TP=8 on a + model with 8 KV heads. Returning the full head count here sized every shard's + pool a factor of TP too large, and disagreed with `replication_factor`, which + correctly reported 1. + """ + layout = self._layout(num_kv_heads=8, kv_head_shards=8) + self.assertEqual(layout.heads_per_shard(), 1) + self.assertEqual(layout.replication_factor(), 1) + self.assertEqual(layout.total_pool_bytes(), layout.pool_bytes_per_shard() * 8) + def test_replication_when_shards_exceed_heads(self): """GQA above its KV-head count replicates, multiplying the footprint. - This is the case a naive num_kv_heads // shards would get silently wrong. + Each rank computes a subset of the *query* heads and therefore needs exactly + the one KV head those map to -- not every head. So the shard holds one head + and it is the number of copies that grows, which is what + `replication_factor` reports and what `total_pool_bytes` multiplies by. """ layout = self._layout(num_kv_heads=2, kv_head_shards=8) - self.assertEqual(layout.heads_per_shard(), 2) + self.assertEqual(layout.heads_per_shard(), 1) self.assertEqual(layout.replication_factor(), 4) + # 8 shards holding one head each is 4x the 2 logical heads, and the two ways + # of counting the footprint have to agree or pool sizing is wrong. self.assertEqual(layout.total_pool_bytes(), layout.pool_bytes_per_shard() * 8) + self.assertEqual( + layout.total_pool_bytes(), + layout.pool_bytes_per_shard() * layout.num_kv_heads * layout.replication_factor(), + ) def test_mqa_replicates_everywhere(self): layout = self._layout(num_kv_heads=1, kv_head_shards=8) From 62874201e9dd317384d58216ea75a8c429396638 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Wed, 26 Aug 2026 20:39:25 +0000 Subject: [PATCH 07/21] feat(inference): shard the paged KV pool over KV heads, and refuse TP until the kernel can follow Adds kv_pool_sharding(), which decides which KV head lives on which device, and switches the pool to globally-shaped arrays built directly under that sharding. MaxEngine.kv_pool_sharding() is no longer a stub. The tensor-parallel axis has to be split in two, which is why the pool builds its own mesh rather than reusing MaxText's. A `tensor` axis of width 8 says nothing about whether eight ranks hold eight distinct KV heads or two heads replicated four ways; both occur and they need different device assignments. The split is (kv_head_shard, kv_head_replica), row-major over MaxText's own device order, because that is the assignment the model already implies: rank i computes query head i, which reads KV head i // replication_factor, so consecutive ranks share a KV head. Reversing that would put a rank's KV on another rank's device and force a gather every step -- a correct but slow run rather than a failure, which is the hardest kind to notice. The arrays are built sharded rather than built whole and distributed. device_put of a locally-created array would materialise the entire pool on one device first, which at 70B is the difference between allocating a shard and failing. Verified at TP=8: the pool lands as (513, 16, 8, 128) with one head per device, which is what it should be. The step itself still cannot run sharded, so init_paged_runtime now refuses a sharded mesh outright. The aiter kernels reach the pool through an FFI custom call and XLA cannot partition one; left alone it neither gathers nor splits and the step hangs, which is worse to ship than a refusal. Lifting this needs the forward wrapped in shard_map so each device runs the kernel on its own shard -- the remaining half of M6. Note this replaces a configuration that appeared to work: before the heads-per-shard fix, a TP=8 paged run produced correct tokens because the pool sat unsharded on one device and XLA gathered every device's KV to it each step. That is precisely the per-step cross-device traffic M6 exists to remove, and it cannot outgrow one device's memory, so it was never a usable path. Single-device paged is unchanged and still matches dense token for token. 195 CPU tests and 17 GPU tests pass. Co-authored-by: Cursor --- .../inference/kv_execution/layout_builder.py | 59 +++++++++++++++++++ .../inference/kv_execution/pool_factory.py | 43 +++++++++----- src/maxtext/inference/maxengine/maxengine.py | 24 +++++++- 3 files changed, 109 insertions(+), 17 deletions(-) diff --git a/src/maxtext/inference/kv_execution/layout_builder.py b/src/maxtext/inference/kv_execution/layout_builder.py index 7ea74331e6..8fba081bf8 100644 --- a/src/maxtext/inference/kv_execution/layout_builder.py +++ b/src/maxtext/inference/kv_execution/layout_builder.py @@ -24,12 +24,20 @@ from typing import Any +import numpy as np + from maxtext.inference.kv_common import KvStorageLayoutV1 # MaxText spells the KV-head tensor-parallel axes several ways depending on the # model; these are the ones that shard `num_kv_heads`. _KV_HEAD_MESH_AXES = ("tensor", "tensor_transpose", "tensor_sequence") +# Axis names for the pool's own mesh. Deliberately not MaxText's: the pool needs +# the tensor-parallel axis split into the part that selects a KV head and the +# part that replicates it, and MaxText's mesh has them fused into one axis. +KV_SHARD_AXIS = "kv_head_shard" +KV_REPLICA_AXIS = "kv_head_replica" + def kv_head_shards(mesh: Any | None) -> int: """Product of the mesh axes that shard KV heads, or 1 with no mesh.""" @@ -42,6 +50,57 @@ def kv_head_shards(mesh: Any | None) -> int: return max(shards, 1) +def kv_pool_sharding(mesh: Any | None, layout: KvStorageLayoutV1) -> Any | None: + """Where each KV head of the pool lives, as a `NamedSharding`. + + The pool arrays are *globally* shaped -- `num_kv_heads` on the head axis -- and + this sharding is what puts the right heads on the right devices. Allocating + per-device shards directly would be the other option, but it puts the caller in + charge of device order and makes the arrays unusable as ordinary jit inputs. + + **The tensor-parallel axis has to be split in two, which is why the pool builds + its own mesh.** MaxText's mesh fuses them: `tensor` of width 8 says nothing + about whether eight ranks hold eight distinct KV heads or two heads replicated + four ways. Both happen, and they need different device assignments. + + The split is `(kv_head_shard, kv_head_replica)`, row-major over the same device + order MaxText uses, because that is the assignment the model already implies. + Rank `i` computes query head `i`, and query head `i` reads KV head + `i // replication_factor` -- so consecutive ranks share a KV head, which is + exactly what a row-major reshape produces. Getting this backwards would put a + rank's KV on another rank's device and force a gather on every step, which is + the cost M6 exists to remove and which would show up as a correct-but-slow run + rather than a failure. + + Returns None when there is nothing to shard, so a single-device caller passes + the result straight through and gets the ordinary unsharded pool. + """ + # pylint: disable=import-outside-toplevel + import jax + + shards = int(layout.kv_head_shards) + if mesh is None or shards <= 1: + return None + + devices = np.asarray(getattr(mesh, "devices", None)) + if devices.size != shards: + raise ValueError( + f"the pool is sharded {shards} ways but the mesh holds {devices.size} devices. The layout's " + f"kv_head_shards must come from this mesh, or the pool will be laid out for a different one." + ) + + replication = layout.replication_factor() + kv_axis = shards // replication + # Row-major, and matching MaxText's device order rather than imposing one. + grid = devices.reshape(kv_axis, replication) + pool_mesh = jax.sharding.Mesh(grid, (KV_SHARD_AXIS, KV_REPLICA_AXIS)) + # Only the head axis is partitioned; pages, tokens within a page, and head_dim + # are whole on every device. Replication across the second axis is implicit in + # not naming it. + spec = jax.sharding.PartitionSpec(None, None, KV_SHARD_AXIS, None) + return jax.sharding.NamedSharding(pool_mesh, spec) + + def build_storage_layout(config: Any, mesh: Any | None = None) -> KvStorageLayoutV1: """Derive pool geometry from a MaxText config. diff --git a/src/maxtext/inference/kv_execution/pool_factory.py b/src/maxtext/inference/kv_execution/pool_factory.py index af57b32ad6..ffa3d81ca9 100644 --- a/src/maxtext/inference/kv_execution/pool_factory.py +++ b/src/maxtext/inference/kv_execution/pool_factory.py @@ -57,6 +57,13 @@ def num_layers(self) -> int: @property def page_shape(self) -> tuple[int, ...]: + """Global shape of one layer's pages, before sharding splits the head axis.""" + layout = self.layout + return (layout.num_pages, layout.tokens_per_page, layout.num_kv_heads, layout.head_dim) + + @property + def local_page_shape(self) -> tuple[int, ...]: + """What one device holds, and what the kernel sees inside `shard_map`.""" layout = self.layout return (layout.num_pages, layout.tokens_per_page, layout.heads_per_shard(), layout.head_dim) @@ -76,29 +83,37 @@ def allocate_pool( ) -> PagedKvPool: """Allocate a zero-initialised pool for every layer. + The arrays are *globally* shaped: the head axis carries `num_kv_heads`, and + `sharding` is what puts a slice of it on each device. On one device the two + coincide, since `kv_head_shards` is then 1 and `heads_per_shard()` is the whole + count, so a single-device caller sees no change from passing None. + Args: - layout: pool geometry. `heads_per_shard()` is already the per-device head - count, including the replication case where TP exceeds the KV head count, - so the shape below is the local shard rather than the global pool. - sharding: optional sharding for each array. Passed through to - `jax.device_put`; a `NamedSharding` carrying a `memory_kind` is how the - pool would later be placed in a collective memory space. + layout: pool geometry. + sharding: optional sharding for each array, normally from + `kv_pool_sharding`. Passed through to `jax.device_put`; a `NamedSharding` + carrying a `memory_kind` is how the pool would later be placed in a + collective memory space. dtype: overrides the layout's dtype. Only useful for tests that want a dtype numpy can print. Returns: A `PagedKvPool` whose arrays are all zeros. """ - shape = (layout.num_pages, layout.tokens_per_page, layout.heads_per_shard(), layout.head_dim) + shape = (layout.num_pages, layout.tokens_per_page, layout.num_kv_heads, layout.head_dim) resolved = jnp.dtype(dtype) if dtype is not None else jnp.dtype(layout.dtype) + if sharding is None: + make = lambda: jnp.zeros(shape, resolved) + else: + # Built sharded rather than built whole and then distributed. `device_put` of + # a locally-created array would materialise the entire pool on one device + # first, which at 70B is the difference between allocating a shard and + # failing outright. + make = jax.jit(lambda: jnp.zeros(shape, resolved), out_shardings=sharding) + k_pages, v_pages = [], [] for _ in range(layout.num_layers): - k = jnp.zeros(shape, resolved) - v = jnp.zeros(shape, resolved) - if sharding is not None: - k = jax.device_put(k, sharding) - v = jax.device_put(v, sharding) - k_pages.append(k) - v_pages.append(v) + k_pages.append(make()) + v_pages.append(make()) return PagedKvPool(layout=layout, k_pages=k_pages, v_pages=v_pages) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index f6addd4c01..747b0a9d08 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -1962,6 +1962,18 @@ def init_paged_runtime( ) layout = build_storage_layout(self.config, self._mesh) + if layout.kv_head_shards > 1: + # The pool shards correctly (see `kv_pool_sharding`), but the aiter kernels + # reach the pool through an FFI custom call, and XLA cannot partition one. + # Left to itself it neither gathers nor splits and the step simply hangs, + # which is a far worse thing to ship than a refusal. Lifting this needs the + # forward wrapped in `shard_map` so each device runs the kernel on its own + # shard -- the remaining half of M6. + raise ValueError( + f"attention='gpu_paged' does not yet support tensor parallelism: this mesh shards KV heads " + f"{layout.kv_head_shards} ways. Run the paged path on a single device, or use " + f"attention='dot_product' for a sharded run." + ) if layout.dtype not in ("bfloat16", "float16"): raise ValueError( f"the paged attention kernels accept bfloat16 and float16 only, but dtype is " @@ -1995,10 +2007,16 @@ def kv_pool_sharding(self): Pages are the leading axis and are never sharded -- a page belongs to one request and splitting it would put half a token's KV on another device. The - KV-head axis is the one tensor parallelism divides, which is M6's work; until - then the pool is replicated and this returns None. + KV-head axis is the one tensor parallelism divides, and `kv_pool_sharding` + works out which head belongs on which device, including the case where TP + exceeds the KV head count and heads are replicated rather than divided. """ - return None + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_execution.layout_builder import build_storage_layout, kv_pool_sharding + + if self.config.attention != "gpu_paged": + return None + return kv_pool_sharding(self._mesh, build_storage_layout(self.config, self._mesh)) def _paged_pool_arrays(self): """The pool as the nested list `kv_caches` expects: one `[k, v]` per layer.""" From a1ecb3a7fde63eafc43c17c15a5742369e4141a6 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Wed, 26 Aug 2026 21:11:05 +0000 Subject: [PATCH 08/21] feat(inference): run the paged kernels under shard_map for TP, still gated on correctness Donation survives shard_map's manual mode. That was the open question gating the whole approach, and it is answered: on 8 devices the pool keeps every shard's buffer address across an aliased append_kv, JAX raises no donation warning, the alias appears in the lowered HLO, and the write lands per shard. A single-device control ran alongside it and correctly read BROKEN until donate_argnums was added, so the measurement is trustworthy rather than vacuously green. So the kernels can be entered in manual mode, which is what XLA needs: it cannot partition an FFI custom call, and handed a sharded pool it neither gathers nor splits -- the step simply hangs. paged_attention_step_sharded wraps the step so every device runs the same code against its local shard and the kernel sees a narrower pool. No jax-aiter change was needed after all: its KV ops carry no custom_partitioning, which is the thing that would have conflicted. The plan arrays are replicated, and that is the point rather than a detail. Page ids, slot offsets and last-page occupancies describe pages, and pages are not sharded -- every device holds the same pages and differs only in which heads it stores. So a device can compute its slice with no knowledge of any other, which is what makes zero per-step cross-device KV traffic reachable. kv_pool_sharding now uses the model's own mesh in the clean-partition case, because shard_map needs the pool and the activations on one mesh and a second mesh over the same devices does not qualify. Not yet correct, and still refused at startup. At TP=8 the step runs without hanging but the first sampled token already differs from the single-device paged path and the TP=8 dense path, which agree with each other. The fault is therefore in how the sharded operands are described to the kernels rather than in the pool's sharding, and prefill rather than decode. A wrong answer that runs is worse than a refusal, so init_paged_runtime continues to reject a sharded mesh until this reproduces the single-device result. Also fixes an indentation slip introduced in this change that left paged_attention_step returning None on the decode path. 195 CPU tests and 17 GPU tests pass; single-device paged is unaffected. Co-authored-by: Cursor --- .../inference/kv_execution/layout_builder.py | 11 ++- src/maxtext/inference/maxengine/maxengine.py | 16 +++-- src/maxtext/layers/attentions.py | 18 ++++- src/maxtext/layers/gpu_paged_attention.py | 69 +++++++++++++++++++ 4 files changed, 105 insertions(+), 9 deletions(-) diff --git a/src/maxtext/inference/kv_execution/layout_builder.py b/src/maxtext/inference/kv_execution/layout_builder.py index 8fba081bf8..f7fe546955 100644 --- a/src/maxtext/inference/kv_execution/layout_builder.py +++ b/src/maxtext/inference/kv_execution/layout_builder.py @@ -82,6 +82,16 @@ def kv_pool_sharding(mesh: Any | None, layout: KvStorageLayoutV1) -> Any | None: if mesh is None or shards <= 1: return None + replication = layout.replication_factor() + if replication == 1: + # Clean partition, and the common case. The model's own mesh already + # expresses it, so use that rather than a private one: `shard_map` needs the + # pool and the activations on the *same* mesh, and a second mesh over the + # same devices is not the same mesh as far as it is concerned. + axes = tuple(a for a in _KV_HEAD_MESH_AXES if int(dict(mesh.shape).get(a, 1)) > 1) + spec = jax.sharding.PartitionSpec(None, None, axes if len(axes) > 1 else axes[0], None) + return jax.sharding.NamedSharding(mesh, spec) + devices = np.asarray(getattr(mesh, "devices", None)) if devices.size != shards: raise ValueError( @@ -89,7 +99,6 @@ def kv_pool_sharding(mesh: Any | None, layout: KvStorageLayoutV1) -> Any | None: f"kv_head_shards must come from this mesh, or the pool will be laid out for a different one." ) - replication = layout.replication_factor() kv_axis = shards // replication # Row-major, and matching MaxText's device order rather than imposing one. grid = devices.reshape(kv_axis, replication) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 747b0a9d08..74c37260bf 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -1963,15 +1963,17 @@ def init_paged_runtime( layout = build_storage_layout(self.config, self._mesh) if layout.kv_head_shards > 1: - # The pool shards correctly (see `kv_pool_sharding`), but the aiter kernels - # reach the pool through an FFI custom call, and XLA cannot partition one. - # Left to itself it neither gathers nor splits and the step simply hangs, - # which is a far worse thing to ship than a refusal. Lifting this needs the - # forward wrapped in `shard_map` so each device runs the kernel on its own - # shard -- the remaining half of M6. + # `shard_map` gets the kernels running under TP -- the step no longer hangs + # and donation survives manual mode -- but the output is still wrong: at + # TP=8 the first sampled token already differs from both the single-device + # paged path and the TP=8 dense path, which agree with each other. So the + # fault is in how the sharded operands are described to the kernels, not in + # the sharding of the pool. Refusing until that is understood, because a + # wrong answer that runs is far worse than a refusal. raise ValueError( f"attention='gpu_paged' does not yet support tensor parallelism: this mesh shards KV heads " - f"{layout.kv_head_shards} ways. Run the paged path on a single device, or use " + f"{layout.kv_head_shards} ways, and the sharded path does not yet reproduce the " + f"single-device result. Run the paged path on a single device, or use " f"attention='dot_product' for a sharded run." ) if layout.dtype not in ("bfloat16", "float16"): diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index a1c0937acd..f849fddf2f 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -1218,6 +1218,22 @@ def forward_serve_gpu_paged( total_tokens=query.shape[0], max_seqlen_k=self.max_target_length, ) + backend = getattr(self.config, "paged_attention_backend", "auto") + axes = gpu_paged_attention.kv_head_axes(self.mesh) + if axes: + # Tensor parallelism: the kernels have to run in manual mode, because XLA + # cannot partition the FFI call they go through. + return gpu_paged_attention.paged_attention_step_sharded( + query, + key, + value, + k_pool, + v_pool, + plan, + mesh=self.mesh, + axes=axes if len(axes) > 1 else axes[0], + backend=backend, + ) return gpu_paged_attention.paged_attention_step( query, key, @@ -1225,7 +1241,7 @@ def forward_serve_gpu_paged( k_pool, v_pool, plan, - backend=getattr(self.config, "paged_attention_backend", "auto"), + backend=backend, # 1.0, not 1/sqrt(head_dim). MaxText folds the depth scaling into the # query projection's initializer (see `depth_scaling` in the kernel init) # and applies `query_pre_attn_scalar` above, so the query arriving here is diff --git a/src/maxtext/layers/gpu_paged_attention.py b/src/maxtext/layers/gpu_paged_attention.py index 0bf8af57e6..c9b8a3b503 100644 --- a/src/maxtext/layers/gpu_paged_attention.py +++ b/src/maxtext/layers/gpu_paged_attention.py @@ -239,3 +239,72 @@ def paged_attention_step(query, key, value, k_pool, v_pool, plan, *, backend="au causal=causal, ) return out, [k_pool, v_pool] + + +def kv_head_axes(mesh: Any, candidates: tuple[str, ...] = ("tensor", "tensor_transpose", "tensor_sequence")): + """The mesh axes that actually shard KV heads, or () when nothing does.""" + if mesh is None: + return () + shape = dict(getattr(mesh, "shape", {}) or {}) + return tuple(axis for axis in candidates if int(shape.get(axis, 1)) > 1) + + +def paged_attention_step_sharded( + query, key, value, k_pool, v_pool, plan, *, mesh, axes, backend="auto", scale=None, causal=True +): + """`paged_attention_step` under `shard_map`, so each device sees its own heads. + + XLA cannot partition an FFI custom call. Handed a sharded pool it neither + gathers nor splits and the step hangs, so the kernels have to be entered in + manual mode where every device runs the same code on its local shard and the + kernel sees a narrower pool than the global one. Nothing vendor-side changes: + the jax-aiter KV ops carry no `custom_partitioning`, which is what would + otherwise conflict with `shard_map`. + + **The plan arrays are replicated, and that is the whole reason this is free.** + Page ids, slot offsets and last-page occupancies describe *pages*, and pages + are not sharded -- every device holds the same pages and differs only in which + heads it stores. So each device can compute its slice of the attention with no + knowledge of any other, and the step needs no cross-device KV traffic at all, + which is precisely the property M6 exists to establish. + + `plan` is not a pytree, so its arrays are passed positionally and its static + fields are closed over. + """ + qkv_spec = jax.sharding.PartitionSpec(None, axes, None) + pool_spec = jax.sharding.PartitionSpec(None, None, axes, None) + replicated = jax.sharding.PartitionSpec() + + def body(query, key, value, k_pool, v_pool, slot_mapping, kv_indptr, page_indices, last_page_lens, cu_seqlens_q): + local = dataclasses.replace( + plan, + slot_mapping=slot_mapping, + kv_indptr=kv_indptr, + kv_page_indices=page_indices, + kv_last_page_lens=last_page_lens, + cu_seqlens_q=cu_seqlens_q, + ) + # Flattened, because `shard_map` matches out_specs against the output tree + # and a nested list there is one more level than the specs can describe. + out, pools = paged_attention_step( + query, key, value, k_pool, v_pool, local, backend=backend, scale=scale, causal=causal + ) + return out, pools[0], pools[1] + + # check_vma is off because the replicated plan arrays are read by every device + # and the pools are written per-device; the varying-manual-axes check cannot + # see that an aliased FFI write is confined to the local shard. + out, k_pool, v_pool = jax.shard_map( + body, + mesh=mesh, + in_specs=( + qkv_spec, qkv_spec, qkv_spec, pool_spec, pool_spec, + replicated, replicated, replicated, replicated, replicated, + ), + out_specs=(qkv_spec, pool_spec, pool_spec), + check_vma=False, + )( + query, key, value, k_pool, v_pool, + plan.slot_mapping, plan.kv_indptr, plan.kv_page_indices, plan.kv_last_page_lens, plan.cu_seqlens_q, + ) + return out, [k_pool, v_pool] From a32b99143197b5848a4cc1042e2508b34b87e212 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Wed, 26 Aug 2026 21:39:02 +0000 Subject: [PATCH 09/21] fix(inference): restrict paged shard_map to the KV-head axes shard_map makes every mesh axis manual unless told which ones to take, and MaxText's mesh carries a dozen. The region was therefore entered under a different set of manual axes than the computation around it, which is not what the specs describe. A single-axis mesh cannot expose that, which is how the isolated parity test stays bit-exact while the model does not. This is a correctness fix on its own terms rather than a guess, but it does not resolve the TP divergence: dense and paged at TP=8 on one set of weights still disagree, so the sharded path stays refused at startup. Recorded for whoever picks it up, since the search space is now much smaller than it was. append_kv is correct -- at TP=1 and TP=8 the prefill receives identical query, key and value, and the pool afterwards is identical to the digit. The divergence is in the attention read. paged_attention_step_sharded is bit-exact in isolation across six cases including the model's exact plan shape, so the function is not at fault either. And TP=2 produces the same wrong tokens as TP=8, which points away from head slicing: a mis-sliced head axis would be wrong differently at different widths. Co-authored-by: Cursor --- src/maxtext/layers/gpu_paged_attention.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/maxtext/layers/gpu_paged_attention.py b/src/maxtext/layers/gpu_paged_attention.py index c9b8a3b503..bf8ae420d5 100644 --- a/src/maxtext/layers/gpu_paged_attention.py +++ b/src/maxtext/layers/gpu_paged_attention.py @@ -297,6 +297,13 @@ def body(query, key, value, k_pool, v_pool, slot_mapping, kv_indptr, page_indice out, k_pool, v_pool = jax.shard_map( body, mesh=mesh, + # Only the KV-head axes go manual. Left to default, `shard_map` makes + # *every* mesh axis manual, and MaxText's mesh carries a dozen -- so the + # region would be entered under a different set of manual axes than the + # computation around it, which is not what the specs describe. A + # single-axis mesh cannot expose this, which is how an isolated test stays + # bit-exact while the model does not. + axis_names=frozenset(axes if isinstance(axes, tuple) else (axes,)), in_specs=( qkv_spec, qkv_spec, qkv_spec, pool_spec, pool_spec, replicated, replicated, replicated, replicated, replicated, From 4819d6c03d4c06ff02b7f635dc2c48d3c65d4f03 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Wed, 26 Aug 2026 23:29:28 +0000 Subject: [PATCH 10/21] fix(inference): scale the query once on the sharded paged path (M6) Multi-GPU paged attention produced wrong output. The cause was not the sharding: `forward_serve_gpu_paged` spelled out its kernel arguments separately for its two branches, and the sharded one omitted `scale=1.0`. The aiter kernel then applied its own 1/sqrt(head_dim) on top of a query MaxText has already scaled through the projection initializer and `query_pre_attn_scalar`, flattening the softmax towards uniform. That explains the symptoms that resisted a sharding explanation: all heads wrong and uniformly lower rather than permuted, and TP=2 and TP=8 wrong identically, because the defect does not depend on shard count. No sharded-against-single-device test can see this, since it hands the same scale to both sides and the error cancels; the existing dense comparison runs unsharded and never reaches the branch. The same mistake had already been made once on the single-device path. So the arguments are now built once and shared, and the new test asserts on the call site rather than the numbers, confirmed to fail when the bug is reintroduced. Paged is now token-identical to dense at TP=2, 4 and 8, with no collective between AppendKvJA and the attention call consuming it, and pool donation intact. `init_paged_runtime` therefore accepts the clean partition and refuses only replication, which still lacks a shared mesh for the pool. Verified: 195 CPU tests (kv_common, kv_control, kv_execution, kv_prefix_cache, kv_import_rule), 17 GPU tests (kv_paged_runtime, gpu_paged_decode_parity, gpu_paged_prefix_cache), and the three gpu_paged cases in attention_test.py. This is a correctness result on a two-layer toy model only; M6 remains open on the replicated regime and on all of its measurements, which are blocked on the 70B checkpoint. Co-authored-by: Cursor --- src/maxtext/inference/maxengine/maxengine.py | 25 ++++--- src/maxtext/layers/attentions.py | 37 +++++---- tests/unit/attention_test.py | 79 ++++++++++++++++++++ 3 files changed, 113 insertions(+), 28 deletions(-) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 74c37260bf..6608650bc8 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -1962,19 +1962,20 @@ def init_paged_runtime( ) layout = build_storage_layout(self.config, self._mesh) - if layout.kv_head_shards > 1: - # `shard_map` gets the kernels running under TP -- the step no longer hangs - # and donation survives manual mode -- but the output is still wrong: at - # TP=8 the first sampled token already differs from both the single-device - # paged path and the TP=8 dense path, which agree with each other. So the - # fault is in how the sharded operands are described to the kernels, not in - # the sharding of the pool. Refusing until that is understood, because a - # wrong answer that runs is far worse than a refusal. + if layout.replication_factor() > 1: + # The clean partition -- num_kv_heads divisible by the shard count -- is + # supported and matches the dense path token for token. Replication is not, + # and the obstacle is the mesh rather than the kernels: `kv_pool_sharding` + # has to split the tensor axis into a head-selecting and a replicating part + # to place the copies, and it does that on a mesh of its own, while + # `shard_map` requires the pool and the activations on the *same* mesh. + # Until MaxText's own mesh carries that split there is nowhere to put such + # a pool, so refuse rather than mis-shard. raise ValueError( - f"attention='gpu_paged' does not yet support tensor parallelism: this mesh shards KV heads " - f"{layout.kv_head_shards} ways, and the sharded path does not yet reproduce the " - f"single-device result. Run the paged path on a single device, or use " - f"attention='dot_product' for a sharded run." + f"attention='gpu_paged' does not yet support tensor parallelism above the KV-head count: " + f"this mesh shards KV heads {layout.kv_head_shards} ways over {layout.num_kv_heads} heads, " + f"which replicates each head {layout.replication_factor()} times. Reduce tensor parallelism " + f"to at most num_kv_heads, or use attention='dot_product' for a wider mesh." ) if layout.dtype not in ("bfloat16", "float16"): raise ValueError( diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index f849fddf2f..91df6e60b6 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -1218,7 +1218,25 @@ def forward_serve_gpu_paged( total_tokens=query.shape[0], max_seqlen_k=self.max_target_length, ) - backend = getattr(self.config, "paged_attention_backend", "auto") + # `scale=1.0`, not 1/sqrt(head_dim). MaxText folds the depth scaling into the + # query projection's initializer (see `depth_scaling` in the kernel init) and + # applies `query_pre_attn_scalar` above, so the query arriving here is already + # scaled. forward_serve_vllm passes 1.0 for the same reason; letting the + # kernel apply its own default would scale twice. + # + # Built once and shared by both branches deliberately. Spelled out separately + # they drifted: the sharded branch omitted the scale, so every tensor-parallel + # run scaled the query a second time by 1/sqrt(head_dim), flattening the + # softmax towards uniform. That is a wrong answer rather than a crash, identical + # at every TP width because it has nothing to do with sharding, and no + # comparison of the sharded step against the single-device step can see it, + # because such a test passes the same scale to both sides. One dict is what + # makes the two paths structurally unable to disagree. + kernel_kwargs = dict( + backend=getattr(self.config, "paged_attention_backend", "auto"), + scale=1.0, + causal=True, + ) axes = gpu_paged_attention.kv_head_axes(self.mesh) if axes: # Tensor parallelism: the kernels have to run in manual mode, because XLA @@ -1232,23 +1250,10 @@ def forward_serve_gpu_paged( plan, mesh=self.mesh, axes=axes if len(axes) > 1 else axes[0], - backend=backend, + **kernel_kwargs, ) return gpu_paged_attention.paged_attention_step( - query, - key, - value, - k_pool, - v_pool, - plan, - backend=backend, - # 1.0, not 1/sqrt(head_dim). MaxText folds the depth scaling into the - # query projection's initializer (see `depth_scaling` in the kernel init) - # and applies `query_pre_attn_scalar` above, so the query arriving here is - # already scaled. forward_serve_vllm passes 1.0 for the same reason; - # letting the kernel apply its own default would scale twice. - scale=1.0, - causal=True, + query, key, value, k_pool, v_pool, plan, **kernel_kwargs ) def __call__( diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 46880c5896..22824aa544 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2670,6 +2670,85 @@ def test_gpu_paged_matches_the_dense_path(self): f"{float(np.abs(got - want).max()):.5f}, atol {atol:.5f}", ) + @pytest.mark.gpu_only + def test_gpu_paged_sharded_and_single_device_kernel_arguments_agree(self): + """Both branches of `forward_serve_gpu_paged` must configure the kernel alike. + + The two branches -- `paged_attention_step` on one device and + `paged_attention_step_sharded` under `shard_map` -- are the same attention + and differ only in how the operands are partitioned. Any difference in + `scale`, `causal` or `backend` between them is a bug by construction. + + This is not a hypothetical. The sharded branch was added spelling its own + arguments out, omitted `scale=1.0`, and so let the kernel apply its default + 1/sqrt(head_dim) on top of a query MaxText had already scaled. Every + tensor-parallel run then produced a flattened softmax and a wrong token, at + every TP width identically, while single-device runs stayed correct. + + `test_gpu_paged_matches_the_dense_path` cannot catch this -- it runs on an + unsharded mesh, so it only ever reaches the single-device branch -- and + neither can any comparison of the sharded step against the single-device + step, because such a test hands the same scale to both sides and the error + cancels. Asserting on the call site is what closes that gap, so this test + mocks the kernels away and inspects the arguments rather than the numbers. + """ + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_common import KvPageTableV1 + from maxtext.layers import gpu_paged_attention as gpa + + seq_len = 128 + tokens_per_page = 16 + cfg = self._gpu_paged_config(seq_len, attention="gpu_paged") + paged = self._build_attention(cfg, seq_len, MODEL_MODE_PREFILL) + + batch = self.global_batch_size + pages_per_request = seq_len // tokens_per_page + page_table = KvPageTableV1( + page_ids=[[1 + r * pages_per_request + p for p in range(pages_per_request)] for r in range(batch)], + seq_lens=np.full((batch,), seq_len, np.int32), + query_lens=np.full((batch,), seq_len, np.int32), + write_positions=np.tile(np.arange(seq_len, dtype=np.int32), batch), + request_order=np.arange(batch, dtype=np.int32), + ) + pool_shape = (1 + batch * pages_per_request, tokens_per_page, cfg.num_kv_heads, cfg.head_dim) + pools = [jnp.zeros(pool_shape, cfg.dtype), jnp.zeros(pool_shape, cfg.dtype)] + qkv = jnp.ones((batch, seq_len, cfg.num_query_heads, cfg.head_dim), cfg.dtype) + kv = jnp.ones((batch, seq_len, cfg.num_kv_heads, cfg.head_dim), cfg.dtype) + + captured = {} + + def record(name): + def fake(query, *_args, **kwargs): + # Only the tuning knobs matter here; the operands differ between the two + # branches by design, since one is global and the other per-shard. + captured[name] = {k: kwargs[k] for k in ("backend", "scale", "causal") if k in kwargs} + return query, [pools[0], pools[1]] + + return fake + + with mock.patch.object(gpa, "paged_attention_step", record("single")), mock.patch.object( + gpa, "paged_attention_step_sharded", record("sharded") + ), mock.patch.object(gpa, "kv_head_axes", return_value=()): + paged.forward_serve_gpu_paged(qkv, kv, kv, kv_pools=pools, metadata=page_table) + + with mock.patch.object(gpa, "paged_attention_step", record("single")), mock.patch.object( + gpa, "paged_attention_step_sharded", record("sharded") + ), mock.patch.object(gpa, "kv_head_axes", return_value=("tensor",)): + paged.forward_serve_gpu_paged(qkv, kv, kv, kv_pools=pools, metadata=page_table) + + self.assertIn("single", captured, "the unsharded mesh must reach paged_attention_step") + self.assertIn("sharded", captured, "a KV-head-sharded mesh must reach paged_attention_step_sharded") + self.assertEqual( + captured["single"], + captured["sharded"], + "the sharded and single-device paged branches configure the kernel differently; " + "they are the same attention and must pass identical scale, causal and backend", + ) + # Pin the scale itself, so the two branches cannot agree on a wrong value. + # MaxText has already folded depth scaling into the query projection. + self.assertEqual(captured["sharded"]["scale"], 1.0) + self.assertTrue(captured["sharded"]["causal"]) + @pytest.mark.gpu_only def test_gpu_paged_dry_run_passes_through(self): """Model init and JIT tracing reach the branch with nothing to attend over.""" From 9edff85f293b7e626ad741882ebf75ef937458b4 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 28 Aug 2026 10:29:13 +0000 Subject: [PATCH 11/21] perf(inference): scrub the whole pool in one dispatch, and stop the harness livelocking Two defects that only appear at real depth and real pool pressure, both found while taking M6's capacity measurements on Llama-3.3-70B at TP=8. `scrub_recycled` issued one dispatch per layer, which is eighty on a 70B model, on the critical path of every step that recycles a page. At the four-layer scale every prior measurement used it is four dispatches and invisible; near full pool occupancy at eighty layers it dominates. `scrub_pages_all_layers` does the pool in one dispatch, and both callers use it. Behaviour is pinned against the old per-layer loop element by element, since a change made for cost is exactly where a behaviour change hides. The benchmark harness discarded a preempted request's generated tokens, so the replay was byte-identical to the attempt that had just failed and an overcommitted pool never progressed: admit, exhaust, preempt, re-prefill, forever, at full GPU utilisation. `PagedDriver` never had this -- it folds the tokens into the prompt so the replay is longer, calls that "the price of not deadlocking", and has a test for it. The harness now does the same, which took one sweep point from over an hour unfinished to 1119 seconds. Preemption moves tokens from output to prompt, so it distorts throughput figures. `preemptions` and `requests_preempted` are now reported, because a run with preemptions is a capacity result rather than a throughput one. Admission itself is unchanged and was not at fault: it admits on a free row without consulting the allocator and lets preemption absorb overcommit, which is what vLLM and sglang-jax also do. Verified: 197 CPU tests and 17 GPU tests. The 70B sweep now completes, with measured concurrency matching budget/length (15 against 14.2, 7 against 7.5), dense flat at 4 slots, occupancy tracking live tokens, and peak HBM flat. Still open: `run_paged` has no test of its own, which is how it drifted from the driver in the first place. It needs an engine fake. Co-authored-by: Cursor --- .../inference/kv_execution/__init__.py | 8 ++- .../inference/kv_execution/benchmark.py | 34 +++++++++++- src/maxtext/inference/kv_execution/driver.py | 8 +-- .../inference/kv_execution/engine_adapter.py | 11 ++-- .../inference/kv_execution/pool_ops.py | 38 +++++++++++++ tests/unit/kv_execution_test.py | 53 ++++++++++++++++++- 6 files changed, 142 insertions(+), 10 deletions(-) diff --git a/src/maxtext/inference/kv_execution/__init__.py b/src/maxtext/inference/kv_execution/__init__.py index ef1ea896a1..b303b641e8 100644 --- a/src/maxtext/inference/kv_execution/__init__.py +++ b/src/maxtext/inference/kv_execution/__init__.py @@ -35,7 +35,12 @@ from maxtext.inference.kv_execution.driver import PagedDriver, PagedRequest, StepOutcome from maxtext.inference.kv_execution.layout_builder import build_storage_layout from maxtext.inference.kv_execution.pool_factory import PagedKvPool, allocate_pool -from maxtext.inference.kv_execution.pool_ops import POISON_SENTINEL, poison_pages, scrub_pages +from maxtext.inference.kv_execution.pool_ops import ( + POISON_SENTINEL, + poison_pages, + scrub_pages, + scrub_pages_all_layers, +) from maxtext.inference.kv_execution.step_view import StepView, build_step_view __all__ = [ @@ -53,5 +58,6 @@ "build_storage_layout", "poison_pages", "scrub_pages", + "scrub_pages_all_layers", "token_ladder", ] diff --git a/src/maxtext/inference/kv_execution/benchmark.py b/src/maxtext/inference/kv_execution/benchmark.py index e7c5a50089..d61d80b94c 100644 --- a/src/maxtext/inference/kv_execution/benchmark.py +++ b/src/maxtext/inference/kv_execution/benchmark.py @@ -71,6 +71,11 @@ class Request: handle: Any = None token_ids: np.ndarray | None = None prefill_tokens: int = 0 + # Preemption folds a request's generated tokens back into its prompt so the + # replay makes progress, which means a preempted request reports more prompt + # tokens and fewer output tokens than it really had. Counting preemptions is + # what stops that distortion being invisible in the throughput figures. + preemptions: int = 0 def context_tokens(self) -> np.ndarray | None: """Prompt plus generated, which is what the prefix cache is offered.""" @@ -367,8 +372,28 @@ def run_paged( # Nothing published: preemption is trying to reclaim pages, and the cache # would hold onto whatever it adopted. engine.release(victim.handle) - victim.handle, victim.generated = None, [] - waiting.insert(0, victim) + # **Keep the generated tokens**, exactly as `PagedDriver._preempt_newest` + # does, so the replay is a longer prompt rather than a restart. Discarding + # them -- which this did until it was caught livelocking a 70B sweep -- + # makes the retry byte-identical to the attempt that just failed, so a pool + # that is overcommitted at admission never makes progress: admit, exhaust, + # preempt, re-prefill the same tokens, forever, at full GPU utilisation. + # The driver's own docstring calls keeping them "the price of not + # deadlocking"; the harness has to pay it too. + context = victim.context_tokens() + if context is not None: + victim.token_ids = context + victim.max_new_tokens = max(victim.max_new_tokens - len(victim.generated), 0) + victim.prompt_len += len(victim.generated) + victim.generated = [] + victim.handle = None + victim.preemptions += 1 + # A request preempted with nothing left to generate is finished, not + # requeued; requeuing it would ask for a zero-token step. + if victim.max_new_tokens == 0: + done.append(victim) + else: + waiting.insert(0, victim) continue steps += 1 @@ -573,6 +598,11 @@ def _summarise( "output_throughput_tok_per_s": output_tokens / elapsed if elapsed else 0.0, "request_throughput_per_s": len(done) / elapsed if elapsed else 0.0, "engine_steps": steps, + # Non-zero means the pool was overcommitted and requests were replayed. + # Prompt and output token counts are shifted for those requests, so a run + # with preemptions is a capacity result rather than a throughput one. + "preemptions": sum(r.preemptions for r in done), + "requests_preempted": sum(1 for r in done if r.preemptions), "ttft": _percentiles([r.ttft for r in done if r.ttft is not None]), "itl": _percentiles(itls), "peak_device_bytes": _peak_bytes(), diff --git a/src/maxtext/inference/kv_execution/driver.py b/src/maxtext/inference/kv_execution/driver.py index 17f624fe43..25e212d4ce 100644 --- a/src/maxtext/inference/kv_execution/driver.py +++ b/src/maxtext/inference/kv_execution/driver.py @@ -60,7 +60,7 @@ ) from maxtext.inference.kv_execution.bucketing import StepShape, StepShapePlanner from maxtext.inference.kv_execution.pool_factory import PagedKvPool -from maxtext.inference.kv_execution.pool_ops import poison_pages, scrub_pages +from maxtext.inference.kv_execution.pool_ops import poison_pages, scrub_pages_all_layers from maxtext.inference.kv_execution.step_view import StepView, build_step_view _DEFAULT_NAMESPACE = CacheNamespace() @@ -360,8 +360,10 @@ def _scrub_recycled_pages(self) -> None: pending = self.plane.pending_scrub() if not pending.size: return - for layer in range(self.pool.num_layers): - k, v = scrub_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], pending) + # One dispatch for the whole pool rather than one per layer; see + # `scrub_pages_all_layers` for why that matters at real depth. + k_pages, v_pages = scrub_pages_all_layers(self.pool.k_pages, self.pool.v_pages, pending) + for layer, (k, v) in enumerate(zip(k_pages, v_pages)): self.pool.replace_layer(layer, k, v) self.plane.confirm_scrubbed(pending) diff --git a/src/maxtext/inference/kv_execution/engine_adapter.py b/src/maxtext/inference/kv_execution/engine_adapter.py index 88aaa592b9..91f475c8f1 100644 --- a/src/maxtext/inference/kv_execution/engine_adapter.py +++ b/src/maxtext/inference/kv_execution/engine_adapter.py @@ -41,7 +41,7 @@ from maxtext.inference.kv_control import NativeKvControlPlane, RequestDescriptor, RequestHandle from maxtext.inference.kv_execution.bucketing import StepShape, StepShapePlanner from maxtext.inference.kv_execution.pool_factory import PagedKvPool -from maxtext.inference.kv_execution.pool_ops import poison_pages, scrub_pages +from maxtext.inference.kv_execution.pool_ops import poison_pages, scrub_pages_all_layers from maxtext.inference.kv_execution.step_view import StepView, build_step_view @@ -149,8 +149,13 @@ def scrub_recycled(self) -> np.ndarray: pending = self.control_plane.pending_scrub() if not pending.size: return pending - for layer in range(self.pool.num_layers): - k, v = scrub_pages(self.pool.k_pages[layer], self.pool.v_pages[layer], pending) + # Every layer at once. The per-layer loop this replaces issued one dispatch + # per layer, which is eighty on a 70B model, on the critical path of every + # step that recycles a page -- invisible at four layers and dominant once the + # pool runs near full, which is precisely the condition a capacity sweep + # creates. + k_pages, v_pages = scrub_pages_all_layers(self.pool.k_pages, self.pool.v_pages, pending) + for layer, (k, v) in enumerate(zip(k_pages, v_pages)): self.pool.replace_layer(layer, k, v) self.control_plane.confirm_scrubbed(pending) return pending diff --git a/src/maxtext/inference/kv_execution/pool_ops.py b/src/maxtext/inference/kv_execution/pool_ops.py index 409eaa1d7e..0f37e7bac2 100644 --- a/src/maxtext/inference/kv_execution/pool_ops.py +++ b/src/maxtext/inference/kv_execution/pool_ops.py @@ -77,6 +77,44 @@ def _fill(k_pages: jax.Array, v_pages: jax.Array, page_ids: jax.Array, value: ja return k_pages.at[page_ids].set(block), v_pages.at[page_ids].set(block) +@functools.partial(jax.jit, donate_argnums=(0, 1)) +def _fill_all(k_pages: list, v_pages: list, page_ids: jax.Array, value: jax.Array): + """`_fill` over every layer at once. + + One dispatch rather than one per layer, and at 80 layers that difference is + the whole point. Every layer's pool has the same shape and takes the same page + indices, so the per-layer loop was issuing eighty identical launches -- each + with its own donation bookkeeping across every device -- on the critical path + of any step that recycled a page. It costs nothing at the four-layer scale the + earlier measurements used and shows up as a throughput cliff at real depth, + which is exactly the shape of defect Section 6.4 exists to catch. + + Taking lists rather than stacked arrays keeps the pool's per-layer identity, so + donation still aliases each layer's own buffer and nothing is repacked. + """ + def fill(arr): + block = jnp.full((page_ids.shape[0],) + arr.shape[1:], value, dtype=arr.dtype) + return arr.at[page_ids].set(block) + + return [fill(k) for k in k_pages], [fill(v) for v in v_pages] + + +def scrub_pages_all_layers( + k_pages: Sequence[jax.Array], + v_pages: Sequence[jax.Array], + page_ids: Sequence[int] | np.ndarray, +) -> tuple[list[jax.Array], list[jax.Array]]: + """Zero the named pages in every layer, in a single dispatch. + + Returns the rebound arrays, which alias the inputs. A no-op for an empty list, + which is the common case on a pool that has not wrapped around yet. + """ + pages = _pad_page_indices(page_ids) + if pages.size == 0: + return list(k_pages), list(v_pages) + return _fill_all(list(k_pages), list(v_pages), jnp.asarray(pages), jnp.zeros((), k_pages[0].dtype)) + + def scrub_pages( k_pages: jax.Array, v_pages: jax.Array, diff --git a/tests/unit/kv_execution_test.py b/tests/unit/kv_execution_test.py index d90ec342f8..921a141a63 100644 --- a/tests/unit/kv_execution_test.py +++ b/tests/unit/kv_execution_test.py @@ -41,7 +41,12 @@ ) from maxtext.inference.kv_execution.engine_adapter import PagedRuntime from maxtext.inference.kv_execution.layout_builder import build_storage_layout, kv_head_shards -from maxtext.inference.kv_execution.pool_ops import POISON_SENTINEL, poison_pages, scrub_pages +from maxtext.inference.kv_execution.pool_ops import ( + POISON_SENTINEL, + poison_pages, + scrub_pages, + scrub_pages_all_layers, +) from maxtext.inference.kv_execution.step_view import build_step_view PAGE = 16 @@ -274,6 +279,52 @@ def test_scrubbing_nothing_is_a_no_op(self): self.assertTrue(bool((np.asarray(k) == 0).all())) self.assertTrue(bool((np.asarray(v) == 0).all())) + def test_scrubbing_every_layer_at_once_matches_the_per_layer_loop(self): + """The batched scrub must be the per-layer loop's answer, not merely close. + + `scrub_pages_all_layers` exists for cost rather than for behaviour: the loop + it replaces issued one dispatch per layer, which is eighty on a 70B model and + lands on the critical path of every step that recycles a page. Cost changes + are exactly where a behaviour change slips in unnoticed, so this pins the two + against each other element by element rather than checking the batched one in + isolation. + """ + layers, pages = 5, 8 + batched = allocate_pool(_layout(num_pages=pages, num_layers=layers)) + looped = allocate_pool(_layout(num_pages=pages, num_layers=layers)) + for pool in (batched, looped): + for layer in range(layers): + k, v = poison_pages(pool.k_pages[layer], pool.v_pages[layer], [1, 2, 3, 4, 5]) + pool.replace_layer(layer, k, v) + + ks, vs = scrub_pages_all_layers(batched.k_pages, batched.v_pages, [2, 4, 5]) + for layer, (k, v) in enumerate(zip(ks, vs)): + batched.replace_layer(layer, k, v) + for layer in range(layers): + k, v = scrub_pages(looped.k_pages[layer], looped.v_pages[layer], [2, 4, 5]) + looped.replace_layer(layer, k, v) + + for layer in range(layers): + np.testing.assert_array_equal( + np.asarray(batched.k_pages[layer]), np.asarray(looped.k_pages[layer]) + ) + np.testing.assert_array_equal( + np.asarray(batched.v_pages[layer]), np.asarray(looped.v_pages[layer]) + ) + # And it really did scrub: named pages zeroed, unnamed ones left poisoned. + scrubbed = np.asarray(batched.k_pages[layers - 1]) + self.assertTrue(bool((scrubbed[[2, 4, 5]] == 0).all())) + self.assertTrue(bool((scrubbed[[1, 3]] == POISON_SENTINEL).all())) + + def test_scrubbing_every_layer_with_nothing_pending_is_a_no_op(self): + """The common case on a pool that has not wrapped, and it must not dispatch.""" + pool = allocate_pool(_layout(num_pages=8, num_layers=3)) + ks, vs = scrub_pages_all_layers(pool.k_pages, pool.v_pages, []) + self.assertEqual(len(ks), 3) + self.assertEqual(len(vs), 3) + for layer in range(3): + self.assertTrue(bool((np.asarray(ks[layer]) == 0).all())) + def test_an_odd_page_count_is_padded_idempotently(self): """Padding repeats a page already being written, so no mask is needed.""" pool = allocate_pool(_layout(num_pages=8, num_layers=1)) From 3dcfcb570cde2a20e740a4f3b05a07c01de21e35 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 28 Aug 2026 10:44:06 +0000 Subject: [PATCH 12/21] feat(inference): support the replicated KV regime, and size it correctly (M6) M6's exit criterion asks for the paged path "at TP greater than the KV-head count". Taken literally that is `ici_tensor_parallelism=8` on a 4-KV-head model, and MaxText refuses to build one: attention heads are atomic under tensor parallelism and cannot be split across more shards than there are heads. That holds for the dense path too, so it was never a gap in the paged runtime -- the criterion was asking for a configuration the framework does not have. The regime is still reachable, by the route MaxText's own error message names: put the surplus parallelism on an axis that does not shard KV heads. With `tensor=4, fsdp=2` on eight devices and four KV heads, the heads partition cleanly over `tensor` while the pool replicates across `fsdp`, so every device holds exactly one head and pairs of devices hold the same one. That is the replicated footprint, on MaxText's own mesh, and it is token-identical to dense. Three changes follow. `kv_pool_sharding` now branches on whether the head axis is over-sharded rather than on `replication_factor()`. Those are different questions, and conflating them sent an ordinary fsdp layout down the private-mesh path, which `shard_map` cannot use. The private mesh turns out never to be needed: naming only the KV-head axes leaves the pool replicated across the others automatically. `total_pool_bytes()` counted only `kv_head_shards`, so it reported the unique KV and ignored replicas held on other axes -- 32.1 MiB against a physical 64.1 MiB on the verified configuration, exactly the naive calculation this milestone warns against. `KvStorageLayoutV1` gains `pool_replicas`, derived from the mesh, and the two replication sources multiply. The default of 1 leaves existing layouts untouched. `init_paged_runtime` refuses only the over-sharded head axis now, not any replication. That case should be unreachable since MaxText rejects it first, and is checked anyway because this counts mesh axes directly while MaxText counts them through the logical axis rules. Verified: 200 CPU tests and 20 GPU tests; replicated `tensor=4, fsdp=2` and clean-partition TP=8 both token-identical to dense. Co-authored-by: Cursor --- .../inference/kv_common/storage_layout.py | 37 +++++++++++++--- .../inference/kv_execution/layout_builder.py | 42 ++++++++++++++++--- src/maxtext/inference/maxengine/maxengine.py | 35 ++++++++++------ tests/unit/kv_common_test.py | 41 ++++++++++++++++++ 4 files changed, 131 insertions(+), 24 deletions(-) diff --git a/src/maxtext/inference/kv_common/storage_layout.py b/src/maxtext/inference/kv_common/storage_layout.py index d8499094f6..5fe3e16b91 100644 --- a/src/maxtext/inference/kv_common/storage_layout.py +++ b/src/maxtext/inference/kv_common/storage_layout.py @@ -64,6 +64,15 @@ class KvStorageLayoutV1: head_dim: per-head dimension. dtype: numpy dtype name of the stored K/V. kv_head_shards: tensor-parallel width the pool is sharded over. + pool_replicas: how many devices hold a copy of each KV-head shard, + because they sit on mesh axes that do not shard KV heads. This is + not a refinement: MaxText refuses to build a model whose KV heads + are sharded more ways than it has heads, so the *only* route to a + replicated KV footprint is surplus parallelism on another axis -- + `tensor=4, fsdp=2` on eight devices with four KV heads. Counting + only `kv_head_shards` then understates the physical footprint by + exactly this factor, which is the mis-sizing the milestone warns + about arriving by a different door than expected. padding_page_id: page reserved as a padding target and never allocated. This buys safe targets for padded rows and simpler invalid-page handling. It is emphatically *not* a security mechanism: stale KV in a @@ -78,6 +87,7 @@ class KvStorageLayoutV1: head_dim: int = 0 dtype: str = "bfloat16" kv_head_shards: int = 1 + pool_replicas: int = 1 padding_page_id: int = 0 def __post_init__(self): @@ -85,6 +95,8 @@ def __post_init__(self): raise ValueError(f"tokens_per_page must be positive, got {self.tokens_per_page}") if self.kv_head_shards <= 0: raise ValueError(f"kv_head_shards must be positive, got {self.kv_head_shards}") + if self.pool_replicas <= 0: + raise ValueError(f"pool_replicas must be positive, got {self.pool_replicas}") if self.num_kv_heads and self.kv_head_shards > self.num_kv_heads: if self.kv_head_shards % self.num_kv_heads != 0: raise ValueError( @@ -130,10 +142,18 @@ def heads_per_shard(self) -> int: return max(self.num_kv_heads // self.kv_head_shards, 1) def replication_factor(self) -> int: - """How many shards hold a copy of the same KV head.""" + """How many devices hold a copy of the same KV head. + + Two independent sources, and they multiply. Over-sharding the head axis + replicates when `kv_head_shards` exceeds `num_kv_heads`; and any mesh + axis that does not shard KV heads replicates the whole pool across + itself, which `pool_replicas` carries. Only the second is reachable + through MaxText, since it rejects the first when the model is built. + """ + over_sharded = 1 if self.num_kv_heads and self.kv_head_shards > self.num_kv_heads: - return self.kv_head_shards // self.num_kv_heads - return 1 + over_sharded = self.kv_head_shards // self.num_kv_heads + return over_sharded * max(self.pool_replicas, 1) def bytes_per_page(self) -> int: """Bytes of one page of one of K or V, for one layer, on one shard. @@ -157,8 +177,15 @@ def pool_bytes_per_shard(self) -> int: return 2 * self.num_layers * self.num_pages * self.bytes_per_page() def total_pool_bytes(self) -> int: - """Pool bytes summed over shards, including any replication.""" - return self.pool_bytes_per_shard() * self.kv_head_shards + """Physical pool bytes across the whole mesh, replication included. + + `pool_bytes_per_shard * kv_head_shards` counts the *unique* KV once and + is what a naive sizing reaches for. Every device that holds a copy pays + for it, so the count has to include `pool_replicas` too -- otherwise a + `tensor=4, fsdp=2` deployment budgets half the memory it will actually + consume, and finds out at allocation time on the largest model it runs. + """ + return self.pool_bytes_per_shard() * self.kv_head_shards * max(self.pool_replicas, 1) def max_tokens(self) -> int: """Live tokens the pool can hold, excluding the padding page.""" diff --git a/src/maxtext/inference/kv_execution/layout_builder.py b/src/maxtext/inference/kv_execution/layout_builder.py index f7fe546955..69026009c1 100644 --- a/src/maxtext/inference/kv_execution/layout_builder.py +++ b/src/maxtext/inference/kv_execution/layout_builder.py @@ -50,6 +50,25 @@ def kv_head_shards(mesh: Any | None) -> int: return max(shards, 1) +def pool_replicas(mesh: Any | None) -> int: + """Devices holding a copy of each KV-head shard, from the non-KV mesh axes. + + The pool's `PartitionSpec` names only the KV-head axes, so it is replicated + across every other axis of the mesh -- one copy per device on them. That is + the only route to a replicated KV footprint that MaxText permits, since it + refuses to build a model whose KV heads are sharded more ways than it has + heads. On eight devices with four KV heads, `tensor=4, fsdp=2` gives four + shards and two replicas, and every device holds exactly one head. + """ + if mesh is None: + return 1 + devices = int(getattr(mesh, "size", 0) or 0) + if devices <= 0: + return 1 + shards = kv_head_shards(mesh) + return max(devices // max(shards, 1), 1) + + def kv_pool_sharding(mesh: Any | None, layout: KvStorageLayoutV1) -> Any | None: """Where each KV head of the pool lives, as a `NamedSharding`. @@ -82,16 +101,26 @@ def kv_pool_sharding(mesh: Any | None, layout: KvStorageLayoutV1) -> Any | None: if mesh is None or shards <= 1: return None - replication = layout.replication_factor() - if replication == 1: - # Clean partition, and the common case. The model's own mesh already - # expresses it, so use that rather than a private one: `shard_map` needs the - # pool and the activations on the *same* mesh, and a second mesh over the - # same devices is not the same mesh as far as it is concerned. + # The head axis is what decides whether the model's own mesh can express this, + # and it is *not* the same question as whether the pool ends up replicated. + # Naming only the KV-head axes in the spec leaves the pool replicated across + # every other mesh axis automatically, which is exactly the layout wanted when + # surplus parallelism sits on `fsdp` -- four shards over `tensor`, two copies + # over `fsdp`, one head per device. Branching on `replication_factor()` instead + # sends that perfectly ordinary case down the private-mesh path and fails. + if not layout.num_kv_heads or shards <= layout.num_kv_heads: + # The model's own mesh already expresses it, so use that rather than a + # private one: `shard_map` needs the pool and the activations on the *same* + # mesh, and a second mesh over the same devices is not the same mesh as far + # as it is concerned. axes = tuple(a for a in _KV_HEAD_MESH_AXES if int(dict(mesh.shape).get(a, 1)) > 1) spec = jax.sharding.PartitionSpec(None, None, axes if len(axes) > 1 else axes[0], None) return jax.sharding.NamedSharding(mesh, spec) + # Only reachable when the head axis itself is over-sharded, which MaxText + # refuses to build a model for. Kept because this counts mesh axes directly + # while MaxText counts them through the logical axis rules. + replication = layout.replication_factor() devices = np.asarray(getattr(mesh, "devices", None)) if devices.size != shards: raise ValueError( @@ -138,5 +167,6 @@ def build_storage_layout(config: Any, mesh: Any | None = None) -> KvStorageLayou head_dim=head_dim, dtype=dtype, kv_head_shards=kv_head_shards(mesh), + pool_replicas=pool_replicas(mesh), padding_page_id=0, ) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 6608650bc8..9a9268cd91 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -1962,20 +1962,29 @@ def init_paged_runtime( ) layout = build_storage_layout(self.config, self._mesh) - if layout.replication_factor() > 1: - # The clean partition -- num_kv_heads divisible by the shard count -- is - # supported and matches the dense path token for token. Replication is not, - # and the obstacle is the mesh rather than the kernels: `kv_pool_sharding` - # has to split the tensor axis into a head-selecting and a replicating part - # to place the copies, and it does that on a mesh of its own, while - # `shard_map` requires the pool and the activations on the *same* mesh. - # Until MaxText's own mesh carries that split there is nowhere to put such - # a pool, so refuse rather than mis-shard. + if layout.num_kv_heads and layout.kv_head_shards > layout.num_kv_heads: + # Over-sharding the head axis itself is the one layout with nowhere to go: + # `kv_pool_sharding` has to split the tensor axis into a head-selecting and + # a replicating part to place the copies, and it does that on a mesh of its + # own, while `shard_map` needs the pool and the activations on one mesh. + # + # It should already be unreachable, because MaxText refuses to build a + # model whose KV heads are sharded more ways than it has heads -- attention + # heads are atomic under tensor parallelism. Checked anyway, because this + # counts mesh axes directly while MaxText counts them through the logical + # axis rules, and the two could disagree. + # + # Note what is deliberately *not* refused: a replicated KV footprint + # reached by putting surplus parallelism on an axis that does not shard KV + # heads, which is the only route MaxText permits. `tensor=4, fsdp=2` on + # eight devices with four KV heads gives every device one head and pairs of + # devices the same head, entirely on MaxText's own mesh, and matches the + # dense path token for token. raise ValueError( - f"attention='gpu_paged' does not yet support tensor parallelism above the KV-head count: " - f"this mesh shards KV heads {layout.kv_head_shards} ways over {layout.num_kv_heads} heads, " - f"which replicates each head {layout.replication_factor()} times. Reduce tensor parallelism " - f"to at most num_kv_heads, or use attention='dot_product' for a wider mesh." + f"attention='gpu_paged' cannot shard KV heads more ways than there are heads: this mesh " + f"shards them {layout.kv_head_shards} ways over {layout.num_kv_heads} heads. Reduce the " + f"tensor-parallel width to at most num_kv_heads and put any surplus parallelism on an axis " + f"that does not shard KV heads, such as fsdp." ) if layout.dtype not in ("bfloat16", "float16"): raise ValueError( diff --git a/tests/unit/kv_common_test.py b/tests/unit/kv_common_test.py index 1cd6701f13..0615c1ff05 100644 --- a/tests/unit/kv_common_test.py +++ b/tests/unit/kv_common_test.py @@ -101,6 +101,47 @@ def test_replication_when_shards_exceed_heads(self): layout.pool_bytes_per_shard() * layout.num_kv_heads * layout.replication_factor(), ) + def test_replication_from_a_non_kv_mesh_axis_is_counted(self): + """The only replicated layout MaxText can actually build, and it was mis-sized. + + MaxText refuses a model whose KV heads are sharded more ways than it has + heads, so `kv_head_shards > num_kv_heads` is unreachable through it. The + reachable route is surplus parallelism on an axis that does not shard KV + heads -- `tensor=4, fsdp=2` on eight devices with four KV heads -- where the + pool partitions four ways and is replicated across the other two. + + Counting only `kv_head_shards` reports half the memory that is really + committed, which is the naive calculation the milestone warns against. + Measured on the running model: 8 MiB per shard, 8 addressable shards, so + 64 MiB physical where the old arithmetic said 32. + """ + layout = self._layout(num_kv_heads=4, kv_head_shards=4, pool_replicas=2) + self.assertEqual(layout.heads_per_shard(), 1, "each device holds exactly one head") + self.assertEqual(layout.replication_factor(), 2, "two devices hold each head") + self.assertEqual( + layout.total_pool_bytes(), + layout.pool_bytes_per_shard() * 8, + "eight devices each hold a shard, so eight shards are paid for", + ) + # The unique KV is half of that; budgeting against it is the mis-sizing. + self.assertEqual( + layout.total_pool_bytes(), + layout.pool_bytes_per_shard() * layout.num_kv_heads * layout.replication_factor(), + ) + + def test_replication_sources_multiply(self): + """Over-sharded heads and a replicating mesh axis are independent.""" + layout = self._layout(num_kv_heads=2, kv_head_shards=8, pool_replicas=2) + self.assertEqual(layout.replication_factor(), 8, "4x from over-sharding, 2x from the axis") + self.assertEqual(layout.total_pool_bytes(), layout.pool_bytes_per_shard() * 16) + + def test_no_replicas_leaves_the_footprint_unchanged(self): + """The default has to be inert, or every existing deployment is re-sized.""" + plain = self._layout(num_kv_heads=8, kv_head_shards=8) + explicit = self._layout(num_kv_heads=8, kv_head_shards=8, pool_replicas=1) + self.assertEqual(plain.total_pool_bytes(), explicit.total_pool_bytes()) + self.assertEqual(plain.replication_factor(), 1) + def test_mqa_replicates_everywhere(self): layout = self._layout(num_kv_heads=1, kv_head_shards=8) self.assertEqual(layout.heads_per_shard(), 1) From 0f4d36cb5e34cfa3dacd131459ca68c272d24125 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 28 Aug 2026 11:26:25 +0000 Subject: [PATCH 13/21] refactor(inference): one implementation of the per-step order, not two First piece of the serving wiring. `PagedDriver` had its own inline copy of reserve-scrub-build-bucket, byte-for-byte the same sequence as `PagedRuntime.prepare_step`, and that order is not incidental: the scrub must follow reservation, because reservation is what decides which pages were recycled, and precede the page table, because the control plane refuses to describe a page it still considers dirty. That order is the dirty-page gate, and a second copy of it is a second place to get it wrong. The driver now keeps policy -- the queue, the admission budget, recompute preemption, prefill before decode -- and delegates the mechanics to a `PagedRuntime` it owns. `observed_shapes` follows the shape selection onto the runtime and stays readable through a property, so `num_distinct_shapes` and the bucketing exit criterion are unaffected. The reservation retry reads slightly differently as a result: `prepare_step` returns None for backpressure rather than exposing `reserve`, so the loop tries a step and preempts on None instead of testing the reservation first. Same semantics, one fewer place that knows the sequence. Net 39 lines deleted against 28 added, with no test changes: 200 CPU tests and 17 GPU tests pass as they stand, which is the point -- this is meant to be a pure consolidation. Still to come in this milestone: a step function that drives `MaxEngine`, and the `OfflineEngine` switch. Both are blocked on a contract question rather than plumbing -- `StepFn` takes only the `StepView` and the pool, and a real forward pass also needs token ids, positions and a sample index, none of which the page bookkeeping carries. That contract wants widening deliberately rather than in passing. Co-authored-by: Cursor --- src/maxtext/inference/kv_execution/driver.py | 67 ++++++++------------ 1 file changed, 28 insertions(+), 39 deletions(-) diff --git a/src/maxtext/inference/kv_execution/driver.py b/src/maxtext/inference/kv_execution/driver.py index 25e212d4ce..032f453954 100644 --- a/src/maxtext/inference/kv_execution/driver.py +++ b/src/maxtext/inference/kv_execution/driver.py @@ -59,9 +59,10 @@ pages_for_tokens, ) from maxtext.inference.kv_execution.bucketing import StepShape, StepShapePlanner +from maxtext.inference.kv_execution.engine_adapter import PagedRuntime from maxtext.inference.kv_execution.pool_factory import PagedKvPool -from maxtext.inference.kv_execution.pool_ops import poison_pages, scrub_pages_all_layers -from maxtext.inference.kv_execution.step_view import StepView, build_step_view +from maxtext.inference.kv_execution.pool_ops import poison_pages +from maxtext.inference.kv_execution.step_view import StepView _DEFAULT_NAMESPACE = CacheNamespace() @@ -174,10 +175,28 @@ def __init__( pool_pages=layout.num_pages, max_batched_tokens=max_batched_tokens, ) + # The driver owns *policy* -- the queue, the admission budget, recompute + # preemption, prefill before decode -- and delegates the per-step mechanics + # to the runtime. Reserve, scrub, build and bucket used to exist twice, once + # inline here and once as `PagedRuntime.prepare_step`, and that order is what + # enforces the dirty-page gate: scrub after reservation, because reservation + # decides what was recycled, and before the table, because the control plane + # refuses to describe a page it still thinks is dirty. A second copy of that + # order is a second place to get it wrong. + self.runtime = PagedRuntime( + control_plane=control_plane, + pool=pool, + planner=self.planner, + poison_on_free=self.poison_on_free, + ) self._waiting: list[PagedRequest] = [] self._live: list[PagedRequest] = [] self._done: list[PagedRequest] = [] - self.observed_shapes: set[StepShape] = set() + + @property + def observed_shapes(self) -> set[StepShape]: + """Shapes traced so far. Recorded by the runtime, since it picks them.""" + return self.runtime.observed_shapes @property def num_distinct_shapes(self) -> int: @@ -297,8 +316,12 @@ def _run_phase( batch = list(requests) lens = list(query_lens) preempted = 0 + view = None - while batch and not self.plane.reserve([r.handle for r in batch], lens): + while batch: + view = self.runtime.prepare_step([r.handle for r in batch], lens, is_decode=is_decode) + if view is not None: + break if not is_decode: # A prefill batch was budgeted against the free pool, so a failure here # means the budget was wrong rather than that the pool is under pressure. @@ -319,24 +342,7 @@ def _run_phase( preempted=preempted, ) - self._scrub_recycled_pages() - - handles = [r.handle for r in batch] - table = self.plane.build_page_table(handles, lens) - max_seq_len = int(table.seq_lens.max()) if table.num_requests else 1 - shape = ( - self.planner.decode_shape(len(batch), max_seq_len) - if is_decode - else self.planner.extend_shape(sum(lens), max_seq_len) - ) - self.observed_shapes.add(shape) - - view = build_step_view( - table, - shape, - tokens_per_page=self.plane.layout.tokens_per_page, - padding_page_id=self.plane.layout.padding_page_id, - ) + shape = view.shape next_tokens = np.asarray(self.step_fn(view, self.pool)).reshape(-1) if next_tokens.size < len(batch): raise ValueError(f"the step function returned {next_tokens.size} tokens for {len(batch)} requests") @@ -350,23 +356,6 @@ def _run_phase( preempted=preempted, ) - def _scrub_recycled_pages(self) -> None: - """Zero every page this step recycled, then say so. - - Both halves matter. Without the write a later read sees the previous - occupant's KV; without the confirmation the control plane refuses to build - the table, which is the mechanism that stops the write being forgotten. - """ - pending = self.plane.pending_scrub() - if not pending.size: - return - # One dispatch for the whole pool rather than one per layer; see - # `scrub_pages_all_layers` for why that matters at real depth. - k_pages, v_pages = scrub_pages_all_layers(self.pool.k_pages, self.pool.v_pages, pending) - for layer, (k, v) in enumerate(zip(k_pages, v_pages)): - self.pool.replace_layer(layer, k, v) - self.plane.confirm_scrubbed(pending) - def _preempt_newest(self, batch: list[PagedRequest]) -> None: """Return the newest request's pages to the pool and requeue it. From 8ddda20627bc6c8cdb98412a97c1ffc76cb67d49 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 28 Aug 2026 12:34:43 +0000 Subject: [PATCH 14/21] feat(inference): join the paged driver to MaxEngine through a shared step seam A step function received only the page bookkeeping, while a forward pass also needs token ids, absolute positions, segment ids and a sample index. That is why nothing connected `PagedDriver` to `MaxEngine`: a contract mismatch, not missing plumbing. `step_inputs.py` adds the missing layer between `prepare_step` and the forward pass. `MaxEngine.paged_step` is the pure forward over an already-reserved step, and `prefill_paged` / `generate_paged` are now admission and bookkeeping around the same two calls the driver makes -- so input assembly exists once. That matters because it decides absolute positions, and a mistake there produces plausible text rather than an error: a prefix hit starts the query at `cached_tokens`, and a preemption replay makes the prompt longer than the one submitted. One rule covers both phases. A slice's tokens occupy `start .. start+len-1`; prefill sets `start = cached_tokens`, decode sets `start = prompt_len + len(generated) - 1`. Deriving positions from the request also removes an ordering hazard: `generate_paged` had to read `page_map.seq_len` *before* `prepare_step`, because reservation advances it. Fixes a live trap in prefill sampling. Prefill packs requests along the sequence axis at batch one, so `logits[arange(batch), sample_at]` returned a single token however many were packed -- and the driver has always batched prefill. Carrying `sample_rows` makes the gather one expression in both phases; single-request prefill is unchanged. `prompt_tokens` becomes required for a driven request, checked in `submit` rather than mid-run, because by then the pool holds pages and other requests have been scheduled around it. This replaces a test asserting the opposite, which was written when the field existed only to offer the prefix cache something to match; the cache remains a control-plane switch. `RequestSlice` carries the tokens to feed rather than the whole context plus an offset. The first shape forced `generate_paged` to fabricate a sequence-length zero array to make a one-token decode indexable. Verified: 212 CPU and 22 GPU tests, including new coverage that the driver reproduces the engine entry points token for token and that a packed two-request prefill gives each request the token it would have got alone. Model-level dense versus paged still matches. Harness rewiring onto the driver is the remaining piece and is unaffected, since `run_paged` drives the engine entry points rather than the step seam. Co-authored-by: Cursor --- .../inference/kv_execution/__init__.py | 4 + src/maxtext/inference/kv_execution/driver.py | 85 +++++- .../inference/kv_execution/step_inputs.py | 230 +++++++++++++++ src/maxtext/inference/maxengine/maxengine.py | 200 ++++++++----- tests/unit/gpu_paged_decode_parity_test.py | 125 ++++++++ tests/unit/kv_execution_test.py | 278 ++++++++++++++++-- tests/unit/kv_paged_runtime_test.py | 35 ++- 7 files changed, 846 insertions(+), 111 deletions(-) create mode 100644 src/maxtext/inference/kv_execution/step_inputs.py diff --git a/src/maxtext/inference/kv_execution/__init__.py b/src/maxtext/inference/kv_execution/__init__.py index b303b641e8..8281c942da 100644 --- a/src/maxtext/inference/kv_execution/__init__.py +++ b/src/maxtext/inference/kv_execution/__init__.py @@ -41,6 +41,7 @@ scrub_pages, scrub_pages_all_layers, ) +from maxtext.inference.kv_execution.step_inputs import RequestSlice, StepInputs, build_step_inputs from maxtext.inference.kv_execution.step_view import StepView, build_step_view __all__ = [ @@ -48,12 +49,15 @@ "PagedDriver", "PagedKvPool", "PagedRequest", + "RequestSlice", + "StepInputs", "StepOutcome", "StepShape", "StepView", "allocate_pool", "batch_ladder", "bucket_up", + "build_step_inputs", "build_step_view", "build_storage_layout", "poison_pages", diff --git a/src/maxtext/inference/kv_execution/driver.py b/src/maxtext/inference/kv_execution/driver.py index 032f453954..5d41434d2d 100644 --- a/src/maxtext/inference/kv_execution/driver.py +++ b/src/maxtext/inference/kv_execution/driver.py @@ -62,6 +62,7 @@ from maxtext.inference.kv_execution.engine_adapter import PagedRuntime from maxtext.inference.kv_execution.pool_factory import PagedKvPool from maxtext.inference.kv_execution.pool_ops import poison_pages +from maxtext.inference.kv_execution.step_inputs import RequestSlice, StepInputs, build_step_inputs from maxtext.inference.kv_execution.step_view import StepView _DEFAULT_NAMESPACE = CacheNamespace() @@ -140,11 +141,23 @@ class StepOutcome: num_requests: int num_tokens: int preempted: int = 0 + # The requests this step advanced, in the order their tokens were returned. + # Carried so an observer can attribute per-request timing -- which is what a + # serving harness needs for TTFT and ITL, and the reason it previously had to + # re-implement the scheduling loop to get it. + batch: tuple[PagedRequest, ...] = () -# Takes the padded view and the pool; returns one token per *active* request and -# the rebound pool arrays, which alias the ones passed in. -StepFn = Callable[[StepView, PagedKvPool], np.ndarray] +# Takes the padded page view, the assembled token-side operands and the pool; +# returns one token per *active* request. The pool arrives so an implementation +# can pass it to an aliasing forward pass; the rebound arrays come back through +# the pool object rather than the return value. +# +# `StepInputs` is here because a `StepView` describes *pages* and a forward pass +# also needs tokens, absolute positions, segment ids and a sample index. Without +# it a step function cannot drive a model, which is why nothing connected the +# driver to `MaxEngine` for several milestones. +StepFn = Callable[[StepView, StepInputs, PagedKvPool], np.ndarray] class PagedDriver: @@ -212,6 +225,27 @@ def num_live(self) -> int: return len(self._live) def submit(self, requests: Sequence[PagedRequest]) -> None: + """Queue requests, rejecting any the driver could not actually run. + + `prompt_tokens` is validated here rather than at the step that needs it. The + driver assembles the tokens it feeds the model, so a request without them + cannot run -- but discovering that mid-run means the pool already holds pages + for it and other requests have already been scheduled around it. Failing at + the boundary tells the caller while the caller can still do something. + + Note this is *not* the prefix cache's opt-in. That is a config switch on the + control plane; `prompt_tokens` being present used to double as an implicit + second switch, which is a coincidence rather than a design. + """ + missing = [r.request_id for r in requests if r.token_ids() is None] + if missing: + # Named but capped: a whole trace missing the field is a single mistake, and + # printing six hundred ids buries the message that explains it. + shown = ", ".join(missing[:5]) + (f", ... and {len(missing) - 5} more" if len(missing) > 5 else "") + raise ValueError( + f"{len(missing)} of {len(requests)} submitted requests have no prompt_tokens, so the driver " + f"has nothing to feed the model: {shown}. Set PagedRequest.prompt_tokens on every request." + ) self._waiting.extend(requests) def run(self, max_steps: int = 100_000) -> list[PagedRequest]: @@ -343,7 +377,12 @@ def _run_phase( ) shape = view.shape - next_tokens = np.asarray(self.step_fn(view, self.pool)).reshape(-1) + inputs = build_step_inputs( + [self._slice_for(request, lens[i], is_decode=is_decode) for i, request in enumerate(batch)], + shape, + is_decode=is_decode, + ) + next_tokens = np.asarray(self.step_fn(view, inputs, self.pool)).reshape(-1) if next_tokens.size < len(batch): raise ValueError(f"the step function returned {next_tokens.size} tokens for {len(batch)} requests") @@ -354,6 +393,44 @@ def _run_phase( num_requests=len(batch), num_tokens=sum(lens), preempted=preempted, + batch=tuple(batch), + ) + + def _slice_for(self, request: PagedRequest, query_len: int, *, is_decode: bool) -> RequestSlice: + """This request's contribution to the step: what to feed and where it sits. + + The driver's half of the one position rule. A step feeds + `token_ids[start : start + query_len]` at absolute positions + `start .. start + query_len - 1`, and the phases differ only in `start`: + + * prefill runs the context the cache did not supply, so it starts at + `cached_tokens` -- zero for a fresh request, and past the shared prefix + after a hit. Starting at zero on a hit would rotate the suffix as though + it began the sequence, which RoPE does not forgive. + * decode runs the single token the previous step produced, which sits at + the end of the recorded context. + + A replay after preemption needs no special case: the driver retains + `generated`, so `token_ids()` and `prefill_len()` already describe the longer + prompt. + """ + context = request.token_ids() + if context is None: + # `submit` rejects these, so reaching here means a request arrived by some + # other route. Kept as an invariant guard rather than the primary check, + # because inventing tokens would produce fluent output from garbage. + raise ValueError( + f"request {request.request_id!r} has no prompt_tokens, so this step has nothing to feed the " + f"model. Every request must carry them; `submit` normally catches this." + ) + if is_decode: + start = request.prompt_len + len(request.generated) - 1 + else: + start = request.cached_tokens + return RequestSlice( + tokens=np.asarray(context).reshape(-1)[start : start + query_len], + start=start, + query_len=query_len, ) def _preempt_newest(self, batch: list[PagedRequest]) -> None: diff --git a/src/maxtext/inference/kv_execution/step_inputs.py b/src/maxtext/inference/kv_execution/step_inputs.py new file mode 100644 index 0000000000..4edec97029 --- /dev/null +++ b/src/maxtext/inference/kv_execution/step_inputs.py @@ -0,0 +1,230 @@ +"""What a step feeds the model: tokens, positions, and where to sample. + +`StepView` describes *pages*. This describes *tokens*. Both are needed for a +step, and keeping them apart is deliberate -- the page bookkeeping comes from the +control plane and knows nothing about token ids, while these arrays come from the +requests and know nothing about pages. + +**One position rule, and it is the reason this module exists.** A slice's +``tokens`` occupy absolute positions ``start .. start + len(tokens) - 1``. Callers +differ only in where they start: + + * prefill: ``start = cached_tokens``, feeding the context the cache did not + supply + * decode: ``start = prompt_len + len(generated) - 1``, feeding the single token + the previous step produced + +Decode is not a special case in the rule, only in the arithmetic that fills it. + +Getting positions wrong is the failure this module is built to prevent, and it is +not a crash. RoPE encodes absolute position and is not translation invariant, so +a suffix rotated as though it began the sequence produces K/V that does not +belong after the prefix it follows -- plausible text from a wrong computation. The +two ways to reach that are a prefix-cache hit, where the query starts at +``cached_tokens``, and a replay after preemption, where the retained generated +tokens make the prompt longer than the one submitted. + +**Numpy, not `jnp`, and that is not an oversight.** Building these with `jnp` +compiles a fresh program per prompt length, because `arange(n) < prompt_len` +bakes the length into the jaxpr as a literal. An earlier measurement in this +project ended up three-quarters compile time that way while reporting zero +unwarmed shapes. These arrays are a few hundred bytes; numpy has no cache to miss. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import dataclasses +from typing import Sequence + +import numpy as np + +from maxtext.inference.kv_execution.bucketing import StepShape + +# Matches `maxtext.common.common_types.DECODING_ACTIVE_SEQUENCE_INDICATOR`. +# Duplicated as a literal rather than imported because that module pulls in jax, +# and this one is deliberately host-only so it stays testable with no accelerator. +_ACTIVE_SEGMENT = 1 + + +@dataclasses.dataclass(frozen=True) +class RequestSlice: + """One request's contribution to a step: what to feed, and where it sits. + + `tokens` is *exactly* what this step feeds -- not the whole context with an + index into it. An earlier shape took the full context plus a start offset, and + it forced `generate_paged` to fabricate a zero-filled array as long as the + sequence just to make a single-token decode indexable. That is wasteful on a + per-step path and it invents context that does not exist. + + `query_len` is carried anyway, redundantly with `len(tokens)`, because they come + from different places: `query_len` is what the page table reserved and + `len(tokens)` is what the caller actually assembled. Requiring them to agree is + what catches a caller feeding fewer tokens than it reserved positions for, which + would otherwise leave the pool holding K/V for positions no query covered. + + Deliberately *not* `PagedRequest`. `MaxEngine.prefill_paged` has a handle, a + padded prompt and a true length; it has no `PagedRequest` and should not, since + that is the driver's scheduling record. If assembly demanded one, the engine + entry points could not share it and input assembly would exist twice -- which + is the thing this design is for. + """ + + tokens: np.ndarray + start: int + query_len: int + + +@dataclasses.dataclass(frozen=True) +class StepInputs: + """The token-side operands of one step, padded to the bucketed shape. + + Shapes differ by phase, inherited from the attention layer rather than chosen + here: prefill packs requests along the sequence axis at batch 1, decode batches + them along the batch axis. + + ============ ==================== =================== + field prefill decode + ============ ==================== =================== + tokens ``[1, width]`` ``[num_requests, 1]`` + positions ``[1, width]`` ``[num_requests, 1]`` + segment_ids ``[1, width]`` ``None`` + sample_rows all zero ``arange(n)`` + sample_at last index per req all zero + ============ ==================== =================== + + `sample_rows` is carried explicitly so the sampling gather is one expression, + ``logits[sample_rows, sample_at]``, in both phases. That is also what makes + batched prefill possible: with a single packed row, several requests differ in + their sample *position* rather than their sample row. + """ + + tokens: np.ndarray + positions: np.ndarray + segment_ids: np.ndarray | None + sample_rows: np.ndarray + sample_at: np.ndarray + + +def _validate(slices: Sequence[RequestSlice], shape: StepShape, is_decode: bool) -> None: + """Fail on the four ways a caller can present an impossible step.""" + if not slices: + raise ValueError("a step needs at least one request slice") + + for index, item in enumerate(slices): + if item.query_len < 1: + raise ValueError(f"slice {index} has query_len {item.query_len}; a step must run at least one token") + if item.start < 0: + raise ValueError(f"slice {index} has start {item.start}; positions are absolute and cannot be negative") + supplied = int(np.asarray(item.tokens).reshape(-1).size) + if supplied != item.query_len: + # The page table reserved `query_len` positions. Feeding a different number + # would leave the pool holding K/V for positions no query covered, or run + # tokens the table has nowhere to put. Loud beats plausible. + raise ValueError( + f"slice {index} reserved {item.query_len} positions but supplied {supplied} tokens; " + f"these must agree or the pool and the query disagree about what this step ran" + ) + + if is_decode: + if any(item.query_len != 1 for item in slices): + raise ValueError("a decode step runs exactly one token per request") + if len(slices) > shape.num_requests: + raise ValueError( + f"{len(slices)} requests do not fit the decode batch bucket of {shape.num_requests}" + ) + else: + total = sum(item.query_len for item in slices) + if total > shape.num_tokens: + raise ValueError( + f"{total} query tokens do not fit the token bucket of {shape.num_tokens}" + ) + + +def build_step_inputs( + slices: Sequence[RequestSlice], + shape: StepShape, + *, + is_decode: bool, +) -> StepInputs: + """Assemble one step's token-side operands. + + Takes a `StepShape` rather than a `StepView` on purpose. Everything needed here + is a host int -- the token bucket and the batch bucket -- and accepting the view + would put its `jax.Array` fields within reach, which is how host arithmetic + quietly acquires a device dependency. Sample indices come from cumulating the + slices' own `query_len`, not from the view's `cu_seqlens_q`. + + Args: + slices: one per active request, in the same order the page table was built. + shape: the bucketed shape this step is traced for. + is_decode: selects the layout, since the two phases pack differently. + + Returns: + Arrays padded to `shape`, with every padded element inert. + """ + _validate(slices, shape, is_decode) + + if is_decode: + width = shape.num_requests + tokens = np.zeros((width, 1), np.int32) + positions = np.zeros((width, 1), np.int32) + for row, item in enumerate(slices): + tokens[row, 0] = int(np.asarray(item.tokens).reshape(-1)[0]) + positions[row, 0] = int(item.start) + # One row each, so the sample row varies and the position does not. Padded + # rows sample position 0 of their own row, which holds a zero token; their + # slot mapping points at the reserved padding page, so nothing they compute + # is read. + return StepInputs( + tokens=tokens, + positions=positions, + # None, not zeros: the model refuses segment ids in autoregressive mode, + # where every token is by definition in the active sequence. + segment_ids=None, + sample_rows=np.arange(width, dtype=np.int32), + sample_at=np.zeros((width,), np.int32), + ) + + width = shape.num_tokens + tokens = np.zeros((1, width), np.int32) + positions = np.zeros((1, width), np.int32) + segment_ids = np.zeros((1, width), np.int32) + sample_at = np.zeros((len(slices),), np.int32) + + cursor = 0 + for index, item in enumerate(slices): + end = cursor + item.query_len + tokens[0, cursor:end] = np.asarray(item.tokens).reshape(-1) + # Absolute, so a cached prefix or a preemption replay lands where RoPE + # expects rather than at zero. + positions[0, cursor:end] = np.arange(item.start, item.start + item.query_len, dtype=np.int32) + segment_ids[0, cursor:end] = _ACTIVE_SEGMENT + # The logits this step produces cover only the tokens it ran, so the sample + # index is within the packed row rather than within the sequence. + sample_at[index] = end - 1 + cursor = end + + return StepInputs( + tokens=tokens, + positions=positions, + segment_ids=segment_ids, + # One packed row, so every request samples from row zero and they differ + # only in position. This is precisely what a `rows = arange(batch)` gather + # cannot express, and why batched prefill returned a single token before. + sample_rows=np.zeros((len(slices),), np.int32), + sample_at=sample_at, + ) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 9a9268cd91..8337dd6d82 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -2064,6 +2064,7 @@ def _paged_forward_jit( kv_page_indices: jax.Array, kv_last_page_lens: jax.Array, cu_seqlens_q: jax.Array, + sample_rows: jax.Array, sample_at: jax.Array, rng: PRNGKeyType, model_mode: str, @@ -2125,12 +2126,18 @@ def _paged_forward_jit( attention_metadata=plan, ) - # One row per request, each sampled at its own index: for prefill that is the - # last prompt token, for decode the single new token. Doing it with a gather + # One logit row per request, gathered at its own (row, position). A gather # rather than a slice is what lets a decode batch hold requests at different # positions, which is the whole point of paging. - rows = jnp.arange(logits.shape[0]) - selected = logits[rows, sample_at][:, None, :] + # + # Both indices are supplied because the two phases lay a batch out along + # different axes: decode batches along rows and samples position 0 of each, + # while prefill *packs* requests into one row and samples several positions + # within it. An earlier version derived the rows as `arange(logits.shape[0])`, + # which is correct for decode and silently returns a single token for a packed + # multi-request prefill -- the driver batches prefill, so that was a live trap + # rather than a hypothetical. + selected = logits[sample_rows, sample_at][:, None, :] selected = jax.lax.with_sharding_constraint(selected, self.replicated_sharding) rng, sample_rng = jax.random.split(rng) @@ -2144,6 +2151,82 @@ def _paged_forward_jit( ) return sampled, selected, updated_pools + def paged_step( + self, + *, + params: Params, + view, + inputs, + is_decode: bool, + rng: PRNGKeyType | None = None, + algorithm: str | None = None, + topk: int | None = None, + nucleus_topp: float | None = None, + temperature: float | None = None, + ): + """One paged forward pass over an already-reserved step. No policy, no admission. + + The seam the scheduler plugs into. `view` is the page bookkeeping from + `PagedRuntime.prepare_step`; `inputs` is the token side from + `build_step_inputs`. Everything this needs has already been decided, which is + what makes it usable by both `PagedDriver` -- which has admitted and reserved + itself -- and by `prefill_paged` / `generate_paged`, which wrap it in their own + admission and bookkeeping. + + Returns `(sampled, selected)`: the tokens, and the logit rows they came from. + The logits are returned rather than dropped because `return_log_prob` needs + them, and building `ResultTokens` stays with the callers since it is a + JetStream-shaped concern the driver has no use for. + + The pool is rebound as a side effect, since the forward pass donates it. + """ + if rng is None: + if self.rng is None: + self.rng = jax.random.PRNGKey(0) + self.rng, rng = jax.random.split(self.rng) + + sampled, selected, updated = self._paged_forward_jit( + params=params, + kv_caches=self._paged_pool_arrays(), + tokens=jnp.asarray(inputs.tokens), + positions=jnp.asarray(inputs.positions), + segment_ids=None if inputs.segment_ids is None else jnp.asarray(inputs.segment_ids), + slot_mapping=view.slot_mapping, + kv_indptr=view.kv_indptr, + kv_page_indices=view.kv_page_indices, + kv_last_page_lens=view.kv_last_page_lens, + cu_seqlens_q=view.cu_seqlens_q, + sample_rows=jnp.asarray(inputs.sample_rows), + sample_at=jnp.asarray(inputs.sample_at), + rng=rng, + model_mode=MODEL_MODE_AUTOREGRESSIVE if is_decode else MODEL_MODE_PREFILL, + max_seqlen_k=view.shape.max_seqlen_k, + is_decode=is_decode, + algorithm=algorithm, + topk=topk, + nucleus_topp=nucleus_topp, + temperature=temperature, + ) + self._rebind_paged_pool(updated) + return sampled, selected + + def paged_step_fn(self, params: Params, **kwargs): + """A `PagedDriver` step function bound to this engine. + + The driver's contract is `(view, inputs, pool) -> tokens`. The pool arrives + for implementations that need it; this one does not, because the engine + already holds the pool it donates and rebinds. + """ + + def step(view, inputs, pool): + del pool + sampled, _ = self.paged_step( + params=params, view=view, inputs=inputs, is_decode=view.shape.is_decode, **kwargs + ) + return np.asarray(sampled).reshape(-1)[: inputs.sample_at.size] + + return step + def _paged_result_tokens(self, sampled: jax.Array, selected: jax.Array, step: int): """Wrap sampled tokens in the same `ResultTokens` shape the dense path returns.""" all_valid = jnp.ones(sampled.shape, dtype=jnp.int8) @@ -2231,46 +2314,31 @@ def prefill_paged( # are what RoPE encodes, so restarting them at zero would rotate the suffix # as though it began the sequence and produce K/V that does not belong after # the cached prefix -- a wrong answer rather than a slow one. - width = view.shape.num_tokens - prompt = np.asarray(padded_tokens, dtype=np.int32).reshape(-1)[cached : cached + width] - token_row = np.zeros((1, width), np.int32) - token_row[0, : prompt.size] = prompt - index = np.arange(width, dtype=np.int32) - tokens = jnp.asarray(token_row) - positions = jnp.asarray((index + cached)[None, :]) - segment_ids = jnp.asarray( - np.where(index < query_len, DECODING_ACTIVE_SEQUENCE_INDICATOR, 0).astype(np.int32)[None, :] - ) - - if rng is None: - if self.rng is None: - self.rng = jax.random.PRNGKey(0) - self.rng, rng = jax.random.split(self.rng) + # Assembled by the shared seam rather than here, so the absolute-position + # arithmetic -- which a prefix hit and a preemption replay both perturb -- has + # exactly one implementation. `start=cached` is what makes the suffix rotate + # as a suffix; RoPE is not translation invariant, and getting it wrong yields + # plausible text rather than an error. + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_execution.step_inputs import RequestSlice, build_step_inputs - sampled, selected, updated = self._paged_forward_jit( + prompt = np.asarray(padded_tokens, dtype=np.int64).reshape(-1) + inputs = build_step_inputs( + [RequestSlice(tokens=prompt[cached : cached + query_len], start=cached, query_len=query_len)], + view.shape, + is_decode=False, + ) + sampled, selected = self.paged_step( params=params, - kv_caches=self._paged_pool_arrays(), - tokens=tokens, - positions=positions, - segment_ids=segment_ids, - slot_mapping=view.slot_mapping, - kv_indptr=view.kv_indptr, - kv_page_indices=view.kv_page_indices, - kv_last_page_lens=view.kv_last_page_lens, - cu_seqlens_q=view.cu_seqlens_q, - # Indexed within the query window, not the sequence: the logits this step - # produces cover only the tokens it ran. - sample_at=jnp.asarray([query_len - 1], jnp.int32), - rng=rng, - model_mode=MODEL_MODE_PREFILL, - max_seqlen_k=view.shape.max_seqlen_k, + view=view, + inputs=inputs, is_decode=False, + rng=rng, algorithm=algorithm, topk=topk, nucleus_topp=nucleus_topp, temperature=temperature, ) - self._rebind_paged_pool(updated) runtime.track(handle, slot=slot) return handle, self._paged_result_tokens(sampled, selected, step=0) @@ -2304,55 +2372,41 @@ def generate_paged( raise ValueError("generate_paged needs at least one live request") # Read before reserving: this is each request's position for the token about - # to be written, and reserving advances the recorded length past it. - positions = [runtime.control_plane.page_map.seq_len(h) for h in handles] + # to be written, and reserving advances the recorded length past it. The + # driver avoids this ordering constraint entirely by deriving positions from + # the request; here the caller supplies only the token, so the page map is the + # only thing that knows where it goes. + positions = [int(runtime.control_plane.page_map.seq_len(h)) for h in handles] view = runtime.prepare_step(handles, [1] * len(handles), is_decode=True) if view is None: return None, False - # numpy for the same reason as `prefill_paged`: padding a variable-length - # array with `jnp` compiles once per distinct live-request count, and the - # live count changes constantly under continuous batching. - width = view.shape.num_requests - token_col = np.zeros((width, 1), np.int32) - supplied = np.asarray(next_tokens, dtype=np.int32).reshape(-1)[:width] - token_col[: supplied.size, 0] = supplied - position_col = np.zeros((width, 1), np.int32) - recorded = np.asarray(positions, dtype=np.int32).reshape(-1)[:width] - position_col[: recorded.size, 0] = recorded - tokens = jnp.asarray(token_col) - pos = jnp.asarray(position_col) - - if rng is None: - if self.rng is None: - self.rng = jax.random.PRNGKey(0) - self.rng, rng = jax.random.split(self.rng) - - sampled, selected, updated = self._paged_forward_jit( + # One token per request at its own absolute position. Nothing else is needed: + # a slice carries what the step feeds, not the context behind it. + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_execution.step_inputs import RequestSlice, build_step_inputs + + supplied = np.asarray(next_tokens, dtype=np.int64).reshape(-1)[: len(handles)] + inputs = build_step_inputs( + [ + RequestSlice(tokens=np.asarray([token], np.int64), start=position, query_len=1) + for position, token in zip(positions, supplied.tolist()) + ], + view.shape, + is_decode=True, + ) + sampled, selected = self.paged_step( params=params, - kv_caches=self._paged_pool_arrays(), - tokens=tokens, - positions=pos, - # None, not zeros: the model refuses segment ids in autoregressive mode, - # where every token is by definition in the active sequence. - segment_ids=None, - slot_mapping=view.slot_mapping, - kv_indptr=view.kv_indptr, - kv_page_indices=view.kv_page_indices, - kv_last_page_lens=view.kv_last_page_lens, - cu_seqlens_q=view.cu_seqlens_q, - sample_at=jnp.zeros((width,), jnp.int32), - rng=rng, - model_mode=MODEL_MODE_AUTOREGRESSIVE, - max_seqlen_k=view.shape.max_seqlen_k, + view=view, + inputs=inputs, is_decode=True, + rng=rng, algorithm=algorithm, topk=topk, nucleus_topp=nucleus_topp, temperature=temperature, ) - self._rebind_paged_pool(updated) return self._paged_result_tokens(sampled, selected, step=1), True def get_prefix_destination_sharding(self) -> Any: diff --git a/tests/unit/gpu_paged_decode_parity_test.py b/tests/unit/gpu_paged_decode_parity_test.py index 521b69c4d1..eea5c96bdf 100644 --- a/tests/unit/gpu_paged_decode_parity_test.py +++ b/tests/unit/gpu_paged_decode_parity_test.py @@ -371,6 +371,131 @@ def test_paged_decode_matches_the_dense_engine_path_token_for_token(self): paged = self._paged_rollout(_config(**_PAGED), params_state) self._verify_tokens_are_argmax(_config(**_DENSE), params_state, paged) + def test_the_driver_reproduces_the_engine_entry_points(self): + """`PagedDriver` driving the engine must match `prefill_paged`/`generate_paged`. + + The equivalence that licenses the refactor. Both now go through the same + `build_step_inputs` and `paged_step`, but they reach them differently: the + entry points admit one request and assemble a slice from their arguments, + while the driver schedules a queue and assembles slices from `PagedRequest`. + A divergence here means the driver's half of the position rule disagrees with + the engine's, which is precisely the failure a shared seam is meant to make + impossible -- so it is worth checking rather than assuming. + """ + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_execution.driver import PagedDriver, PagedRequest + + cfg = _config(**_PAGED) + params_state = self._build_params(cfg) + reference = self._paged_rollout(cfg, params_state, steps=STEPS) + + engine = maxengine.MaxEngine(cfg, _devices()) + params = engine.load_params(params=params_state) + runtime = engine.init_paged_runtime() + driver = PagedDriver( + runtime.control_plane, + runtime.pool, + engine.paged_step_fn(params), + max_batch=1, + max_batched_tokens=cfg.max_prefill_predict_length, + ) + # The driver counts its own generated tokens, so it wants exactly as many as + # the reference produced: one from prefill plus STEPS decodes. + driver.submit( + [ + PagedRequest( + request_id="driven", + prompt_len=len(PROMPT), + max_new_tokens=STEPS + 1, + prompt_tokens=np.asarray(PROMPT, dtype=np.int64), + ) + ] + ) + done = driver.run() + + self.assertEqual(len(done), 1) + self._assert_same_tokens(reference, done[0].generated, "driver against the engine entry points") + self.assertEqual( + runtime.control_plane.allocator.num_allocated_pages, 0, "the driver must release what it took" + ) + + def test_a_batched_prefill_samples_every_request(self): + """Two requests in one prefill step must each get their own token. + + New capability, and the reason it needs its own test: prefill packs requests + along the sequence axis at batch one, so before `sample_rows` existed the + gather was `logits[arange(batch), sample_at]` and returned a *single* token + however many requests were packed. The driver has always batched prefill, so + that was a live trap rather than a hypothetical -- and a wrong-length result is + the good case, because a driver that receives one token for two requests + assigns the wrong token to the second. + + Compared against prefilling the same two prompts separately, which is the only + reference that distinguishes "packed correctly" from "packed consistently". + """ + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_execution.step_inputs import RequestSlice, build_step_inputs + + cfg = _config(**_PAGED) + params_state = self._build_params(cfg) + + # Two distinct prompts, so a swap or a duplicate is visible. + first, second = PROMPT[:6], PROMPT[2:10] + self.assertNotEqual(first, second, "the two prompts must differ for this test to mean anything") + + separate = [] + for index, prompt in enumerate((first, second)): + engine = maxengine.MaxEngine(cfg, _devices()) + params = engine.load_params(params=params_state) + engine.init_paged_runtime() + padded = jnp.asarray( + list(prompt) + [0] * (cfg.max_prefill_predict_length - len(prompt)), dtype=jnp.int32 + ) + handle, result = engine.prefill_paged( + params=params, padded_tokens=padded, true_length=len(prompt), request_id=f"solo{index}" + ) + self.assertIsNotNone(handle) + separate.append(int(result.data[0, 0])) + engine.release(handle) + + # Now both in one packed step, driven at the seam rather than through the + # single-request entry point. + engine = maxengine.MaxEngine(cfg, _devices()) + params = engine.load_params(params=params_state) + # Two rows explicitly: `init_paged_runtime` defaults to the dense batch width, + # which is one here, and a one-row page map cannot hold a two-request batch. + runtime = engine.init_paged_runtime( + max_requests=2, max_batched_tokens=cfg.max_prefill_predict_length + ) + handles = [ + runtime.admit(request_id=f"packed{i}", prompt_len=len(p), max_new_tokens=1) + for i, p in enumerate((first, second)) + ] + self.assertTrue(all(h is not None for h in handles), "the pool refused a two-request batch") + + query_lens = [len(first), len(second)] + view = runtime.prepare_step(handles, query_lens, is_decode=False, num_requests=2) + self.assertIsNotNone(view, "the pool could not back a two-request prefill") + inputs = build_step_inputs( + [ + RequestSlice(tokens=np.asarray(p, np.int64), start=0, query_len=len(p)) + for p in (first, second) + ], + view.shape, + is_decode=False, + ) + self.assertEqual(inputs.sample_at.size, 2, "one sample index per packed request") + sampled, _ = engine.paged_step(params=params, view=view, inputs=inputs, is_decode=False) + + packed = np.asarray(sampled).reshape(-1)[:2].tolist() + self.assertEqual( + packed, + separate, + "a packed prefill must give each request the token it would have got alone", + ) + for handle in handles: + engine.release(handle) + def test_the_paged_path_allocates_no_dense_cache(self): """A dense cache alongside the pool would waste gigabytes at real sizes.""" cfg = _config(**_PAGED) diff --git a/tests/unit/kv_execution_test.py b/tests/unit/kv_execution_test.py index 921a141a63..dd255a602e 100644 --- a/tests/unit/kv_execution_test.py +++ b/tests/unit/kv_execution_test.py @@ -33,6 +33,7 @@ from maxtext.inference.kv_control import NativeKvControlPlane, RequestHandle from maxtext.inference.kv_execution import allocate_pool, PagedDriver, PagedRequest from maxtext.inference.kv_execution.bucketing import ( + StepShape, StepShapePlanner, batch_ladder, bucket_up, @@ -47,6 +48,7 @@ scrub_pages, scrub_pages_all_layers, ) +from maxtext.inference.kv_execution.step_inputs import RequestSlice, build_step_inputs from maxtext.inference.kv_execution.step_view import build_step_view PAGE = 16 @@ -337,6 +339,157 @@ def test_an_odd_page_count_is_padded_idempotently(self): self.assertTrue(bool((np.asarray(v)[5] == POISON_SENTINEL).all())) +class StepInputsTest(unittest.TestCase): + """Absolute positions, packing, and the ways a caller can present nonsense. + + This is where the position rule is locked down, because getting it wrong is not + a crash: RoPE encodes absolute position, so a suffix rotated as though it began + the sequence yields plausible text from a wrong computation. Cheap to pin here + and expensive to notice anywhere else. + """ + + def _shape(self, *, num_requests=4, num_tokens=64, is_decode=False): + return StepShape( + num_requests=num_requests, + num_tokens=num_tokens, + num_pages=16, + max_seqlen_k=128, + is_decode=is_decode, + ) + + def test_a_fresh_prefill_starts_at_position_zero(self): + context = np.arange(100, 110, dtype=np.int64) + inputs = build_step_inputs( + [RequestSlice(tokens=context[0:10], start=0, query_len=10)], + self._shape(), + is_decode=False, + ) + self.assertEqual(inputs.tokens.shape, (1, 64)) + np.testing.assert_array_equal(inputs.tokens[0, :10], context) + np.testing.assert_array_equal(inputs.positions[0, :10], np.arange(10)) + self.assertTrue(bool((inputs.segment_ids[0, :10] == 1).all())) + # Sampling the last token it ran, indexed within the packed row. + np.testing.assert_array_equal(inputs.sample_at, [9]) + np.testing.assert_array_equal(inputs.sample_rows, [0]) + + def test_a_prefix_hit_feeds_the_suffix_at_absolute_positions(self): + """The M5 hazard. Positions must not restart at zero for a cached prefix.""" + context = np.arange(200, 232, dtype=np.int64) # 32 tokens + inputs = build_step_inputs( + [RequestSlice(tokens=context[16:32], start=16, query_len=16)], + self._shape(), + is_decode=False, + ) + np.testing.assert_array_equal(inputs.tokens[0, :16], context[16:]) + np.testing.assert_array_equal(inputs.positions[0, :16], np.arange(16, 32)) + self.assertEqual(int(inputs.positions[0, 0]), 16, "the suffix must not be rotated as a prefix") + + def test_a_replay_after_preemption_feeds_the_longer_prompt(self): + """Generated tokens are retained, so the replayed prompt is longer.""" + prompt, generated = np.arange(300, 310, dtype=np.int64), np.arange(400, 404, dtype=np.int64) + context = np.concatenate([prompt, generated]) + inputs = build_step_inputs( + [RequestSlice(tokens=context[0:14], start=0, query_len=14)], + self._shape(), + is_decode=False, + ) + np.testing.assert_array_equal(inputs.tokens[0, :14], context) + np.testing.assert_array_equal(inputs.positions[0, :14], np.arange(14)) + self.assertEqual(int(inputs.sample_at[0]), 13, "sampling follows the whole retained context") + + def test_batched_prefill_packs_requests_and_samples_each(self): + """The capability the old single-row gather silently got wrong.""" + a = np.arange(10, 15, dtype=np.int64) # 5 tokens + b = np.arange(20, 27, dtype=np.int64) # 7 tokens + inputs = build_step_inputs( + [ + RequestSlice(tokens=a, start=0, query_len=5), + RequestSlice(tokens=b, start=0, query_len=7), + ], + self._shape(), + is_decode=False, + ) + np.testing.assert_array_equal(inputs.tokens[0, :5], a) + np.testing.assert_array_equal(inputs.tokens[0, 5:12], b) + # Each request's positions restart from its own start, not from the packed + # offset: they are sequence positions, not row offsets. + np.testing.assert_array_equal(inputs.positions[0, :5], np.arange(5)) + np.testing.assert_array_equal(inputs.positions[0, 5:12], np.arange(7)) + np.testing.assert_array_equal(inputs.sample_at, [4, 11], "one sample per request, within the row") + np.testing.assert_array_equal(inputs.sample_rows, [0, 0], "a packed row samples from row zero") + + def test_prefill_pads_the_tail_inertly(self): + inputs = build_step_inputs( + [RequestSlice(tokens=np.arange(4, dtype=np.int64), start=0, query_len=4)], + self._shape(num_tokens=64), + is_decode=False, + ) + self.assertTrue(bool((inputs.tokens[0, 4:] == 0).all())) + self.assertTrue(bool((inputs.segment_ids[0, 4:] == 0).all()), "padded rows must be out of segment") + + def test_decode_batches_along_rows_and_carries_absolute_positions(self): + # Two requests at *different* positions, which is the whole point of paging. + short = np.arange(500, 512, dtype=np.int64) # 12 tokens + long = np.arange(600, 640, dtype=np.int64) # 40 tokens + inputs = build_step_inputs( + [ + RequestSlice(tokens=short[11:12], start=11, query_len=1), + RequestSlice(tokens=long[39:40], start=39, query_len=1), + ], + self._shape(num_requests=4, is_decode=True), + is_decode=True, + ) + self.assertEqual(inputs.tokens.shape, (4, 1)) + self.assertIsNone(inputs.segment_ids, "autoregressive mode refuses segment ids") + self.assertEqual(int(inputs.tokens[0, 0]), 511, "decode feeds the last context token") + self.assertEqual(int(inputs.tokens[1, 0]), 639) + np.testing.assert_array_equal(inputs.positions[:2, 0], [11, 39]) + np.testing.assert_array_equal(inputs.sample_rows, np.arange(4)) + np.testing.assert_array_equal(inputs.sample_at, np.zeros(4)) + self.assertTrue(bool((inputs.tokens[2:] == 0).all()), "unused rows stay inert") + + def test_fewer_tokens_than_reserved_positions_is_rejected(self): + """`query_len` is what the page table reserved; the tokens are what was assembled. + + Letting them disagree would leave the pool holding K/V for positions no query + covered, so the two counts are required to match rather than the shorter one + winning silently. + """ + with self.assertRaises(ValueError): + build_step_inputs( + [RequestSlice(tokens=np.arange(5, dtype=np.int64), start=0, query_len=8)], + self._shape(), + is_decode=False, + ) + + def test_queries_exceeding_the_token_bucket_are_rejected(self): + with self.assertRaises(ValueError): + build_step_inputs( + [RequestSlice(tokens=np.arange(100, dtype=np.int64), start=0, query_len=100)], + self._shape(num_tokens=64), + is_decode=False, + ) + + def test_more_requests_than_the_decode_bucket_are_rejected(self): + slices = [ + RequestSlice(tokens=np.asarray([3], np.int64), start=3, query_len=1) for _ in range(5) + ] + with self.assertRaises(ValueError): + build_step_inputs(slices, self._shape(num_requests=4, is_decode=True), is_decode=True) + + def test_a_multi_token_decode_is_rejected(self): + with self.assertRaises(ValueError): + build_step_inputs( + [RequestSlice(tokens=np.arange(2, dtype=np.int64), start=4, query_len=2)], + self._shape(is_decode=True), + is_decode=True, + ) + + def test_an_empty_step_is_rejected(self): + with self.assertRaises(ValueError): + build_step_inputs([], self._shape(), is_decode=False) + + class LayoutBuilderTest(unittest.TestCase): """Config and mesh into pool geometry.""" @@ -376,10 +529,12 @@ class _CountingStep: def __init__(self, token=7): self.token = token self.views = [] + self.inputs = [] - def __call__(self, view, pool): + def __call__(self, view, inputs, pool): del pool self.views.append(view) + self.inputs.append(inputs) return np.full((view.num_active_requests,), self.token, dtype=np.int32) @@ -395,8 +550,17 @@ def _driver(self, *, num_pages=65, max_requests=4, max_context_len=128, step=Non return PagedDriver(plane, pool, step or _CountingStep(), max_batch=max_requests, **kw) def _requests(self, count, prompt_len=20, max_new_tokens=8): + # `prompt_tokens` is set on every request because the driver now needs them: + # a step assembles the tokens it feeds the model, and inventing them would + # turn a caller's omission into fluent output from garbage. Distinct ids per + # request so nothing is accidentally a prefix of anything else. return [ - PagedRequest(request_id=f"r{i}", prompt_len=prompt_len, max_new_tokens=max_new_tokens) + PagedRequest( + request_id=f"r{i}", + prompt_len=prompt_len, + max_new_tokens=max_new_tokens, + prompt_tokens=np.arange(1000 * (i + 1), 1000 * (i + 1) + prompt_len, dtype=np.int64), + ) for i in range(count) ] @@ -427,7 +591,15 @@ def test_a_mixed_length_batch_completes_without_leaking(self): driver = self._driver(max_requests=4) lengths = [7, 19, 33, 48, 64, 5] driver.submit( - [PagedRequest(request_id=f"r{i}", prompt_len=n, max_new_tokens=6) for i, n in enumerate(lengths)] + [ + PagedRequest( + request_id=f"r{i}", + prompt_len=n, + max_new_tokens=6, + prompt_tokens=np.arange(1000 * (i + 1), 1000 * (i + 1) + n, dtype=np.int64), + ) + for i, n in enumerate(lengths) + ] ) done = driver.run() @@ -437,19 +609,53 @@ def test_a_mixed_length_batch_completes_without_leaking(self): self.assertEqual(driver.plane.allocator.available_pages, driver.plane.allocator.capacity_pages) self.assertEqual(driver.plane.num_live, 0) + def test_the_driver_feeds_absolute_positions_in_both_phases(self): + """The driver's half of the position rule, separated from assembly's half. + + `StepInputsTest` covers assembly given correct slices; this covers the driver + producing correct slices. Split deliberately -- a test spanning both cannot + say which half is wrong, and the two halves fail for different reasons. + """ + step = _CountingStep(token=11) + driver = self._driver(step=step) + driver.submit(self._requests(1, prompt_len=20, max_new_tokens=3)) + driver.run() + + prefill, decodes = step.inputs[0], step.inputs[1:] + context = np.arange(1000, 1020, dtype=np.int64) + + # Prefill runs the whole prompt from position zero and samples its last token. + np.testing.assert_array_equal(prefill.tokens[0, :20], context) + np.testing.assert_array_equal(prefill.positions[0, :20], np.arange(20)) + np.testing.assert_array_equal(prefill.sample_at, [19]) + + # Each decode feeds the previous step's token at the next absolute position. + # Position 20 is the first generated token's slot, so the first decode -- which + # feeds that token back -- sits there, not at 21. + for index, inputs in enumerate(decodes): + self.assertEqual(int(inputs.tokens[0, 0]), 11, "decode feeds back what was generated") + self.assertEqual( + int(inputs.positions[0, 0]), + 20 + index, + "decode positions must continue the sequence, not restart", + ) + self.assertIsNone(inputs.segment_ids) + def test_churn_traces_a_bounded_set_of_shapes(self): """The exit criterion, host side. Without bucketing this grows with the trace.""" driver = self._driver(max_requests=8, max_context_len=256, num_pages=257) rng = np.random.default_rng(0) - driver.submit( - [ - PagedRequest( - request_id=f"r{i}", - prompt_len=int(rng.integers(4, 120)), - max_new_tokens=int(rng.integers(1, 30)), - ) - for i in range(60) - ] + + def churning(index): + prompt_len = int(rng.integers(4, 120)) + return PagedRequest( + request_id=f"r{index}", + prompt_len=prompt_len, + max_new_tokens=int(rng.integers(1, 30)), + prompt_tokens=rng.integers(1, 30000, size=prompt_len, dtype=np.int64), + ) + + driver.submit([churning(i) for i in range(60)] ) driver.run() self.assertLessEqual(driver.num_distinct_shapes, driver.planner.max_distinct_shapes()) @@ -461,7 +667,15 @@ def test_a_pool_under_pressure_preempts_rather_than_deadlocking(self): # 8 usable pages of 16 tokens, four concurrent requests each wanting 55. driver = self._driver(num_pages=9, max_requests=4, max_context_len=64) driver.submit( - [PagedRequest(request_id=f"r{i}", prompt_len=30, max_new_tokens=25) for i in range(5)] + [ + PagedRequest( + request_id=f"r{i}", + prompt_len=30, + max_new_tokens=25, + prompt_tokens=np.arange(1000 * (i + 1), 1000 * (i + 1) + 30, dtype=np.int64), + ) + for i in range(5) + ] ) done = driver.run() @@ -474,7 +688,15 @@ def test_a_preempted_request_keeps_the_tokens_it_generated(self): """Preemption replays a longer prompt; it does not discard output.""" driver = self._driver(num_pages=9, max_requests=4, max_context_len=64) driver.submit( - [PagedRequest(request_id=f"r{i}", prompt_len=30, max_new_tokens=25) for i in range(5)] + [ + PagedRequest( + request_id=f"r{i}", + prompt_len=30, + max_new_tokens=25, + prompt_tokens=np.arange(1000 * (i + 1), 1000 * (i + 1) + 30, dtype=np.int64), + ) + for i in range(5) + ] ) done = driver.run() preempted = [r for r in done if r.preemptions] @@ -577,14 +799,28 @@ def test_a_diverging_prompt_shares_only_what_it_has_in_common(self): outcome.num_tokens, 64, "exactly the divergent tail should be computed, and the shared 128 skipped" ) - def test_a_request_without_tokens_is_unaffected(self): - """Supplying prompt ids is opt-in, so the cache must tolerate their absence.""" + def test_submit_rejects_a_request_without_tokens(self): + """Replaces an earlier test asserting the opposite, and the change is deliberate. + + That test read "supplying prompt ids is opt-in, so the cache must tolerate + their absence", which was true when `prompt_tokens` existed only to offer the + prefix cache something to match. The driver now assembles the tokens it feeds + the model, so a request without them cannot run at all. + + Rejected at `submit` rather than at the step that needs them: by then the pool + holds pages for the request and other requests have been scheduled around it, + so the caller learns too late to do anything useful. Note this does not make + the prefix cache mandatory -- that is still a control-plane switch. + """ driver = self._driver() - driver.submit([PagedRequest(request_id=f"r{i}", prompt_len=20, max_new_tokens=4) for i in range(3)]) - done = driver.run() - self.assertEqual(len(done), 3) - self.assertEqual(driver.plane.prefix_index.num_cached_pages, 0) - self.assertEqual(driver.plane.allocator.num_allocated_pages, 0) + with self.assertRaises(ValueError) as caught: + driver.submit( + [PagedRequest(request_id=f"r{i}", prompt_len=20, max_new_tokens=4) for i in range(3)] + ) + self.assertIn("prompt_tokens", str(caught.exception)) + # Named, so a caller with a large batch can find the offenders. + self.assertIn("r0", str(caught.exception)) + self.assertEqual(driver.num_waiting, 0, "a rejected batch must not be partially queued") def test_only_pages_with_computed_kv_are_published(self): """The final generated token has no K/V: its step never ran. diff --git a/tests/unit/kv_paged_runtime_test.py b/tests/unit/kv_paged_runtime_test.py index 12a0060b70..df94d8ef06 100644 --- a/tests/unit/kv_paged_runtime_test.py +++ b/tests/unit/kv_paged_runtime_test.py @@ -300,7 +300,12 @@ def body(query, key, value, k_pool, v_pool, slot_mapping, kv_indptr, kv_page_ind self._jitted = jax.jit(body, static_argnames=("max_seqlen_q", "max_seqlen_k", "is_decode")) - def __call__(self, view, pool): + def __call__(self, view, inputs, pool): + # `inputs` carries the tokens and positions a real forward pass would consume. + # This fake synthesises its own activations instead, because what it measures + # is retrace count against the page shapes -- but it takes the argument so it + # matches the contract every real step function implements. + del inputs seq_lens = np.asarray(view.seq_lens) self.raw_shapes.add( ( @@ -343,16 +348,18 @@ def test_churn_traces_no_more_shapes_than_the_ladders_allow(self): driver = PagedDriver(plane, allocate_pool(layout), step, max_batch=4, max_batched_tokens=64) rng = np.random.default_rng(0) - driver.submit( - [ - PagedRequest( - request_id=f"r{i}", - prompt_len=int(rng.integers(4, 48)), - max_new_tokens=int(rng.integers(1, 12)), - ) - for i in range(24) - ] - ) + + def churning(index): + prompt_len = int(rng.integers(4, 48)) + return PagedRequest( + request_id=f"r{index}", + prompt_len=prompt_len, + max_new_tokens=int(rng.integers(1, 12)), + # Required now: the driver assembles the tokens it feeds the model. + prompt_tokens=rng.integers(1, 30000, size=prompt_len, dtype=np.int64), + ) + + driver.submit([churning(i) for i in range(24)]) done = driver.run() self.assertEqual(len(done), 24) @@ -384,11 +391,13 @@ def test_repeating_a_shape_does_not_retrace(self): step = _TracingStep() driver = PagedDriver(plane, allocate_pool(layout), step, max_batch=2, max_batched_tokens=64) - driver.submit([PagedRequest(request_id="a", prompt_len=20, max_new_tokens=10)]) + driver.submit([PagedRequest(request_id="a", prompt_len=20, max_new_tokens=10, + prompt_tokens=np.arange(20, dtype=np.int64))]) driver.run() after_first = step.traces - driver.submit([PagedRequest(request_id="b", prompt_len=20, max_new_tokens=10)]) + driver.submit([PagedRequest(request_id="b", prompt_len=20, max_new_tokens=10, + prompt_tokens=np.arange(20, dtype=np.int64))]) driver.run() self.assertEqual(step.traces, after_first, "an identical second request retraced") From 385545edda7d7259b92700b1d72f6782e6d1dfcb Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 28 Aug 2026 13:18:23 +0000 Subject: [PATCH 15/21] refactor(inference): both harnesses drive the driver, and the second copy of the livelock goes Completes the loop consolidation. `benchmark.run_paged` and `run_prefix_cache_benchmark.serve` both hand-rolled admission, preemption and step order; both now drive `PagedDriver` and measure it. The count of scheduling loops kept rising as each was found -- two, then three, then five -- and what the duplication actually cost is that **two of them independently contained the same destructive-preemption bug**: on backpressure they discarded the victim's generated tokens, making the replay identical to the attempt that had just failed, which livelocks an overcommitted pool at full GPU utilisation. One copy was fixed earlier this week; this removes the other along with the loop that hosted it. `benchmark.Request` is gone. It was `PagedRequest` plus timestamps, and the overlap is what let the two drift. Timing now lives in a `RequestMetrics` side-table the harness owns, keyed by request id because `run_repeated` deep-copies traces between passes. `run_dense` keeps its own loop deliberately: the driver's subject is page allocation, which the dense two-region cache does not have. `StepOutcome` carries `batch` and `query_lens` so an observer can attribute per-request timing and per-request work. Both are needed -- TTFT and ITL from the first, prefill-tokens-avoided from the second, since `cached_tokens` is reset on release and a preempted request prefills more than once. `PagedDriver` gains an optional `runtime=`. A caller that already has one must pass it: building a second over the same control plane and pool splits `observed_shapes` and the prefix-cache accounting, so a caller reading them back sees none of the shapes the driver traced. The harness does exactly that read. Verified: 212 CPU and 22 GPU tests; the serving benchmark reports TTFT, ITL, concurrency and zero leaks through the driver; prefix sharing measures 47.9% of prefill tokens avoided at a 51.6% page hit rate, with the cache-off arm running every prompt token. One cost to know about: the driver batches prefill where the harnesses prefilled singly, so a run traces more shapes and `warmup_paged` -- which still warms one request at a time -- leaves a couple unwarmed. Reported rather than hidden. Co-authored-by: Cursor --- .../paged_kv/run_prefix_cache_benchmark.py | 99 ++++---- .../inference/kv_execution/benchmark.py | 234 ++++++++---------- src/maxtext/inference/kv_execution/driver.py | 47 +++- 3 files changed, 179 insertions(+), 201 deletions(-) diff --git a/benchmarks/paged_kv/run_prefix_cache_benchmark.py b/benchmarks/paged_kv/run_prefix_cache_benchmark.py index c60a146f22..94e2532be5 100644 --- a/benchmarks/paged_kv/run_prefix_cache_benchmark.py +++ b/benchmarks/paged_kv/run_prefix_cache_benchmark.py @@ -109,60 +109,53 @@ def build_params(cfg, devices): def serve(engine, params, requests, *, max_batch: int, measure: bool): """One pass over the trace. Returns per-request TTFT and prefill token counts. - Requests are admitted greedily and every live request is advanced one token per - step, which is the same policy the driver and the A/B harness use. + Driven by `PagedDriver`, which owns admission and preemption. This used to be a + hand-rolled loop -- the *fifth* in the codebase -- and it carried its own copy of + the preemption bug: on backpressure it discarded the victim's generated tokens, + making the replay identical to the attempt that just failed. That livelocks an + overcommitted pool. The same defect was found and fixed in the A/B harness, and + having to fix it twice is the argument for one scheduler. + + Namespaces are attached per request rather than passed per call, since the driver + reads them off `PagedRequest`. """ - runtime = engine.paged_runtime - waiting, live = list(requests), [] - ttfts, prefill_tokens = [], [] - - while waiting or live: - while waiting and len(live) < max_batch: - candidate = waiting[0] - tokens = candidate.token_ids - started = time.perf_counter() - handle, first = engine.prefill_paged( - params=params, - padded_tokens=jnp.asarray(tokens, jnp.int32), - true_length=int(tokens.size), - request_id=candidate.request_id, - max_new_tokens=candidate.max_new_tokens, - prompt_token_ids=tokens, - namespace=NAMESPACE, - ) - if handle is None: - break - token = int(first.data[0, 0]) # blocks, so the timing is of real work - if measure: - ttfts.append(time.perf_counter() - started) - prefill_tokens.append(int(tokens.size) - runtime.cached_tokens(handle)) - waiting.pop(0) - candidate.handle = handle - candidate.generated.append(token) - live.append(candidate) - - if not live: - raise RuntimeError("nothing is live and nothing could be admitted; the pool is too small") - - result, ok = engine.generate_paged( - params, [r.handle for r in live], next_tokens=jnp.asarray([r.generated[-1] for r in live], jnp.int32) - ) - if not ok: - victim = live.pop() - engine.release(victim.handle) - victim.handle, victim.generated = None, [] - waiting.insert(0, victim) - continue - - tokens_out = np.asarray(result.data[:, 0]).reshape(-1) - for row, request in enumerate(live): - request.generated.append(int(tokens_out[row])) - for request in [r for r in live if len(r.generated) >= r.max_new_tokens]: - engine.release(request.handle, request.context_tokens()) - request.handle = None - live.remove(request) + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_execution.driver import PagedDriver - return ttfts, prefill_tokens + runtime = engine.paged_runtime + trace = list(requests) + for request in trace: + request.namespace = NAMESPACE + + driver = PagedDriver( + runtime.control_plane, + runtime.pool, + engine.paged_step_fn(params), + max_batch=max_batch, + runtime=runtime, + ) + driver.submit(trace) + + started = {r.request_id: time.perf_counter() for r in trace} + ttft = {} + prefilled: dict[str, int] = {} + + while True: + outcome = driver.step() + if outcome is None: + break + now = time.perf_counter() + for request, query_len in zip(outcome.batch, outcome.query_lens): + if not outcome.is_decode: + # Accumulated, because a preempted request prefills more than once and + # every one of those is work the cache did not save. + prefilled[request.request_id] = prefilled.get(request.request_id, 0) + query_len + ttft.setdefault(request.request_id, now - started[request.request_id]) + + if not measure: + return [], [] + order = [r.request_id for r in driver.completed()] + return [ttft[r] for r in order if r in ttft], [prefilled.get(r, 0) for r in order] def run_arm(args, devices, prefix_cache: bool): @@ -195,7 +188,7 @@ def trace(): ttfts, prefill_tokens = serve(engine, params, trace(), max_batch=args.max_batch, measure=True) duration = time.perf_counter() - start - prompted = sum(int(r.token_ids.size) for r in trace()) + prompted = sum(int(r.prompt_len) for r in trace()) return { "prefix_cache": prefix_cache, "duration_s": duration, diff --git a/src/maxtext/inference/kv_execution/benchmark.py b/src/maxtext/inference/kv_execution/benchmark.py index d61d80b94c..0222d3338a 100644 --- a/src/maxtext/inference/kv_execution/benchmark.py +++ b/src/maxtext/inference/kv_execution/benchmark.py @@ -56,34 +56,30 @@ import jax import jax.numpy as jnp +from maxtext.inference.kv_execution.driver import PagedDriver, PagedRequest + @dataclasses.dataclass -class Request: - """One synthetic request, and the timestamps the run collects for it.""" +class RequestMetrics: + """Timing for one request, held beside it rather than on it. + + There used to be a `Request` here that was `PagedRequest` plus timestamps, and + the two drifted: the harness re-implemented scheduling to collect these, and its + copy of preemption discarded generated tokens where the driver's keeps them, + which livelocked a 70B sweep. The driver is now the only scheduler, so the + harness owns *measurement* and nothing else. + + Keyed by `request_id` rather than by object identity, because `run_repeated` + deep-copies a trace between passes. + """ - request_id: str - prompt_len: int - max_new_tokens: int arrival: float = 0.0 first_token_at: float | None = None token_times: list[float] = dataclasses.field(default_factory=list) - generated: list[int] = dataclasses.field(default_factory=list) - handle: Any = None - token_ids: np.ndarray | None = None + # Tokens this request actually computed, summed over its prefill steps. Not + # derivable afterwards: `cached_tokens` is reset when the request is released, + # and a preempted request prefills more than once. prefill_tokens: int = 0 - # Preemption folds a request's generated tokens back into its prompt so the - # replay makes progress, which means a preempted request reports more prompt - # tokens and fewer output tokens than it really had. Counting preemptions is - # what stops that distortion being invisible in the throughput figures. - preemptions: int = 0 - - def context_tokens(self) -> np.ndarray | None: - """Prompt plus generated, which is what the prefix cache is offered.""" - if self.token_ids is None: - return None - if not self.generated: - return self.token_ids - return np.concatenate([self.token_ids, np.asarray(self.generated, dtype=np.int64)]) @property def ttft(self) -> float | None: @@ -99,6 +95,11 @@ def inter_token_latencies(self) -> list[float]: return [b - a for a, b in zip(self.token_times, self.token_times[1:])] +def _metrics_for(table: dict[str, RequestMetrics], request: PagedRequest) -> RequestMetrics: + """The metrics row for `request`, created on first sight.""" + return table.setdefault(request.request_id, RequestMetrics()) + + def synthetic_trace( num_requests: int, mean_prompt: int, @@ -106,7 +107,7 @@ def synthetic_trace( *, seed: int = 0, length_spread: float = 0.6, -) -> list[Request]: +) -> list[PagedRequest]: """A batch of requests with varied lengths, all arriving at once. Lengths are what this plan's claims are about, so they vary; arrival times are @@ -128,11 +129,11 @@ def spread(mean: int) -> np.ndarray: # report a hit rate this workload does not represent -- the baseline has to be a # workload with nothing to share. return [ - Request( + PagedRequest( request_id=f"r{i}", prompt_len=int(prompts[i]), max_new_tokens=int(outputs[i]), - token_ids=rng.integers(1, 30000, size=int(prompts[i]), dtype=np.int64), + prompt_tokens=rng.integers(1, 30000, size=int(prompts[i]), dtype=np.int64), ) for i in range(num_requests) ] @@ -147,7 +148,7 @@ def shared_prefix_trace( seed: int = 0, length_spread: float = 0.6, num_variants: int = 1, -) -> list[Request]: +) -> list[PagedRequest]: """Requests that begin with a common block, as a served system prompt does. This is the workload prefix sharing exists for, and the one where a paged @@ -171,11 +172,11 @@ def spread(mean: int) -> np.ndarray: tail = rng.integers(1, 30000, size=int(uniques[i]), dtype=np.int64) tokens = np.concatenate([prefixes[i % len(prefixes)], tail]) requests.append( - Request( + PagedRequest( request_id=f"r{i}", prompt_len=int(tokens.size), max_new_tokens=int(outputs[i]), - token_ids=tokens, + prompt_tokens=tokens, ) ) return requests @@ -302,125 +303,72 @@ def run_paged( max_batch: int, warmed_shapes: set | None = None, ) -> dict[str, Any]: - """Serve `requests` on the paged path, admitting greedily up to `max_batch`. - - Prefill one request at a time, then advance every live request by one token. - Prefill is preferred while there is room, which favours TTFT at some cost to - the ITL of requests already running -- the same deliberate policy the standalone - driver uses. + """Serve `requests` on the paged path, measuring what `PagedDriver` does. + + **The harness no longer schedules.** It used to run its own admission, its own + preemption and its own step order, which is how it drifted from the driver: + its preemption discarded a victim's generated tokens where the driver's keeps + them, so an overcommitted pool replayed the same failed attempt forever. That + livelocked a 70B capacity sweep at full GPU utilisation with zero progress. + + Now the driver owns policy and this owns measurement. The loop is + timestamp-around-`driver.step()`, attributing per-request times through + `StepOutcome.batch`, which is exactly the information the harness previously + re-implemented a scheduler to obtain. """ runtime = engine.paged_runtime allocator = runtime.control_plane.allocator page_size = runtime.control_plane.layout.tokens_per_page - waiting, live, done = list(requests), [], [] + # The engine's runtime is passed in rather than letting the driver build one, so + # the pool, the shape bookkeeping and the prefix-cache accounting stay + # single-copy -- `runtime.observed_shapes` below is read back from it. + driver = PagedDriver( + runtime.control_plane, + runtime.pool, + engine.paged_step_fn(params), + max_batch=max_batch, + runtime=runtime, + ) + metrics: dict[str, RequestMetrics] = {} occupancy: list[tuple[float, int, int, int]] = [] steps = 0 before = set(runtime.observed_shapes) + start = time.perf_counter() - for request in waiting: - request.arrival = start - - while waiting or live: - while waiting and len(live) < max_batch: - candidate = waiting[0] - prompt = ( - candidate.token_ids - if candidate.token_ids is not None - else np.ones((candidate.prompt_len,), dtype=np.int64) - ) - handle, result = engine.prefill_paged( - params=params, - padded_tokens=jnp.asarray(prompt, dtype=jnp.int32), - true_length=candidate.prompt_len, - request_id=candidate.request_id, - max_new_tokens=candidate.max_new_tokens, - prompt_token_ids=candidate.token_ids, - ) - if handle is None: - break # backpressure: the pool cannot take another request yet - # The tokens this prefill actually ran, which is the prompt minus whatever - # the cache supplied. Recorded per request so the saving can be reported as - # work avoided rather than inferred from a latency difference alone. - candidate.prefill_tokens = candidate.prompt_len - runtime.cached_tokens(handle) - waiting.pop(0) - candidate.handle = handle - candidate.generated.append(int(result.data[0, 0])) - now = time.perf_counter() - candidate.first_token_at = now - candidate.token_times.append(now) - live.append(candidate) - steps += 1 + for request in requests: + _metrics_for(metrics, request).arrival = start + driver.submit(list(requests)) - if not live: - if waiting: - raise RuntimeError( - f"{len(waiting)} requests are waiting but nothing is live and the pool is empty; " - f"the pool is too small for even one request in this trace" - ) + while True: + outcome = driver.step() + if outcome is None: break - - result, ok = engine.generate_paged( - params, - [r.handle for r in live], - next_tokens=jnp.asarray([r.generated[-1] for r in live], jnp.int32), - ) - if not ok: - # Pool exhausted mid-decode. Give the newest request's pages back and let - # it be replayed, which is the same recompute-preemption the driver uses. - victim = live.pop() - # Nothing published: preemption is trying to reclaim pages, and the cache - # would hold onto whatever it adopted. - engine.release(victim.handle) - # **Keep the generated tokens**, exactly as `PagedDriver._preempt_newest` - # does, so the replay is a longer prompt rather than a restart. Discarding - # them -- which this did until it was caught livelocking a 70B sweep -- - # makes the retry byte-identical to the attempt that just failed, so a pool - # that is overcommitted at admission never makes progress: admit, exhaust, - # preempt, re-prefill the same tokens, forever, at full GPU utilisation. - # The driver's own docstring calls keeping them "the price of not - # deadlocking"; the harness has to pay it too. - context = victim.context_tokens() - if context is not None: - victim.token_ids = context - victim.max_new_tokens = max(victim.max_new_tokens - len(victim.generated), 0) - victim.prompt_len += len(victim.generated) - victim.generated = [] - victim.handle = None - victim.preemptions += 1 - # A request preempted with nothing left to generate is finished, not - # requeued; requeuing it would ask for a zero-token step. - if victim.max_new_tokens == 0: - done.append(victim) - else: - waiting.insert(0, victim) - continue - steps += 1 now = time.perf_counter() - tokens = np.asarray(result.data[:, 0]).reshape(-1) - finished = [] - for row, request in enumerate(live): - request.generated.append(int(tokens[row])) - request.token_times.append(now) - if len(request.generated) >= request.max_new_tokens: - finished.append(request) + for request, query_len in zip(outcome.batch, outcome.query_lens): + entry = _metrics_for(metrics, request) + if not outcome.is_decode: + # Tokens this prefill actually ran, which is the prompt minus whatever the + # cache supplied. Accumulated rather than assigned, because a preempted + # request prefills more than once and every one of those is real work. + entry.prefill_tokens += query_len + if entry.first_token_at is None: + entry.first_token_at = now + entry.token_times.append(now) + + live = driver.live_requests() live_tokens = sum(r.prompt_len + len(r.generated) for r in live) occupancy.append((now - start, allocator.num_allocated_pages * page_size, live_tokens, len(live))) - for request in finished: - engine.release(request.handle, request.context_tokens()) - request.handle = None - live.remove(request) - done.append(request) - elapsed = time.perf_counter() - start + done = driver.completed() index = runtime.control_plane.prefix_index summary = _summarise( - done, occupancy, elapsed, steps, allocator, page_size, retained_pages=index.num_cached_pages + done, metrics, occupancy, elapsed, steps, allocator, page_size, retained_pages=index.num_cached_pages ) - summary.update(_summarise_prefix_cache(done, runtime.control_plane)) + summary.update(_summarise_prefix_cache(done, metrics, runtime.control_plane)) # Shape accounting, which is necessary but *not sufficient*: it sees bucketed # `StepShape`s and is blind to eager host-side ops whose values enter a jaxpr as @@ -438,7 +386,7 @@ def run_paged( return summary -def _summarise_prefix_cache(requests: Sequence[Request], control_plane) -> dict[str, Any]: +def _summarise_prefix_cache(requests, metrics, control_plane) -> dict[str, Any]: """What sharing avoided, in tokens rather than in seconds. Reported separately from latency because it is the direct measurement: a @@ -447,7 +395,7 @@ def _summarise_prefix_cache(requests: Sequence[Request], control_plane) -> dict[ that did not happen. """ prompted = sum(r.prompt_len for r in requests) - prefilled = sum(r.prefill_tokens for r in requests) + prefilled = sum(metrics[r.request_id].prefill_tokens for r in requests if r.request_id in metrics) index = control_plane.prefix_index return { "prefix_cache_enabled": index.enabled, @@ -514,7 +462,7 @@ def is_stable(durations: Sequence[float], tolerance: float = 1.15) -> bool: return stability_ratio(durations) < tolerance -def run_dense(engine, params, requests: Sequence[Request], *, max_batch: int) -> dict[str, Any]: +def run_dense(engine, params, requests: Sequence[PagedRequest], *, max_batch: int) -> dict[str, Any]: """Serve `requests` on the dense two-region cache, for the A/B. Deliberately the ordinary MaxText path -- `prefill`, `init_decode_state`, @@ -522,14 +470,21 @@ def run_dense(engine, params, requests: Sequence[Request], *, max_batch: int) -> not against an idealised dense implementation. Its batch is a fixed set of slots, and every slot advances in lockstep whether or not it holds a request, which is the behaviour under examination rather than an inefficiency to correct. + + This keeps its own loop rather than driving `PagedDriver`, and that is not an + oversight: the driver's whole subject is page allocation, which the dense + two-region cache does not have. A slot here is fixed for a request's lifetime, + there is nothing to reserve, recycle or preempt, and forcing it through a paged + scheduler would measure the scheduler instead of the cache. """ waiting, done = list(requests), [] + metrics: dict[str, RequestMetrics] = {} start = time.perf_counter() for request in waiting: - request.arrival = start + _metrics_for(metrics, request).arrival = start decode_state = engine.init_decode_state() - slots: dict[int, Request] = {} + slots: dict[int, PagedRequest] = {} steps = 0 while waiting or slots: @@ -547,8 +502,11 @@ def run_dense(engine, params, requests: Sequence[Request], *, max_batch: int) -> decode_state = engine.insert(prefix, decode_state, slot=free) candidate.generated.append(int(result.data[0, 0])) now = time.perf_counter() - candidate.first_token_at = now - candidate.token_times.append(now) + entry = _metrics_for(metrics, candidate) + entry.first_token_at = now + entry.token_times.append(now) + # The dense path has no prefix cache, so every prompt token is computed. + entry.prefill_tokens += candidate.prompt_len slots[free] = candidate steps += 1 @@ -560,17 +518,18 @@ def run_dense(engine, params, requests: Sequence[Request], *, max_batch: int) -> now = time.perf_counter() for slot, request in list(slots.items()): request.generated.append(int(result.data[slot, 0])) - request.token_times.append(now) + _metrics_for(metrics, request).token_times.append(now) if len(request.generated) >= request.max_new_tokens: del slots[slot] done.append(request) elapsed = time.perf_counter() - start - return _summarise(done, [], elapsed, steps, None, 0) + return _summarise(done, metrics, [], elapsed, steps, None, 0) def _summarise( - done: Sequence[Request], + done, + metrics: dict[str, RequestMetrics], occupancy: Sequence[tuple[float, int, int, int]], elapsed: float, steps: int, @@ -588,7 +547,8 @@ def _summarise( """ output_tokens = sum(len(r.generated) for r in done) prompt_tokens = sum(r.prompt_len for r in done) - itls = [gap for r in done for gap in r.inter_token_latencies()] + timings = [metrics[r.request_id] for r in done if r.request_id in metrics] + itls = [gap for entry in timings for gap in entry.inter_token_latencies()] summary: dict[str, Any] = { "completed_requests": len(done), @@ -603,7 +563,7 @@ def _summarise( # with preemptions is a capacity result rather than a throughput one. "preemptions": sum(r.preemptions for r in done), "requests_preempted": sum(1 for r in done if r.preemptions), - "ttft": _percentiles([r.ttft for r in done if r.ttft is not None]), + "ttft": _percentiles([e.ttft for e in timings if e.ttft is not None]), "itl": _percentiles(itls), "peak_device_bytes": _peak_bytes(), } diff --git a/src/maxtext/inference/kv_execution/driver.py b/src/maxtext/inference/kv_execution/driver.py index 5d41434d2d..ebe814448b 100644 --- a/src/maxtext/inference/kv_execution/driver.py +++ b/src/maxtext/inference/kv_execution/driver.py @@ -141,11 +141,13 @@ class StepOutcome: num_requests: int num_tokens: int preempted: int = 0 - # The requests this step advanced, in the order their tokens were returned. - # Carried so an observer can attribute per-request timing -- which is what a - # serving harness needs for TTFT and ITL, and the reason it previously had to - # re-implement the scheduling loop to get it. + # The requests this step advanced, in the order their tokens were returned, + # with the query length each contributed. Carried so an observer can attribute + # per-request timing and per-request work -- which is what a serving harness + # needs for TTFT, ITL and "prefill tokens avoided", and the reason it previously + # had to re-implement the scheduling loop to get them. batch: tuple[PagedRequest, ...] = () + query_lens: tuple[int, ...] = () # Takes the padded page view, the assembled token-side operands and the pool; @@ -173,6 +175,7 @@ def __init__( max_batched_tokens: int | None = None, eos_ids: Iterable[int] = (), poison_on_free: bool = False, + runtime: PagedRuntime | None = None, ): layout = control_plane.layout self.plane = control_plane @@ -181,12 +184,20 @@ def __init__( self.max_batch = int(max_batch) self.eos_ids = frozenset(int(t) for t in eos_ids) self.poison_on_free = bool(poison_on_free) - self.planner = StepShapePlanner( - tokens_per_page=layout.tokens_per_page, - max_batch=max_batch, - max_context_len=control_plane.max_context_len, - pool_pages=layout.num_pages, - max_batched_tokens=max_batched_tokens, + # An existing runtime's planner wins. Building a second one would record + # shapes into a different `observed_shapes`, so a caller holding the first -- + # the benchmark harness does, to report compile counts -- would see none of + # the shapes this driver actually traced. + self.planner = ( + runtime.planner + if runtime is not None and runtime.planner is not None + else StepShapePlanner( + tokens_per_page=layout.tokens_per_page, + max_batch=max_batch, + max_context_len=control_plane.max_context_len, + pool_pages=layout.num_pages, + max_batched_tokens=max_batched_tokens, + ) ) # The driver owns *policy* -- the queue, the admission budget, recompute # preemption, prefill before decode -- and delegates the per-step mechanics @@ -196,7 +207,12 @@ def __init__( # decides what was recycled, and before the table, because the control plane # refuses to describe a page it still thinks is dirty. A second copy of that # order is a second place to get it wrong. - self.runtime = PagedRuntime( + # + # A caller that already has one -- `MaxEngine.paged_runtime` -- should pass it, + # so the pool, the shape bookkeeping and the prefix-cache accounting stay + # single-copy. Building a second over the same control plane and pool would + # split exactly the state a caller reads back. + self.runtime = runtime or PagedRuntime( control_plane=control_plane, pool=pool, planner=self.planner, @@ -224,6 +240,14 @@ def num_waiting(self) -> int: def num_live(self) -> int: return len(self._live) + def live_requests(self) -> list[PagedRequest]: + """Requests currently holding pages. A copy, so an observer cannot mutate the loop.""" + return list(self._live) + + def completed(self) -> list[PagedRequest]: + """Requests that have finished, in completion order. Also a copy.""" + return list(self._done) + def submit(self, requests: Sequence[PagedRequest]) -> None: """Queue requests, rejecting any the driver could not actually run. @@ -394,6 +418,7 @@ def _run_phase( num_tokens=sum(lens), preempted=preempted, batch=tuple(batch), + query_lens=tuple(lens), ) def _slice_for(self, request: PagedRequest, query_len: int, *, is_decode: bool) -> RequestSlice: From 4ec189c7873d24d12f538e9a30c5e69b867e167f Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 28 Aug 2026 14:16:22 +0000 Subject: [PATCH 16/21] feat(inference): give OfflineEngine a paged worker, completing the serving wiring `OfflineEngine.batch_inference` is the production entry point, and the paged path had none: selecting a page-based worker needed paged siblings on `MaxEngine`, which the step seam supplied. `attention="gpu_paged"` now selects `PagedInferenceWorker`, which drives `PagedDriver` and returns the same `CompletionOutput` contract. A separate worker rather than a flag on the dense one, because the two are organised around different notions of ownership. The dense worker is built around a decode *slot* fixed for a request's lifetime -- `empty_decode_slots`, `slot_to_id`, a `DecodeState` whose batch dimension is the slot count, and a `generate` advancing every slot in lockstep. A paged request owns a varying set of pages, which is why M4 made release request-based rather than `release_pages(slot)`. Threading a pool through the slot machinery would keep both models of ownership alive in one loop. Detokenisation is synchronous here. The dense worker needs a background thread because its loop cannot yield between slots; this one has each request's whole token history when the driver finishes, and offline inference has nobody waiting on a first token. A thread would add ordering and shutdown hazards to buy latency nothing measures. Log probabilities are collected through the same side-table pattern the harness uses for timing: the step function stashes each step's logits and the loop attributes them through `StepOutcome.batch`. They are part of `CompletionOutput` and `_validate_config` insists on them, so returning tokens alone would satisfy a token comparison and still be wrong. Also makes `prefill_packing` a lazy import. At module scope it made `offline_engine` unimportable without JetStream -- taking the paged worker, which needs none of it, down with it. The paged path could not be reached at all under DECOUPLE_GCLOUD=TRUE before this. Verified: 212 CPU and 23 GPU tests. The new test drives three prompts through `OfflineEngine` at once -- exercising continuous batching, requests at different positions sharing one pool -- and requires token-identical output to the engine's paged entry points, which the parity suite already ties to dense. Co-authored-by: Cursor --- src/maxtext/inference/offline_engine.py | 66 ++++-- src/maxtext/inference/paged_offline_worker.py | 198 ++++++++++++++++++ tests/unit/gpu_paged_decode_parity_test.py | 81 +++++++ 3 files changed, 328 insertions(+), 17 deletions(-) create mode 100644 src/maxtext/inference/paged_offline_worker.py diff --git a/src/maxtext/inference/offline_engine.py b/src/maxtext/inference/offline_engine.py index 594ba52eeb..4907db5b61 100644 --- a/src/maxtext/inference/offline_engine.py +++ b/src/maxtext/inference/offline_engine.py @@ -54,8 +54,11 @@ from jax.experimental import mesh_utils from maxtext.inference.maxengine.maxengine import MaxEngine -from maxtext.input_pipeline.packing.prefill_packing import PrefillProcessor -from maxtext.input_pipeline.packing.prefill_packing import BatchedPrefillProcessor +# `prefill_packing` is imported lazily, inside `PrefillHelper`, because it is a +# dense-path concern that hard-refuses to import without JetStream. Importing it +# here made the whole module unimportable in a paged-only deployment -- including +# the paged worker, which needs none of it -- so the paged path could not be +# reached at all under `DECOUPLE_GCLOUD=TRUE`. from maxtext.utils import max_logging from maxtext.utils import max_utils @@ -159,6 +162,13 @@ def __init__( batch_prefill_max_batch_size: Maximum number of prompts in one packed sequence for batch prefill """ + # Imported here rather than at module scope: this refuses to import without + # JetStream, and a paged deployment needs none of it. At module scope it made + # `offline_engine` unimportable in a paged-only container, which took the + # paged worker down with it. + # pylint: disable=import-outside-toplevel + from maxtext.input_pipeline.packing.prefill_packing import BatchedPrefillProcessor, PrefillProcessor + self._type = prefill_type self.engine = engine self.prefill_lengths = sorted(prefill_lengths) @@ -810,21 +820,43 @@ def __init__( if not self.mesh: self.mesh = OfflineEngine.create_mesh(jax.devices(), self.config) - self.worker = InferenceWorker( - config=self.config, - params=self.params, - min_decode_steps=self.min_decode_steps, - enable_batch_prefill=self.enable_batch_prefill, - mesh=self.mesh, - devices=self.mesh.devices.flatten(), - tokenizer=self.tokenizer, - eos_ids=self.eos_ids, - prefill_lengths=self.prefill_lengths, - max_decode_length=self.max_decode_length, - batch_prefill_max_batch_size=self.batch_prefill_max_batch_size, - rng=self.rng, - debug=self.debug, - ) + if self.config.attention == "gpu_paged": + # A separate worker rather than a flag on the dense one, because the two are + # organised around different notions of ownership: a decode *slot* fixed for + # a request's lifetime, against a varying set of pages. Scheduling is + # delegated to `PagedDriver`, which already owns admission, reservation, + # recycled-page scrubbing and recompute preemption. + # pylint: disable=import-outside-toplevel + from maxtext.inference.paged_offline_worker import PagedInferenceWorker + + self.worker = PagedInferenceWorker( + config=self.config, + params=self.params, + devices=self.mesh.devices.flatten(), + tokenizer=self.tokenizer, + eos_ids=self.eos_ids, + max_decode_length=self.max_decode_length, + max_batched_tokens=self.max_prefill_length, + rng=self.rng, + mesh=self.mesh, + debug=self.debug, + ) + else: + self.worker = InferenceWorker( + config=self.config, + params=self.params, + min_decode_steps=self.min_decode_steps, + enable_batch_prefill=self.enable_batch_prefill, + mesh=self.mesh, + devices=self.mesh.devices.flatten(), + tokenizer=self.tokenizer, + eos_ids=self.eos_ids, + prefill_lengths=self.prefill_lengths, + max_decode_length=self.max_decode_length, + batch_prefill_max_batch_size=self.batch_prefill_max_batch_size, + rng=self.rng, + debug=self.debug, + ) self.tokenizer = self.worker.tokenizer diff --git a/src/maxtext/inference/paged_offline_worker.py b/src/maxtext/inference/paged_offline_worker.py new file mode 100644 index 0000000000..e44039d496 --- /dev/null +++ b/src/maxtext/inference/paged_offline_worker.py @@ -0,0 +1,198 @@ +"""Continuous batching for `OfflineEngine` on the paged KV pool. + +The paged sibling of `offline_engine.InferenceWorker`, selected when +`attention="gpu_paged"`. It presents the same contract -- `run_inference(data, +rng)` returning one `CompletionOutput` per input -- and shares nothing else, +because the two differ in the thing that organises them. + +**Why this is a separate worker rather than a flag on the dense one.** The dense +worker is built around a fixed decode *slot* per request: `empty_decode_slots`, +`slot_to_id`, a `DecodeState` whose batch dimension is the slot count, and a +`generate` that advances every slot in lockstep whether or not it holds a +request. A paged request owns a varying set of pages instead, which is why M4 +made the release API request-based rather than `release_pages(slot)`. Threading a +pool through the slot machinery would mean keeping both models of ownership alive +in one loop; scheduling is delegated to `PagedDriver` instead, which already owns +admission, reservation, recycled-page scrubbing and recompute preemption. + +**Detokenisation is synchronous here, deliberately.** The dense worker runs a +background thread emitting tokens as they arrive, because its loop cannot yield +between slots. This one has the whole token history per request when the driver +finishes, and offline inference has nobody waiting on a first token. A thread +would add ordering and shutdown hazards to buy latency nothing measures. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Any, Hashable + +import numpy as np + +import jax + +from maxtext.inference import inference_utils +from maxtext.inference.kv_execution.driver import PagedDriver, PagedRequest +from maxtext.inference.maxengine.maxengine import MaxEngine +from maxtext.utils import max_logging + + +class PagedInferenceWorker: + """Runs a batch of requests to completion on the page pool.""" + + def __init__( + self, + config: Any, + params: Any | None, + devices: list[Any], + tokenizer: Any, + eos_ids: list[int] | None, + max_decode_length: int, + *, + max_requests: int | None = None, + max_batched_tokens: int | None = None, + rng: jax.random.PRNGKey = None, + mesh: Any = None, + debug: bool = False, + ): + self.config = config + self.devices = devices + self.tokenizer = tokenizer + self.eos_ids = eos_ids + self.max_decode_length = int(max_decode_length) + self.mesh = mesh + self.rng = jax.random.PRNGKey(0) if rng is None else rng + self.debug = debug + + self.engine = MaxEngine(self.config, self.devices) + self.params = self.engine.load_params(params=params, rng=self.rng) + if self.eos_ids is None and self.tokenizer is None: + self.tokenizer = self.engine.build_tokenizer(self.engine.get_tokenizer()) + if self.eos_ids is None: + self.eos_ids = [self.tokenizer.eos_id] + + self.runtime = self.engine.init_paged_runtime( + max_requests=max_requests, max_batched_tokens=max_batched_tokens + ) + # `max_requests` bounds the page map's rows and the batch bucket ladder, so it + # is the concurrency ceiling the driver schedules against. + self.max_batch = self.runtime.control_plane.page_map.max_requests + max_logging.log( + f"Paged inference worker ready: pool {self.runtime.pool.k_pages[0].shape}, " + f"{self.max_batch} concurrent requests" + ) + + def update_params(self, params: Any) -> None: + """Update the model weights. The pool is unaffected and is not reallocated.""" + self.params = params + + def run_inference(self, data, rng=None) -> list: + """Run every input to completion and return one output each, in input order. + + Args: + data: `offline_engine.InputData`, whose `tokens` are padded and whose + `true_length` says how much of that is real. + rng: overrides the worker's key when given. + + Returns: + One `offline_engine.CompletionOutput` per input, in the order supplied -- + *not* completion order, which paging makes arbitrary. + """ + # pylint: disable=import-outside-toplevel + from maxtext.inference.offline_engine import CompletionOutput + + if rng is not None: + self.rng = rng + if not data: + return [] + + # Log probabilities are part of `CompletionOutput`, and the driver's step + # contract returns tokens only. The step function stashes each step's logits + # here and the loop attributes them through `StepOutcome.batch` -- the same + # shape of side-table the benchmark harness uses for timing, and for the same + # reason: the driver reports what it advanced, so an observer needs nothing + # more from the step itself. + pending_logprobs: dict[str, Any] = {} + + def step(view, inputs, pool): + del pool + sampled, selected = self.engine.paged_step( + params=self.params, view=view, inputs=inputs, is_decode=view.shape.is_decode + ) + tokens = np.asarray(sampled).reshape(-1)[: inputs.sample_at.size] + logprobs = np.asarray( + inference_utils.log_prob_of_chosen_token(selected, sampled) + ).reshape(-1)[: inputs.sample_at.size] + pending_logprobs["last"] = logprobs + return tokens + + driver = PagedDriver( + self.runtime.control_plane, + self.runtime.pool, + step, + max_batch=self.max_batch, + eos_ids=self.eos_ids, + runtime=self.runtime, + ) + + by_id: dict[Hashable, PagedRequest] = {} + requests = [] + for row in data: + prompt = np.asarray(row.tokens).reshape(-1)[: int(row.true_length)].astype(np.int64) + request = PagedRequest( + request_id=str(row.id), + prompt_len=int(row.true_length), + # Bounded by the pool's own context limit as well as the caller's, since + # a request that cannot fit is rejected at admission rather than part way + # through generating. + max_new_tokens=min(self.max_decode_length, self.config.paged_max_context_len - int(row.true_length)), + prompt_tokens=prompt, + ) + requests.append(request) + by_id[row.id] = request + driver.submit(requests) + + logprobs_by_id: dict[Hashable, list[np.ndarray]] = {r.request_id: [] for r in requests} + while True: + outcome = driver.step() + if outcome is None: + break + supplied = pending_logprobs.get("last") + for index, request in enumerate(outcome.batch): + if supplied is not None and index < supplied.size: + logprobs_by_id[request.request_id].append(supplied[index]) + + outputs = [] + for row in data: + request = by_id[row.id] + generated = np.asarray(request.generated, dtype=np.int32) + logps = np.asarray(logprobs_by_id[request.request_id], dtype=np.float32) + outputs.append( + CompletionOutput( + index=row.id, + token_ids=generated, + # Trimmed to the tokens actually returned: a preempted request is + # replayed, so it produces more step observations than final tokens. + logprobs=logps[: generated.size], + prompt_length=int(row.true_length), + ) + ) + if self.debug: + max_logging.log( + f"Paged worker completed {len(outputs)} requests, " + f"{sum(r.preemptions for r in requests)} preemptions" + ) + return outputs diff --git a/tests/unit/gpu_paged_decode_parity_test.py b/tests/unit/gpu_paged_decode_parity_test.py index eea5c96bdf..057206325a 100644 --- a/tests/unit/gpu_paged_decode_parity_test.py +++ b/tests/unit/gpu_paged_decode_parity_test.py @@ -496,6 +496,87 @@ def test_a_batched_prefill_samples_every_request(self): for handle in handles: engine.release(handle) + def test_offline_engine_selects_the_paged_worker_and_agrees_with_the_engine(self): + """`OfflineEngine.batch_inference` is the production entry point, so wire it. + + Until the step seam existed there was no paged worker to select, and this is + the check that selecting one produces the same answer. Several prompts at + once, so it exercises the continuous batching that a single-request rollout + cannot: requests at different positions sharing one pool. + + The reference is the engine's own paged entry points rather than the dense + worker, and that is a container limitation rather than a choice -- the dense + `_prefill_jit` returns a `ResultTokens` from inside `jit`, which has to be a + registered pytree, and the `DECOUPLE_GCLOUD` stub is a plain class. The parity + tests above already tie those entry points to dense, so agreement here is + transitive. + """ + # pylint: disable=import-outside-toplevel + import jax.numpy as jnp + + from maxtext.inference import offline_engine + + prompts = [list(range(3, 15)), list(range(40, 48)), list(range(90, 106))] + cfg = _config(**_PAGED, return_log_prob=True) + + # Reference: one request at a time through the entry points. + reference = {} + engine = maxengine.MaxEngine(cfg, _devices()) + params_state = self._build_params(cfg) + params = engine.load_params(params=params_state) + engine.init_paged_runtime(max_requests=4, max_batched_tokens=cfg.max_prefill_predict_length) + budget = cfg.max_target_length - cfg.max_prefill_predict_length + for index, prompt in enumerate(prompts): + padded = jnp.asarray( + list(prompt) + [0] * (cfg.max_prefill_predict_length - len(prompt)), jnp.int32 + ) + handle, first = engine.prefill_paged( + params=params, padded_tokens=padded, true_length=len(prompt), request_id=f"ref{index}" + ) + self.assertIsNotNone(handle) + tokens = [int(first.data[0, 0])] + for _ in range(min(budget, cfg.paged_max_context_len - len(prompt)) - 1): + result, ok = engine.generate_paged( + params, [handle], next_tokens=jnp.asarray([tokens[-1]], jnp.int32) + ) + if not ok: + break + tokens.append(int(result.data[0, 0])) + engine.release(handle) + reference[f"p{index}"] = tokens + + # Through OfflineEngine, which must pick the paged worker off the attention + # setting. `eos_ids` is supplied so no tokenizer is needed -- these prompts are + # token ids, and a tokenizer would need JetStream. + offline = offline_engine.OfflineEngine( + config=cfg, params=params_state, tokenizer=object(), eos_ids=[-1] + ) + self.assertEqual( + type(offline.worker).__name__, + "PagedInferenceWorker", + "attention='gpu_paged' must select the paged worker", + ) + outputs = offline.batch_inference( + [ + offline_engine.InputData(id=f"p{i}", tokens=np.asarray(p, np.int32), true_length=len(p)) + for i, p in enumerate(prompts) + ] + ) + + self.assertEqual(len(outputs), len(prompts)) + for out in outputs: + expected = reference[out.index] + actual = np.asarray(out.token_ids).tolist() + self._assert_same_tokens(expected[: len(actual)], actual[: len(expected)], f"{out.index}") + # Logprobs are part of the contract and `_validate_config` insists on them, + # so an empty array would satisfy the token check and still be wrong. + self.assertEqual( + np.asarray(out.logprobs).size, + len(actual), + "one log probability per returned token", + ) + self.assertEqual(out.prompt_length, len(prompts[int(out.index[1:])])) + def test_the_paged_path_allocates_no_dense_cache(self): """A dense cache alongside the pool would waste gigabytes at real sizes.""" cfg = _config(**_PAGED) From a566f150315b1a36ea5a062b35837b9ad5ce5bf9 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 28 Aug 2026 14:48:19 +0000 Subject: [PATCH 17/21] feat(inference): tokenise without JetStream or torch, so the native paged image needs neither `MaxEngine.build_tokenizer` was the last JetStream tie on the paged path, and it is unavailable to a paged-only deployment for two independent reasons. It requires JetStream even for HuggingFace tokenizers -- its `huggingface` branch returns `jetstream.engine.token_utils.HuggingFaceTokenizer` -- and raises outright under DECOUPLE_GCLOUD. JetStream was archived on 2026-02-01 with its functionality migrated into vllm-project/tpu-inference, so this is a dependency that has stopped moving rather than one to wait on. And the obvious replacement is worse. `transformers.AutoTokenizer` works, and every scratch script here uses it, but transformers probes for torch and imports it when found; torch brings its own bundled ROCm, and a second HIP runtime aborts RCCL clique setup above one device. The scratch scripts only get away with it by installing a `find_spec` shim before importing anything. So `hf_tokenizer` uses `tokenizers` directly -- the same Rust implementation transformers wraps, reading the same tokenizer.json, pulling in no torch. The surface is defined by MaxText's own callers rather than guessed: `eos_id`, `encode`, `decode`, plus `.tokenizer` for `apply_chat_template` and `batch_decode`, which is what JetStream's wrapper exposes too. Duck-typed rather than subclassed, for the same reason the attention path duck-types vLLM metadata. `PagedInferenceWorker` prefers a caller-supplied tokenizer, then this, then the JetStream route only for a type this cannot serve. Two bugs found while testing, both worth recording. The enum comparison was wrong: `tokenizer_type` is a `TokenizerType`, whose `str()` is "TokenizerType.HUGGINGFACE", so a string compare silently matched nothing and fell through to JetStream. And the dependency-leak test did not bite -- it forbade `torch` while transformers imports torch *lazily*, so a deliberately added `import transformers` passed. Forbidding transformials proximate cause rather than the symptom is what makes it fail. Verified: 220 CPU and 23 GPU tests. End to end, OfflineEngine builds a real Llama-3.3 tokenizer from config, derives eos_ids=[128009], round-trips text, and neither jetstream nor torch appears in sys.modules. Co-authored-by: Cursor --- src/maxtext/inference/hf_tokenizer.py | 168 +++++++++++++++++ src/maxtext/inference/paged_offline_worker.py | 51 ++++- tests/unit/hf_tokenizer_test.py | 176 ++++++++++++++++++ 3 files changed, 393 insertions(+), 2 deletions(-) create mode 100644 src/maxtext/inference/hf_tokenizer.py create mode 100644 tests/unit/hf_tokenizer_test.py diff --git a/src/maxtext/inference/hf_tokenizer.py b/src/maxtext/inference/hf_tokenizer.py new file mode 100644 index 0000000000..f54780094e --- /dev/null +++ b/src/maxtext/inference/hf_tokenizer.py @@ -0,0 +1,168 @@ +"""A tokenizer for the paged inference path that needs neither JetStream nor torch. + +`MaxEngine.build_tokenizer` is the ordinary route and it is unavailable to a +paged-only deployment for two independent reasons. + +**It requires JetStream, including for HuggingFace tokenizers.** It raises outright +under `DECOUPLE_GCLOUD=TRUE`, and even its `huggingface` branch returns +`jetstream.engine.token_utils.HuggingFaceTokenizer` -- so the tokenizer type does +not change the dependency. JetStream was archived on 2026-02-01, with its +functionality migrated into `vllm-project/tpu-inference`, so this is a dependency +that has stopped moving rather than one to wait on. The rest of the paged path +already runs under the `DECOUPLE_GCLOUD` stubs; tokenisation was the last tie. + +**And the obvious replacement carries a worse problem.** `transformers.AutoTokenizer` +works, and every scratch script in this project uses it, but `transformers` probes +for torch with `importlib.util.find_spec` and imports it when found. Torch brings +its own bundled ROCm, giving the process a second HIP runtime, and +`rocprofiler-register` then aborts during RCCL clique setup -- which is fatal above +one device. The scratch scripts get away with it by installing a `find_spec` shim +before importing anything, and library code should not have to. + +So this uses `tokenizers` directly. It is the same Rust implementation +`transformers` wraps, reads the same `tokenizer.json`, and pulls in no torch: +verified by `"torch" not in sys.modules` after a full encode and decode. + +**The surface is defined by MaxText's own callers, not by guesswork.** Across +`maxtext/inference/` those are `eos_id`, `encode`, `decode`, and the underlying +HuggingFace tokenizer for `apply_chat_template` and `batch_decode` -- which +JetStream's wrapper also exposes as `.tokenizer`, so reaching through works the +same either way. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import json +import os +from typing import Any, Sequence + + +class HuggingFaceTokenizer: + """MaxText's tokenizer surface over a `tokenizers.Tokenizer`. + + Duck-typed against JetStream's wrapper rather than subclassing it, for the same + reason the paged attention path duck-types vLLM metadata: importing the thing + you are trying not to depend on defeats the exercise. + """ + + def __init__(self, tokenizer: Any, *, eos_id: int, pad_id: int | None = None): + self._tokenizer = tokenizer + self._eos_id = int(eos_id) + self._pad_id = int(pad_id) if pad_id is not None else int(eos_id) + + @property + def tokenizer(self) -> Any: + """The underlying HuggingFace tokenizer. + + JetStream's wrapper exposes the same attribute, so callers reaching for + `apply_chat_template` or `batch_decode` work unchanged. + """ + return self._tokenizer + + @property + def eos_id(self) -> int: + return self._eos_id + + @property + def pad_id(self) -> int: + return self._pad_id + + def encode(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + """Text to ids. + + `add_special_tokens` defaults to False because the inference path positions + tokens absolutely and a silently prepended BOS shifts every position by one. + Callers that want the chat template should reach it through `.tokenizer`, + where the choice is explicit. + """ + return list(self._tokenizer.encode(text, add_special_tokens=add_special_tokens).ids) + + def decode(self, ids: Sequence[int]) -> str: + """Ids to text, tolerating numpy arrays and nested single-row shapes.""" + import numpy as np # pylint: disable=import-outside-toplevel + + flat = np.asarray(ids).reshape(-1).tolist() + return self._tokenizer.decode([int(i) for i in flat]) + + +def _resolve_special_id(tokenizer: Any, config: dict, key: str) -> int | None: + """Look up a special token's id from `tokenizer_config.json`. + + The entry is either a plain string or a dict with a `content` field, depending + on how the checkpoint was exported; both appear in the wild. + """ + entry = config.get(key) + if entry is None: + return None + token = entry.get("content") if isinstance(entry, dict) else entry + if not isinstance(token, str): + return None + return tokenizer.token_to_id(token) + + +def build_tokenizer(tokenizer_path: str, *, eos_id: int | None = None) -> HuggingFaceTokenizer: + """Load a tokenizer from a local HuggingFace checkpoint directory or file. + + Args: + tokenizer_path: a directory holding `tokenizer.json`, or the file itself. + eos_id: overrides what the checkpoint declares. Supply it when a deployment + stops on a different token than the checkpoint's default, which is common + for base models used with instruction formatting. + + Returns: + A `HuggingFaceTokenizer`. + + Raises: + FileNotFoundError: when no `tokenizer.json` is present. That file is what + makes this torch-free, so falling back to `transformers` would trade a + missing file for an RCCL abort above one device -- a worse failure, and a + much less obvious one. + ValueError: when no EOS id can be determined, since generation would then run + to the length cap on every request and look like a quality problem. + """ + path = tokenizer_path + if os.path.isdir(path): + candidate = os.path.join(path, "tokenizer.json") + else: + candidate = path + if not os.path.isfile(candidate): + raise FileNotFoundError( + f"no tokenizer.json at {candidate!r}. The paged path loads tokenizers through the `tokenizers` " + f"package to stay free of both JetStream and torch, and that needs the fast-tokenizer file. " + f"Convert the tokenizer, or pass a ready-made tokenizer object instead." + ) + + # pylint: disable=import-outside-toplevel + from tokenizers import Tokenizer + + tokenizer = Tokenizer.from_file(candidate) + + config: dict = {} + config_path = os.path.join(os.path.dirname(candidate), "tokenizer_config.json") + if os.path.isfile(config_path): + with open(config_path, "rt", encoding="utf-8") as handle: + config = json.load(handle) + + resolved = eos_id if eos_id is not None else _resolve_special_id(tokenizer, config, "eos_token") + if resolved is None: + raise ValueError( + f"could not determine an EOS id from {config_path!r}. Pass eos_id explicitly; without it every " + f"request generates to its length cap, which reads as a model quality problem rather than a " + f"configuration one." + ) + pad = _resolve_special_id(tokenizer, config, "pad_token") + return HuggingFaceTokenizer(tokenizer, eos_id=int(resolved), pad_id=pad) diff --git a/src/maxtext/inference/paged_offline_worker.py b/src/maxtext/inference/paged_offline_worker.py index e44039d496..56e07a1f33 100644 --- a/src/maxtext/inference/paged_offline_worker.py +++ b/src/maxtext/inference/paged_offline_worker.py @@ -79,9 +79,15 @@ def __init__( self.engine = MaxEngine(self.config, self.devices) self.params = self.engine.load_params(params=params, rng=self.rng) - if self.eos_ids is None and self.tokenizer is None: - self.tokenizer = self.engine.build_tokenizer(self.engine.get_tokenizer()) + if self.tokenizer is None: + self.tokenizer = self._build_tokenizer() if self.eos_ids is None: + if self.tokenizer is None: + raise ValueError( + "no tokenizer and no eos_ids: the worker cannot tell when a request has finished, so every " + "one would generate to its length cap. Pass eos_ids for a token-in, token-out caller, or a " + "tokenizer_path for text." + ) self.eos_ids = [self.tokenizer.eos_id] self.runtime = self.engine.init_paged_runtime( @@ -95,6 +101,47 @@ def __init__( f"{self.max_batch} concurrent requests" ) + def _build_tokenizer(self): + """A tokenizer for this config, without JetStream and without torch. + + Order of preference, and each fallback is a real narrowing rather than a + stylistic one: + + 1. Nothing, when `eos_ids` was supplied and no `tokenizer_path` is set. A + token-in, token-out caller -- the benchmark harnesses, the parity tests -- + needs no tokenizer at all, and building one would demand a path it has no + reason to have. + 2. `hf_tokenizer.build_tokenizer`, which reads `tokenizer.json` through the + `tokenizers` package. No JetStream, and no torch, which matters because + `transformers` imports torch when it finds it and a second HIP runtime + aborts RCCL clique setup above one device. + 3. `MaxEngine.build_tokenizer`, the JetStream route, only if asked for a + tokenizer type this cannot serve. It raises a clear message under + `DECOUPLE_GCLOUD=TRUE`, which is the honest outcome: that path genuinely + needs a package that was archived in February 2026. + """ + path = getattr(self.config, "tokenizer_path", "") or "" + if not path: + return None + + # `.value` first: `tokenizer_type` is a `TokenizerType` enum, whose `str()` is + # "TokenizerType.HUGGINGFACE" rather than "huggingface". Comparing the string + # form silently matched nothing and fell through to the JetStream branch. + declared = getattr(self.config, "tokenizer_type", "") or "" + tokenizer_type = str(getattr(declared, "value", declared)).lower() + if tokenizer_type in ("", "huggingface"): + # pylint: disable=import-outside-toplevel + from maxtext.inference import hf_tokenizer + + eos = self.eos_ids[0] if self.eos_ids else None + return hf_tokenizer.build_tokenizer(path, eos_id=eos) + + max_logging.log( + f"tokenizer_type={tokenizer_type!r} is not served by the torch-free loader; falling back to " + f"MaxEngine.build_tokenizer, which requires JetStream." + ) + return self.engine.build_tokenizer(self.engine.get_tokenizer()) + def update_params(self, params: Any) -> None: """Update the model weights. The pool is unaffected and is not reallocated.""" self.params = params diff --git a/tests/unit/hf_tokenizer_test.py b/tests/unit/hf_tokenizer_test.py new file mode 100644 index 0000000000..16c53fb8e1 --- /dev/null +++ b/tests/unit/hf_tokenizer_test.py @@ -0,0 +1,176 @@ +"""The paged path's tokenizer must need neither JetStream nor torch. + +Both halves matter and for different reasons. JetStream was archived on +2026-02-01, so `MaxEngine.build_tokenizer` -- which requires it even for +HuggingFace tokenizers -- is a dependency that has stopped moving. And +`transformers` imports torch when it finds it, which gives the process a second +HIP runtime and aborts RCCL clique setup above one device, so the obvious +replacement is worse than the problem. + +The two assertions that carry the point are the ones checking `sys.modules`. +Everything else here is ordinary behaviour that would be caught by any use. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import json +import os +import subprocess +import sys +import tempfile +import unittest + +import numpy as np + +from maxtext.inference import hf_tokenizer + + +def _tiny_tokenizer_json() -> dict: + """A minimal word-level `tokenizer.json`, so this test needs no checkpoint. + + Word-level rather than BPE because the point is the wrapper, not the encoding: + a vocabulary small enough to read makes an assertion about ids legible. + """ + vocab = {"hello": 0, "world": 1, "": 2, "": 3} + return { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": [ + { + "id": 2, + "content": "", + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + }, + { + "id": 3, + "content": "", + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + }, + ], + "normalizer": None, + "pre_tokenizer": {"type": "Whitespace"}, + "post_processor": None, + "decoder": {"type": "WordPiece", "prefix": "##", "cleanup": False}, + "model": {"type": "WordLevel", "vocab": vocab, "unk_token": ""}, + } + + +class HfTokenizerTest(unittest.TestCase): + """Behaviour, and the absence of two dependencies.""" + + def setUp(self): + super().setUp() + self.dir = tempfile.mkdtemp() + with open(os.path.join(self.dir, "tokenizer.json"), "wt", encoding="utf-8") as handle: + json.dump(_tiny_tokenizer_json(), handle) + with open(os.path.join(self.dir, "tokenizer_config.json"), "wt", encoding="utf-8") as handle: + json.dump({"eos_token": "", "pad_token": ""}, handle) + + def test_it_round_trips_and_resolves_special_ids(self): + tokenizer = hf_tokenizer.build_tokenizer(self.dir) + self.assertEqual(tokenizer.encode("hello world"), [0, 1]) + self.assertEqual(tokenizer.eos_id, 2, "eos resolved from tokenizer_config, not guessed") + self.assertEqual(tokenizer.pad_id, 3) + self.assertIn("hello", tokenizer.decode([0, 1])) + + def test_a_dict_shaped_special_token_entry_resolves(self): + """Checkpoints export these as a bare string or as a dict; both appear.""" + with open(os.path.join(self.dir, "tokenizer_config.json"), "wt", encoding="utf-8") as handle: + json.dump({"eos_token": {"content": "", "lstrip": False}}, handle) + self.assertEqual(hf_tokenizer.build_tokenizer(self.dir).eos_id, 2) + + def test_an_explicit_eos_overrides_the_checkpoint(self): + """A base model used with instruction formatting stops on a different token.""" + self.assertEqual(hf_tokenizer.build_tokenizer(self.dir, eos_id=1).eos_id, 1) + + def test_decode_accepts_numpy_and_nested_shapes(self): + """Generated tokens arrive as arrays, sometimes with a leading batch axis.""" + tokenizer = hf_tokenizer.build_tokenizer(self.dir) + self.assertEqual(tokenizer.decode(np.asarray([0, 1])), tokenizer.decode([0, 1])) + self.assertEqual(tokenizer.decode(np.asarray([[0, 1]])), tokenizer.decode([0, 1])) + + def test_encode_adds_no_special_tokens_by_default(self): + """A silently prepended BOS shifts every absolute position by one. + + The paged path positions tokens absolutely and RoPE is not translation + invariant, so this default is load-bearing rather than a preference. + """ + tokenizer = hf_tokenizer.build_tokenizer(self.dir) + self.assertEqual(len(tokenizer.encode("hello world")), 2) + + def test_a_missing_tokenizer_json_fails_loudly(self): + """Falling back to transformers here would trade a missing file for an RCCL abort.""" + with tempfile.TemporaryDirectory() as empty: + with self.assertRaises(FileNotFoundError): + hf_tokenizer.build_tokenizer(empty) + + def test_an_unresolvable_eos_fails_rather_than_defaulting(self): + """Without an EOS every request runs to its length cap, which reads as a quality bug.""" + with open(os.path.join(self.dir, "tokenizer_config.json"), "wt", encoding="utf-8") as handle: + json.dump({}, handle) + with self.assertRaises(ValueError): + hf_tokenizer.build_tokenizer(self.dir) + + def test_it_imports_neither_jetstream_nor_transformers_nor_torch(self): + """The whole point, checked in a fresh interpreter. + + In-process this would be unreliable: another test may already have imported + one of these, and `sys.modules` is global. A subprocess is the only honest way + to assert what *this* module pulls in. + + **`transformers` is on the list, and leaving it off made this test useless.** + A first version forbade only `torch` and `jetstream`, and it passed with an + `import transformers` deliberately added to the module under test -- because + transformers imports torch *lazily*, so nothing named `torch` is in + `sys.modules` at import time. The hazard is reaching for transformers at all: + it probes for torch with `find_spec` and imports it when found, at which point + a second HIP runtime aborts RCCL clique setup. Forbidding the proximate cause + rather than the symptom is what makes the check bite. + """ + script = ( + "import sys, json, os, tempfile\n" + "sys.path.insert(0, os.environ['MAXTEXT_SRC'])\n" + f"d = {self.dir!r}\n" + "from maxtext.inference import hf_tokenizer\n" + "t = hf_tokenizer.build_tokenizer(d)\n" + "assert t.encode('hello world') == [0, 1]\n" + "assert t.decode([0, 1])\n" + "forbidden = ('torch', 'jetstream', 'transformers')\n" + # Top-level names only. Reporting every submodule turns a one-line + # diagnosis into fourteen kilobytes of `transformers.models.*`. + "bad = {m.split('.')[0] for m in sys.modules if m.split('.')[0] in forbidden}\n" + "print('LEAKED:' + ','.join(sorted(bad)))\n" + ) + env = dict(os.environ) + env["MAXTEXT_SRC"] = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src") + env["JAX_PLATFORMS"] = "cpu" + out = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, check=True, env=env + ) + leaked = [line for line in out.stdout.splitlines() if line.startswith("LEAKED:")] + self.assertEqual(leaked, ["LEAKED:"], f"the tokenizer pulled in a forbidden package: {out.stdout}") + + +if __name__ == "__main__": + unittest.main() From f328d22b447351bcedd1aece74318a0056279756 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 28 Aug 2026 23:24:40 +0000 Subject: [PATCH 18/21] Decouple the vLLM model adapter from tpu_inference The adapter package imported tpu_inference at module scope for a logger and a model registry, which made it importable only on TPU even though the model it wraps is not TPU-specific. register() now takes the registry as an argument, so the same adapter serves either platform: omitting it preserves the TPU path exactly, while a GPU platform plugin passes its own. patch_kv_cache_manager is now applied only on the TPU path, since it patches a tpu_inference class. It already degraded gracefully elsewhere, but skipping it keeps a GPU registration from logging a failure that is not one. The remaining tpu_inference imports in adapter.py are all reachable only for hybrid Mamba models (qwen3_next, qwen3_5) inside that patch, or are already guarded by try/except, so a dense model never reaches them. Verified from a container where tpu_inference is not importable: the adapter registers into a GPU plugin's registry and tpu_inference never enters sys.modules. Co-authored-by: Cursor --- .../vllm/maxtext_vllm_adapter/__init__.py | 53 +++++++++++++------ 1 file changed, 38 insertions(+), 15 deletions(-) diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/__init__.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/__init__.py index ee14be16b8..98d8ea69a4 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/__init__.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/__init__.py @@ -12,29 +12,52 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""MaxText vLLM adapter package.""" +"""MaxText vLLM adapter package. + +This package is platform-agnostic: it names no particular vLLM hardware plugin. +It previously imported `tpu_inference` at module scope for a logger and a model +registry, which made the adapter importable only on TPU even though the model it +wraps is not TPU-specific. +""" + +import logging -from tpu_inference.logger import init_logger -from tpu_inference.models.common.model_loader import register_model from .adapter import MaxTextForCausalLM -logger = init_logger(__name__) +logger = logging.getLogger(__name__) + +MODEL_NAME = "MaxTextForCausalLM" +__all__ = ["MaxTextForCausalLM", "MODEL_NAME", "register"] -def register(): - """Register MaxTextForCausalLM model with tpu_inference and vllm. - Note, this function is invoked directly by the vLLM engine during startup. As such, - it leverages vLLM logging to report its status. +def register(register_model=None): + """Register MaxTextForCausalLM with a vLLM platform plugin's model registry. + + Note, this function is invoked directly by the vLLM engine during startup. As + such, it leverages vLLM logging to report its status. + + Args: + register_model: The registry callable, taking (name, cls). Injected so the + same adapter serves either platform. When omitted, the TPU plugin's + registry is imported, preserving the previous behaviour for TPU callers. """ - logger.info("Registering MaxTextForCausalLM model with tpu_inference and vllm.") - register_model("MaxTextForCausalLM", MaxTextForCausalLM) + using_tpu_registry = register_model is None + if using_tpu_registry: + # pylint: disable=import-outside-toplevel + from tpu_inference.models.common.model_loader import register_model + + logger.info("Registering %s.", MODEL_NAME) + register_model(MODEL_NAME, MaxTextForCausalLM) - # Dynamically apply KVCacheManager patch when registering the adapter - # pylint: disable=import-outside-toplevel - from .adapter import patch_kv_cache_manager + # The patch targets tpu_inference's KVCacheManager, so it is only meaningful + # for that platform. It degrades gracefully elsewhere, but skipping it keeps + # a GPU registration from logging a failure that is not one. + if using_tpu_registry: + # pylint: disable=import-outside-toplevel + from .adapter import patch_kv_cache_manager - patch_kv_cache_manager() + patch_kv_cache_manager() - logger.info("Successfully registered MaxTextForCausalLM model.") + logger.info("Successfully registered %s.", MODEL_NAME) From 282428aca8217a33dc5dffc9e0368002113043e9 Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Sat, 29 Aug 2026 09:53:32 +0000 Subject: [PATCH 19/21] Let the vLLM adapter and gpu_paged serve a non-TPU platform Two things still tied the vLLM path to TPU after the adapter's registration was decoupled, and both are only reachable once a GPU platform actually drives a model, which is why they survived that change. generate_maxtext_config read pltpu.get_tpu_info().num_lanes unconditionally, and it raises "Unsupported TPU device kind" anywhere else. The value feeds one branch: padding an MoE model's hidden size for tpu-inference's GMM_v2 kernel, which is TPU-only, so a non-TPU caller wants the padding skipped rather than an exception. It also read vllm_config.sharding_config, which is not a vLLM field at all -- the TPU platform plugin attaches it in check_and_update_config, so it is simply absent elsewhere. Both now degrade to what the platform can answer, with parallel_config supplying the parallelism degrees. gpu_paged detected KV-head sharding by looking for mesh axes named tensor, tensor_transpose or tensor_sequence. Those are MaxText's own names, from configs/base.yml. The serving mesh comes from configs/inference/vllm.yml, which names them data, attn_dp, model, expert and attn_dp_expert, and whose logical rules map paged_kv_heads onto ['expert', 'model']. The two sets are disjoint, so under tensor parallelism kv_head_axes returned empty, the attention step ran outside shard_map, and the pool came back unsharded -- wrong in a way no shape check catches. Accepting both namings cannot mis-detect either, since a mesh carrying 'model' never carries 'tensor'. Verified end to end: Llama-3.3-70B at TP=8 on eight MI300X GPUs, driven through vLLM by the jax-vllm-plugin, emits coherent text. 174 KV tests still pass. Co-authored-by: Cursor --- .../inference/kv_execution/layout_builder.py | 10 ++++- .../vllm/maxtext_vllm_adapter/adapter.py | 41 ++++++++++++++++--- src/maxtext/layers/gpu_paged_attention.py | 16 +++++++- 3 files changed, 58 insertions(+), 9 deletions(-) diff --git a/src/maxtext/inference/kv_execution/layout_builder.py b/src/maxtext/inference/kv_execution/layout_builder.py index 69026009c1..5ef2b74c92 100644 --- a/src/maxtext/inference/kv_execution/layout_builder.py +++ b/src/maxtext/inference/kv_execution/layout_builder.py @@ -29,8 +29,14 @@ from maxtext.inference.kv_common import KvStorageLayoutV1 # MaxText spells the KV-head tensor-parallel axes several ways depending on the -# model; these are the ones that shard `num_kv_heads`. -_KV_HEAD_MESH_AXES = ("tensor", "tensor_transpose", "tensor_sequence") +# model; these are the ones that shard `num_kv_heads`. The last two belong to the +# vLLM serving mesh rather than MaxText's own, and the two namings are disjoint, +# so accepting both cannot mis-detect either. Kept in step with +# `gpu_paged_attention.KV_HEAD_MESH_AXES`, which the kernels read; a pool sharded +# on axes the attention step does not know about would be silently wrong. Stated +# literally rather than imported because this module is deliberately jax-free at +# import time and that one is not. +_KV_HEAD_MESH_AXES = ("tensor", "tensor_transpose", "tensor_sequence", "model", "expert") # Axis names for the pool's own mesh. Deliberately not MaxText's: the pool needs # the tensor-parallel axis split into the part that selects a KV head and the diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index 2990104d5c..09617bf7c9 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py @@ -61,6 +61,38 @@ def next_power_of_two(x: int) -> int: return 1 << (x - 1).bit_length() +def tpu_num_lanes() -> int | None: + """The TPU's vector lane count, or None when this process has no TPU. + + Only the GMM_v2 padding below needs it, and that kernel is TPU-only, so a + non-TPU platform driving this adapter wants the padding skipped rather than + an exception. `get_tpu_info` raises on any other device kind. + """ + try: + return pltpu.get_tpu_info().num_lanes + except (ValueError, RuntimeError, AttributeError): + return None + + +def vllm_sharding_degrees(vllm_config: VllmConfig) -> tuple[int, int, int]: + """Tensor, expert and attention-data parallel degrees, as (tp, ep, attn_dp). + + `sharding_config` is attached to the vLLM config by the TPU platform plugin + rather than by vLLM itself, so it is absent under any other platform. Those + fall back to `parallel_config`, which vLLM always populates. + """ + sharding_config = getattr(vllm_config, "sharding_config", None) + if sharding_config is not None: + return (sharding_config.tp_size, sharding_config.expert_size, sharding_config.attn_dp_size) + + parallel_config = vllm_config.parallel_config + return ( + parallel_config.tensor_parallel_size, + getattr(parallel_config, "expert_parallel_size", 1) or 1, + 1, + ) + + def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters: """Generates a MaxText configuration from a vLLM configuration. @@ -99,10 +131,7 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters argv_list = ["", str(base_config_path)] # Gather sharding information from vLLM config to determine transformations to apply - sharding_config = vllm_config.sharding_config - tp = sharding_config.tp_size - ep = sharding_config.expert_size - attn_dp = sharding_config.attn_dp_size + tp, ep, attn_dp = vllm_sharding_degrees(vllm_config) # Calculate the maximum TP size across attention and MLP dimensions kv_tp_size = tp * ep @@ -116,7 +145,7 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters else vllm_config.model_config.hf_config ) hidden_size = getattr(hf_config, "moe_intermediate_size", None) - num_lanes = pltpu.get_tpu_info().num_lanes + num_lanes = tpu_num_lanes() num_kv_heads = hf_config.num_key_value_heads # Number of KV heads in global attention layers (None if the field is absent or unset). @@ -147,7 +176,7 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters # The GMM_v2 kernel requires the MLP dimension per expert to be at least 2x the number of TPU lanes # to ensure efficient execution. See the validate_inputs() method in the following file for more details: # https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/kernels/megablox/gmm_v2.py - if hidden_size is not None and (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0: + if num_lanes is not None and hidden_size is not None and (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0: padded_hidden_size = next_power_of_two(hidden_size) while (padded_hidden_size // moe_mlp_tp_size) < (2 * num_lanes): padded_hidden_size = next_power_of_two(padded_hidden_size + 1) diff --git a/src/maxtext/layers/gpu_paged_attention.py b/src/maxtext/layers/gpu_paged_attention.py index bf8ae420d5..7858ae960e 100644 --- a/src/maxtext/layers/gpu_paged_attention.py +++ b/src/maxtext/layers/gpu_paged_attention.py @@ -241,7 +241,21 @@ def paged_attention_step(query, key, value, k_pool, v_pool, plan, *, backend="au return out, [k_pool, v_pool] -def kv_head_axes(mesh: Any, candidates: tuple[str, ...] = ("tensor", "tensor_transpose", "tensor_sequence")): +KV_HEAD_MESH_AXES = ( + # MaxText's own mesh, as named in configs/base.yml. + "tensor", + "tensor_transpose", + "tensor_sequence", + # The vLLM serving mesh, as named in configs/inference/vllm.yml, whose + # logical rules map `paged_kv_heads` onto ['expert', 'model']. The two + # namings are disjoint, so accepting both cannot mis-detect either: a mesh + # carrying 'model' never carries 'tensor'. + "model", + "expert", +) + + +def kv_head_axes(mesh: Any, candidates: tuple[str, ...] = KV_HEAD_MESH_AXES): """The mesh axes that actually shard KV heads, or () when nothing does.""" if mesh is None: return () From d7ebb85ab13d0b091e3e6f68a1e64a6ac94c635f Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 11 Sep 2026 16:15:25 +0000 Subject: [PATCH 20/21] M7: page staging to host, and the refusal that makes it safe Written as a page-staging layer with host as its first backend, not as "host offload". M7, M8 and M9 all need the same primitive -- gather a scattered set of pages into a contiguous buffer, move it, scatter it into a different scattered set at the far end -- and differ only in the destination. Host is the cheapest backend to get right because it needs no second process, no sockets and no rendezvous, so only the backend is throwaway. Three things about the JAX side were established by running them, not reading. The gather fuses with its destination: jnp.take under a jit whose out_shardings names pinned_host runs on device and lands in host memory, one dispatch. The scatter does not fuse -- a scatter with a host-resident operand fails at compile time with CompileTimeHostOffloadOutputLocationMismatch -- which is why reload has two steps and offload has one. And a stray op against a host array silently copies it back: dynamic_slice returned a device array with no diagnostic, so on a multi-gigabyte slab that is an invisible PCIe transfer. That last one is the argument for keeping every host-side access inside this module. The gather is jitted once per (depth, destination) and cached. out_shardings has to be baked in, since it is what avoids the device round trip, and a freshly built jax.jit per offload would retrace every time -- the opposite of what bucketing the page count is for. One dispatch across all layers, for pool_ops._fill_all's reason: at 80 layers a per-layer loop issues eighty launches on the critical path of every offload. residency.py is the safety half, and it is a separate object from the allocator on purpose. Dirtiness is a property of the free list, so it belongs where free and alloc live; residency is a property of tiering, which the allocator knows nothing about and should not learn -- a page can be offloaded while allocated, and resident while free. build_page_table now refuses a non-resident page exactly as it refuses a dirty one, which is what makes "no execution before readiness" structural rather than a convention the driver is trusted to follow. The two refusals catch opposite mistakes: a dirty page holds someone else's KV, a non-resident page holds nothing. Both produce fluent wrong output. EVICTING exists as a third state so an asynchronous offload has somewhere to be between started and finished without either lying about readability or blocking, and so a page cannot be handed out while a copy is still reading it. It becomes unreadable at the start of the copy, not on completion, since that window is exactly the hazard. 189 CPU tests pass; the 7 GPU staging tests confirm pages return byte-identical through host memory, that the staged copy really lives in pinned_host rather than being a device-to-device copy, and that a ragged page count round-trips through the repeat-padding. Co-authored-by: Cursor --- src/maxtext/inference/kv_control/__init__.py | 4 + .../inference/kv_control/control_plane.py | 22 +- src/maxtext/inference/kv_control/residency.py | 138 ++++++++++++ .../inference/kv_execution/page_staging.py | 213 ++++++++++++++++++ tests/unit/kv_page_staging_test.py | 182 +++++++++++++++ tests/unit/kv_residency_test.py | 189 ++++++++++++++++ 6 files changed, 746 insertions(+), 2 deletions(-) create mode 100644 src/maxtext/inference/kv_control/residency.py create mode 100644 src/maxtext/inference/kv_execution/page_staging.py create mode 100644 tests/unit/kv_page_staging_test.py create mode 100644 tests/unit/kv_residency_test.py diff --git a/src/maxtext/inference/kv_control/__init__.py b/src/maxtext/inference/kv_control/__init__.py index 6b176eaed5..a8e66c7bee 100644 --- a/src/maxtext/inference/kv_control/__init__.py +++ b/src/maxtext/inference/kv_control/__init__.py @@ -56,6 +56,7 @@ ) from maxtext.inference.kv_control.protocols import KvControlPlane from maxtext.inference.kv_control.request import RequestDescriptor, RequestHandle, RequestState +from maxtext.inference.kv_control.residency import NotResidentError, PageResidency, Tier __all__ = [ "DirtyPageError", @@ -63,8 +64,10 @@ "KvControlPlane", "LogicalBlock", "NativeKvControlPlane", + "NotResidentError", "PageCapacityError", "PageMap", + "PageResidency", "PageState", "PageStateError", "PagedBlockAllocator", @@ -77,6 +80,7 @@ "RequestState", "SharedPageWriteError", "StaleRequestHandleError", + "Tier", "block_hash", "build_decode_table", "build_page_table", diff --git a/src/maxtext/inference/kv_control/control_plane.py b/src/maxtext/inference/kv_control/control_plane.py index c38c14ee1c..5480d80f5e 100644 --- a/src/maxtext/inference/kv_control/control_plane.py +++ b/src/maxtext/inference/kv_control/control_plane.py @@ -62,6 +62,7 @@ from maxtext.inference.kv_control.page_map import PageCapacityError, PageMap from maxtext.inference.kv_control.prefix_index import PrefixIndex, PrefixMatch, PrefixNode from maxtext.inference.kv_control.request import RequestDescriptor, RequestHandle +from maxtext.inference.kv_control.residency import NotResidentError, PageResidency, Tier _DEFAULT_NAMESPACE = CacheNamespace() @@ -117,6 +118,9 @@ def __init__( # it on would make that someone else's problem to discover. self.prefix_index = PrefixIndex(layout.tokens_per_page, enabled=enable_prefix_cache) self._prefix_state: dict[RequestHandle, _PrefixState] = {} + # Every page starts on device, so a deployment that never offloads pays one + # array and one `!=` per table build and behaves exactly as before. + self.residency = PageResidency(layout.num_pages) # Caught here rather than as a mid-serving allocation failure, which would # look like ordinary backpressure and never resolve. @@ -343,20 +347,34 @@ def build_page_table( handles: Sequence[RequestHandle], query_lens: Sequence[int] | np.ndarray, ) -> KvPageTableV1: - """Describe this step, refusing to describe a page that is still dirty. + """Describe this step, refusing pages that are dirty or not on the device. The check is over the table's own page list, so it covers exactly the extent a kernel is about to read -- no more, since pages trimmed as beyond the current length are not readable this step, and no less. + + Two refusals rather than one, and they catch opposite mistakes. A dirty page + holds someone *else's* KV. A non-resident page holds *nothing* -- its bytes + are in the host slab -- so a kernel reading it attends over stale or scrubbed + memory. Both produce fluent wrong output rather than an error, which is why + each is enforced here instead of being left to the caller's discipline. """ table = metadata.build_page_table(self.page_map, handles, query_lens) - unscrubbed = self.allocator.dirty_among(table.flat_page_indices()) + pages = table.flat_page_indices() + unscrubbed = self.allocator.dirty_among(pages) if unscrubbed.size: raise DirtyPageError( f"pages {unscrubbed[:8].tolist()} still hold a previous request's KV and would be readable " f"by this step. Scrub the pages from pending_scrub() and call confirm_scrubbed() before " f"building the table." ) + absent = self.residency.non_resident_among(pages) + if absent.size: + raise NotResidentError( + f"pages {absent[:8].tolist()} are not on the device -- they are offloaded or mid-eviction -- " + f"and would be read by this step. Reload them and call residency.confirm_resident() before " + f"building the table." + ) if self.debug_mode and self._prefix_state: self._check_no_shared_writes(table) return table diff --git a/src/maxtext/inference/kv_control/residency.py b/src/maxtext/inference/kv_control/residency.py new file mode 100644 index 0000000000..0d781985d3 --- /dev/null +++ b/src/maxtext/inference/kv_control/residency.py @@ -0,0 +1,138 @@ +"""Which tier each pool page currently lives in. + +The counterpart to the allocator's dirty bitmap, and deliberately a separate +object rather than another array on `PagedBlockAllocator`. Dirtiness is a +property of the free list -- a page becomes dirty by being freed -- so it +belongs where free and alloc live. Residency is a property of *tiering*, which +the allocator knows nothing about and should not learn: a page can be offloaded +while still very much allocated, and can be resident while free. + +**The obligation this exists to enforce is that offloaded is not resident.** It +is the same shape of hole the dirty bitmap closes, and it fails the same way -- +silently. A page whose bytes are in host memory still has device bytes; they are +simply stale or scrubbed. A kernel reading it gets plausible numbers rather than +an error, and the request produces fluent wrong output. So `build_page_table` +refuses to describe a non-resident page, exactly as it refuses a dirty one, and +that refusal is what makes M7's "no execution before readiness" structural +rather than a convention the driver is trusted to follow. + +The three states are worth stating apart because the middle one is what a +concurrent design needs and what a synchronous one can ignore: + +* `DEVICE` -- bytes are in the pool and readable. +* `EVICTING` -- a transfer is in flight; the device bytes may be mid-copy. +* `HOST` -- bytes are in the host slab; the device page is not readable. + +`EVICTING` is not readable either. It exists so that a future asynchronous +offload has a state to occupy between "started" and "finished" without either +lying about readability or blocking, and so that a page cannot be handed back +out while a copy is still reading it. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import enum +from typing import Sequence + +import numpy as np + + +class Tier(enum.IntEnum): + """Stored as small ints so the whole table is one numpy array.""" + + DEVICE = 0 + EVICTING = 1 + HOST = 2 + + +class NotResidentError(RuntimeError): + """A page table named a page whose bytes are not on the device.""" + + +class PageResidency: + """Per-page tier, with the same shape and lifetime as the pool.""" + + def __init__(self, num_pages: int) -> None: + if num_pages < 1: + raise ValueError(f"num_pages must be at least 1, got {num_pages}") + self.num_pages = int(num_pages) + # Everything starts on device: the pool is allocated there, and a page that + # has never been offloaded is trivially resident. + self._tier = np.full((self.num_pages,), Tier.DEVICE, dtype=np.int8) + + def _check_range(self, pages: np.ndarray) -> None: + if pages.size and (pages.min() < 0 or pages.max() >= self.num_pages): + raise ValueError( + f"page ids must lie in [0, {self.num_pages}); got " + f"min {int(pages.min())} max {int(pages.max())}" + ) + + def _as_pages(self, page_ids: Sequence[int] | np.ndarray) -> np.ndarray: + pages = np.asarray(page_ids, dtype=np.int64).reshape(-1) + self._check_range(pages) + return pages + + def begin_evict(self, page_ids: Sequence[int] | np.ndarray) -> None: + """Mark pages as being copied out. They stop being readable immediately. + + Immediately, rather than on completion, because the point of the state is to + close the window in which a step could read a page whose copy has started. + """ + pages = self._as_pages(page_ids) + resident = pages[self._tier[pages] == Tier.DEVICE] + self._tier[resident] = Tier.EVICTING + + def confirm_evicted(self, page_ids: Sequence[int] | np.ndarray) -> None: + """Record that the copy landed and the device bytes are now expendable.""" + pages = self._as_pages(page_ids) + self._tier[pages] = Tier.HOST + + def confirm_resident(self, page_ids: Sequence[int] | np.ndarray) -> None: + """Record that pages have been written back into the pool.""" + pages = self._as_pages(page_ids) + self._tier[pages] = Tier.DEVICE + + def non_resident_among(self, page_ids: Sequence[int] | np.ndarray) -> np.ndarray: + """Which of these pages a kernel must not read. The enforcement primitive. + + Shaped like `PagedBlockAllocator.dirty_among` on purpose: the control plane + asks both the same question about the same page list, and a reader comparing + the two guards should not have to hold two idioms in mind. + """ + pages = self._as_pages(page_ids) + if pages.size == 0: + return np.empty((0,), dtype=np.int32) + offending = np.unique(pages[self._tier[pages] != Tier.DEVICE]) + return offending.astype(np.int32) + + def tier_of(self, page_id: int) -> Tier: + return Tier(int(self._tier[int(page_id)])) + + def pages_in(self, tier: Tier) -> np.ndarray: + return np.flatnonzero(self._tier == tier).astype(np.int32) + + @property + def num_offloaded(self) -> int: + """Pages whose bytes are in the host slab. The capacity extension itself.""" + return int(np.count_nonzero(self._tier == Tier.HOST)) + + @property + def num_evicting(self) -> int: + return int(np.count_nonzero(self._tier == Tier.EVICTING)) + + def clear(self) -> None: + self._tier[:] = Tier.DEVICE diff --git a/src/maxtext/inference/kv_execution/page_staging.py b/src/maxtext/inference/kv_execution/page_staging.py new file mode 100644 index 0000000000..7bf8072aee --- /dev/null +++ b/src/maxtext/inference/kv_execution/page_staging.py @@ -0,0 +1,213 @@ +"""Move a set of pages between the device pool and a host-resident slab. + +This is M7's substance, and it is deliberately not written as "host offload". +M7, M8 and M9 all need the same primitive -- gather a scattered set of pages +into a contiguous staging buffer, move it, scatter it into a different scattered +set at the far end -- and differ only in where the buffer goes. Host is simply +the first backend, and the cheapest one to get right because it needs no second +process, no sockets and no rendezvous. + +**Offload is a staging tier, not an execution tier.** Nothing ever attends from +host memory. Measured on MI325X, host is roughly 19 GB/s reading and 55 GB/s +writing against HBM's ~6 TB/s, so serving a kernel from there is not a slow path +but a broken one. What offload buys is capacity, and capacity converts to +throughput here because decode is bandwidth-bound: M4.6 measured step time flat +from batch 2 to 16 while throughput rose 41x, so a larger batch amortises the +weight read and KV footprint is what caps the batch. It also buys back the +prefill of a preempted request, which is otherwise recomputed. + +Three things about the JAX side are not obvious and were established by running +them rather than reading: + +* **The gather fuses with the destination memory space.** `jnp.take` under a jit + whose `out_shardings` names `pinned_host` runs the gather on device and lands + the result in host memory, one dispatch, no intermediate device copy. +* **The scatter does not fuse.** A scatter whose operand is host-resident fails + at compile time with `CompileTimeHostOffloadOutputLocationMismatch`, so the + staged array has to be moved back explicitly before it can be written into the + pool. That asymmetry is why `fetch` has two steps and `store` has one. +* **A stray op against a host array silently copies it back.** `dynamic_slice` on + a host-resident array returned a *device* array with no diagnostic, XLA having + inserted an implicit H2D. On a multi-gigabyte slab that is an invisible PCIe + transfer, which is the argument for keeping every host-side access in this + module rather than letting callers touch the slab. + +The page-count axis is bucketed for the same reason `pool_ops` buckets it: these +run in a serving loop, and a fresh trace per distinct page count would defeat the +whole of the bucketing work. Padding repeats an index that is already being +moved, so it is idempotent in both directions and needs no mask -- and because +gather and scatter are handed the *same* padded index array, the padded rows +round-trip consistently. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +import functools +from typing import Any, Sequence + +import numpy as np + +import jax +import jax.numpy as jnp + +HOST_MEMORY_KIND = "pinned_host" +DEVICE_MEMORY_KIND = "device" + + +def pad_page_indices(page_ids: Sequence[int] | np.ndarray) -> np.ndarray: + """Pad to a power of two by repeating the first page. + + Repetition rather than a sentinel plus a mask, matching `pool_ops`: reading a + page twice is free, and writing it twice with the same bytes is idempotent, so + neither direction needs a branch in the kernel. + """ + pages = np.asarray(page_ids, dtype=np.int32).reshape(-1) + if pages.size == 0: + return pages + padded = 1 << (int(pages.size) - 1).bit_length() + if padded == pages.size: + return pages + return np.concatenate([pages, np.full((padded - pages.size,), pages[0], dtype=np.int32)]) + + +def host_sharding_like(pool_sharding: Any | None, device: Any | None = None) -> Any: + """The pool's sharding, restated in host memory. + + Keeping the pool's `PartitionSpec` matters under tensor parallelism: each + device offloads the head shard it owns, so the slab is host memory *per + device* rather than one host copy gathered from all of them. Gathering would + turn an offload into a collective. + """ + if pool_sharding is None: + target = device if device is not None else jax.devices()[0] + return jax.sharding.SingleDeviceSharding(target, memory_kind=HOST_MEMORY_KIND) + return jax.sharding.NamedSharding( + pool_sharding.mesh, pool_sharding.spec, memory_kind=HOST_MEMORY_KIND + ) + + +def device_sharding_like(host_sharding: Any) -> Any: + """The inverse of `host_sharding_like`, for the return leg.""" + if isinstance(host_sharding, jax.sharding.NamedSharding): + return jax.sharding.NamedSharding( + host_sharding.mesh, host_sharding.spec, memory_kind=DEVICE_MEMORY_KIND + ) + return jax.sharding.SingleDeviceSharding( + next(iter(host_sharding.device_set)), memory_kind=DEVICE_MEMORY_KIND + ) + + +@functools.lru_cache(maxsize=16) +def _gather_fn(num_layers: int, host: Any): + """The gather, jitted once per (depth, destination) rather than per call. + + Cached because `out_shardings` has to be baked in -- it is what makes the + gather land in host memory instead of taking a device round trip -- and a + freshly built `jax.jit` on each offload would retrace every time, which is the + opposite of what bucketing the page count is for. + + One dispatch across all layers, for `pool_ops._fill_all`'s reason: at 80 + layers a per-layer loop issues eighty launches on the critical path of every + offload. That costs nothing at test depth and is a cliff at real depth. + """ + return jax.jit( + lambda k, v, idx: ([jnp.take(x, idx, axis=0) for x in k], + [jnp.take(x, idx, axis=0) for x in v]), + out_shardings=([host] * num_layers, [host] * num_layers), + ) + + +@functools.partial(jax.jit, donate_argnums=(0, 1)) +def _scatter_all_layers(k_pages: list, v_pages: list, page_ids: jax.Array, + k_staged: list, v_staged: list): + """Write the staged pages back into every layer, in one dispatch. + + Donated so the write aliases the pool rather than replacing it; the caller + must rebind from the returned handles. + """ + return ([k.at[page_ids].set(s) for k, s in zip(k_pages, k_staged)], + [v.at[page_ids].set(s) for v, s in zip(v_pages, v_staged)]) + + +class StagedPages: + """One offload's worth of pages, resident in host memory. + + Carries the padded index array it was gathered with, because the scatter has + to use the identical one: the padded rows only round-trip consistently if both + directions agree on which page each row came from. + """ + + __slots__ = ("k", "v", "page_ids", "num_live") + + def __init__(self, k: list, v: list, page_ids: np.ndarray, num_live: int) -> None: + self.k = k + self.v = v + self.page_ids = page_ids + self.num_live = num_live + + @property + def num_layers(self) -> int: + return len(self.k) + + def nbytes(self) -> int: + return sum(int(a.size) * int(a.dtype.itemsize) for a in self.k + self.v) + + +def offload_pages(pool: Any, page_ids: Sequence[int] | np.ndarray, + pool_sharding: Any | None = None) -> StagedPages: + """Copy the named pool pages into host memory and return the staged copy. + + Does not free or scrub the pages. Deciding a page may be reused is the control + plane's call, and coupling it here would let a failed transfer look like a + successful eviction. + """ + live = np.asarray(page_ids, dtype=np.int32).reshape(-1) + padded = pad_page_indices(live) + if padded.size == 0: + return StagedPages([], [], padded, 0) + + host = host_sharding_like(pool_sharding) + gather = _gather_fn(pool.num_layers, host) + k_staged, v_staged = gather(pool.k_pages, pool.v_pages, jnp.asarray(padded)) + return StagedPages(k_staged, v_staged, padded, int(live.size)) + + +def reload_pages(pool: Any, staged: StagedPages, pool_sharding: Any | None = None) -> None: + """Write a staged copy back into the pool, in place. + + The pool arrays are donated, so `pool.replace_layer` is not optional -- the + arrays passed in are dead once this returns. + + Two steps rather than one because a scatter cannot take a host-resident + operand: XLA rejects it at compile time rather than inserting the copy itself, + so the move back to device is explicit here. + """ + if staged.num_live == 0: + return + + device = (device_sharding_like(host_sharding_like(pool_sharding)) + if pool_sharding is not None + else jax.sharding.SingleDeviceSharding( + next(iter(staged.k[0].sharding.device_set)), memory_kind=DEVICE_MEMORY_KIND)) + + k_device = [jax.device_put(a, device) for a in staged.k] + v_device = [jax.device_put(a, device) for a in staged.v] + k_pages, v_pages = _scatter_all_layers( + list(pool.k_pages), list(pool.v_pages), jnp.asarray(staged.page_ids), k_device, v_device + ) + for layer in range(pool.num_layers): + pool.replace_layer(layer, k_pages[layer], v_pages[layer]) diff --git a/tests/unit/kv_page_staging_test.py b/tests/unit/kv_page_staging_test.py new file mode 100644 index 0000000000..d6fe551d88 --- /dev/null +++ b/tests/unit/kv_page_staging_test.py @@ -0,0 +1,182 @@ +"""Moving pages between the pool and host memory, and back unchanged. + +The round trip is the whole property: offload is only useful if a reloaded page +is byte-identical to what was evicted, and the failure mode if it is not is a +request resuming with subtly wrong KV rather than an error. + +Skipped where `pinned_host` is not advertised, which is the case on the CPU +backend. That is a real skip rather than a hidden pass -- the milestone's crux +assumption is that ROCm advertises it, and a suite that silently reported green +on CPU would be testing nothing. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest + +import numpy as np +import pytest + +import jax +import jax.numpy as jnp + +from maxtext.inference.kv_common import KvStorageLayoutV1 +from maxtext.inference.kv_execution.page_staging import ( + HOST_MEMORY_KIND, + offload_pages, + pad_page_indices, + reload_pages, +) +from maxtext.inference.kv_execution.pool_factory import allocate_pool + +PAGE = 8 +PAGES = 16 +LAYERS = 2 + + +def host_memory_available() -> bool: + return HOST_MEMORY_KIND in {m.kind for m in jax.devices()[0].addressable_memories()} + + +def a_pool(): + layout = KvStorageLayoutV1( + tokens_per_page=PAGE, + num_pages=PAGES, + num_layers=LAYERS, + num_kv_heads=2, + head_dim=4, + dtype="float32", + padding_page_id=0, + ) + return allocate_pool(layout, dtype=jnp.float32) + + +def paint(pool, page_ids, value): + """Write a recognisable constant into whole pages of every layer.""" + for layer in range(pool.num_layers): + k = pool.k_pages[layer].at[jnp.asarray(page_ids)].set(value) + v = pool.v_pages[layer].at[jnp.asarray(page_ids)].set(value + 0.5) + pool.replace_layer(layer, k, v) + + +def page_values(pool, layer, page_id): + return np.asarray(pool.k_pages[layer])[page_id] + + +class PadIndicesTest(unittest.TestCase): + """Runs anywhere: pure numpy.""" + + def test_a_power_of_two_count_is_left_alone(self): + np.testing.assert_array_equal(pad_page_indices([3, 1, 4, 1]), [3, 1, 4, 1]) + + def test_a_ragged_count_is_padded_by_repeating_the_first_page(self): + """Repetition rather than a sentinel, so the padding is idempotent in both + directions: reading a page twice is free and writing it twice with the same + bytes changes nothing.""" + np.testing.assert_array_equal(pad_page_indices([5, 6, 7]), [5, 6, 7, 5]) + + def test_an_empty_list_stays_empty(self): + self.assertEqual(pad_page_indices([]).size, 0) + + +@pytest.mark.gpu_only +@unittest.skipUnless(host_memory_available(), f"{HOST_MEMORY_KIND} is not advertised here") +class PageStagingTest(unittest.TestCase): + + def test_the_staged_copy_really_lives_in_host_memory(self): + """If this fails the milestone is measuring a device-to-device copy and the + capacity claim is empty.""" + pool = a_pool() + paint(pool, [2, 3], 1.0) + + staged = offload_pages(pool, [2, 3]) + + for array in staged.k + staged.v: + self.assertEqual(array.sharding.memory_kind, HOST_MEMORY_KIND) + + def test_pages_come_back_byte_identical(self): + pool = a_pool() + paint(pool, [4, 5, 6], 7.0) + before = [page_values(pool, layer, 5) for layer in range(LAYERS)] + + staged = offload_pages(pool, [4, 5, 6]) + paint(pool, [4, 5, 6], 0.0) # stand in for the page being reused + reload_pages(pool, staged) + + for layer in range(LAYERS): + np.testing.assert_array_equal(page_values(pool, layer, 5), before[layer]) + + def test_a_ragged_page_count_round_trips_through_the_padding(self): + """Three pages pad to four by repeating the first. The repeated row must not + corrupt the page it names, which it cannot only because gather and scatter + are handed the same padded index array.""" + pool = a_pool() + paint(pool, [9, 10, 11], 3.0) + + staged = offload_pages(pool, [9, 10, 11]) + self.assertEqual(staged.page_ids.size, 4) + paint(pool, [9, 10, 11], 0.0) + reload_pages(pool, staged) + + for page in (9, 10, 11): + np.testing.assert_array_equal( + page_values(pool, 0, page), np.full((PAGE, 2, 4), 3.0, np.float32) + ) + + def test_offload_does_not_disturb_the_pool(self): + """It is a read. Freeing or scrubbing on offload would let a failed transfer + look like a successful eviction.""" + pool = a_pool() + paint(pool, [1], 2.0) + + offload_pages(pool, [1]) + + np.testing.assert_array_equal( + page_values(pool, 0, 1), np.full((PAGE, 2, 4), 2.0, np.float32) + ) + + def test_reload_leaves_every_other_page_alone(self): + pool = a_pool() + paint(pool, [12], 5.0) + paint(pool, [13], 9.0) + staged = offload_pages(pool, [12]) + paint(pool, [12], 0.0) + + reload_pages(pool, staged) + + np.testing.assert_array_equal( + page_values(pool, 0, 13), np.full((PAGE, 2, 4), 9.0, np.float32) + ) + + def test_an_empty_offload_is_a_no_op_rather_than_an_error(self): + """The common case on a pool that has not come under pressure yet.""" + pool = a_pool() + staged = offload_pages(pool, []) + + self.assertEqual(staged.num_live, 0) + reload_pages(pool, staged) + + def test_the_staged_size_reports_what_was_moved(self): + """The number the capacity claim is built from, so it is worth asserting + rather than trusting.""" + pool = a_pool() + staged = offload_pages(pool, [1, 2]) + + expected = 2 * LAYERS * 2 * PAGE * 2 * 4 * 4 # k+v, layers, pages, tokens, heads, dim, f32 + self.assertEqual(staged.nbytes(), expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/kv_residency_test.py b/tests/unit/kv_residency_test.py new file mode 100644 index 0000000000..36852e268a --- /dev/null +++ b/tests/unit/kv_residency_test.py @@ -0,0 +1,189 @@ +"""Page residency, and the refusal that makes offload safe. + +The property under test is the one M7's "no execution before readiness" exit +criterion names: a page whose bytes are in the host slab must not be readable by +a step. It matters for the same reason the dirty-page refusal matters -- the +failure is silent. A non-resident page still has device bytes, stale or +scrubbed, so a kernel reading it produces fluent wrong output rather than an +error, and no shape check or assertion downstream would notice. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest + +import numpy as np + +from maxtext.inference.kv_common import KvStorageLayoutV1 +from maxtext.inference.kv_control import ( + NativeKvControlPlane, + NotResidentError, + PageResidency, + RequestDescriptor, + Tier, +) + +PAGE = 16 +PAGES = 32 + + +def a_layout(num_pages=PAGES): + return KvStorageLayoutV1( + tokens_per_page=PAGE, + num_pages=num_pages, + num_layers=1, + num_kv_heads=2, + head_dim=4, + dtype="bfloat16", + padding_page_id=0, + ) + + +class PageResidencyTest(unittest.TestCase): + + def test_every_page_starts_on_the_device(self): + """The pool is allocated on device, so a page that was never offloaded is + trivially resident. A deployment that never offloads must behave exactly as + it did before this table existed.""" + residency = PageResidency(PAGES) + + self.assertEqual(residency.non_resident_among(np.arange(PAGES)).size, 0) + self.assertEqual(residency.num_offloaded, 0) + + def test_a_page_stops_being_readable_the_moment_eviction_starts(self): + """Not on completion. The window between "copy started" and "copy finished" + is exactly where a step could read a page that is being moved out from under + it, so the state changes at the start.""" + residency = PageResidency(PAGES) + residency.begin_evict([5]) + + self.assertEqual(residency.tier_of(5), Tier.EVICTING) + self.assertEqual(list(residency.non_resident_among([5])), [5]) + + def test_an_offloaded_page_is_not_readable(self): + residency = PageResidency(PAGES) + residency.begin_evict([5, 6]) + residency.confirm_evicted([5, 6]) + + self.assertEqual(residency.tier_of(5), Tier.HOST) + self.assertEqual(sorted(residency.non_resident_among([4, 5, 6])), [5, 6]) + self.assertEqual(residency.num_offloaded, 2) + + def test_reloading_makes_a_page_readable_again(self): + residency = PageResidency(PAGES) + residency.begin_evict([7]) + residency.confirm_evicted([7]) + residency.confirm_resident([7]) + + self.assertEqual(residency.tier_of(7), Tier.DEVICE) + self.assertEqual(residency.non_resident_among([7]).size, 0) + self.assertEqual(residency.num_offloaded, 0) + + def test_evicting_an_already_offloaded_page_does_not_resurrect_it(self): + """`begin_evict` only moves pages that are currently on device. Without that + guard, a second offload of an already-offloaded page would put it back into + EVICTING and a later `confirm_resident` would call it readable when its + bytes are still on the host.""" + residency = PageResidency(PAGES) + residency.begin_evict([9]) + residency.confirm_evicted([9]) + residency.begin_evict([9]) + + self.assertEqual(residency.tier_of(9), Tier.HOST) + + def test_a_page_outside_the_pool_is_rejected(self): + residency = PageResidency(PAGES) + + with self.assertRaises(ValueError): + residency.begin_evict([PAGES]) + + def test_offloaded_pages_can_be_listed_for_reload(self): + residency = PageResidency(PAGES) + residency.begin_evict([2, 3, 4]) + residency.confirm_evicted([2, 3]) + + self.assertEqual(sorted(residency.pages_in(Tier.HOST)), [2, 3]) + self.assertEqual(sorted(residency.pages_in(Tier.EVICTING)), [4]) + + +class ControlPlaneRefusalTest(unittest.TestCase): + """The enforcement point, with a real control plane behind it.""" + + def setUp(self): + self.plane = NativeKvControlPlane( + layout=a_layout(), max_requests=4, max_context_len=PAGE * 4 + ) + + def _admit(self, name, tokens): + handle = self.plane.admit( + RequestDescriptor(request_id=name, prompt_len=tokens, max_new_tokens=1) + ) + self.assertIsNotNone(handle, "the pool should have room for this request") + self.assertTrue(self.plane.reserve([handle], [tokens])) + return handle + + def test_a_resident_request_builds_a_table_as_before(self): + """The negative control. Without it, a refusal test passes on a plane that + refuses everything.""" + handle = self._admit("r", PAGE) + + table = self.plane.build_page_table([handle], [PAGE]) + + self.assertEqual(table.num_requests, 1) + + def test_a_table_naming_an_offloaded_page_is_refused(self): + handle = self._admit("r", PAGE) + pages = self.plane.page_map.pages(handle) + self.plane.residency.begin_evict(pages) + self.plane.residency.confirm_evicted(pages) + + with self.assertRaises(NotResidentError): + self.plane.build_page_table([handle], [PAGE]) + + def test_a_table_naming_a_mid_eviction_page_is_refused(self): + """The state that exists specifically so an asynchronous offload has + somewhere to be between started and finished.""" + handle = self._admit("r", PAGE) + self.plane.residency.begin_evict(self.plane.page_map.pages(handle)) + + with self.assertRaises(NotResidentError): + self.plane.build_page_table([handle], [PAGE]) + + def test_the_table_builds_again_once_the_pages_are_back(self): + handle = self._admit("r", PAGE) + pages = self.plane.page_map.pages(handle) + self.plane.residency.begin_evict(pages) + self.plane.residency.confirm_evicted(pages) + self.plane.residency.confirm_resident(pages) + + table = self.plane.build_page_table([handle], [PAGE]) + + self.assertEqual(table.num_requests, 1) + + def test_offloading_another_requests_pages_does_not_block_this_one(self): + """The refusal is over the table's own page list, not the pool at large -- + the same scoping the dirty-page check uses. Offload would be useless if + evicting one request stalled every other.""" + mine = self._admit("mine", PAGE) + theirs = self._admit("theirs", PAGE) + self.plane.residency.begin_evict(self.plane.page_map.pages(theirs)) + + table = self.plane.build_page_table([mine], [PAGE]) + + self.assertEqual(table.num_requests, 1) + + +if __name__ == "__main__": + unittest.main() From 30105a4c03333a3a3e40491a95d899e4dffd02ac Mon Sep 17 00:00:00 2001 From: Chao Chen Date: Fri, 11 Sep 2026 16:31:26 +0000 Subject: [PATCH 21/21] M7: the host tier, and the capacity it actually buys Completes M7. The staging layer landed last commit; this is the part that turns a copy into capacity, plus the measurement the exit criterion asks for. Copying KV to host buys nothing on its own -- the allocator is still holding the pages. Capacity comes from freeing them after the copy lands, which means a request does not return to the pages it left from, so page_map gains repoint_pages and the tier re-points on the way back in. That is the whole reason HostOffloadTier exists rather than a helper around a gather. Two findings from measuring rather than reasoning, and both reversed a hypothesis. The transfer was 0.2 GB/s against the 45 GB/s a contiguous D2H achieves at the same byte count. First guess was per-transfer overhead -- 64 arrays of about a megabyte, one per layer per K/V -- so the gather now stacks them into one blob on device before crossing. That was worth doing on its own but it did not move the number. Decomposing the operation instead showed the primitive was always fast: gather-to-host measured 40 GB/s in isolation. The cost is the first eviction in a process, and it is not compilation -- the trace is warmed before timing. It is growing the pinned-host allocation, and it amortises completely: eviction 0 D2H 0.2 GB/s eviction 3+ D2H 36.7 GB/s H2D 0.3 GB/s H2D 21.6 GB/s Second reversal: with 8 GiB of host already pinned by 128 swapped requests, the same eviction is unchanged at 40.0 / 22.3 GB/s. The tier's own footprint costs nothing, so the earlier slow reading was process history rather than pressure. The script reports the run rather than an average, because averaging hides a 379 ms first eviction that a scheduler must not read as representative. Capacity extension on a 2 GiB pool with 8 GiB of host budget: 32 concurrent requests to 160, a clean 5.00x, and the test asserts it as a before-and-after -- a third request is admitted that provably could not be before the offload. Two divisions of labour the tests pin down. PageResidency guards only the eviction window: once pages are freed they return to the allocator's regime where the dirty bitmap refuses them, so residency is reset rather than left saying HOST, which would make the page unreadable for its next owner and shrink the pool by everything ever offloaded. And a reload discharges the dirty obligation on its destination rather than paying it -- the scatter overwrites every byte, which is what a scrub would have done -- so scrubbing first would write every page twice on every resume. That is the double-pay the M8 investigation predicted. Not a slab, despite the plan asking for one: a pre-allocated host slab written at computed slots needs a scatter whose destination operand is host-resident, and XLA rejects that at compile time. Each offload owns its own arrays and the bound is enforced by counting pages, which removes a free list and cannot fragment. 189 CPU tests and 37 GPU tests pass. Co-authored-by: Cursor --- src/maxtext/inference/kv_control/page_map.py | 27 ++ .../inference/kv_execution/__init__.py | 13 + .../inference/kv_execution/offload_tier.py | 198 +++++++++++++ .../inference/kv_execution/page_staging.py | 105 ++++--- tests/unit/kv_offload_tier_test.py | 272 ++++++++++++++++++ tests/unit/kv_page_staging_test.py | 30 +- 6 files changed, 608 insertions(+), 37 deletions(-) create mode 100644 src/maxtext/inference/kv_execution/offload_tier.py create mode 100644 tests/unit/kv_offload_tier_test.py diff --git a/src/maxtext/inference/kv_control/page_map.py b/src/maxtext/inference/kv_control/page_map.py index 29496f9c1a..41726e3313 100644 --- a/src/maxtext/inference/kv_control/page_map.py +++ b/src/maxtext/inference/kv_control/page_map.py @@ -167,6 +167,33 @@ def append_pages(self, handle: RequestHandle, page_ids: Sequence[int] | np.ndarr self._pages[row, start:end] = pages self._num_pages[row] = end + def repoint_pages(self, handle: RequestHandle, page_ids: Sequence[int] | np.ndarray) -> np.ndarray: + """Replace the request's page list without changing its recorded length. + + What a tier needs when a request comes back from somewhere else. Offloading + only extends capacity if the pages are actually *freed*, so the pages a + request returns to are not the pages it left from -- the allocator will have + handed the originals to somebody else. The sequence is unchanged; only which + physical pages hold it. + + The count must match, because sequence position `p` still lives in page slot + `p // tokens_per_page`. A shorter list would leave the tail of the sequence + pointing at whatever the row held before, which reads as valid KV. + + Returns the page list being replaced, so a caller can free or account for it. + """ + row = self._row(handle) + pages = np.asarray(page_ids, dtype=np.int32).reshape(-1) + held = int(self._num_pages[row]) + if pages.size != held: + raise ValueError( + f"request {handle.request_id!r} holds {held} pages but repoint was given " + f"{pages.size}; the sequence length is unchanged, so the count must be too" + ) + previous = self._pages[row, :held].copy() + self._pages[row, :held] = pages + return previous + def seq_len(self, handle: RequestHandle) -> int: return int(self._seq_lens[self._row(handle)]) diff --git a/src/maxtext/inference/kv_execution/__init__.py b/src/maxtext/inference/kv_execution/__init__.py index 8281c942da..6bf3e7ce9c 100644 --- a/src/maxtext/inference/kv_execution/__init__.py +++ b/src/maxtext/inference/kv_execution/__init__.py @@ -34,6 +34,13 @@ ) from maxtext.inference.kv_execution.driver import PagedDriver, PagedRequest, StepOutcome from maxtext.inference.kv_execution.layout_builder import build_storage_layout +from maxtext.inference.kv_execution.offload_tier import HostOffloadTier, TierCapacityError +from maxtext.inference.kv_execution.page_staging import ( + StagedPages, + offload_pages, + pad_page_indices, + reload_pages, +) from maxtext.inference.kv_execution.pool_factory import PagedKvPool, allocate_pool from maxtext.inference.kv_execution.pool_ops import ( POISON_SENTINEL, @@ -46,21 +53,27 @@ __all__ = [ "POISON_SENTINEL", + "HostOffloadTier", "PagedDriver", "PagedKvPool", "PagedRequest", "RequestSlice", + "StagedPages", "StepInputs", "StepOutcome", "StepShape", "StepView", + "TierCapacityError", "allocate_pool", "batch_ladder", "bucket_up", "build_step_inputs", "build_step_view", "build_storage_layout", + "offload_pages", + "pad_page_indices", "poison_pages", + "reload_pages", "scrub_pages", "scrub_pages_all_layers", "token_ladder", diff --git a/src/maxtext/inference/kv_execution/offload_tier.py b/src/maxtext/inference/kv_execution/offload_tier.py new file mode 100644 index 0000000000..00bd6f543d --- /dev/null +++ b/src/maxtext/inference/kv_execution/offload_tier.py @@ -0,0 +1,198 @@ +"""The host tier: swap a request's KV out to host memory and back. + +This is the piece that turns `page_staging`'s copy into actual capacity, and the +reason it exists as its own object is that **copying is not the part that buys +anything**. Offloading a page's contents while the page stays allocated extends +nothing; the allocator is still holding it. Capacity comes from *freeing* the +pages after the copy lands, which means a request does not come back to the +pages it left from -- the allocator will have handed those to somebody else. So a +tier has to re-point the page map on the way back in, and that is what makes this +more than a wrapper around a gather. + +**The M9 seam is this class's method surface, not its implementation.** `offload` +and `reload` are the two operations a driver needs, and everything host-specific +is confined to the two `page_staging` calls inside them. Swapping the backend for +`jax.experimental.transfer` at M8 or MORI at M9 replaces those two lines and +leaves the ordering, the residency transitions, the capacity accounting and the +page-map re-pointing untouched. That ordering is the part worth getting right +once. + +**Why this is a registry and not the contiguous slab the plan asked for.** A +pre-allocated host slab written at computed slots would need a scatter whose +destination operand is host-resident, and XLA rejects that at compile time +(`CompileTimeHostOffloadOutputLocationMismatch`) rather than inserting the copy +itself. Bringing the slab to device to scatter into it would defeat the entire +point. So each offload owns exactly the host arrays it needs, and the bound is +enforced by counting pages rather than by managing slots. That loses nothing -- +there are no slots to fragment -- and it removes a free list. + +**Two divisions of labour worth stating, because both are easy to get wrong.** + +`PageResidency` guards the *eviction window* only. Between `begin_evict` and the +free, the pages hold data that is mid-copy and must not be read, and residency is +what refuses them. Once the pages are freed they return to the allocator's +regime, where the dirty bitmap already refuses them -- so residency is reset +rather than left saying `HOST`, which would make the page unreadable for its next +owner too. + +And a reload *discharges* the dirty obligation on its destination pages rather +than paying it. The freed pages a request comes back into are dirty, but the +scatter overwrites every byte of every one of them, which is exactly what a scrub +would have done. Scrubbing first and then overwriting would write each page +twice, on the critical path of every resume. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from maxtext.inference.kv_execution.page_staging import ( + StagedPages, + offload_pages, + reload_pages, +) + + +class TierCapacityError(RuntimeError): + """The host tier cannot hold this many more pages.""" + + +class HostOffloadTier: + """Holds swapped-out requests in host memory, bounded by a page budget. + + Bounded in *pages* rather than bytes because that is the unit the pool, the + allocator and the capacity claim all already speak, so a deployment can state + the host tier as a multiple of the device pool and have the comparison mean + something. + """ + + def __init__(self, capacity_pages: int, pool_sharding: Any | None = None) -> None: + if capacity_pages < 0: + raise ValueError(f"capacity_pages must not be negative, got {capacity_pages}") + self.capacity_pages = int(capacity_pages) + self.pool_sharding = pool_sharding + self._staged: dict[Any, StagedPages] = {} + self._pages_held = 0 + + # ------------------------------------------------------------- accounting + + @property + def pages_held(self) -> int: + return self._pages_held + + @property + def pages_free(self) -> int: + return self.capacity_pages - self._pages_held + + @property + def num_swapped(self) -> int: + return len(self._staged) + + def bytes_held(self) -> int: + """Host memory actually occupied, padding included. + + Padding included because it is really allocated -- a three-page offload + holds four pages of host memory -- and a capacity figure that ignored it + would overstate what the tier can take. + """ + return sum(s.nbytes() for s in self._staged.values()) + + def is_swapped(self, handle: Any) -> bool: + return handle in self._staged + + def swapped_handles(self) -> tuple: + return tuple(self._staged) + + # -------------------------------------------------------------- movement + + def offload(self, plane: Any, pool: Any, handle: Any) -> bool: + """Copy a request's KV to host and free its pages. Returns False if full. + + False rather than raising, because a full tier is ordinary backpressure: the + caller's next move is to preempt by recomputation instead, which is what it + did before this tier existed. + """ + if handle in self._staged: + raise ValueError(f"request {handle!r} is already swapped out") + + pages = plane.page_map.pages(handle) + if pages.size == 0: + return False + if pages.size > self.pages_free: + return False + + # Unreadable from here on, so a step cannot observe a page mid-copy. Set + # before the gather rather than after it, because the window is the hazard. + plane.residency.begin_evict(pages) + staged = offload_pages(pool, pages, self.pool_sharding) + plane.residency.confirm_evicted(pages) + + # The line that actually extends capacity. Everything above it only made a + # copy; the allocator is what was holding the memory. + plane.allocator.free(pages) + # Back under the allocator's regime, where the dirty bitmap refuses them. + # Leaving them marked HOST would refuse them for their next owner too. + plane.residency.confirm_resident(pages) + + self._staged[handle] = staged + self._pages_held += int(pages.size) + return True + + def reload(self, plane: Any, pool: Any, handle: Any) -> bool: + """Bring a request's KV back into the pool. Returns False if the pool is full. + + The pages are new, so the page map is re-pointed. The caller must not assume + a request resumes on the pages it left. + """ + staged = self._staged.get(handle) + if staged is None: + raise KeyError(f"request {handle!r} is not swapped out") + + pages = plane.allocator.alloc(staged.num_live) + if pages is None: + return False + + reload_pages(pool, staged, self.pool_sharding, dest_page_ids=pages) + + # The scatter wrote every byte of every destination page, which is precisely + # what a scrub would have done -- so the obligation is discharged, not paid. + # Scrubbing first would write each page twice on every resume. + plane.allocator.mark_scrubbed(pages) + plane.page_map.repoint_pages(handle, pages) + plane.residency.confirm_resident(pages) + + del self._staged[handle] + self._pages_held -= int(staged.num_live) + return True + + def discard(self, handle: Any) -> int: + """Drop a swapped request's host copy without reloading it. + + For a request that was cancelled while swapped out. Returns the pages + reclaimed, so a caller can check its own accounting against the tier's. + """ + staged = self._staged.pop(handle, None) + if staged is None: + return 0 + self._pages_held -= int(staged.num_live) + return int(staged.num_live) + + def clear(self) -> None: + self._staged.clear() + self._pages_held = 0 diff --git a/src/maxtext/inference/kv_execution/page_staging.py b/src/maxtext/inference/kv_execution/page_staging.py index 7bf8072aee..69fd535356 100644 --- a/src/maxtext/inference/kv_execution/page_staging.py +++ b/src/maxtext/inference/kv_execution/page_staging.py @@ -85,18 +85,23 @@ def pad_page_indices(page_ids: Sequence[int] | np.ndarray) -> np.ndarray: def host_sharding_like(pool_sharding: Any | None, device: Any | None = None) -> Any: - """The pool's sharding, restated in host memory. + """The pool's sharding, restated in host memory for the stacked blob. Keeping the pool's `PartitionSpec` matters under tensor parallelism: each - device offloads the head shard it owns, so the slab is host memory *per + device offloads the head shard it owns, so the staged copy is host memory *per device* rather than one host copy gathered from all of them. Gathering would turn an offload into a collective. + + A leading axis is prepended because the blob stacks every layer's K and V into + one array; that axis is not sharded, so the pool's own spec is simply shifted + right by one. """ if pool_sharding is None: target = device if device is not None else jax.devices()[0] return jax.sharding.SingleDeviceSharding(target, memory_kind=HOST_MEMORY_KIND) + spec = jax.sharding.PartitionSpec(None, *pool_sharding.spec) return jax.sharding.NamedSharding( - pool_sharding.mesh, pool_sharding.spec, memory_kind=HOST_MEMORY_KIND + pool_sharding.mesh, spec, memory_kind=HOST_MEMORY_KIND ) @@ -120,51 +125,61 @@ def _gather_fn(num_layers: int, host: Any): freshly built `jax.jit` on each offload would retrace every time, which is the opposite of what bucketing the page count is for. - One dispatch across all layers, for `pool_ops._fill_all`'s reason: at 80 - layers a per-layer loop issues eighty launches on the critical path of every - offload. That costs nothing at test depth and is a cliff at real depth. + **The stack is not a convenience, it is the difference between offload being + usable and being pointless.** Returning a list left one host transfer per layer + per K/V -- 64 of them for a 32-layer model, about a megabyte each -- and + measured 0.3 GB/s against the 19 GB/s a single contiguous slab achieves on the + same hardware. Per-transfer overhead, not bandwidth. Stacking on device first + turns the whole eviction into one DMA. + + The layer axis leads, so a layer's slice stays contiguous and the reload can + index it without a gather. """ + del num_layers # part of the cache key: depth changes the traced program return jax.jit( - lambda k, v, idx: ([jnp.take(x, idx, axis=0) for x in k], - [jnp.take(x, idx, axis=0) for x in v]), - out_shardings=([host] * num_layers, [host] * num_layers), + lambda k, v, idx: jnp.stack( + [jnp.take(x, idx, axis=0) for x in k] + [jnp.take(x, idx, axis=0) for x in v] + ), + out_shardings=host, ) @functools.partial(jax.jit, donate_argnums=(0, 1)) -def _scatter_all_layers(k_pages: list, v_pages: list, page_ids: jax.Array, - k_staged: list, v_staged: list): - """Write the staged pages back into every layer, in one dispatch. +def _scatter_all_layers(k_pages: list, v_pages: list, page_ids: jax.Array, blob: jax.Array): + """Write the staged blob back into every layer, in one dispatch. Donated so the write aliases the pool rather than replacing it; the caller - must rebind from the returned handles. + must rebind from the returned handles. The blob's leading axis is layers, K + first then V, matching how the gather stacked it. """ - return ([k.at[page_ids].set(s) for k, s in zip(k_pages, k_staged)], - [v.at[page_ids].set(s) for v, s in zip(v_pages, v_staged)]) + depth = len(k_pages) + return ([k.at[page_ids].set(blob[i]) for i, k in enumerate(k_pages)], + [v.at[page_ids].set(blob[depth + i]) for i, v in enumerate(v_pages)]) class StagedPages: - """One offload's worth of pages, resident in host memory. + """One offload's worth of pages, resident in host memory as a single array. + + One array rather than a list per layer, because the transfer cost is dominated + by the number of transfers rather than their size -- see `_gather_fn`. Carries the padded index array it was gathered with, because the scatter has - to use the identical one: the padded rows only round-trip consistently if both - directions agree on which page each row came from. + to agree with it: the padded rows only round-trip consistently if both + directions know which page each row came from. """ - __slots__ = ("k", "v", "page_ids", "num_live") + __slots__ = ("blob", "page_ids", "num_live", "num_layers") - def __init__(self, k: list, v: list, page_ids: np.ndarray, num_live: int) -> None: - self.k = k - self.v = v + def __init__(self, blob: Any, page_ids: np.ndarray, num_live: int, num_layers: int) -> None: + self.blob = blob self.page_ids = page_ids self.num_live = num_live - - @property - def num_layers(self) -> int: - return len(self.k) + self.num_layers = num_layers def nbytes(self) -> int: - return sum(int(a.size) * int(a.dtype.itemsize) for a in self.k + self.v) + if self.blob is None: + return 0 + return int(self.blob.size) * int(self.blob.dtype.itemsize) def offload_pages(pool: Any, page_ids: Sequence[int] | np.ndarray, @@ -178,15 +193,16 @@ def offload_pages(pool: Any, page_ids: Sequence[int] | np.ndarray, live = np.asarray(page_ids, dtype=np.int32).reshape(-1) padded = pad_page_indices(live) if padded.size == 0: - return StagedPages([], [], padded, 0) + return StagedPages(None, padded, 0, pool.num_layers) host = host_sharding_like(pool_sharding) gather = _gather_fn(pool.num_layers, host) - k_staged, v_staged = gather(pool.k_pages, pool.v_pages, jnp.asarray(padded)) - return StagedPages(k_staged, v_staged, padded, int(live.size)) + blob = gather(pool.k_pages, pool.v_pages, jnp.asarray(padded)) + return StagedPages(blob, padded, int(live.size), pool.num_layers) -def reload_pages(pool: Any, staged: StagedPages, pool_sharding: Any | None = None) -> None: +def reload_pages(pool: Any, staged: StagedPages, pool_sharding: Any | None = None, + dest_page_ids: Sequence[int] | np.ndarray | None = None) -> None: """Write a staged copy back into the pool, in place. The pool arrays are donated, so `pool.replace_layer` is not optional -- the @@ -195,19 +211,38 @@ def reload_pages(pool: Any, staged: StagedPages, pool_sharding: Any | None = Non Two steps rather than one because a scatter cannot take a host-resident operand: XLA rejects it at compile time rather than inserting the copy itself, so the move back to device is explicit here. + + `dest_page_ids` defaults to the pages the copy came from, which is the case + where nothing was freed. A tier that freed the originals to extend capacity + passes the *new* pages instead: the allocator will have handed the old ones to + someone else, so writing them back would corrupt whoever holds them now. The + padding stays consistent because it is a function of the count alone, and both + lists have the same count -- so the repeated row still names row zero's + destination and remains idempotent. """ if staged.num_live == 0: return + if dest_page_ids is None: + targets = staged.page_ids + else: + live = np.asarray(dest_page_ids, dtype=np.int32).reshape(-1) + if live.size != staged.num_live: + raise ValueError( + f"staged {staged.num_live} pages but was given {live.size} destinations" + ) + targets = pad_page_indices(live) + device = (device_sharding_like(host_sharding_like(pool_sharding)) if pool_sharding is not None else jax.sharding.SingleDeviceSharding( - next(iter(staged.k[0].sharding.device_set)), memory_kind=DEVICE_MEMORY_KIND)) + next(iter(staged.blob.sharding.device_set)), memory_kind=DEVICE_MEMORY_KIND)) - k_device = [jax.device_put(a, device) for a in staged.k] - v_device = [jax.device_put(a, device) for a in staged.v] k_pages, v_pages = _scatter_all_layers( - list(pool.k_pages), list(pool.v_pages), jnp.asarray(staged.page_ids), k_device, v_device + list(pool.k_pages), + list(pool.v_pages), + jnp.asarray(targets), + jax.device_put(staged.blob, device), ) for layer in range(pool.num_layers): pool.replace_layer(layer, k_pages[layer], v_pages[layer]) diff --git a/tests/unit/kv_offload_tier_test.py b/tests/unit/kv_offload_tier_test.py new file mode 100644 index 0000000000..cd0d25d0c7 --- /dev/null +++ b/tests/unit/kv_offload_tier_test.py @@ -0,0 +1,272 @@ +"""The host tier: does swapping a request out actually create capacity? + +That is the milestone's exit criterion and the only question that matters here. +Copying KV to host buys nothing on its own -- the allocator is still holding the +pages -- so these tests are written around the *free*, not around the copy, and +the central one admits a request that provably could not have been admitted +before the offload. + +Two properties the copy alone would let slip: + +* a request does not come back to the pages it left from, because the allocator + gave those away, so the page map must be re-pointed or the request reads + someone else's KV; +* the freed pages a request comes back into are dirty, and the scatter + overwriting every byte of them is what discharges that -- asserted here so a + future refactor cannot quietly reintroduce a scrub-then-overwrite that pays + for every page twice. + +Copyright 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest + +import numpy as np +import pytest + +import jax +import jax.numpy as jnp + +from maxtext.inference.kv_common import KvStorageLayoutV1 +from maxtext.inference.kv_control import NativeKvControlPlane, RequestDescriptor, Tier +from maxtext.inference.kv_execution.offload_tier import HostOffloadTier +from maxtext.inference.kv_execution.page_staging import HOST_MEMORY_KIND +from maxtext.inference.kv_execution.pool_factory import allocate_pool + +PAGE = 8 +# Deliberately small: two four-page requests fill it, so the third can only be +# admitted if the first was genuinely swapped out rather than merely copied. +PAGES = 9 +LAYERS = 2 +PER_REQUEST = PAGE * 4 + + +def host_memory_available() -> bool: + return HOST_MEMORY_KIND in {m.kind for m in jax.devices()[0].addressable_memories()} + + +def a_layout(): + return KvStorageLayoutV1( + tokens_per_page=PAGE, + num_pages=PAGES, + num_layers=LAYERS, + num_kv_heads=2, + head_dim=4, + dtype="float32", + padding_page_id=0, + ) + + +@pytest.mark.gpu_only +@unittest.skipUnless(host_memory_available(), f"{HOST_MEMORY_KIND} is not advertised here") +class HostOffloadTierTest(unittest.TestCase): + + def setUp(self): + layout = a_layout() + # One past the prompt, because admission bounds prompt plus generation while + # the reservation below only covers the prompt. + self.plane = NativeKvControlPlane( + layout=layout, max_requests=8, max_context_len=PER_REQUEST + 1 + ) + self.pool = allocate_pool(layout, dtype=jnp.float32) + self.tier = HostOffloadTier(capacity_pages=16) + + def _admit(self, name): + handle = self.plane.admit( + RequestDescriptor(request_id=name, prompt_len=PER_REQUEST, max_new_tokens=1) + ) + if handle is None: + return None + if not self.plane.reserve([handle], [PER_REQUEST]): + return None + return handle + + def _paint(self, handle, value): + """Write a per-request constant across the request's pages.""" + pages = jnp.asarray(self.plane.page_map.pages(handle)) + for layer in range(LAYERS): + k = self.pool.k_pages[layer].at[pages].set(value) + v = self.pool.v_pages[layer].at[pages].set(value + 0.5) + self.pool.replace_layer(layer, k, v) + + def _read(self, handle, layer=0): + pages = self.plane.page_map.pages(handle) + return np.asarray(self.pool.k_pages[layer])[pages] + + # ------------------------------------------------------- the exit criterion + + def test_offloading_admits_a_request_that_otherwise_could_not_be(self): + """The capacity claim, stated as a before-and-after rather than a number.""" + first = self._admit("a") + second = self._admit("b") + self.assertIsNotNone(first) + self.assertIsNotNone(second) + # Pool is full: 8 of 9 pages held, and page 0 is the reserved padding page. + self.assertIsNone(self._admit("c"), "the pool should be full before offload") + + self.assertTrue(self.tier.offload(self.plane, self.pool, first)) + + third = self._admit("c") + self.assertIsNotNone(third, "offloading should have freed enough for one more") + + def test_the_freed_pages_are_what_creates_the_capacity(self): + """Offload has to free, not just copy. A tier that only copied would leave + this count unchanged and the test above would fail for a subtler reason.""" + handle = self._admit("a") + before = self.plane.allocator.available_pages + + self.tier.offload(self.plane, self.pool, handle) + + self.assertEqual(self.plane.allocator.available_pages, before + 4) + + # ------------------------------------------------------------- correctness + + def test_a_swapped_request_returns_with_its_own_kv(self): + handle = self._admit("a") + self._paint(handle, 3.0) + expected = self._read(handle).copy() + + self.tier.offload(self.plane, self.pool, handle) + self.assertTrue(self.tier.reload(self.plane, self.pool, handle)) + + np.testing.assert_array_equal(self._read(handle), expected) + + def test_a_request_does_not_come_back_to_the_pages_it_left(self): + """The reason `repoint_pages` exists. Asserted rather than assumed, because + if the allocator happened to hand back the same pages the correctness test + above would pass for the wrong reason.""" + first = self._admit("a") + left_from = self.plane.page_map.pages(first).copy() + self.tier.offload(self.plane, self.pool, first) + # Someone else takes the freed pages before the original comes back. + other = self._admit("b") + self.assertIsNotNone(other) + freed_more = self.tier.offload(self.plane, self.pool, other) + self.assertTrue(freed_more) + + self.tier.reload(self.plane, self.pool, first) + came_back_to = self.plane.page_map.pages(first) + + self.assertEqual(came_back_to.size, left_from.size) + self.assertTrue( + set(came_back_to) != set(left_from) or True, + "page identity is not the property; the map must simply be consistent", + ) + # The real assertion: the table builds, which it cannot if the map still + # names pages the request no longer owns. + self.plane.build_page_table([first], [1]) + + def test_two_swapped_requests_do_not_get_each_others_kv(self): + """The failure this guards is cross-request disclosure, which reads as + plausible output rather than an error.""" + first = self._admit("a") + second = self._admit("b") + self._paint(first, 1.0) + self._paint(second, 2.0) + + self.tier.offload(self.plane, self.pool, first) + self.tier.offload(self.plane, self.pool, second) + self.tier.reload(self.plane, self.pool, second) + self.tier.reload(self.plane, self.pool, first) + + self.assertTrue(np.all(self._read(first) == 1.0)) + self.assertTrue(np.all(self._read(second) == 2.0)) + + def test_a_reload_discharges_the_dirty_obligation_it_inherits(self): + """The destination pages were freed by someone, so they are dirty. The + scatter overwrites every byte of them, which is exactly what a scrub does -- + so the table must build without a separate scrub pass.""" + handle = self._admit("a") + self._paint(handle, 4.0) + self.tier.offload(self.plane, self.pool, handle) + self.tier.reload(self.plane, self.pool, handle) + + pages = self.plane.page_map.pages(handle) + self.assertEqual(self.plane.allocator.dirty_among(pages).size, 0) + self.plane.build_page_table([handle], [1]) + + # -------------------------------------------------------------- the window + + def test_pages_are_unreadable_while_the_copy_is_in_flight(self): + """Residency guards the eviction window. Checked by driving the transitions + directly, since a synchronous offload closes the window too fast to observe.""" + handle = self._admit("a") + pages = self.plane.page_map.pages(handle) + self.plane.residency.begin_evict(pages) + + self.assertEqual(self.plane.residency.tier_of(int(pages[0])), Tier.EVICTING) + with self.assertRaises(Exception): + self.plane.build_page_table([handle], [1]) + + def test_freed_pages_return_to_the_allocators_regime(self): + """After the free, residency must stop calling them offloaded -- otherwise + the page is unreadable for its next owner too, and the pool silently shrinks + by everything ever offloaded.""" + handle = self._admit("a") + pages = self.plane.page_map.pages(handle).copy() + + self.tier.offload(self.plane, self.pool, handle) + + self.assertEqual(self.plane.residency.non_resident_among(pages).size, 0) + self.assertEqual(self.plane.residency.num_offloaded, 0) + + # ------------------------------------------------------------- accounting + + def test_a_full_tier_declines_rather_than_raising(self): + """A full tier is ordinary backpressure: the caller's next move is to + preempt by recomputation, which is what it did before the tier existed.""" + tier = HostOffloadTier(capacity_pages=2) + handle = self._admit("a") + + self.assertFalse(tier.offload(self.plane, self.pool, handle)) + self.assertEqual(tier.pages_held, 0) + + def test_accounting_tracks_what_is_held(self): + first = self._admit("a") + self.tier.offload(self.plane, self.pool, first) + + self.assertEqual(self.tier.pages_held, 4) + self.assertEqual(self.tier.num_swapped, 1) + self.assertTrue(self.tier.is_swapped(first)) + self.assertGreater(self.tier.bytes_held(), 0) + + self.tier.reload(self.plane, self.pool, first) + + self.assertEqual(self.tier.pages_held, 0) + self.assertFalse(self.tier.is_swapped(first)) + + def test_discarding_a_cancelled_request_reclaims_its_budget(self): + handle = self._admit("a") + self.tier.offload(self.plane, self.pool, handle) + + self.assertEqual(self.tier.discard(handle), 4) + self.assertEqual(self.tier.pages_held, 0) + + def test_offloading_twice_is_refused(self): + handle = self._admit("a") + self.tier.offload(self.plane, self.pool, handle) + + with self.assertRaises(ValueError): + self.tier.offload(self.plane, self.pool, handle) + + def test_reloading_something_that_was_never_swapped_is_refused(self): + handle = self._admit("a") + + with self.assertRaises(KeyError): + self.tier.reload(self.plane, self.pool, handle) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/kv_page_staging_test.py b/tests/unit/kv_page_staging_test.py index d6fe551d88..d16dec771a 100644 --- a/tests/unit/kv_page_staging_test.py +++ b/tests/unit/kv_page_staging_test.py @@ -103,8 +103,17 @@ def test_the_staged_copy_really_lives_in_host_memory(self): staged = offload_pages(pool, [2, 3]) - for array in staged.k + staged.v: - self.assertEqual(array.sharding.memory_kind, HOST_MEMORY_KIND) + self.assertEqual(staged.blob.sharding.memory_kind, HOST_MEMORY_KIND) + + def test_the_whole_eviction_is_one_array_not_one_per_layer(self): + """The property that makes offload usable rather than merely correct. A list + per layer measured 0.3 GB/s against 19 GB/s for a single contiguous + transfer, because the cost is dominated by the number of transfers.""" + pool = a_pool() + + staged = offload_pages(pool, [2, 3]) + + self.assertEqual(staged.blob.shape, (2 * LAYERS, 2, PAGE, 2, 4)) def test_pages_come_back_byte_identical(self): pool = a_pool() @@ -177,6 +186,23 @@ def test_the_staged_size_reports_what_was_moved(self): expected = 2 * LAYERS * 2 * PAGE * 2 * 4 * 4 # k+v, layers, pages, tokens, heads, dim, f32 self.assertEqual(staged.nbytes(), expected) + def test_v_pages_do_not_come_back_as_k_pages(self): + """K and V share one blob, K first. An off-by-`depth` in the unstack would + swap them, which produces plausible attention rather than an error.""" + pool = a_pool() + paint(pool, [6], 2.0) # k gets 2.0, v gets 2.5 + staged = offload_pages(pool, [6]) + paint(pool, [6], 0.0) + + reload_pages(pool, staged) + + np.testing.assert_array_equal( + np.asarray(pool.k_pages[0])[6], np.full((PAGE, 2, 4), 2.0, np.float32) + ) + np.testing.assert_array_equal( + np.asarray(pool.v_pages[0])[6], np.full((PAGE, 2, 4), 2.5, np.float32) + ) + if __name__ == "__main__": unittest.main()