ATOM: hybrid KV tables and single-block PD (#207, #231, #236)

opt_prime/docs/validation/batching/ATOM.md

@@ -100,3 +100,13 @@ shared (load average 75–135 from other users' jobs), so absolute values differ
 With one cache block the host tier quadruples the reused prefix (2048 → 11264 tokens). Under that run's host load the
 reloads did not turn into throughput: host staging is CPU-side work. PD mode (KV prefilled elsewhere, ATOM decodes
 only) is measured in `PD.md`.
+
+## Hybrid attention (#207)
+
+Gemma3 standalone external-cache serving uses global block ownership and lane-local sliding cache ownership; a fresh
+prefill resets the owning lane's decode attention mask. PD supports two cases: fresh full-prefix hybrid PD (the remote
+producer returns the whole prompt KV from position zero; global layers go to the lane's block, the latest sliding
+window to the lane-local sliding cache), and batch-1 single-block PD with prefix and host caching off (the sole block is
+used for the request and released afterwards; no prefix entries are retained). Prefix caching and host caching remain
+rejected for hybrid models until sliding-state restoration is implemented (#207), and partial remote prefixes are
+rejected for both cases.

opt_prime/opt_prime/accel/platforms/atom/batching.py

@@ -26,6 +26,11 @@
   - a host-DRAM tier: device entries evicted for space are copied to host memory (per-layer `get_kv_cache_for_layer`
     through a 2 MiB-aligned staging buffer that mirrors the layer's block layout) and written back on a later hit;
   - KV events (BlockStored / BlockRemoved, device and CPU tiers) and prefix-hit counters, as for the GPU pool.
+Hybrid fresh PD imports global layers into their assigned block and the latest sliding window into the request
+lane. Source dtype metadata may differ from physical storage, so staging validates allocation size. Hybrid prefix
+and host-cache restoration remain unsupported; remote hybrid requests must return their full prompt KV.
+Static B1 cache-off PD may use its sole block: it never decodes an idle lane alongside an active lane.
+This path requires full prompt KV and releases the block after each request; it retains no prefix entries.
 Idle lanes then decode on a reserved scratch block. Measured on ATOM (Llama-3.2-1B, 600-token prompt): full prefill
 85.5 ms; device copy of a 512-token prefix 0.8 ms + 17.5 ms suffix prefill; host reload 8.9 ms; tokens identical.
 """
@@ -39,6 +44,10 @@
 from opt_prime.accel.pd.sampling import select_token
 
 
+class InvalidKVImport(ValueError):
+    """Invalid remote KV rejected before any device write; the lane remains reusable."""
+
+
 _DLFLOAT = None
 
 
@@ -74,20 +83,23 @@ def __init__(self, model, lanes, *, prefix_cache=True, host_bytes=0, medium_devi
         from opt_prime.accel.pd.paged import block_hashes
         self._hashes = block_hashes
         rc = model.rbln_config
+        self.hybrid_cache = getattr(rc, "cache_impl", None) == "hybrid"
         self.unit, self.lanes = rc.prefill_chunk_size, lanes
         self.blocks = rc.kvcache_num_blocks
-        if self.blocks < lanes + 1:
+        self.single_block = self.blocks == lanes == 1 and not (prefix_cache or host_bytes or self.hybrid_cache)
+        if self.blocks < lanes + 1 and not self.single_block:
             raise ValueError("own block tables need kv_blocks >= lanes + 1 (a scratch block for idle lanes)")
         if prefix_cache and self.blocks < lanes + 2:     # else a prefix hit may find only its own source to evict
             raise ValueError("prefix caching needs kv_blocks >= lanes + 2 (a scratch block and an evictable entry)")
         self.scratch = self.blocks - 1
-        self.free = list(range(self.blocks - 1))
+        # B1 cache-off PD never decodes an idle lane beside an active one.
+        self.free = list(range(self.blocks if self.single_block else self.blocks - 1))
         self.lane_block = [None] * lanes
         self.prefix_cache = prefix_cache
         self.entries = OrderedDict()             # block -> prompt unit hashes (unreferenced, least recent first)
         self.index = {}                          # unit hash -> a block holding that prefix
         self.copy_rt, self.io_rt = model.model[0], model.model[1]
-        self.metas = [m if isinstance(m, dict) else {k: getattr(m, k) for k in ("name", "shape", "dtype")
+        self.metas = [m if isinstance(m, dict) else {k: getattr(m, k) for k in ("name", "shape", "dtype", "layer_type")
                                                      if hasattr(m, k)} for m in rc.kvcache_metas]
         self.host_bytes, self.host_used = host_bytes, 0
         self.host = OrderedDict()                # last unit hash -> (hashes, tokens, [per-layer arrays])
@@ -106,7 +118,13 @@ def _staging(self):
             self._stage = []
             for m in self.metas:
                 n = self.io_rt.get_kv_cache_size_for_layer(m["name"])
-                dt = np.uint16 if m.get("dtype", "bfloat16") == "bfloat16" else np.float32
+                count = int(np.prod(m["shape"]))
+                if n == count * 2:
+                    dt = np.uint16  # physical DLFloat16, even when source metadata says float32
+                elif n == count * 4:
+                    dt = np.float32
+                else:
+                    raise ValueError("unsupported physical KV element size")
                 self._stage.append(aligned_tensor(n // np.dtype(dt).itemsize, dtype=dt, alignment=2 << 20,
                                                   tensor_type="np").reshape(m["shape"]))
         return self._stage
@@ -229,17 +247,37 @@ 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"""
         import numpy as np
+        if keys.ndim != 4 or values.ndim != 4:
+            raise InvalidKVImport("KV export must have layer/head/token/dimension axes")
         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:
+            raise InvalidKVImport("KV export geometry does not match compiled cache")
         stage = self._staging()
+        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):
+            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"
+            if self.hybrid_cache and m.get("layer_type") not in ("sliding_attention", "full_attention"):
+                raise ValueError("unknown hybrid KV layer type")
+            dest, offset, kept = (lane, 0, min(n, buf.shape[2])) if local else (block, start, n)
+            if dest is None or not 0 <= dest < buf.shape[0] or offset < 0 or offset + kept > buf.shape[2]:
+                raise InvalidKVImport("KV import exceeds assigned cache capacity")
+            destinations.append((dest, offset, kept))
         # all layers converted at once ([layers, 2, kv_heads, n, head_dim], K then V), then one copy per KV tensor
         kv = torch.stack([keys, values], dim=1).detach().to("cpu")
         if stage[0].dtype == np.uint16:                      # the device holds DLFloat16
             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)):
