[Draft] kv paged kernels - #93
Draft
i-chaochen wants to merge 10 commits into
Draft
Conversation
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 <cursoragent@cursor.com>
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.
… 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 <cursoragent@cursor.com>
… 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 <cursoragent@cursor.com>
…enchmark 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 <cursoragent@cursor.com>
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 <cursoragent@cursor.com>
… 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 <cursoragent@cursor.com>
…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 <cursoragent@cursor.com>
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 <cursoragent@cursor.com>
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 <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
toy model result
https://github.com/AMD-ROCm-Internal/rocm-jax-xla-architecture/pull/1
it's still WIP
llama3.1-8b result
Step time is flat from batch 2 to 16 — 11.20 ms to 10.65 ms — while throughput rises 41× (1502/36.3) in one gpu setup.
Prefix cache