atom: prefill one queued request during B1 remote-only decode (#253)

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)")

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,10 +537,26 @@ 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
         if self.kv_events is not None:
             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

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

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)