-            buf[block, :, start:start + n] = data[i // 2, i % 2]
-            self.io_rt.set_kv_cache_for_layer(m["name"], block, start, n, buf)
+            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)
 
     def drain_events(self):
         out, self.events = self.events, []
@@ -262,6 +300,9 @@ def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib=
         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)."""
         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
         decoders = getattr(model, "decoders", None) or {}
         if self.lanes not in decoders:
@@ -288,7 +329,10 @@ def __init__(self, model, *, pad_token_id=None, prefix_cache=False, host_kv_gib=
         self._remote_s = self._local_s = None    # moving averages (seconds per prompt)
         self.local_prefill = local_prefill
         self.blocks = None
-        if (getattr(rc, "kvcache_num_blocks", self.lanes) or self.lanes) > self.lanes:
+        blocks = getattr(rc, "kvcache_num_blocks", self.lanes) or self.lanes
+        single_pd = blocks == self.lanes == 1 and prefill_worker is not None and not (
+            prefix_cache or host_kv_gib or self.hybrid_cache)
+        if blocks > self.lanes or single_pd:
             self.blocks = AtomKVBlocks(model, self.lanes, prefix_cache=prefix_cache,
                                        host_bytes=int(host_kv_gib * 2 ** 30))
             self.block_size = self.blocks.unit
@@ -313,9 +357,18 @@ def kv_usage(self):
     def _prefill(self, lane, prompt, start=0):
         kw = {} if self.blocks is None else dict(block_tables=torch.tensor([self.blocks.lane_block[lane]],
                                                                            dtype=torch.int16))
-        return self.model.prefill_decoder(input_ids=torch.tensor([list(prompt[start:])], dtype=torch.int64),
-                                          cache_position=torch.arange(start, len(prompt), dtype=torch.int32)[None],
-                                          batch_idx=lane, **kw)
+        if self.hybrid_cache and self.blocks is not None:
+            if start:
+                raise ValueError("hybrid sliding cache requires a fresh prefill")
+            kw["local_block_tables"] = torch.tensor([lane], dtype=torch.int16)
+        out = self.model.prefill_decoder(input_ids=torch.tensor([list(prompt[start:])], dtype=torch.int64),
+                                         cache_position=torch.arange(start, len(prompt), dtype=torch.int32)[None],
+                                         batch_idx=lane, **kw)
+        if self.hybrid_cache and self.blocks is not None:
+            # External tables bypass the SDK's prompt mask reset; decode only sets new bits.
+            self.decoder.dec_attn_mask[lane].zero_()
+            self.decoder.dec_attn_mask[lane, :len(prompt)] = 1
+        return out
 
     @torch.inference_mode()
     def step(self):
@@ -367,6 +420,8 @@ def _decode_step(self):
             tokens[r.slot], positions[r.slot] = r.tokens[-1], self.kv.length[r.slot]
         kw = {} if self.blocks is None else dict(block_tables=self.blocks.table([r.slot for r in rows]))
         cache_position = torch.tensor(positions, dtype=torch.int32)[:, None].contiguous()
+        if self.hybrid_cache and self.blocks is not None:
+            kw["local_block_tables"] = torch.arange(self.lanes, dtype=torch.int16)[:, None]
         if getattr(self.model.rbln_config, "use_position_ids", False):     # e.g. Gemma 3 (sliding attention)
             kw["position_ids"] = cache_position.clone()
         out = self.decoder(input_ids=torch.tensor(tokens, dtype=torch.int64)[:, None].contiguous(),
@@ -432,15 +487,25 @@ def _pd_step(self):
                     self.prefilling.remove(r)
                 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)
                     continue
                 r.pending = None
                 ids = r.request.input_ids
                 t0 = time.perf_counter()
-                if out.keys.shape[2]:
-                    self.blocks.import_span(r.slot, out.start, out.keys, out.values)
+                try:
+                    if 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)
+                    continue
+                if self.hybrid_cache:
+                    self.decoder.dec_attn_mask[r.slot].zero_()
+                    self.decoder.dec_attn_mask[r.slot, :len(ids)] = 1
                 # remote cost per prompt: the prefill itself, its transfer and the import here
                 self._remote_s = _ema(self._remote_s, (out.prefill_s or 0.0) + (out.transfer_s or 0.0) +
                                       time.perf_counter() - t0)

opt_prime/tests/test_atom_batching.py

@@ -180,7 +180,7 @@ def _decode(self, input_ids, cache_position, block_tables):
         logits = torch.zeros(B, 1, self.config.vocab_size)
         for lane in range(B):
             block, pos = int(block_tables[lane, 0]), int(cache_position[lane, 0])
-            if block == self.rbln_config.kvcache_num_blocks - 1:
+            if self.rbln_config.kvcache_num_blocks > 1 and block == self.rbln_config.kvcache_num_blocks - 1:
                 continue                                       # scratch block: idle lane
             with torch.inference_mode():
                 out = self.m(input_ids[lane:lane + 1], past_key_values=self._cache(block, pos), use_cache=True,
@@ -351,3 +351,86 @@ def test_pd_remote_only_with_host_tier_events_and_metrics(fake_rebel):
     flat = [e for batch in events.batches for e in batch[1]]
     assert any(e["type"] == "BlockStored" and e["medium"] == "CPU" for e in flat)
     worker.close()
+
+
+@pytest.mark.parametrize("partial", [False, True])
+def test_single_block_pd_reuses_full_prefix_without_cache(fake_rebel, partial):
+    import copy
+    from opt_prime.accel.pd.pd_batching import PrefillWorker, TorchPrefillEngine
+    model = _llama()
+    fake = _blocks_model(fake_rebel, model, batch=1, blocks=1)
+    engine = TorchPrefillEngine(copy.deepcopy(model), max_context=128, chunk=16, page_size=16)
+    original = engine.prefill
+    calls = 0
+    def prefill(request, send_from=0):
+        nonlocal calls
+        out = original(request, send_from)
+        calls += 1
+        if partial and calls == 1:
+            out.start = 1
+            out.keys, out.values = out.keys[:, :, 1:], out.values[:, :, 1:]
+        return out
+    engine.prefill = prefill
+    worker = PrefillWorker(engine)
+    try:
+        b = AtomBatcher(fake, prefix_cache=False, prefill_worker=worker, local_prefill=False)
+        prompts = _prompts(4, 20, 70, seed=236)
+        fs = [b.submit(PDRequest(str(i), p, 5)) for i,p in enumerate(prompts)]
+        b.run_until_idle()
+        for i,(p,f) in enumerate(zip(prompts,fs)):
+            if partial and i == 0:
+                with pytest.raises(ValueError, match="full prompt"):
+                    f.result()
+            else:
+                assert f.result() == _greedy(model,p,5)
+        assert b.stats_pd == {"remote":4,"local":0}
+        assert b.blocks.free == [0] and b.blocks.lane_block == [None]
+        assert not fake.calls and not b.by_id and not b.poisoned
+    finally:
+        worker.close()
+
+
+@pytest.mark.parametrize("batch,blocks,prefix,host", [(1,1,True,0),(1,1,False,1),(2,2,False,0)])
+def test_single_block_pd_keeps_cache_and_multilane_guards(fake_rebel,batch,blocks,prefix,host):
+    model = _llama()
+    fake = _blocks_model(fake_rebel,model,batch=batch,blocks=blocks)
+    with pytest.raises(ValueError):
+        AtomBatcher(fake,prefix_cache=prefix,host_kv_gib=host,prefill_worker=object(),local_prefill=False)
+
+
+def test_single_block_pd_cancel_late_result_and_reuse(fake_rebel):
+    import copy
+    from concurrent.futures import Future
+    from opt_prime.accel.pd.pd_batching import TorchPrefillEngine
+
+    class Worker:
+        def __init__(self):
+            self.pending = []
+
+        def submit(self, request, send_from=0, slot=None):
+            assert send_from == slot == 0
+            f = Future()
+            self.pending.append((request, f))
+            return f
+
+    model = _llama()
+    fake = _blocks_model(fake_rebel, model, batch=1, blocks=1)
+    worker = Worker()
+    engine = TorchPrefillEngine(copy.deepcopy(model), max_context=128, chunk=16, page_size=16)
+    b = AtomBatcher(fake, prefix_cache=False, prefill_worker=worker, local_prefill=False)
+    p,q = _prompts(2,20,35)
+    gone = b.submit(PDRequest('gone',p,5))
+    b.step()
+    assert b.cancel('gone')
+    b.step()
+    with pytest.raises(RequestCancelled):
+        gone.result()
+    good = b.submit(PDRequest('next',q,5))
+    b.step()
+    for request,f in worker.pending:
+        if not f.done():
+            f.set_result(engine.prefill(request))
+    b.run_until_idle()
+    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