From 86d487b212b4b8db639f2784016a14688aa32f45 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 30 Oct 2026 10:00:34 +0900 Subject: [PATCH] atom: prefill one queued request during B1 remote-only decode (#253) AtomBatcher(remote_prefetch=True) / serve_batched.py --pd-prefetch-next: with a single decode lane and remote-only PD, submit the next queued prompt to the producer while the current request decodes, so its KV is ready when the lane frees. One prefetch at a time; the source slot stays reserved until the owner has imported the KV (also after cancel or producer error). Off by default; rejected unless single-lane, remote-only, with a prefill worker. With a prefix cache the prefetched prompt is re-sent in full and overwrites the restored prefix (correct, but forfeits cache savings). From #255 (3e16c33), code and tests only. Refs #253 Co-Authored-By: Claude Opus 5.5 --- opt_prime/examples/pd/serve_batched.py | 9 ++- .../accel/platforms/atom/batching.py | 46 ++++++++++- opt_prime/tests/test_atom_batching.py | 77 +++++++++++++++++++ opt_prime/tests/test_serving_prefill_cache.py | 15 ++++ 4 files changed, 142 insertions(+), 5 deletions(-) diff --git a/opt_prime/examples/pd/serve_batched.py b/opt_prime/examples/pd/serve_batched.py index 5f3f38cf..07e774c8 100644 --- a/opt_prime/examples/pd/serve_batched.py +++ b/opt_prime/examples/pd/serve_batched.py @@ -66,6 +66,10 @@ def remote_tier(a, model): def build(a): + prefetch = getattr(a, "pd_prefetch_next", False) + if prefetch and (a.backend != "atom" or a.slots != 1 or not a.pd_remote_only or + not (a.pd_prefill_device or a.pd_prefill_python)): + raise ValueError("next prefill requires single-lane ATOM remote-only PD") tok = AutoTokenizer.from_pretrained(a.model) events = None if a.kv_events: @@ -87,7 +91,7 @@ def build(a): engine.start(a.model) b, model = engine.batcher(prefix_cache=a.prefix_cache, host_kv_gib=a.kv_host_memory_gb, kv_events=events, prefill_worker=prefill_worker(a, a.slots, engine.model.config), - local_prefill=not a.pd_remote_only), engine.model + local_prefill=not a.pd_remote_only, remote_prefetch=prefetch), engine.model else: model = AutoModelForCausalLM.from_pretrained(a.model, dtype=getattr(torch, a.dtype)).to(a.device).eval() kw = dict(slots=a.slots, max_context=a.max_context, kv=a.kv, page_size=a.page_size, @@ -132,6 +136,9 @@ def build(a): ap.add_argument("--pd-prefill-pythonpath", default=None) 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", + help="B1 ATOM remote-only PD: prefill one queued prompt while the current request decodes " + "(with --prefix-cache it re-sends the full prompt, forfeiting cache savings)") ap.add_argument("--remote-kv", default=None, help="JSON with Mooncake store setup arguments (shared KV, #128)") ap.add_argument("--remote-kv-keys", default="lmcache", choices=["lmcache", "optimus"]) ap.add_argument("--remote-kv-model", default=None, help="model name in the keys (the vLLM/LMCache side's name)") diff --git a/opt_prime/opt_prime/accel/platforms/atom/batching.py b/opt_prime/opt_prime/accel/platforms/atom/batching.py index 1fa67d0b..7b233274 100644 --- a/opt_prime/opt_prime/accel/platforms/atom/batching.py +++ b/opt_prime/opt_prime/accel/platforms/atom/batching.py @@ -288,7 +288,7 @@ class AtomBatcher(ContinuousBatcher): supports_sampling = True # the decode graph returns logits; sampling runs on the host def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib=0.0, kv_events=None, - prefill_worker=None, remote_inflight=None, local_prefill=True): + prefill_worker=None, remote_inflight=None, local_prefill=True, remote_prefetch=False): """prefix_cache / host_kv_gib need a model compiled with kv_blocks > batch_size (own block tables; prefix_cache needs kv_blocks >= batch_size + 2). prefill_worker (PD, #120): prompts are prefilled elsewhere (`pd.pd_batching.PrefillWorker` and engines) and their @@ -298,12 +298,19 @@ def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib= (in-flight remote prefills + 1) x remote time per prompt (prefill + KV import) against local (prompts queued here + 1) x ATOM prefill time; local prefills run one per iteration, as in aggregated serving. A slow prefill device is therefore used only for overflow, and a fast one takes most prompts. - local_prefill=False keeps every prefill remote (decode never stalls on a prompt; TPOT first).""" + local_prefill=False keeps every prefill remote (decode never stalls on a prompt; TPOT first). + remote_prefetch=True (B1 remote-only): keep one full-prompt result ahead of decode. Cache admission still + happens on the owner after the lane is free; early prefills cannot omit a cached prefix from transfer. + With prefix_cache the prefetched prompt is re-sent in full and overwrites the restored prefix: correct, + but it forfeits the cache's transfer savings.""" rc = model.rbln_config self.hybrid_cache = getattr(rc, "cache_impl", None) == "hybrid" if self.hybrid_cache and (prefix_cache or host_kv_gib): raise ValueError("hybrid sliding cache restoration is not implemented for prefix/host cache") self.lanes = rc.batch_size + if remote_prefetch and (self.lanes != 1 or prefill_worker is None or local_prefill): + raise ValueError("remote_prefetch requires single-lane remote-only PD") + self.remote_prefetch, self._prefetched = remote_prefetch, None decoders = getattr(model, "decoders", None) or {} if self.lanes not in decoders: raise ValueError(f"no compiled decode graph for batch {self.lanes} (have {sorted(decoders)})") @@ -343,6 +350,10 @@ def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib= for lane in range(self.lanes): # every lane owns a page before the B-lane graph runs self._prefill(lane, [self.pad]) + @property + def idle(self): + return super().idle and self._prefetched is None + @property def hit_tokens(self): return self.blocks.hit_tokens if self.blocks else 0 @@ -432,6 +443,8 @@ def _decode_step(self): self._decoded(rows, [self.kv.length[r.slot] for r in rows], logits) def _finish(self, r, error): + if self._prefetched is r and (r.pending is None or r.pending.done()): + self._prefetched = None if self.blocks is not None and r.slot >= 0: self.blocks.release(r.slot, r.request.input_ids if getattr(r, "admitted", False) else None) super()._finish(r, error) @@ -440,10 +453,14 @@ def _finish(self, 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 admitted, local = [], None remote = sum(1 for r in self.prefilling if r.pending is not None) queued = 0 # prompts that will be prefilled here before this one while self.free and self.waiting: + if self._prefetched is not None and self.waiting[0] is not self._prefetched: + break # a cancelled producer still owns its source buffer go_remote = remote < self.remote_inflight or not self.local_prefill if not self.local_prefill: pass @@ -468,14 +485,17 @@ def _pd_step(self): 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 - r.pending = self.worker.submit(r.request, send_from=start, slot=r.slot) - self.stats_pd["remote"] += 1 + if r.pending is None: + r.pending = self.worker.submit(r.request, send_from=start, slot=r.slot) + self.stats_pd["remote"] += 1 if local is not None: self.stats_pd["local"] += 1 t0 = time.perf_counter() self._prefill_request(local) self._local_s = _ema(self._local_s, time.perf_counter() - t0) pending = [r for r in self.prefilling if r.pending is not None] + 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) for r in pending: @@ -503,6 +523,8 @@ def _pd_step(self): with self._lock: self._finish(r, error) continue + if self._prefetched is r: + self._prefetched = None # only release after the owner has imported the KV if self.hybrid_cache: self.decoder.dec_attn_mask[r.slot].zero_() self.decoder.dec_attn_mask[r.slot, :len(ids)] = 1 @@ -515,6 +537,8 @@ def _pd_step(self): self.decoding.append(r) if self.decoding: self._decode_step() + if self.remote_prefetch: + self._prefetch_next() except BaseException: self.poisoned = True raise @@ -522,3 +546,17 @@ def _pd_step(self): events = self.blocks.drain_events() if events: self.kv_events.publish(events) + + def _prefetch_next(self): + from concurrent.futures import Future + with self._lock: + if self._prefetched is not None or not self.decoding or not self.waiting: + return + r = self._prefetched = self.waiting[0] + try: + # Cache admission waits for the decode lane; reserve one full-prompt result in the existing source slot. + r.pending = self.worker.submit(r.request, send_from=0, slot=0) + except BaseException as error: + r.pending = Future() + r.pending.set_exception(error) + self.stats_pd["remote"] += 1 diff --git a/opt_prime/tests/test_atom_batching.py b/opt_prime/tests/test_atom_batching.py index 2434a808..dd3523d0 100644 --- a/opt_prime/tests/test_atom_batching.py +++ b/opt_prime/tests/test_atom_batching.py @@ -434,3 +434,80 @@ def submit(self, request, send_from=0, slot=None): assert good.result() == _greedy(model,q,5) assert b.blocks.free == [0] and b.blocks.lane_block == [None] assert not b.by_id and not b.poisoned + + +@pytest.mark.parametrize('outcome', ['success', 'cancel', 'error']) +@pytest.mark.parametrize('cache', [False, True]) +def test_b1_prefetch_overlaps_decode_and_holds_source_until_import(fake_rebel, outcome, cache): + """A reserved prefill must neither wait for the decode lane nor reuse a still-written source slot.""" + import copy + from concurrent.futures import Future + from opt_prime.accel.pd.pd_batching import TorchPrefillEngine + + class ControlledProducer: + def __init__(self): + self.calls = [] + + def submit(self, request, send_from=0, slot=None): + assert send_from == slot == 0 + future = Future() + self.calls.append((request, future)) + return future + + model = _llama() + fake = _blocks_model(fake_rebel, model, batch=1, blocks=3 if cache else 1) + producer = ControlledProducer() + engine = TorchPrefillEngine(copy.deepcopy(model), max_context=128, chunk=16, page_size=16) + b = AtomBatcher(fake, prefix_cache=cache, host_kv_gib=1.0 if cache else 0, + prefill_worker=producer, local_prefill=False, remote_prefetch=True) + p, q = _prompts(2, 40, 41, seed=253) + first = b.submit(PDRequest('first', p, 8)) + second = b.submit(PDRequest('second', p, 5)) + third = b.submit(PDRequest('third', q, 5)) + b.step() + request, pending = producer.calls[0] + pending.set_result(engine.prefill(request)) + b.step() + assert not first.done() and b.decoding + assert len(producer.calls) == 2, 'next producer must start before the first decode finishes' + request, pending = producer.calls[1] + if outcome == 'success': + pending.set_result(engine.prefill(request)) + elif outcome == 'error': + pending.set_exception(RuntimeError('producer failed')) + else: + assert b.cancel('second') + assert b.cancel('third') + b.step() + with pytest.raises(RequestCancelled): + second.result() + with pytest.raises(RequestCancelled): + third.result() + assert not pending.cancelled(), 'cancellation must not imply that the remote writer has stopped' + while not first.done(): + assert len(producer.calls) == 2, 'one retained prefill must bound producer reservations' + b.step() + assert first.result() == _greedy(model, p, 8) + if outcome == 'cancel': + assert not b.idle, 'the owner must keep draining an orphaned producer after every client has left' + third = b.submit(PDRequest('third', q, 5)) + b.step() + assert len(producer.calls) == 2, 'cancelled writer still owns its shared source slot' + pending.set_result(engine.prefill(request)) + for _ in range(30): + b.step() + if len(producer.calls) == 3: + break + assert len(producer.calls) == 3 + request, pending = producer.calls[2] + pending.set_result(engine.prefill(request)) + b.run_until_idle() + if outcome == 'success': + assert second.result() == _greedy(model, p, 5) + if cache: + assert b.hit_tokens >= 32 + elif outcome == 'error': + with pytest.raises(RuntimeError, match='producer failed'): + second.result() + assert third.result() == _greedy(model, q, 5) + assert b.free == [0] and not b.by_id and not b.poisoned diff --git a/opt_prime/tests/test_serving_prefill_cache.py b/opt_prime/tests/test_serving_prefill_cache.py index d2d4f8af..5e858acf 100644 --- a/opt_prime/tests/test_serving_prefill_cache.py +++ b/opt_prime/tests/test_serving_prefill_cache.py @@ -72,3 +72,18 @@ def observe(model, kv, slot, ids, start, end, device): args.pd_prefill_chunk = chunk with pytest.raises(ValueError, match="positive"): factory.prefill_worker(args, 1, model.config) + + +@pytest.mark.parametrize('backend,slots,remote,worker', [('torch',1,True,True),('atom',2,True,True), + ('atom',1,False,True),('atom',1,True,False)]) +def test_next_prefill_rejects_unsupported_route_before_load(monkeypatch,backend,slots,remote,worker): + path=Path(__file__).parents[1]/'examples/pd/serve_batched.py' + spec=importlib.util.spec_from_file_location('serve_prefetch_test',path) + factory=importlib.util.module_from_spec(spec);spec.loader.exec_module(factory) + args=SimpleNamespace(model='must-not-load',backend=backend,slots=slots,pd_prefetch_next=True, + pd_remote_only=remote,pd_prefill_python='python' if worker else None,pd_prefill_device=None) + def unexpected_load(*args,**kwargs): + pytest.fail('unsupported prefetch route reached model loading') + monkeypatch.setattr(factory.AutoTokenizer,'from_pretrained',unexpected_load) + with pytest.raises(ValueError,match='single-lane'): + factory.build(args)