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_prefix_cache_benchmark.py b/benchmarks/paged_kv/run_prefix_cache_benchmark.py new file mode 100644 index 0000000000..94e2532be5 --- /dev/null +++ b/benchmarks/paged_kv/run_prefix_cache_benchmark.py @@ -0,0 +1,262 @@ +"""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. + + 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`. + """ + # pylint: disable=import-outside-toplevel + from maxtext.inference.kv_execution.driver import PagedDriver + + 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): + # 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.prompt_len) 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 new file mode 100644 index 0000000000..74c6ab0d13 --- /dev/null +++ b/benchmarks/paged_kv/run_serving_benchmark.py @@ -0,0 +1,310 @@ +"""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["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) + + +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( + "--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). " + "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) + 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(): + 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 + # 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 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"] + 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/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..a11122add0 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -598,7 +598,59 @@ 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." + ), + ) + 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." + ), + ) + 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." @@ -3158,6 +3210,45 @@ 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": + 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") def set_derived_and_validate_values(self) -> "MaxTextConfig": """ 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/kv_common/__init__.py b/src/maxtext/inference/kv_common/__init__.py new file mode 100644 index 0000000000..5fa881272a --- /dev/null +++ b/src/maxtext/inference/kv_common/__init__.py @@ -0,0 +1,42 @@ +"""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.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, + KvStorageLayoutV1, +) + +__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_common/page_table.py b/src/maxtext/inference/kv_common/page_table.py new file mode 100644 index 0000000000..9dd38d5438 --- /dev/null +++ b/src/maxtext/inference/kv_common/page_table.py @@ -0,0 +1,163 @@ +"""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] 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) + 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..5fe3e16b91 --- /dev/null +++ b/src/maxtext/inference/kv_common/storage_layout.py @@ -0,0 +1,192 @@ +"""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. + 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 + 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 + pool_replicas: 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.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( + 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, 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 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 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: + 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. + + 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: + """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.""" + return max(self.num_pages - 1, 0) * self.tokens_per_page diff --git a/src/maxtext/inference/kv_control/__init__.py b/src/maxtext/inference/kv_control/__init__.py new file mode 100644 index 0000000000..6b176eaed5 --- /dev/null +++ b/src/maxtext/inference/kv_control/__init__.py @@ -0,0 +1,88 @@ +"""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, + SharedPageWriteError, +) +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.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 + +__all__ = [ + "DirtyPageError", + "DoubleFreeError", + "KvControlPlane", + "LogicalBlock", + "NativeKvControlPlane", + "PageCapacityError", + "PageMap", + "PageState", + "PageStateError", + "PagedBlockAllocator", + "PrefixIndex", + "PrefixMatch", + "PrefixNode", + "PublishResult", + "RequestDescriptor", + "RequestHandle", + "RequestState", + "SharedPageWriteError", + "StaleRequestHandleError", + "block_hash", + "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..c38c14ee1c --- /dev/null +++ b/src/maxtext/inference/kv_control/control_plane.py @@ -0,0 +1,386 @@ +"""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. + +**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"); +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_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`.""" + + def __init__( + self, + layout: KvStorageLayoutV1, + 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. + 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) + + @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], + 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)) + + needed = int(sum(per_request_pages)) + pages = self.allocator.alloc(needed) + if pages is None: + # 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: + 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 evict_cached(self, num_pages: int) -> np.ndarray: + """Drop `num_pages` unreferenced cached pages and return them to the pool. + + The pages come back dirty, as any freed page does, so a later step that + reserves them will scrub them through the ordinary path. + """ + 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 + + 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." + ) + 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/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/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_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..8281c942da --- /dev/null +++ b/src/maxtext/inference/kv_execution/__init__.py @@ -0,0 +1,67 @@ +"""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, + 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__ = [ + "POISON_SENTINEL", + "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", + "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 new file mode 100644 index 0000000000..0222d3338a --- /dev/null +++ b/src/maxtext/inference/kv_execution/benchmark.py @@ -0,0 +1,647 @@ +"""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 + +from maxtext.inference.kv_execution.driver import PagedDriver, PagedRequest + + +@dataclasses.dataclass +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. + """ + + arrival: float = 0.0 + first_token_at: float | None = None + token_times: list[float] = dataclasses.field(default_factory=list) + # 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 + + @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 _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, + mean_output: int, + *, + seed: int = 0, + length_spread: float = 0.6, +) -> 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 + 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) + # 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 [ + PagedRequest( + request_id=f"r{i}", + prompt_len=int(prompts[i]), + max_new_tokens=int(outputs[i]), + prompt_tokens=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[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 + 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( + PagedRequest( + request_id=f"r{i}", + prompt_len=int(tokens.size), + max_new_tokens=int(outputs[i]), + prompt_tokens=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. + + 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, 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 + + # 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 requests: + _metrics_for(metrics, request).arrival = start + driver.submit(list(requests)) + + while True: + outcome = driver.step() + if outcome is None: + break + steps += 1 + now = time.perf_counter() + + 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))) + + elapsed = time.perf_counter() - start + done = driver.completed() + index = runtime.control_plane.prefix_index + summary = _summarise( + done, metrics, occupancy, elapsed, steps, allocator, page_size, retained_pages=index.num_cached_pages + ) + 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 + # 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 _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 + 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(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, + "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. + + 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. + """ + 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 + ) + 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[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`, + `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. + + 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: + _metrics_for(metrics, request).arrival = start + + decode_state = engine.init_decode_state() + slots: dict[int, PagedRequest] = {} + 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() + 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 + + 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])) + _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, metrics, [], elapsed, steps, None, 0) + + +def _summarise( + done, + metrics: dict[str, RequestMetrics], + occupancy: Sequence[tuple[float, int, int, int]], + elapsed: float, + steps: int, + allocator: Any, + page_size: int, + retained_pages: int = 0, +) -> dict[str, Any]: + """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) + 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), + "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, + # 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([e.ttft for e in timings if e.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_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 + + +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: + 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']} ({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: + 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," + 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..ebe814448b --- /dev/null +++ b/src/maxtext/inference/kv_execution/driver.py @@ -0,0 +1,542 @@ +"""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_common import CacheNamespace +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.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() + + +@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 + # 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: + 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), + ) + + 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: + """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 + # 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; +# 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: + """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, + runtime: PagedRuntime | None = None, + ): + 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) + # 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 + # 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. + # + # 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, + poison_on_free=self.poison_on_free, + ) + self._waiting: list[PagedRequest] = [] + self._live: list[PagedRequest] = [] + self._done: list[PagedRequest] = [] + + @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: + """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 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. + + `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]: + """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.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 + + 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] + 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) + # 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 + # 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], + 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 + view = None + + 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. + 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, + ) + + shape = view.shape + 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") + + 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, + batch=tuple(batch), + query_lens=tuple(lens), + ) + + 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: + """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. + + 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, 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, 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. + + `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 + 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], freed) + self.pool.replace_layer(layer, k, v) + 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. + + 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..91f475c8f1 --- /dev/null +++ b/src/maxtext/inference/kv_execution/engine_adapter.py @@ -0,0 +1,220 @@ +"""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_all_layers +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] = {} + 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.""" + return self.control_plane.admit( + 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], + 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 + # 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 + + 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, 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 + 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. + + `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) + 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], freed) + self.pool.replace_layer(layer, k, v) + self._forget(handle) + return freed + + 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) + 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/kv_execution/layout_builder.py b/src/maxtext/inference/kv_execution/layout_builder.py new file mode 100644 index 0000000000..5ef2b74c92 --- /dev/null +++ b/src/maxtext/inference/kv_execution/layout_builder.py @@ -0,0 +1,178 @@ +"""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 + +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`. 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 +# 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.""" + 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 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`. + + 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 + + # 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( + 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." + ) + + 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. + + `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), + pool_replicas=pool_replicas(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..ffa3d81ca9 --- /dev/null +++ b/src/maxtext/inference/kv_execution/pool_factory.py @@ -0,0 +1,119 @@ +"""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, ...]: + """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) + + 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. + + 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. + 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.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_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/kv_execution/pool_ops.py b/src/maxtext/inference/kv_execution/pool_ops.py new file mode 100644 index 0000000000..0f37e7bac2 --- /dev/null +++ b/src/maxtext/inference/kv_execution/pool_ops.py @@ -0,0 +1,148 @@ +"""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) + + +@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, + 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_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/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..8337dd6d82 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,533 @@ def copy(path, partial_cache, full_cache, annotations): "token_logp": inserted_token_logp, } + 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 + 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. + + `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, token_ids) + 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.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' 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( + 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, + enable_prefix_cache=self.config.paged_enable_prefix_cache, + ) + 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, 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. + """ + # 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.""" + 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_rows: 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 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. + # + # 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) + 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_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) + 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, + prompt_token_ids=None, + namespace=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. + + 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. + """ + 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 + + 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 + + # 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. + # + # 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. + # 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 + + 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, + view=view, + inputs=inputs, + is_decode=False, + rng=rng, + algorithm=algorithm, + topk=topk, + nucleus_topp=nucleus_topp, + temperature=temperature, + ) + 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. 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 + + # 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, + view=view, + inputs=inputs, + is_decode=True, + rng=rng, + algorithm=algorithm, + topk=topk, + nucleus_topp=nucleus_topp, + temperature=temperature, + ) + return self._paged_result_tokens(sampled, selected, step=1), True def get_prefix_destination_sharding(self) -> Any: return { 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..56e07a1f33 --- /dev/null +++ b/src/maxtext/inference/paged_offline_worker.py @@ -0,0 +1,245 @@ +"""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.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( + 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 _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 + + 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/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) 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/attentions.py b/src/maxtext/layers/attentions.py index a2f4b48afd..91df6e60b6 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,82 @@ 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, + ) + # `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 + # 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], + **kernel_kwargs, + ) + return gpu_paged_attention.paged_attention_step( + query, key, value, k_pool, v_pool, plan, **kernel_kwargs + ) + def __call__( self, inputs_q: Array, @@ -1334,6 +1410,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..7858ae960e --- /dev/null +++ b/src/maxtext/layers/gpu_paged_attention.py @@ -0,0 +1,331 @@ +"""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] + + +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 () + 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, + # 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, + ), + 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] 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 e16d0de4aa..22824aa544 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2544,6 +2544,224 @@ 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 + # 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) + + 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_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.""" + 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..ae565aecd9 --- /dev/null +++ b/tests/unit/gpu_paged_attention_test.py @@ -0,0 +1,223 @@ +"""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 _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._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._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_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__": + unittest.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..057206325a --- /dev/null +++ b/tests/unit/gpu_paged_decode_parity_test.py @@ -0,0 +1,628 @@ +"""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_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_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) + 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/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/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() diff --git a/tests/unit/kv_common_test.py b/tests/unit/kv_common_test.py new file mode 100644 index 0000000000..0615c1ff05 --- /dev/null +++ b/tests/unit/kv_common_test.py @@ -0,0 +1,269 @@ +"""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 unittest + +import numpy as np + +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): + """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) + + 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_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. + + 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(), 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_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) + 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) + + +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..dd255a602e --- /dev/null +++ b/tests/unit/kv_execution_test.py @@ -0,0 +1,974 @@ +"""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 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 ( + StepShape, + 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, + 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 + + +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_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)) + 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 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.""" + + 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 = [] + self.inputs = [] + + 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) + + +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): + # `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, + prompt_tokens=np.arange(1000 * (i + 1), 1000 * (i + 1) + prompt_len, dtype=np.int64), + ) + 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, + prompt_tokens=np.arange(1000 * (i + 1), 1000 * (i + 1) + n, dtype=np.int64), + ) + 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_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) + + 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()) + 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, + prompt_tokens=np.arange(1000 * (i + 1), 1000 * (i + 1) + 30, dtype=np.int64), + ) + 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, + 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] + 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 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_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() + 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. + + 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.""" + + 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..00e792f8c5 --- /dev/null +++ b/tests/unit/kv_import_rule_test.py @@ -0,0 +1,205 @@ +"""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", + "heapq", + "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..df94d8ef06 --- /dev/null +++ b/tests/unit/kv_paged_runtime_test.py @@ -0,0 +1,406 @@ +"""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, 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( + ( + 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) + + 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) + 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, + 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, + prompt_tokens=np.arange(20, dtype=np.int64))]) + driver.run() + self.assertEqual(step.traces, after_first, "an identical second request retraced") + + +if __name__ == "__main__": + unittest.main() 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()