[Draft] kv paged kernels - #93
Draft
i-chaochen wants to merge 19 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>
…arness 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 <cursoragent@cursor.com>
…tly (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 <cursoragent@cursor.com>
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 <cursoragent@cursor.com>
…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 <cursoragent@cursor.com>
…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 <cursoragent@cursor.com>
…rving 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 <cursoragent@cursor.com>
…aged 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 <cursoragent@cursor.com>
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 <cursoragent@cursor.com>
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 <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