From c5953604a0ce3e69ea2781749445f1333930f530 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 30 Oct 2026 10:09:32 +0900 Subject: [PATCH 1/3] pd: stream RNGD prefill KV layers into Torch and ATOM decode (opt-in) (#265) Port the library and test parts of #266 onto optimus-rm: - layer_progress.LayerProgress: validated, non-overlapping per-layer token coverage of a reserved slab; finish() requires complete coverage. - PagedKV.write_layer_span: import a model-layer subset into reserved pages without publishing length/prefix/state. - SubprocessPrefillWorker(stream_layers=True): requires producer capability kv_progress=1, accepts {id, kv_span} messages, exposes future.kv_progress. - PDBatcher and AtomBatcher import published spans while the producer runs and admit only after completion and full coverage. - prefill_server: stream_result() writes validated spans into the slab before notifying; _RngdFull.prefill_stream; --engine factory loader. - serve_batched.py --pd-prefill-stream-layers (off by default). From #266 commits 42f59c3, 8e0d271, 952f491, 171e8c4 (prefill_server.py and test_prefill_factory.py only), 33d6cb0, 69180aa. The #254 chunk-copy code (_ChunkCopy, --pd-prefill-copy-overlap and its test) is stripped; streaming_factories.py, STREAMING.md evidence and scripts are not taken. Refs #265 Co-Authored-By: Claude Opus 5.5 --- opt_prime/examples/pd/serve_batched.py | 7 +- .../opt_prime/accel/pd/layer_progress.py | 63 +++++++ opt_prime/opt_prime/accel/pd/paged.py | 30 +++ opt_prime/opt_prime/accel/pd/pd_batching.py | 122 ++++++++++-- .../opt_prime/accel/pd/prefill_server.py | 125 ++++++++++++- .../accel/platforms/atom/batching.py | 89 +++++++-- opt_prime/tests/test_atom_stream_progress.py | 177 ++++++++++++++++++ opt_prime/tests/test_pd_layer_progress.py | 101 ++++++++++ opt_prime/tests/test_pd_stream_admission.py | 98 ++++++++++ opt_prime/tests/test_prefill_factory.py | 87 +++++++++ .../tests/test_prefill_source_progress.py | 140 ++++++++++++++ opt_prime/tests/test_serving_prefill_cache.py | 24 +++ .../tests/test_subprocess_layer_progress.py | 132 +++++++++++++ 13 files changed, 1156 insertions(+), 39 deletions(-) create mode 100644 opt_prime/opt_prime/accel/pd/layer_progress.py create mode 100644 opt_prime/tests/test_atom_stream_progress.py create mode 100644 opt_prime/tests/test_pd_layer_progress.py create mode 100644 opt_prime/tests/test_pd_stream_admission.py create mode 100644 opt_prime/tests/test_prefill_factory.py create mode 100644 opt_prime/tests/test_prefill_source_progress.py create mode 100644 opt_prime/tests/test_subprocess_layer_progress.py diff --git a/opt_prime/examples/pd/serve_batched.py b/opt_prime/examples/pd/serve_batched.py index 07e774c8..d2f97032 100644 --- a/opt_prime/examples/pd/serve_batched.py +++ b/opt_prime/examples/pd/serve_batched.py @@ -29,6 +29,9 @@ def prefill_worker(a, slots, config): """the PD prefill side from the flags, or None (aggregated serving)""" + stream_layers = getattr(a, "pd_prefill_stream_layers", False) + if stream_layers and (a.backend not in ("torch", "atom") or not a.pd_prefill_python or a.pd_prefill_device): + raise ValueError("PD layer streaming requires a subprocess producer and Torch or ATOM decoder") if a.pd_prefill_device: chunk = getattr(a, "pd_prefill_chunk", 2048) if chunk <= 0: @@ -47,7 +50,7 @@ def prefill_worker(a, slots, config): return SubprocessPrefillWorker(a.pd_prefill_python, shlex.split(a.pd_prefill_args) + ["--model", a.model], slots=slots, layers=c.num_hidden_layers, kv_heads=c.num_key_value_heads, head_dim=getattr(c, "head_dim", None) or c.hidden_size // c.num_attention_heads, - max_context=a.max_context, env=env) + max_context=a.max_context, env=env, stream_layers=stream_layers) return None @@ -134,6 +137,8 @@ def build(a): ap.add_argument("--pd-prefill-python", default=None, help="PD: prefill_server under this interpreter") ap.add_argument("--pd-prefill-args", default="--engine torch --device cpu") ap.add_argument("--pd-prefill-pythonpath", default=None) + ap.add_argument("--pd-prefill-stream-layers", action="store_true", + help="experimental: import layer KV during capable subprocess prefill (Torch or ATOM decoder)") ap.add_argument("--pd-remote-only", action="store_true", help="PD on ATOM: every prompt is prefilled by the PD prefill side (no cost routing to ATOM)") ap.add_argument("--pd-prefetch-next", action="store_true", diff --git a/opt_prime/opt_prime/accel/pd/layer_progress.py b/opt_prime/opt_prime/accel/pd/layer_progress.py new file mode 100644 index 00000000..a13fd219 --- /dev/null +++ b/opt_prime/opt_prime/accel/pd/layer_progress.py @@ -0,0 +1,63 @@ +"""Validated, bounded-by-coverage notifications for a reserved KV slab. + +This tracks source writes only. A consumer must finish its own transfers before +publishing the cache or releasing the reservation; source completion is not a fence. +""" +from threading import Lock + + +class LayerProgress: + def __init__(self, layers, start, end): + if any(type(x) is not int for x in (layers,start,end)) or layers<=0 or not 0<=start<=end: + raise ValueError('invalid layer-progress geometry') + self.layers,self.start,self.end=layers,start,end + self._coverage=[[] for _ in range(layers)] + self._pending=[] + self._closed=False + self._lock=Lock() + + def publish(self, layer_start, layer_end, token_start, token_end): + span=(layer_start,layer_end,token_start,token_end) + with self._lock: + if self._closed: + raise ValueError('layer progress is closed') + if (any(type(x) is not int for x in span) + or not 0<=layer_start min(self.max_context, len(self.owned[slot]) * self.page_size): + raise ValueError("partial KV span exceeds reserved pages") + layers = range(layer_start, layer_start + keys.shape[0]) + for layer in layers: + if layer not in self.layer_group: + raise ValueError("partial KV span contains an unknown attention layer") + group = self.groups[self.layer_group[layer][0]] + if keys.shape[1] != group.kv_heads or keys.shape[3] != group.head_dim: + raise ValueError("partial KV geometry does not match receiver group") + idx = self.span_index(slot, start, end) + for offset, layer in enumerate(layers): + pk, pv = self.layer_kv(layer) + pk[:, idx] = keys[offset].to(self.device, pk.dtype) + pv[:, idx] = values[offset].to(self.device, pv.dtype) + def read_span(self, slot, start, end): """[layers, kv_heads, end - start, head_dim] keys and values of `slot` (the prefill side's export); a list per KV group for hybrid models""" diff --git a/opt_prime/opt_prime/accel/pd/pd_batching.py b/opt_prime/opt_prime/accel/pd/pd_batching.py index 57edc7c6..a1e2533c 100644 --- a/opt_prime/opt_prime/accel/pd/pd_batching.py +++ b/opt_prime/opt_prime/accel/pd/pd_batching.py @@ -22,7 +22,7 @@ `platforms.atom.batching.AtomBatcher(prefill_worker=...)` decodes PD requests on ATOM (#123). """ from collections import deque -from concurrent.futures import Future, wait, FIRST_COMPLETED +from concurrent.futures import Future from dataclasses import dataclass import queue import threading @@ -35,6 +35,7 @@ from opt_prime.accel.pd.batching import ContinuousBatcher, prefill_chunk from opt_prime.accel.pd.runtime import RequestCancelled +from opt_prime.accel.pd.layer_progress import LayerProgress from opt_prime.accel.pd.sampling import select_token from opt_prime.accel.pd.torch_models import describe_model @@ -252,13 +253,39 @@ def close(self): self._proc.join(timeout=30) +class _KVProgress: + """Source coverage plus borrowed SHM views, valid for this request's reservation.""" + def __init__(self, slab, start, end): + self.coverage = LayerProgress(slab.shape[1], start, end) + kv = torch.from_numpy(slab[:, :, :, :end - start].view("int16")).view(torch.bfloat16) + self.keys, self.values = kv[0], kv[1] + self.notify = lambda: None + + def publish(self, *span): + self.coverage.publish(*span) + self.notify() + + def drain(self): + start = self.coverage.start + return [(lo, first, self.keys[lo:hi, :, first - start:last - start], + self.values[lo:hi, :, first - start:last - start]) + for lo, hi, first, last in self.coverage.drain()] + + class SubprocessPrefillWorker: """`opt_prime.accel.pd.prefill_server` run by `python` (another environment); one KV slot per decode lane in a shared-memory slab. A result's keys/values are bfloat16 views of the lane's slot: valid until that lane is - submitted again, which the decode side does only after importing them.""" + submitted again, which the decode side does only after importing them. + + Opt-in layer streaming requires ready capability ``kv_progress: 1`` and messages + ``{id, kv_span: [layer_start, layer_end, token_start, token_end]}`` after each CPU write. + Token offsets are absolute; slab storage starts at send_from. Published spans must remain + immutable until completion. Producers must process requests serially, including cancelled + requests, so an old request cannot overwrite a reused slot after its successor starts. + """ def __init__(self, python, engine_args, *, slots, layers, kv_heads, head_dim, max_context, env=None, - startup_timeout=1800): + startup_timeout=1800, stream_layers=False): import json, subprocess import numpy as np from multiprocessing import shared_memory @@ -268,6 +295,7 @@ def __init__(self, python, engine_args, *, slots, layers, kv_heads, head_dim, ma cmd = [python, "-m", "opt_prime.accel.pd.prefill_server", "--shm", self._shm.name, "--slots", str(slots), "--layers", str(layers), "--kv-heads", str(kv_heads), "--head-dim", str(head_dim), "--max-context", str(max_context), *engine_args] + self.stream_layers = stream_layers self._json = json self._futures, self._next, self._lock = {}, 0, threading.Lock() self._write_lock = threading.Lock() @@ -297,6 +325,10 @@ def submit(self, request, send_from=0, slot=None) -> Future: if slot is None or not 0 <= slot < self.slab.shape[0]: raise ValueError("SubprocessPrefillWorker needs a valid decode lane (slot)") f = Future() + if self.stream_layers and (type(send_from) is not int or + not 0 <= send_from <= len(request.input_ids) <= self.slab.shape[-2]): + f.set_exception(ValueError("invalid streamed KV suffix")) + return f s = request.sampling write_error = None # Keep pipe writes ordered without holding the response-map lock across blocking I/O. @@ -305,14 +337,19 @@ def submit(self, request, send_from=0, slot=None) -> Future: error = self._failed key, self._next = self._next, self._next + 1 if error is None: + if self.stream_layers: + f.kv_progress = _KVProgress(self.slab[slot], send_from, len(request.input_ids)) self._futures[key] = (f, slot, send_from) if error is not None: f.set_exception(error) return f try: - line = self._json.dumps(dict(id=key, tokens=list(request.input_ids), send_from=send_from, slot=slot, - sampling=dict(temperature=s.temperature, top_p=s.top_p, top_k=s.top_k, - seed=s.seed, stop_token_ids=list(s.stop_token_ids)))) + message = dict(id=key, tokens=list(request.input_ids), send_from=send_from, slot=slot, + sampling=dict(temperature=s.temperature, top_p=s.top_p, top_k=s.top_k, + seed=s.seed, stop_token_ids=list(s.stop_token_ids))) + if self.stream_layers: + message["stream_layers"] = True + line = self._json.dumps(message) self._proc.stdin.write(line + "\n") self._proc.stdin.flush() except (OSError, ValueError, TypeError) as error: @@ -342,11 +379,15 @@ def _receive(self, ready): if not isinstance(msg, dict): continue if msg.get("ready"): + if self.stream_layers and (type(msg.get("kv_progress")) is not int or msg["kv_progress"] != 1): + raise ValueError("producer does not support KV progress version 1") if not ready.done(): ready.set_result(True) continue if "id" not in msg: continue # structured SDK logs + if type(msg["id"]) is not int: + raise ValueError("invalid prefill response ID") with self._lock: entry = self._futures.get(msg["id"]) failed = self._failed @@ -355,6 +396,13 @@ def _receive(self, ready): if entry is None: raise ValueError("unknown response ID") f, slot, start = entry + progress = getattr(f, "kv_progress", None) + if "kv_span" in msg: + if (progress is None or set(msg) != {"id", "kv_span"} + or not isinstance(msg["kv_span"], list) or len(msg["kv_span"]) != 4): + raise ValueError("invalid KV progress message") + progress.publish(*msg["kv_span"]) + continue error = None if "error" in msg: error = RuntimeError(f"prefill subprocess: {msg['error']}") @@ -362,6 +410,10 @@ def _receive(self, ready): n = msg["n"] if type(n) is not int or not 0 <= n <= self.slab.shape[-2]: raise ValueError("invalid KV span length") + if progress is not None: + if n != progress.coverage.end - start: + raise ValueError("streamed KV completion length differs from requested suffix") + progress.coverage.finish() kv = torch.from_numpy(self.slab[slot, :, :, :, :n].view("int16")).view(torch.bfloat16) out = PrefillOut(msg["token"], start, kv[0], kv[1], msg.get("cached", 0), msg.get("prefill_s", 0.0)) @@ -455,7 +507,8 @@ def __init__(self, model, worker, **kw): raise ValueError("PD batching needs kv='paged'") super().__init__(model, **kw) self.worker = worker - self.stats.update(prefills=0, transferred_tokens=0, decode_cached_tokens=0, prefill_cached_tokens=0, + self._prefill_ready = threading.Event() + self.stats.update(progress_spans=0, prefills=0, transferred_tokens=0, decode_cached_tokens=0, prefill_cached_tokens=0, prefill_s=0.0, transfer_s=0.0) def step(self): @@ -473,6 +526,11 @@ def step(self): r.slot = self.free.pop(0) self.kv.length[r.slot] = r.prefilled = cached r.pending = self.worker.submit(r.request, send_from=cached, slot=r.slot) # cached blocks not sent + r.pending.add_done_callback(lambda _: self._prefill_ready.set()) + progress = getattr(r.pending, "kv_progress", None) + if progress is not None: + progress.notify = self._prefill_ready.set + self._prefill_ready.set() # include progress emitted before binding the notifier self.prefilling.append(r) self._admit_prefilled(block=not self.decoding) if self.decoding: @@ -482,28 +540,60 @@ def step(self): if events: self.kv_events.publish(events) + @torch.inference_mode() + def _import_progress(self, r): + progress = getattr(r.pending, "kv_progress", None) + if progress is None: + return + if self.kv.rstate: + raise ValueError("layer-progress import does not support recurrent state") + touched = False + try: + for layer, start, keys, values in progress.drain(): + touched = True + self.kv.write_layer_span(r.slot, start, layer, keys, values) + self.stats['progress_spans'] += 1 + finally: + # Source slab reads and destination writes must finish before cancellation can reuse the pages. + if touched and torch.device(self.device).type == "cuda": + torch.cuda.current_stream(self.device).synchronize() + def _admit_prefilled(self, block): pending = [r for r in self.prefilling if r.pending is not None] if not pending: return - if block: # nothing to decode: wait for the first prefill - wait([r.pending for r in pending], timeout=0.05, return_when=FIRST_COMPLETED) + if block: + # Clear before inspecting state: a concurrent publication then either remains visible or wakes us. + self._prefill_ready.clear() + if not any(r.pending.done() or (getattr(r.pending, "kv_progress", None) is not None + and r.pending.kv_progress.coverage.pending) for r in pending): + self._prefill_ready.wait(timeout=0.05) for r in pending: - if not r.pending.done(): - continue with self._lock: - if r not in self.prefilling: # cancelled meanwhile: lane already released + if r not in self.prefilling: continue - self.prefilling.remove(r) try: - out = r.pending.result() + if r.pending.done(): + out = r.pending.result() # failures never import remaining notifications + else: + self._import_progress(r) + continue + # Drain after observing completion, including progress published between the last drain and result. + self._import_progress(r) + progress = getattr(r.pending, "kv_progress", None) + if progress is not None and not progress.coverage.finished: + raise RuntimeError("streamed KV coverage was not validated before completion") + with torch.inference_mode(): + if progress is None: + self.kv.write_span(r.slot, out.start, out.keys, out.values) except BaseException as error: with self._lock: + self.prefilling.remove(r) self._finish(r, error) continue + with self._lock: + self.prefilling.remove(r) ids = r.request.input_ids - with torch.inference_mode(): - self.kv.write_span(r.slot, out.start, out.keys, out.values) self.kv.length[r.slot] = r.prefilled = len(ids) self.kv.commit_prompt(r.slot) if self.kv.rstate: diff --git a/opt_prime/opt_prime/accel/pd/prefill_server.py b/opt_prime/opt_prime/accel/pd/prefill_server.py index 5595f51e..0816fb68 100644 --- a/opt_prime/opt_prime/accel/pd/prefill_server.py +++ b/opt_prime/opt_prime/accel/pd/prefill_server.py @@ -100,16 +100,121 @@ def prefill(self, tokens): stage = self.engine.begin_device_prefill(list(tokens)) while not stage.advance(self.engine.ngroups): pass - logits = torch.from_numpy(self.engine._logits(stage.x, n)) + logits = self._last_logits(stage.x, n) # per layer group: keys [NG, kv_heads, head_dim, C], values [NG, kv_heads, C, head_dim] keys = torch.cat([k.cpu()[..., :n] for k in stage.Ks]).transpose(-1, -2) values = torch.cat([v.cpu()[:, :, :n] for v in stage.Vs]) return keys, values, logits + def prefill_stream(self, tokens, emit): + """Export each completed group before scheduling the next group.""" + import torch + n = len(tokens) + if not 0 < n <= self.widths[-1]: + raise ValueError("prompt exceeds compiled capacity or is empty") + stage = self.engine.begin_device_prefill(list(tokens)) + layer = 0 + while True: + done = stage.advance(1) + k, v = stage.Ks[-1], stage.Vs[-1] + keys = k.cpu()[..., :n].transpose(-1, -2) + values = v.cpu()[:, :, :n] + emit(layer, 0, keys, values) + layer += keys.shape[0] + if done: + break + return self._last_logits(stage.x, n) + + def _last_logits(self, hidden, n): + import torch + engine = self.engine + if callable(getattr(engine, "_logits", None)): + return torch.from_numpy(engine._logits(hidden, n)) + # Older rngd-pd exposes the same on-card head only through _next_token. + out = engine.lmh(hidden[n - 1], engine.lm_norm_d, engine.lm_eps_d, + engine.lm_nh_d, engine.wlm_d) + logits = out[0] if isinstance(out, (tuple, list)) else out + return logits.detach().cpu().float().reshape(-1)[:engine.V] + + +def stream_result(engine, req, slab, publish): + """Copy validated source spans into reserved SHM before notifying the receiver. + + Engines call emit(layer_start, token_start, K, V) with dense [L,H,N,D] spans + and return final logits. Calls are synchronous; source buffers may be reused + after emit returns. Absolute token spans are clipped to the requested suffix. + """ + import numpy as np + import torch + from opt_prime.accel.pd.layer_progress import LayerProgress + slot, start, end = req["slot"], req["send_from"], len(req["tokens"]) + if (type(req["id"]) is not int or type(slot) is not int + or not 0 <= slot < slab.shape[0] or type(start) is not int + or not 0 <= start <= end <= slab.shape[-2] or end == 0): + raise ValueError("invalid streamed request reservation") + coverage = LayerProgress(slab.shape[2], start, end) + + def emit(layer, first, keys, values): + if (type(layer) is not int or type(first) is not int + or keys.ndim != 4 or values.shape != keys.shape): + raise ValueError("invalid source KV span") + layers, heads, count, dim = keys.shape + if (layers <= 0 or count <= 0 or not 0 <= layer < layer + layers <= slab.shape[2] + or heads != slab.shape[3] or dim != slab.shape[-1] + or not 0 <= first < first + count <= end): + raise ValueError("invalid source KV geometry") + lo, hi = max(start, first), first + count + if lo >= hi: + return + # Validate before touching any published memory. Failed copies abort this request. + coverage.publish(layer, layer + layers, lo, hi) + arrays = [t[:, :, lo - first:hi - first].to("cpu", torch.bfloat16) + .contiguous().view(torch.int16).numpy().view(np.uint16) for t in (keys, values)] + dst = slab[slot, :, layer:layer + layers, :, lo - start:hi - start] + dst[0], dst[1] = arrays + coverage.drain() + publish(dict(id=req["id"], kv_span=[layer, layer + layers, lo, hi])) + + logits = engine.prefill_stream(list(req["tokens"]), emit) + coverage.finish() + return logits + + +def factory_engine(a): + """Load a trusted CLI factory using the same FP32 checkpoint inputs as research workers.""" + from importlib import import_module + from pathlib import Path + from safetensors import safe_open + path = Path(a.model) + config = json.loads((path / "config.json").read_text()) + width = int(a.width) + head_dim = config.get("head_dim") or config["hidden_size"] // config["num_attention_heads"] + if (not 0 < width <= a.max_context or a.slots <= 0 + or (a.layers, a.kv_heads, a.head_dim) != + (config["num_hidden_layers"], config["num_key_value_heads"], head_dim)): + raise ValueError("factory width or shared-memory geometry does not match checkpoint") + module, sep, name = (a.engine_factory or "").partition(":") + if not sep or not module or not name: + raise ValueError("engine factory must be module:callable") + factory = getattr(import_module(module), name) + if not callable(factory): + raise ValueError("engine factory must be callable") + weights = {} + for shard in sorted(path.glob("*.safetensors")): + with safe_open(shard, framework="pt") as f: + for key in f.keys(): + if key in weights: + raise ValueError("duplicate tensor in checkpoint shards") + weights[key] = f.get_tensor(key).float() + if not weights: + raise ValueError("factory requires local safetensors weights") + return factory(config, weights, width=width, chip=a.chip) + def main(): ap = argparse.ArgumentParser() - ap.add_argument("--engine", required=True, choices=["torch", "rngd-qwen", "rngd-full"]) + ap.add_argument("--engine", required=True, choices=["torch", "rngd-qwen", "rngd-full", "factory"]) + ap.add_argument("--engine-factory", help="trusted local module:callable; used only with --engine factory") ap.add_argument("--model", required=True); ap.add_argument("--model-id", default=None) ap.add_argument("--shm", required=True); ap.add_argument("--slots", type=int, required=True) ap.add_argument("--layers", type=int, required=True); ap.add_argument("--kv-heads", type=int, required=True) @@ -141,9 +246,12 @@ def main(): engine = TorchPrefillEngine(model, max_context=a.max_context, chunk=a.chunk) elif a.engine == "rngd-full": engine = _RngdFull(a) + elif a.engine == "factory": + engine = factory_engine(a) else: engine = _RngdQwen(a) - print(json.dumps({"ready": True}), flush=True) + streaming = callable(getattr(engine, "prefill_stream", None)) + print(json.dumps({"ready": True, **({"kv_progress": 1} if streaming else {})}), flush=True) for line in sys.stdin: req = json.loads(line) try: @@ -151,11 +259,20 @@ def main(): from opt_prime.accel.pd import PDRequest sampling = _sampling(req.get("sampling")) tokens, start = tuple(req["tokens"]), int(req["send_from"]) + if req.get("stream_layers"): + if req["stream_layers"] is not True or not streaming: + raise ValueError("engine does not support requested layer streaming") + logits = stream_result(engine, req, slab, + lambda msg: print(json.dumps(msg), flush=True)) + token = select_token(logits, sampling, len(tokens)) + print(json.dumps(dict(id=req["id"], token=int(token), n=len(tokens) - start, + cached=0, prefill_s=time.perf_counter() - t0)), flush=True) + continue if a.engine == "torch": out = engine.prefill(PDRequest(f"p{req['id']}", tokens, 1, sampling), start) keys, values, token, cached = out.keys, out.values, out.token, out.prefill_cached else: - keys, values, logits = engine.prefill(tokens) + keys, values, logits = engine.prefill(list(tokens)) keys, values = keys[:, :, start:], values[:, :, start:] token, cached = select_token(logits, sampling, len(tokens)), 0 n = keys.shape[2] diff --git a/opt_prime/opt_prime/accel/platforms/atom/batching.py b/opt_prime/opt_prime/accel/platforms/atom/batching.py index 7b233274..2f760192 100644 --- a/opt_prime/opt_prime/accel/platforms/atom/batching.py +++ b/opt_prime/opt_prime/accel/platforms/atom/batching.py @@ -246,19 +246,30 @@ def _reload(self, h, block, n): def import_span(self, lane, start, keys, values): """PD: write KV computed elsewhere ([layers, kv_heads, n, head_dim], any float dtype) into the lane's block at positions start..start+n; the SDK's KV tensors are (key, value) per layer in bfloat16""" + if keys.ndim != 4 or len(self.metas) != keys.shape[0] * 2: + raise InvalidKVImport("KV export geometry does not match compiled cache") + self.import_layer_span(lane, start, 0, keys, values) + + def import_layer_span(self, lane, start, layer_start, keys, values): + """Write a model-layer subset into an already reserved lane; never publish cache state.""" import numpy as np if keys.ndim != 4 or values.ndim != 4: raise InvalidKVImport("KV export must have layer/head/token/dimension axes") + if (any(type(i) is not int for i in (lane, start, layer_start)) + or not 0 <= lane < self.lanes or start < 0 + or not 0 <= layer_start < layer_start + keys.shape[0] <= len(self.metas) // 2): + raise InvalidKVImport("invalid KV layer span or lane") block, n = self.lane_block[lane], keys.shape[2] if (self.hybrid_cache or self.single_block) and start: raise InvalidKVImport("hybrid or single-block PD requires a full prefix from position zero") - if keys.shape != values.shape or len(self.metas) != keys.shape[0] * 2: + if keys.shape != values.shape: raise InvalidKVImport("KV export geometry does not match compiled cache") - stage = self._staging() + lo, hi = 2 * layer_start, 2 * (layer_start + keys.shape[0]) + metas, stage = self.metas[lo:hi], self._staging()[lo:hi] if len({buf.dtype for buf in stage}) != 1: raise ValueError("mixed physical KV formats are unsupported") destinations = [] - for m, buf in zip(self.metas, stage): + for m, buf in zip(metas, stage): if (keys.shape[1], keys.shape[3]) != (buf.shape[1], buf.shape[3]): raise InvalidKVImport("KV export heads/dimensions do not match compiled cache") local = self.hybrid_cache and m.get("layer_type") == "sliding_attention" @@ -274,7 +285,7 @@ def import_span(self, lane, start, keys, values): data = bf16_to_dlfloat(kv.to(torch.bfloat16).contiguous().view(torch.int16).numpy().view(np.uint16)) else: data = kv.float().numpy() - for i, (m, buf) in enumerate(zip(self.metas, stage)): + for i, (m, buf) in enumerate(zip(metas, stage)): dest, offset, kept = destinations[i] buf[dest, :, offset:offset + kept] = data[i // 2, i % 2, :, n - kept:] self.io_rt.set_kv_cache_for_layer(m["name"], dest, offset, kept, buf) @@ -327,7 +338,8 @@ def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib= self.waiting, self.prefilling, self.decoding = deque(), deque(), [] self.by_id = {} self._lock = threading.Lock() - self.stats = dict(iterations=0, prefill_chunks=0, decode_steps=0, decode_rows=0) + self._prefill_ready = threading.Event() + self.stats = dict(iterations=0, prefill_chunks=0, decode_steps=0, decode_rows=0, progress_spans=0) self.poisoned = False self.kv_events = kv_events self.worker = prefill_worker @@ -412,12 +424,12 @@ def step(self): def _prefill_request(self, r): ids = r.request.input_ids start = self.blocks.admit(r.slot, ids) if self.blocks is not None else 0 - r.admitted = True out = self._prefill(r.slot, ids, start) self.stats['prefill_chunks'] += 1 logits = out.logits[0, -1].float() if not torch.isfinite(logits).all(): raise RuntimeError("non-finite ATOM logits") + r.admitted = True padded = int(getattr(out, "padded_cache_lengths", 0) or 0) self.kv.length[r.slot] = r.prefilled = len(ids) + padded self.prefilling.remove(r) # PD: requests waiting on remote prefills may be ahead of it @@ -450,8 +462,36 @@ def _finish(self, r, error): super()._finish(r, error) # ---- PD: prefill elsewhere, decode here ---- + def _watch_prefill(self, r): + r.pending.add_done_callback(lambda _: self._prefill_ready.set()) + progress = getattr(r.pending, "kv_progress", None) + if progress is not None: + progress.notify = self._prefill_ready.set + self._prefill_ready.set() + + def _import_progress(self, r, *, final=False): + progress = getattr(r.pending, "kv_progress", None) + if progress is None: + return + if r.slot < 0 or r not in self.prefilling: + return # prefetch has no destination until the current decode releases it + if self.hybrid_cache and self.lanes > 1 and not final: + return # idle hybrid lanes still own local KV; defer until joining decode + if (progress.coverage.layers * 2 != len(self.blocks.metas) + or progress.coverage.end != len(r.request.input_ids)): + raise InvalidKVImport("streamed coverage does not match the request cache") + for layer, start, keys, values in progress.drain(): + if (self.hybrid_cache or self.blocks.single_block) and (start != 0 or keys.shape[2] != len(r.request.input_ids)): + raise InvalidKVImport("hybrid or single-block layer streaming requires full-prefix layer spans") + self.blocks.import_layer_span(r.slot, start, layer, keys, values) + self.stats['progress_spans'] += 1 + + def _reject_prefill(self, r, error): + with self._lock: + self.prefilling.remove(r) + self._finish(r, error) + def _pd_step(self): - from concurrent.futures import FIRST_COMPLETED, wait with self._lock: if self._prefetched is not None and self._prefetched.future.done() and self._prefetched.pending.done(): self._prefetched = None # cancelled writer has now relinquished the source slot @@ -484,10 +524,10 @@ def _pd_step(self): try: for r in admitted: # decode-side prefix (device copy / host reload) is not sent start = self.blocks.admit(r.slot, r.request.input_ids) - r.admitted = True if r.pending is None: r.pending = self.worker.submit(r.request, send_from=start, slot=r.slot) self.stats_pd["remote"] += 1 + self._watch_prefill(r) if local is not None: self.stats_pd["local"] += 1 t0 = time.perf_counter() @@ -497,32 +537,43 @@ def _pd_step(self): if self._prefetched is not None and self._prefetched.future.done(): pending.append(self._prefetched) # wait for a cancelled writer without busy-spinning if pending and not self.decoding: - wait([r.pending for r in pending], timeout=0.05, return_when=FIRST_COMPLETED) + self._prefill_ready.clear() + if not any(r.pending.done() or (r.slot >= 0 and getattr(r.pending, "kv_progress", None) is not None + and r.pending.kv_progress.coverage.pending) for r in pending): + self._prefill_ready.wait(timeout=0.05) for r in pending: - if not r.pending.done(): - continue with self._lock: if r not in self.prefilling: continue - self.prefilling.remove(r) + if not r.pending.done(): + try: + self._import_progress(r) + except InvalidKVImport as error: + self._reject_prefill(r, error) + continue try: out = r.pending.result() if (self.hybrid_cache or self.blocks.single_block) and (out.start != 0 or out.keys.shape[2] != len(r.request.input_ids)): raise ValueError("hybrid or single-block PD requires the full prompt KV") except BaseException as error: # this request fails; the lanes go on - with self._lock: - self._finish(r, error) + self._reject_prefill(r, error) continue - r.pending = None ids = r.request.input_ids t0 = time.perf_counter() try: - if out.keys.shape[2]: + progress = getattr(r.pending, "kv_progress", None) + if progress is not None: + if not progress.coverage.finished: + raise InvalidKVImport("streamed coverage was not validated before completion") + self._import_progress(r, final=True) + elif out.keys.shape[2]: self.blocks.import_span(r.slot, out.start, out.keys, out.values) except InvalidKVImport as error: - with self._lock: - self._finish(r, error) + self._reject_prefill(r, error) continue + with self._lock: + self.prefilling.remove(r) + r.pending = None if self._prefetched is r: self._prefetched = None # only release after the owner has imported the KV if self.hybrid_cache: @@ -532,6 +583,7 @@ def _pd_step(self): self._remote_s = _ema(self._remote_s, (out.prefill_s or 0.0) + (out.transfer_s or 0.0) + time.perf_counter() - t0) self.kv.length[r.slot] = r.prefilled = len(ids) + r.admitted = True self.stats['prefill_chunks'] += 1 if self._emit(r, out.token): self.decoding.append(r) @@ -559,4 +611,5 @@ def _prefetch_next(self): except BaseException as error: r.pending = Future() r.pending.set_exception(error) + self._watch_prefill(r) self.stats_pd["remote"] += 1 diff --git a/opt_prime/tests/test_atom_stream_progress.py b/opt_prime/tests/test_atom_stream_progress.py new file mode 100644 index 00000000..5f0ed13c --- /dev/null +++ b/opt_prime/tests/test_atom_stream_progress.py @@ -0,0 +1,177 @@ +"""Partial ATOM cache state stays private and never overwrites an active lane.""" +import copy +from concurrent.futures import Future +from dataclasses import replace +import numpy as np +import pytest +import torch +from opt_prime.accel.pd import PDRequest +from opt_prime.accel.pd.pd_batching import TorchPrefillEngine,_KVProgress +from opt_prime.accel.platforms.atom.batching import AtomBatcher,InvalidKVImport +from test_atom_batching import FakeBlocks,_llama,_prompts,fake_rebel + + +class Producer: + def __init__(self,model): + self.engine=TorchPrefillEngine(copy.deepcopy(model),max_context=128,chunk=16,page_size=16) + self.jobs=[] + def submit(self,request,send_from=0,slot=None): + out=self.engine.prefill(request,send_from) + k,v=[(t if torch.is_tensor(t) else torch.cat(t)).to(torch.bfloat16) for t in (out.keys,out.values)] + slab=torch.stack((k,v)).view(torch.int16).numpy().view(np.uint16) + progress=_KVProgress(slab,send_from,len(request.input_ids));f=Future();f.kv_progress=progress + progress.publish(0,1,send_from,len(request.input_ids)) + self.jobs.append((f,progress,replace(out,keys=k,values=v))) + return f + def finish(self,index): + f,p,out=self.jobs[index] + p.publish(1,out.keys.shape[0],p.coverage.start,p.coverage.end);p.coverage.finish();f.set_result(out) + + +def setup(holder,*,cache=False,prefetch=False): + model=_llama();fake=FakeBlocks(model,batch=1,blocks=4 if cache else 1) + holder['model']=fake;worker=Producer(model) + return fake,worker,AtomBatcher(fake,prefix_cache=cache,prefill_worker=worker,local_prefill=False,remote_prefetch=prefetch) + + +def test_partial_import_before_completion_does_not_admit_decode(fake_rebel): + fake,w,b=setup(fake_rebel);prompt=_prompts(1,20,35)[0] + f=b.submit(PDRequest('a',prompt,3));b.step() + torch.testing.assert_close(fake.kv[0][0,:,:len(prompt)],w.jobs[0][2].keys[0].float()) + assert not fake.kv[2].any() and b.kv.length[0]==0 and not b.decoding and not f.done() + w.finish(0);b.run_until_idle();assert len(f.result())==3 and b.idle + + +def test_prefetch_progress_never_overwrites_active_decode(fake_rebel): + fake,w,b=setup(fake_rebel,prefetch=True);p,q=_prompts(2,20,35) + first=b.submit(PDRequest('first',p,12));second=b.submit(PDRequest('second',q,3)) + b.step();w.finish(0);b.step() + assert len(w.jobs)==2 and b._prefetched.slot==-1 and len(b.decoding)==1 + block=b.blocks.lane_block[0];snapshot=[k[block,:,:len(p)].clone() for k in fake.kv] + b.step() + for k,old in zip(fake.kv,snapshot):torch.testing.assert_close(k[block,:,:len(p)],old) + w.finish(1);b.run_until_idle();assert len(first.result())==12 and len(second.result())==3 + assert b.idle + + +def test_failed_prefill_does_not_publish_prefix(fake_rebel): + fake,w,b=setup(fake_rebel,cache=True);p=_prompts(1,32,33)[0] + f=b.submit(PDRequest('bad',p,3));b.step() + w.jobs[0][0].set_exception(RuntimeError('producer failed')) + b.run_until_idle() + with pytest.raises(RuntimeError):f.result() + assert not b.blocks.index and not b.blocks.entries, 'incomplete prompt published as reusable prefix' + assert not b.blocks.drain_events() + + +def test_layer_span_rejects_bad_geometry_before_any_write(fake_rebel): + fake,w,b=setup(fake_rebel);b.blocks.admit(0,(1,2)) + k=torch.ones((1,99,2,4)) + with pytest.raises(InvalidKVImport):b.blocks.import_layer_span(0,0,0,k,k) + assert not any(t.any() for t in fake.kv) + + +def test_layer_span_only_touches_selected_layers(fake_rebel): + fake,w,b=setup(fake_rebel);b.blocks.admit(0,(1,2)) + c=fake.config;k=torch.ones((1,c.num_key_value_heads,2,c.head_dim)) + b.blocks.import_layer_span(0,0,1,k,k*2) + assert not fake.kv[0].any() and not fake.kv[1].any() + assert torch.all(fake.kv[2][0,:,:2]==1) and torch.all(fake.kv[3][0,:,:2]==2) + + +def test_streamed_decode_matches_whole_import_of_identical_kv(fake_rebel): + from types import SimpleNamespace + fake,w,b=setup(fake_rebel);prompt=_prompts(1,20,35)[0] + f=b.submit(PDRequest('stream',prompt,8));b.step();w.finish(0);b.run_until_idle() + control=FakeBlocks(copy.deepcopy(fake.m),batch=1,blocks=1);fake_rebel['model']=control + def submit(*args,**kw): + result=Future();result.set_result(w.jobs[0][2]);return result + baseline=AtomBatcher(control,prefix_cache=False,prefill_worker=SimpleNamespace(submit=submit),local_prefill=False) + g=baseline.submit(PDRequest('whole',prompt,8));baseline.run_until_idle() + assert f.result()==g.result() + + +def test_cancelled_partial_import_never_publishes_or_overwrites_reused_lane(fake_rebel): + from opt_prime.accel.pd.runtime import RequestCancelled + fake,w,b=setup(fake_rebel,cache=True);p,q=_prompts(2,32,33) + old=b.submit(PDRequest('old',p,3));b.step();assert b.cancel('old');b.step() + with pytest.raises(RequestCancelled):old.result() + assert not b.blocks.index + new=b.submit(PDRequest('new',q,3));b.step() + snapshot=[t.clone() for t in fake.kv] + w.finish(0);b.step() + for t,s in zip(fake.kv,snapshot):torch.testing.assert_close(t,s) + w.finish(1);b.run_until_idle();assert len(new.result())==3 and b.idle + + +def test_partial_sdk_failure_poisons_owner(fake_rebel): + fake,w,b=setup(fake_rebel) + def fail(*args):raise RuntimeError('SDK failed during partial write') + b.blocks.io_rt.set_kv_cache_for_layer=fail + b.submit(PDRequest('fail',_prompts(1,20,35)[0],3)) + with pytest.raises(RuntimeError,match='SDK failed'):b.step() + assert b.poisoned and not b.blocks.index + + +def test_hybrid_layer_import_uses_lane_for_sliding_and_block_for_global(fake_rebel): + from types import SimpleNamespace + fake=FakeBlocks(_llama(),batch=2,blocks=4);fake.rbln_config.cache_impl='hybrid' + for i,m in enumerate(fake.rbln_config.kvcache_metas): + m['layer_type']='sliding_attention' if i<2 else 'full_attention' + m['shape']=([2,m['shape'][1],4,m['shape'][3]] if i<2 else m['shape']) + fake_rebel['model']=fake + b=AtomBatcher(fake,prefix_cache=False,prefill_worker=SimpleNamespace(),local_prefill=False) + # First free global block differs from lane1. + b.blocks.admit(1,tuple(range(10))) + c=fake.config;k=torch.arange(10).reshape(1,1,10,1).float().expand(1,c.num_key_value_heads,10,c.head_dim) + writes=[] + b.blocks.io_rt.get_kv_cache_size_for_layer=lambda name: int(np.prod(next(m['shape'] for m in b.blocks.metas if m['name']==name)))*4 + b.blocks.io_rt.set_kv_cache_for_layer=lambda name,block,offset,n,buf:writes.append((name,block,offset,n,buf[block,:,offset:offset+n].copy())) + b.blocks.import_layer_span(1,0,0,k,k) + b.blocks.import_layer_span(1,0,1,k,k) + assert [row[1:4] for row in writes]==[(1,0,4),(1,0,4),(0,0,10),(0,0,10)] + assert writes[0][4][0,:,0].tolist()==[6.,7.,8.,9.] + assert writes[2][4][0,:,0].tolist()==list(map(float,range(10))) + + +def test_multilane_hybrid_defers_partial_cache_until_decode_admission(fake_rebel): + from test_atom_hybrid import HybridFake + fake=HybridFake();fake_rebel['model']=fake + for meta in fake.rbln_config.kvcache_metas[:2]:meta['layer_type']='sliding_attention' + w=Producer(fake.m);b=AtomBatcher(fake,prefill_worker=w,local_prefill=False) + f=b.submit(PDRequest('hybrid',_prompts(1,20,35)[0],3));b.step() + assert b.stats['progress_spans']==0, 'idle hybrid decode lanes can overwrite partial local KV' + w.finish(0);b.run_until_idle();assert len(f.result())==3 and b.stats['progress_spans']==2 + + +@pytest.mark.parametrize('outcome', ['error', 'cancel']) +def test_prefetched_partial_failure_preserves_active_decode_and_source_ownership(fake_rebel, outcome): + from opt_prime.accel.pd.runtime import RequestCancelled + fake, w, b = setup(fake_rebel, prefetch=True) + p, q = _prompts(2, 20, 35) + first = b.submit(PDRequest('first', p, 12)) + second = b.submit(PDRequest('second', q, 3)) + b.step(); w.finish(0); b.step() + assert b.decoding and b._prefetched.slot == -1 + pending, progress, _ = w.jobs[1] + assert progress.coverage.pending and not pending.done() + imported = b.stats['progress_spans'] + if outcome == 'error': + pending.set_exception(RuntimeError('partial source failed')) + else: + assert b.cancel('second') + while not first.done(): + b.step() + assert b.stats['progress_spans'] == imported + if outcome == 'cancel': + assert not pending.done() and not b.idle + w.finish(1) + b.run_until_idle() + with pytest.raises(RuntimeError if outcome == 'error' else RequestCancelled): + second.result() + assert b.idle and not b.poisoned and b.free == [0] + # Compare with the same BF16 source/cache path without the failed prefetch. + _, control_worker, control = setup(fake_rebel) + expected = control.submit(PDRequest('reference', p, 12)) + control.step(); control_worker.finish(0); control.run_until_idle() + assert first.result() == expected.result() diff --git a/opt_prime/tests/test_pd_layer_progress.py b/opt_prime/tests/test_pd_layer_progress.py new file mode 100644 index 00000000..99f5277b --- /dev/null +++ b/opt_prime/tests/test_pd_layer_progress.py @@ -0,0 +1,101 @@ +"""Partial KV must cover every layer/position without overlaps before publication.""" +import importlib.util +from pathlib import Path +import pytest + +path=Path(__file__).parents[1]/'opt_prime/accel/pd/layer_progress.py' + +def tracker(*args): + assert path.exists(), 'layer-progress validator is not implemented' + spec=importlib.util.spec_from_file_location('layer_progress',path) + m=importlib.util.module_from_spec(spec);spec.loader.exec_module(m) + return m.LayerProgress(*args) + + +def test_partial_prefix_and_out_of_order_layers_require_complete_coverage(): + p=tracker(3,5,9) + p.publish(1,3,5,9) + p.publish(0,1,7,9) + assert p.drain()==[(1,3,5,9),(0,1,7,9)] + with pytest.raises(ValueError,match='incomplete'):p.finish() + p.publish(0,1,5,7) + p.finish() + assert p.drain()==[(0,1,5,7)] + assert p.drain()==[] + with pytest.raises(ValueError,match='closed'):p.publish(0,1,5,7) + + +def test_overlap_rejection_is_atomic_across_layers(): + p=tracker(3,0,4) + p.publish(1,2,0,2) + with pytest.raises(ValueError,match='overlap'):p.publish(0,3,1,3) + # If the rejected span mutated layer0, this valid replacement would fail. + p.publish(0,1,0,4) + p.publish(1,2,2,4) + p.publish(2,3,0,4) + p.finish() + assert len(p.drain())==4 + + +@pytest.mark.parametrize('span',[(-1,1,0,4),(0,4,0,4),(0,1,-1,4), + (0,1,0,5),(0,0,0,4),(0,1,2,2),(True,2,0,4),(0,1,0.0,4)]) +def test_invalid_progress_never_enters_queue(span): + p=tracker(3,0,4) + with pytest.raises(ValueError):p.publish(*span) + assert p.drain()==[] + + +def test_empty_uncached_suffix_needs_no_transfer(): + p=tracker(3,8,8) + p.finish() + assert p.drain()==[] + + +def test_partial_hybrid_write_preserves_other_layers_and_positions(): + import torch + from test_pd_dense_hybrid import pool + cache=pool() + assert hasattr(cache,'write_layer_span'), 'partial receiver import is not implemented' + keys=torch.arange(20).reshape(2,1,5,2).float()+1 + cache.write_layer_span(0,3,1,keys,keys+100) + idx=cache.span_index(0,3,8) + for layer in [1,2]: + k,v=cache.layer_kv(layer) + torch.testing.assert_close(k[:,idx],keys[layer-1]) + torch.testing.assert_close(v[:,idx],keys[layer-1]+100) + assert not cache.layer_kv(0)[0].any() + assert cache.length[0]==0 and not cache.cached + + +@pytest.mark.parametrize('slot,start,layer',[(0,10,0),(0,-1,0),(1,0,0),(0,0,2)]) +def test_partial_write_rejects_unreserved_positions_and_unknown_layers(slot,start,layer): + import torch + from test_pd_dense_hybrid import pool + cache=pool() + assert hasattr(cache,'write_layer_span'), 'partial receiver import is not implemented' + keys=torch.ones(2,1,5,2) + with pytest.raises(ValueError):cache.write_layer_span(slot,start,layer,keys,keys) + assert all(not t.any() for t in cache.gk+cache.gv) + + +def test_partial_write_rejects_later_group_geometry_before_mutation(): + import torch + from opt_prime.accel.pd.paged import PagedKV + from opt_prime.accel.pd.torch_models import KVGroup + cache=PagedKV(2,1,1,2,torch.float32,'cpu',pages=4,page_size=4,max_context=16, + groups=(KVGroup('a',None,1,2,(0,)),KVGroup('b',None,2,2,(1,)))) + cache.alloc(0,12) + assert hasattr(cache,'write_layer_span'), 'partial receiver import is not implemented' + keys=torch.ones(2,1,5,2) + with pytest.raises(ValueError):cache.write_layer_span(0,0,0,keys,keys) + assert all(not t.any() for t in cache.gk+cache.gv) + + +def test_partial_write_cannot_overwrite_an_admitted_prefix(): + import torch + from test_pd_dense_hybrid import pool + cache=pool() + cache.length[0]=4 + keys=torch.ones(1,1,5,2) + with pytest.raises(ValueError,match='prefix'):cache.write_layer_span(0,3,0,keys,keys) + assert all(not t.any() for t in cache.gk+cache.gv) diff --git a/opt_prime/tests/test_pd_stream_admission.py b/opt_prime/tests/test_pd_stream_admission.py new file mode 100644 index 00000000..73ebf1eb --- /dev/null +++ b/opt_prime/tests/test_pd_stream_admission.py @@ -0,0 +1,98 @@ +"""Real model/cache, controlled producer completion: never publish partial state.""" +import copy +from concurrent.futures import Future +from dataclasses import replace +import numpy as np +import pytest +import torch +from opt_prime.accel.pd import PDRequest +from opt_prime.accel.pd.pd_batching import PDBatcher, PrefillWorker, TorchPrefillEngine, _KVProgress +from opt_prime.accel.pd.runtime import RequestCancelled +from test_pd_hybrid import _gemma3,_prompts + + +class ControlledProducer: + def __init__(self,model): + self.engine=TorchPrefillEngine(copy.deepcopy(model),max_context=128,chunk=16,page_size=8) + self.jobs=[] + + def submit(self,request,send_from=0,slot=None): + out=self.engine.prefill(request,send_from) + def dense(groups): + mapping={layer:t[i] for g,t in zip(self.engine.kv.groups,groups) for i,layer in enumerate(g.layers)} + return torch.stack([mapping[i] for i in range(len(mapping))]) + keys,values=dense(out.keys),dense(out.values) + slab=torch.stack((keys,values)).view(torch.int16).numpy().view(np.uint16) + progress=_KVProgress(slab,send_from,len(request.input_ids)) + f=Future();f.kv_progress=progress + progress.coverage.publish(0,1,send_from,len(request.input_ids)) + self.jobs.append((f,progress,replace(out,keys=keys,values=values))) + return f + + def finish(self,index=0,complete=True): + f,progress,out=self.jobs[index] + progress.coverage.publish(1,out.keys.shape[0],progress.coverage.start,progress.coverage.end) + if complete:progress.coverage.finish() + f.set_result(out) + + +def setup(device="cpu"): + # CUDA flex-attention dot products require a KV block dimension of at least16. + if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") + model=_gemma3().to(torch.bfloat16) + worker=ControlledProducer(model) + batch=PDBatcher(copy.deepcopy(model).to(device),worker,slots=2,max_context=128,page_size=(16 if device == "cuda" else 8),prefix_cache=False) + return model,worker,batch + + +def baseline(model,prompt,device="cpu"): + worker=PrefillWorker(TorchPrefillEngine(copy.deepcopy(model),max_context=128,chunk=16,page_size=8)) + try: + b=PDBatcher(copy.deepcopy(model).to(device),worker,slots=1,max_context=128,page_size=(16 if device == "cuda" else 8),prefix_cache=False) + f=b.submit(PDRequest('whole',prompt,8));b.run_until_idle();return f.result() + finally:worker.close() + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_hybrid_partial_import_happens_before_completion_without_admission(device): + model,worker,b=setup(device);prompt=_prompts(1,25,26,265)[0] + result=b.submit(PDRequest('stream',prompt,8));b.step() + request=b.by_id['stream'];idx=b.kv.span_index(request.slot,0,len(prompt)) + keys,_=b.kv.layer_kv(0) + torch.testing.assert_close(keys[:,idx].cpu(),worker.jobs[0][2].keys[0]) + assert not b.kv.layer_kv(1)[0].any() + assert b.kv.length[request.slot]==0 and not b.decoding and not result.done() + worker.finish();b.run_until_idle() + assert result.result()==baseline(model,prompt,device) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_cancellation_and_late_old_completion_do_not_corrupt_reused_cache(device): + model,worker,b=setup(device);p,q=_prompts(2,25,26,267) + old=b.submit(PDRequest('old',p,8));b.step();assert b.cancel('old');b.step() + with pytest.raises(RequestCancelled):old.result() + new=b.submit(PDRequest('new',q,8));b.step() + worker.finish(0);worker.finish(1);b.run_until_idle() + assert new.result()==baseline(model,q,device) + assert b.idle and not b.by_id and sorted(b.free)==[0,1] + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_partial_producer_failure_preserves_other_active_decode(device): + model,worker,b=setup(device);p,q=_prompts(2,25,26,269) + kept=b.submit(PDRequest('kept',p,8));b.step();worker.finish(0);b.step() + bad=b.submit(PDRequest('bad',q,8));b.step() + worker.jobs[1][0].set_exception(RuntimeError('producer failed')) + b.run_until_idle() + with pytest.raises(RuntimeError,match='producer failed'):bad.result() + assert kept.result()==baseline(model,p,device) + assert b.idle and not b.by_id + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_unvalidated_stream_completion_cannot_publish_cache(device): + _,worker,b=setup(device);prompt=_prompts(1,25,26,271)[0] + result=b.submit(PDRequest('incomplete',prompt,8));b.step();worker.finish(complete=False) + b.run_until_idle() + with pytest.raises(RuntimeError,match='coverage'):result.result() + assert b.idle and not b.by_id and not b.kv.cached diff --git a/opt_prime/tests/test_prefill_factory.py b/opt_prime/tests/test_prefill_factory.py new file mode 100644 index 00000000..ba6af700 --- /dev/null +++ b/opt_prime/tests/test_prefill_factory.py @@ -0,0 +1,87 @@ +"""Explicit local factories preserve checkpoint inputs and validate slab geometry.""" +import json,sys +from types import SimpleNamespace,ModuleType +import pytest +import torch +from safetensors.torch import save_file +from opt_prime.accel.pd import prefill_server as server + + +def setup(tmp_path,monkeypatch): + config=dict(num_hidden_layers=2,num_key_value_heads=1,num_attention_heads=2,hidden_size=8) + (tmp_path/'config.json').write_text(json.dumps(config)) + save_file({'w':torch.ones(2,dtype=torch.bfloat16)},tmp_path/'model.safetensors') + module=ModuleType('test_local_prefill') + calls=[] + def factory(c,w,**kw): + calls.append((c,w,kw));return SimpleNamespace(prefill=lambda tokens:None) + module.make=factory;monkeypatch.setitem(sys.modules,module.__name__,module) + args=SimpleNamespace(model=str(tmp_path),engine_factory='test_local_prefill:make',width='32',chip=0, + slots=1,layers=2,kv_heads=1,head_dim=4,max_context=64) + return args,calls + + +def test_factory_receives_same_fp32_weights_and_config(tmp_path,monkeypatch): + args,calls=setup(tmp_path,monkeypatch) + engine=server.factory_engine(args) + assert callable(engine.prefill) + assert len(calls)==1 + c,w,kw=calls[0] + assert c['hidden_size']==8 and w['w'].dtype==torch.float32 + assert w['w'].tolist()==[1.,1.] and kw==dict(width=32,chip=0) + + +@pytest.mark.parametrize('field,value',[('head_dim',8),('layers',3),('width','0'),('width','128'),('engine_factory','bad')]) +def test_factory_rejects_bad_setup_before_construction(tmp_path,monkeypatch,field,value): + args,calls=setup(tmp_path,monkeypatch);setattr(args,field,value) + with pytest.raises(ValueError):server.factory_engine(args) + assert not calls + + +@pytest.mark.parametrize('stream',[False,True]) +def test_production_server_loads_factory_over_real_pipes(children,stream): + import time + from opt_prime.accel.pd import PDRequest + make,_,_,gate=children + root=gate.parent + (root/'config.json').write_text(json.dumps(dict(num_hidden_layers=1,num_key_value_heads=1,head_dim=2))) + save_file({'w':torch.ones(1)},root/'model.safetensors') + (root/'child').write_text(f'#!{sys.executable}\n'+''' +import os,sys,time +from pathlib import Path +from types import ModuleType +sys.path.insert(0,os.getcwd()) +import torch +from opt_prime.accel.pd import prefill_server as server +root=Path(os.environ['PREFILL_TEST_GATE']).parent +module=ModuleType('local_fixture') +class Engine: + def __init__(self,config,weights,**kw): + assert weights['w'].tolist()==[1.] and kw['width']==8 + def prefill(self,tokens): + assert isinstance(tokens,list), "factory expects token list" + k=torch.tensor([1.,2.],dtype=torch.bfloat16).reshape(1,1,2,1).expand(1,1,2,2) + return k,k,torch.tensor([0.,1.]) + def prefill_stream(self,tokens,emit): + k,v,logits=self.prefill(tokens) + emit(0,0,k[:,:,:1],v[:,:,:1]) + while not (root/'gate').exists():time.sleep(.005) + emit(0,1,k[:,:,1:],v[:,:,1:]) + return logits +module.make=Engine;sys.modules[module.__name__]=module +sys.argv=[sys.argv[0]]+sys.argv[3:]+['--engine','factory','--engine-factory','local_fixture:make','--model',str(root),'--width','8'] +server.main() +''') + worker=make(stream_layers=stream) + future=worker.submit(PDRequest('factory',(1,2),1),slot=0) + if stream: + deadline=time.monotonic()+5;rows=[] + while not rows and time.monotonic()1: + slab[slot,:,:,:,1:n]=0x4000 + print(json.dumps({'id':msg['id'],'kv_span':[0,1,start+1,start+n]}),flush=True) + print(json.dumps({'id':msg['id'],'n':1 if mode=='short' else n,'token':17}),flush=True) +''' + + +def make_stream(children, mode='ok'): + make,_,_,gate=children + path=gate.parent/'child' + path.write_text(f'#!{sys.executable}\n'+CHILD) + return make(mode,stream_layers=True),gate + + +def first_span(future): + deadline=time.monotonic()+3 + while time.monotonic() Date: Fri, 30 Oct 2026 10:11:21 +0900 Subject: [PATCH 2/3] pd: reject layer streaming for recurrent models at setup (#265) Layer streaming carries attention KV only, not linear-attention state, so PDBatcher imported nothing until the first span and then raised. Reject the combination when it is configured instead: PDBatcher with a stream_layers worker over a recurrent cache, SubprocessPrefillWorker(stream_layers=True, recurrent_state=True) before it allocates the slab or spawns, and serve_batched.py --pd-prefill-stream-layers for a recurrent model before the producer starts. The import-time check stays as a guard for other producers. Refs #265 Co-Authored-By: Claude Opus 5.5 --- opt_prime/examples/pd/serve_batched.py | 4 +++ opt_prime/opt_prime/accel/pd/pd_batching.py | 10 ++++-- opt_prime/tests/test_pd_stream_recurrent.py | 39 +++++++++++++++++++++ 3 files changed, 51 insertions(+), 2 deletions(-) create mode 100644 opt_prime/tests/test_pd_stream_recurrent.py diff --git a/opt_prime/examples/pd/serve_batched.py b/opt_prime/examples/pd/serve_batched.py index d2f97032..28422468 100644 --- a/opt_prime/examples/pd/serve_batched.py +++ b/opt_prime/examples/pd/serve_batched.py @@ -32,6 +32,10 @@ def prefill_worker(a, slots, config): stream_layers = getattr(a, "pd_prefill_stream_layers", False) if stream_layers and (a.backend not in ("torch", "atom") or not a.pd_prefill_python or a.pd_prefill_device): raise ValueError("PD layer streaming requires a subprocess producer and Torch or ATOM decoder") + if stream_layers: + from opt_prime.accel.pd.torch_models import describe_model + if describe_model(config).recurrent: + raise ValueError("PD layer streaming does not support recurrent (linear-attention) models") if a.pd_prefill_device: chunk = getattr(a, "pd_prefill_chunk", 2048) if chunk <= 0: diff --git a/opt_prime/opt_prime/accel/pd/pd_batching.py b/opt_prime/opt_prime/accel/pd/pd_batching.py index a1e2533c..561e9e02 100644 --- a/opt_prime/opt_prime/accel/pd/pd_batching.py +++ b/opt_prime/opt_prime/accel/pd/pd_batching.py @@ -285,10 +285,14 @@ class SubprocessPrefillWorker: """ def __init__(self, python, engine_args, *, slots, layers, kv_heads, head_dim, max_context, env=None, - startup_timeout=1800, stream_layers=False): + startup_timeout=1800, stream_layers=False, recurrent_state=False): + """recurrent_state: the decode model has linear-attention (recurrent) layers; their state is not part of + the layer stream, so stream_layers is rejected for it here, before any process or slab exists.""" import json, subprocess import numpy as np from multiprocessing import shared_memory + if stream_layers and recurrent_state: + raise ValueError("layer streaming does not support recurrent (linear-attention) state") shape = (slots, 2, layers, kv_heads, max_context, head_dim) self._shm = shared_memory.SharedMemory(create=True, size=int(np.prod(shape)) * 2) self.slab = np.ndarray(shape, dtype=np.uint16, buffer=self._shm.buf) @@ -506,6 +510,8 @@ def __init__(self, model, worker, **kw): if kw["kv"] != "paged": raise ValueError("PD batching needs kv='paged'") super().__init__(model, **kw) + if getattr(worker, "stream_layers", False) and self.kv.rstate: + raise ValueError("layer streaming does not support recurrent (linear-attention) state") self.worker = worker self._prefill_ready = threading.Event() self.stats.update(progress_spans=0, prefills=0, transferred_tokens=0, decode_cached_tokens=0, prefill_cached_tokens=0, @@ -545,7 +551,7 @@ def _import_progress(self, r): progress = getattr(r.pending, "kv_progress", None) if progress is None: return - if self.kv.rstate: + if self.kv.rstate: # setup rejects stream_layers workers; guards other producers raise ValueError("layer-progress import does not support recurrent state") touched = False try: diff --git a/opt_prime/tests/test_pd_stream_recurrent.py b/opt_prime/tests/test_pd_stream_recurrent.py new file mode 100644 index 00000000..34fc403d --- /dev/null +++ b/opt_prime/tests/test_pd_stream_recurrent.py @@ -0,0 +1,39 @@ +"""Layer streaming cannot carry recurrent (linear-attention) state: reject it at setup, not at the first import.""" +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +from test_pd_hybrid import _qwen35 + + +def test_pd_batcher_rejects_streaming_worker_for_recurrent_model(): + from opt_prime.accel.pd.pd_batching import PDBatcher + worker = SimpleNamespace(stream_layers=True, submit=lambda *a, **k: pytest.fail('must not submit')) + with pytest.raises(ValueError, match='recurrent'): + PDBatcher(_qwen35(), worker, slots=1, max_context=64, page_size=8) + + +def test_subprocess_worker_rejects_streaming_with_recurrent_state_before_spawn(monkeypatch): + import subprocess + from multiprocessing import shared_memory + from opt_prime.accel.pd.pd_batching import SubprocessPrefillWorker + monkeypatch.setattr(subprocess, 'Popen', lambda *a, **k: pytest.fail('must not spawn')) + monkeypatch.setattr(shared_memory, 'SharedMemory', lambda *a, **k: pytest.fail('must not allocate')) + with pytest.raises(ValueError, match='recurrent'): + SubprocessPrefillWorker('python', [], slots=1, layers=4, kv_heads=2, head_dim=16, max_context=64, + stream_layers=True, recurrent_state=True) + + +def test_cli_rejects_streaming_for_recurrent_model_before_worker(monkeypatch): + import opt_prime.accel.pd.pd_batching as pd + path = Path(__file__).parents[1] / 'examples/pd/serve_batched.py' + spec = importlib.util.spec_from_file_location('stream_recurrent_factory', path) + factory = importlib.util.module_from_spec(spec) + spec.loader.exec_module(factory) + monkeypatch.setattr(pd, 'SubprocessPrefillWorker', lambda *a, **k: pytest.fail('must not start a producer')) + args = SimpleNamespace(pd_prefill_stream_layers=True, backend='torch', pd_prefill_python='python', + pd_prefill_device=None, pd_prefill_pythonpath=None, pd_prefill_args='--engine torch', + model='local', max_context=64) + with pytest.raises(ValueError, match='recurrent'): + factory.prefill_worker(args, 1, _qwen35().config) From 1853ba29ff0cc297ca86c342ee7271b6f3057bff Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 30 Oct 2026 10:11:51 +0900 Subject: [PATCH 3/3] pd: document layer streaming scope and factory-engine trust/memory cost (#265) Add a short STREAMING.md design note: protocol, decode-side invariants, and the measured attribution. The notification-delay control kept most of the completion-time gain (22.7/19.6/2.8/21.3%) while early sending added only 4.6/3.0/0.4/1.7%, so the gain comes from the research factories' slower whole-result export, not overlap. State in --engine-factory help and the factory_engine docstring that it is a trusted research flag that imports arbitrary module:callable code and loads the checkpoint as fp32 on the host. Refs #265 Co-Authored-By: Claude Opus 5.5 --- opt_prime/examples/pd/STREAMING.md | 63 +++++++++++++++++++ .../opt_prime/accel/pd/prefill_server.py | 10 ++- 2 files changed, 71 insertions(+), 2 deletions(-) create mode 100644 opt_prime/examples/pd/STREAMING.md diff --git a/opt_prime/examples/pd/STREAMING.md b/opt_prime/examples/pd/STREAMING.md new file mode 100644 index 00000000..90210619 --- /dev/null +++ b/opt_prime/examples/pd/STREAMING.md @@ -0,0 +1,63 @@ +# RNGD layer streaming (opt-in design note) + +`serve_batched.py --pd-prefill-stream-layers` lets a subprocess prefill producer +(`--pd-prefill-python`) export KV one layer or layer group at a time while a Torch +or ATOM decoder imports each finished span. It is off by default, and whole-KV +transfer stays the default path. + +## Protocol + +- The producer advertises `{"ready": true, "kv_progress": 1}`. A request with + `stream_layers: true` gets `{id, kv_span: [layer_start, layer_end, token_start, token_end]}` + messages, then the normal completion message. +- Each span is copied into the request's reserved shared-memory slot before it is + announced. Published spans must not change until the request completes. + Producers process requests serially, cancelled ones included, so an old request + cannot overwrite a reused slot. +- `LayerProgress` rejects malformed or overlapping spans. At completion, every + layer must cover the requested suffix exactly. + +## Decode side + +- `PagedKV.write_layer_span` writes into reserved pages only. Length, prefix-cache + entries and recurrent state are published after completion and full coverage, + and only once destination copies are fenced. +- `PDBatcher` (Torch) imports spans as they arrive. +- `AtomBatcher` imports only into a reserved lane. A B1 prefetched request + (`--pd-prefetch-next`) may compute ahead, but its spans wait until the current + decode releases the lane. Multi-lane hybrid caches defer imports until admission. +- An invalid source fails only its request. An SDK write failure poisons the owner. + Incomplete prompts never enter the reusable prefix index. +- Recurrent (linear-attention) models are rejected at setup by the CLI, + `SubprocessPrefillWorker` and `PDBatcher`, because their state is not part of + the stream. + +## Producers + +- `prefill_server --engine rngd-full` streams through `_RngdFull.prefill_stream`. +- `--engine factory --engine-factory module:callable` loads a research engine. It + is a trusted flag: the callable is imported and run as-is, and the checkpoint is + first loaded into host memory as fp32 (4 bytes per parameter). +- The research factories used in #266 (`streaming_factories.py`) are not part of + this package, because they depend on example modules that are not on this branch. + +## What the measurements show + +#266 measured end-to-end completion-time reductions against whole-KV transfer: +-20% to -35% for Qwen3-0.6B, Qwen3-4B and KONI-8B, and -2% to -10% for Gemma3-1B. +Most of that gain does **not** come from overlapping transfer with prefill. + +A notification-delay control kept per-layer export but held every progress +message until the source finished: + +| model | gain with delayed notification (deferred vs whole) | extra gain from early sending (early vs deferred) | +|---|---:|---:| +| Qwen3-0.6B | 22.7% | 4.6% | +| Qwen3-4B | 19.6% | 3.0% | +| Gemma3-1B | 2.8% | 0.4% | +| KONI-8B | 21.3% | 1.7% | + +Most of the gain is still there when notification is delayed. It comes from the +research factories' slower whole-result export, which the per-layer export path +avoids, not from overlap. This change makes no claim that overlap improves +performance. diff --git a/opt_prime/opt_prime/accel/pd/prefill_server.py b/opt_prime/opt_prime/accel/pd/prefill_server.py index 0816fb68..6957f652 100644 --- a/opt_prime/opt_prime/accel/pd/prefill_server.py +++ b/opt_prime/opt_prime/accel/pd/prefill_server.py @@ -181,7 +181,10 @@ def emit(layer, first, keys, values): def factory_engine(a): - """Load a trusted CLI factory using the same FP32 checkpoint inputs as research workers.""" + """Load a trusted CLI factory using the same FP32 checkpoint inputs as research workers. + + Research-only: `--engine-factory module:callable` is imported and called without any allow-list, and the whole + checkpoint is materialized on the host as fp32 (4 bytes per parameter) before the factory runs.""" from importlib import import_module from pathlib import Path from safetensors import safe_open @@ -214,7 +217,10 @@ def factory_engine(a): def main(): ap = argparse.ArgumentParser() ap.add_argument("--engine", required=True, choices=["torch", "rngd-qwen", "rngd-full", "factory"]) - ap.add_argument("--engine-factory", help="trusted local module:callable; used only with --engine factory") + ap.add_argument("--engine-factory", + help="--engine factory only: module:callable, imported and called as-is (a trusted research " + "flag: it runs arbitrary code from PYTHONPATH). Loads every *.safetensors tensor of " + "--model into host memory as fp32 first (about 4 bytes per parameter, e.g. ~32 GB for 8B)") ap.add_argument("--model", required=True); ap.add_argument("--model-id", default=None) ap.add_argument("--shm", required=True); ap.add_argument("--slots", type=int, required=True) ap.add_argument("--layers", type=int, required=True); ap.add_argument("--kv-heads", type=int, required=True)