pd: stream RNGD prefill KV layers into Torch and ATOM decode (opt-in) (#265)

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.

opt_prime/examples/pd/serve_batched.py

@@ -29,6 +29,13 @@
 
 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 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:
@@ -47,7 +54,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 +141,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",

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<layer_end<=self.layers
+                or not self.start<=token_start<token_end<=self.end):
+                raise ValueError('invalid layer-progress span')
+            for layer in range(layer_start,layer_end):
+                if any(token_start<b and a<token_end for a,b in self._coverage[layer]):
+                    raise ValueError('overlapping layer-progress span')
+            # Validate the complete message before changing any layer's coverage.
+            for layer in range(layer_start,layer_end):
+                self._coverage[layer].append((token_start,token_end))
+            self._pending.append(span)
+
+    def drain(self):
+        with self._lock:
+            pending,self._pending=self._pending,[]
+            return pending
+
+    @property
+    def pending(self):
+        with self._lock:
+            return bool(self._pending)
+
+    @property
+    def finished(self):
+        with self._lock:
+            return self._closed
+
+    def finish(self):
+        with self._lock:
+            if self._closed:
+                raise ValueError('layer progress is closed')
+            for intervals in self._coverage:
+                position=self.start
+                for start,end in sorted(intervals):
+                    if start!=position:
+                        raise ValueError('incomplete layer-progress coverage')
+                    position=end
+                if position!=self.end:
+                    raise ValueError('incomplete layer-progress coverage')
+            self._closed=True

opt_prime/opt_prime/accel/pd/paged.py

@@ -473,6 +473,36 @@ def write_span(self, slot, start, keys, values):
             pk[:, :, idx] = k.to(self.device, pk.dtype, non_blocking=True)
             pv[:, :, idx] = v.to(self.device, pv.dtype, non_blocking=True)
 
+    def write_layer_span(self, slot, start, layer_start, keys, values):
+        """Import a dense model-layer subset into reserved pages, without publishing it.
+
+        The caller owns the reservation until all transfers finish. Length, prefix
+        admission and recurrent state are deliberately unchanged by partial writes.
+        """
+        if (any(type(x) is not int for x in (slot, start, layer_start))
+                or not 0 <= slot < len(self.owned) or start < 0 or layer_start < 0
+                or not isinstance(keys, torch.Tensor) or not isinstance(values, torch.Tensor)
+                or keys.ndim != 4 or values.shape != keys.shape
+                or keys.shape[0] == 0 or keys.shape[2] == 0):
+            raise ValueError("invalid partial KV span")
+        if start < self.length[slot]:
+            raise ValueError("partial KV span cannot overwrite an admitted prefix")
+        end = start + keys.shape[2]
+        if end > 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"""