feat(PR-A3c): batched decode scheduler — 8.45× served multi-tenant throughput + detailed test report

docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md

@@ -216,10 +216,35 @@ needle, primed then decoded interleaved.
 Each session recalled **its own** needle (`MAPLE-7890`, `IOTA-8961`,
 `THETA-6866`, `IOTA-3281` — note two `IOTA-*` sessions got their *own* numbers),
 proving **per-session KV isolation** through the real served path. So
-multi-tenant serving is end-to-end correct + recall-preserving. (Execution is
-still RPC-serialized on the asyncio loop — a batched scheduler to fuse
-concurrent decodes into §3.5's parallel forward is the remaining throughput
-step; correctness/isolation is done.)
+multi-tenant serving is end-to-end correct + recall-preserving. Execution was
+still RPC-serialized on the asyncio loop; the batched scheduler that fuses the
+cohort into one forward is §3.7.
+
+### 3.7 PR-A3c batched scheduler — fusing concurrent decodes for throughput
+
+§3.6's served path was *correct* multi-tenant but RPC-serialized (each session's
+decode forward ran alone). `BatchedDecodeScheduler`
+(`inference_engine/session/batch_scheduler.py`) takes the cohort of per-session
+adapters from the registry, **stacks their restored caches along the batch dim,
+and runs one verifier forward per step** (dropping finished rows) — the
+served-path realisation of §3.5's parallel decode.
+
+Result (H200, 8 sessions, NIAH ctx≈1238,
+`results/research/k3_served_batched_scheduler_gpu.json`):
+
+| path | aggregate decode tok/s | per-session recall |
+| --- | --- | --- |
+| serialized (§3.6, each session alone) | 26.6 | 1.0 |
+| **batched scheduler (§3.7)** | **224.9** | **1.0** |
+| **speedup** | **8.45×** | — |
+
+So the batched scheduler converts the correct-but-serialized multi-tenant path
+into **8.45× aggregate throughput at 8 sessions, recall preserved** — matching
+the engine-level near-linear scaling (§3.5) now driven through the served
+per-session adapters. Scope: a **fixed-cohort** batcher (synchronized burst —
+the dominant multi-tenant case); dynamic mid-flight arrival + ragged-length
+continuous batching (and the async-gRPC futures glue that lets independent
+`Generate` RPC coroutines feed one batch loop) is the remaining productization.
 
 ## 4. Case 2 — cross-host proposer/verifier (FEASIBILITY VERDICT)
 
@@ -427,6 +452,7 @@ the committed evidence JSON, and the headline result).
 | multi-tenant capacity A/B (ctx2048, model-level) | per-agent KV native 256.9 MB vs **S5 61.1 MB**; **~4.2× more agents** (budget hit 15 vs 32; derived 22 vs 93) — recall-preserving | `results/research/k3_multitenant_pressure_mac.json` |
 | PR-A3c parallel throughput (H200, batched S5) | **8.04× near-linear scaling at N=8** (220 tok/s ≈ AR), **per-session recall 1.0** — per-session binding works | `results/research/k3_cuda_multitenant_parallel_gpu.json` |
 | PR-A3c served path (H200, gRPC + 4 SDK clients) | **true multi-tenant serving end-to-end**: 4 concurrent sessions, **per-session recall 1.0**, isolated | `results/research/k3_grpc_multitenant_e2e_gpu.json` |
+| PR-A3c batched scheduler (H200, 8 sessions) | **8.45× throughput** (26.6 → 224.9 tok/s) fusing cohort into one forward, **recall 1.0** | `results/research/k3_served_batched_scheduler_gpu.json` |
 
 **Case 2 (H200 NVL, Gemma-4-26B + DFlash, fused spec-decode vs AR):**
 

docs/reports/pr-a3c-multitenant-serving-test-report.md

@@ -0,0 +1,130 @@
+# PR-A3c — Multi-tenant serving: detailed test report
+
+Detailed record of the per-session-binding (PR-A3c) work and its end-to-end
+tests, from the single-tenant pressure finding through batched parallel
+throughput, the gRPC served path, and the batched scheduler. Summarized in
+[ADR 0014](../adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md)
+§3.4–3.7; this report carries the full methodology, numbers, and evidence index.
+
+## 1. Environment
+
+| Component | Detail |
+| --- | --- |
+| GPU | NVIDIA **H200 NVL** (143 GB), vast.ai; torch 2.12.0+cu130 |
+| Verifier | `google/gemma-4-26B-A4B-it` (bf16, eager attention) |
+| Drafter | `z-lab/gemma-4-26B-A4B-it-DFlash` |
+| f_θ | `results/research/f_theta_v5_s5_sliding` (S5, 5 exact full-attn layers) |
+| Mac | Mac mini M4 (24 GB) via the git-bus Mac bridge (single-tenant pressure, §3.4 source) |
+| Recall task | NIAH (needle-in-a-haystack); per-session distinct needle |
+| Bottom line | **recall must stay 1.0** — recall-sacrificing configs (pure sink+window) are out of scope |
+
+## 2. Motivation — the single-tenant gap
+
+The first capacity test measured the v0.3 **served** path (`grpc_agent_capacity_loadtest.py`):
+256 concurrent agent **connections** admitted, but v0.3 is **single-tenant** — one
+shared verifier, RPCs serialized on one asyncio loop, no per-session KV isolation.
+So "256" = connections *served*, not parallel inferences; concurrent sessions
+would corrupt each other's KV. PR-A3c (per-session binding) is the fix.
+
+## 3. Tests & results
+
+### 3.1 Multi-tenant memory capacity (model-level A/B) — ADR §3.4
+
+`mlx_multitenant_pressure.py` (Mac M4), per-agent KV at ctx2048, 21 GB budget:
+
+| config | per-agent KV | budget hit at | derived max agents |
+| --- | --- | --- | --- |
+| MLX-native (gemma hybrid) | 256.9 MB | N=15 | ~22 |
+| **Kakeya S5** (recall-preserving) | 61.1 MB | N=32 | ~93 |
+
+→ **~4.2× more concurrent agents** at equal context, recall-preserving (the win
+vs native, which already bounds sliding layers to 1024; pure sink+window's 16.8×
+is excluded — it drops full-attn recall).
+
+### 3.2 Batched parallel throughput (engine level) — ADR §3.5
+
+`k3_cuda_multitenant_parallel_bench.py` (H200), batched restored-S5 decode,
+each row = a session with its own KV-cache row:
+
+| sessions N | restored-S5 agg tok/s | parallel speedup | per-session recall |
+| --- | --- | --- | --- |
+| 1 | 27.4 | 1.00× | 1.0 |
+| 2 | 54.6 | 1.99× | 1.0 |
+| 4 | 111.6 | 4.07× | 1.0 |
+| 8 | 220.4 | **8.04×** | 1.0 |
+
+→ near-linear parallel scaling; restored S5 ≈ native AR (220.4 vs 216.4 @ N=8).
+One batch-1 fix was required: RoPE `cos`/`sin` batch-1 broadcast
+(`restored_attention.py`). Evidence: `k3_cuda_multitenant_parallel_gpu.json`.
+
+### 3.3 gRPC served path — per-session binding (end-to-end) — ADR §3.6
+
+Implementation: `CrossModelRestoredSinkWindowVerifier.spawn()` (fresh per-session
+adapter, shared weights) + `PerSessionVerifierRegistry` (session→adapter; also
+the `SessionStore` cache-inspector + coordinator resolver) + coordinator
+resolver + servicer `on_session_close` cleanup + `start_grpc_runtime_server
+--multi-tenant`. Back-compat: single-tenant unchanged (271 session+server unit
+tests pass; `test_verifier_registry.py` proves interleaved-session isolation).
+
+E2E (`k3_grpc_multitenant_e2e.py`, H200): launch the multi-tenant server, 4
+concurrent SDK clients, each its own session + distinct needle:
+
+| sessions | transport | per-session recall | isolation |
+| --- | --- | --- | --- |
+| 4 concurrent | real gRPC `RuntimeService` + Python SDK | **1.0** | ✓ |
+
+Each recalled its own needle (`MAPLE-7890`/`IOTA-8961`/`THETA-6866`/`IOTA-3281` —
+two `IOTA-*` sessions got their own numbers) → per-session KV isolation through
+the real served path. Evidence: `k3_grpc_multitenant_e2e_gpu.json`.
+
+### 3.4 Batched scheduler — fusing concurrent decodes — ADR §3.7
+
+`BatchedDecodeScheduler` (`inference_engine/session/batch_scheduler.py`) stacks
+the cohort's per-session restored caches along the batch dim and runs one
+verifier forward per step (drops finished rows). `k3_served_batched_scheduler_bench.py`
+(H200, 8 sessions):
+
+| path | aggregate decode tok/s | per-session recall |
+| --- | --- | --- |
+| serialized (each session alone) | 26.6 | 1.0 |
+| **batched scheduler** | **224.9** | **1.0** |
+| **speedup** | **8.45×** | — |
+
+→ the served path goes from correct-but-serialized to **8.45× aggregate
+throughput at 8 sessions, recall preserved**. Evidence:
+`k3_served_batched_scheduler_gpu.json`.
+
+## 4. Net result
+
+PR-A3c delivers, recall-preserving (recall 1.0 throughout), the three multi-tenant
+properties together:
+
+- **Bounded memory** — ~4.2× more concurrent agents per GB (§3.1).
+- **Parallel throughput** — 8.04× engine-level (§3.2) / **8.45× through the
+  served per-session adapters via the batched scheduler** (§3.4).
+- **Correct isolation** — true multi-tenant serving end-to-end through gRPC,
+  per-session recall 1.0 (§3.3).
+
+## 5. Remaining work (productization)
+
+- **Async continuous batching transport**: wire the batched scheduler under the
+  async gRPC streaming `Generate` handlers via per-step futures + a background
+  batch loop, so independent RPC coroutines feed one batch — and support
+  **dynamic mid-flight arrival + ragged-length** cohorts (this report's
+  scheduler is a fixed synchronized cohort, the dominant burst case).
+- **Batched fused spec-decode** (DFlash is batch-1 today).
+- **Mac served path**: the served MLX gemma verifier can't load gemma-4's nested
+  config (`MLXSinkWindowVerifier`) — a v0.4 item; CUDA is the recall-preserving
+  served path today.
+
+## 6. Evidence index
+
+| Test | Script | Evidence JSON |
+| --- | --- | --- |
+| Single-tenant capacity (Mac) | `grpc_agent_capacity_loadtest.py` | `k3_agent_capacity_mac.json`, `k3_agent_capacity_stress_mac.json` |
+| Memory A/B (Mac) | `mlx_multitenant_pressure.py` | `k3_multitenant_pressure_mac.json` |
+| Parallel throughput (H200) | `k3_cuda_multitenant_parallel_bench.py` | `k3_cuda_multitenant_parallel_gpu.json` |
+| Served e2e (H200) | `k3_grpc_multitenant_e2e.py` | `k3_grpc_multitenant_e2e_gpu.json` |
+| Batched scheduler (H200) | `k3_served_batched_scheduler_bench.py` | `k3_served_batched_scheduler_gpu.json` |
+
+(All JSON under `results/research/`.)

inference_engine/session/batch_scheduler.py

@@ -0,0 +1,136 @@
+"""Batched decode scheduler — PR-A3c throughput step.
+
+§3.6 made the served path *correct* multi-tenant (per-session binding), but
+execution was RPC-serialized: N concurrent ``Generate`` calls each ran their own
+verifier forward, one after another. This scheduler **fuses** the decode step of
+a cohort of sessions into **one batched forward** — the served-path realisation
+of the parallel throughput validated at the engine level in ADR 0014 §3.5.
+
+Scope: a **fixed-cohort** batched decoder (the common multi-tenant burst: N
+sessions admitted together, decoded in lockstep). It sources its KV from the
+per-session adapters created by :class:`PerSessionVerifierRegistry`, stacks
+their restored caches along the batch dim, and runs one ``verifier_model``
+forward per step. Sessions that hit EOS/max drop out of the batch; the remainder
+keep batching. Dynamic mid-flight arrival + ragged-length continuous batching is
+a follow-up (this covers the synchronized cohort that dominates burst load).
+
+Recall-preserving only (the per-session adapters are restored S5).
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Dict, List, Optional, Sequence
+
+import torch
+
+
+class BatchedDecodeScheduler:
+    """Fuse a cohort of per-session restored adapters into batched decode.
+
+    Parameters
+    ----------
+    verifier_model
+        The shared HF verifier ``nn.Module`` (the same weights every adapter
+        wraps). One batched forward serves the whole cohort.
+    device
+        Torch device for the batched tensors.
+    """
+
+    def __init__(self, verifier_model: Any, device: Any) -> None:
+        self._model = verifier_model
+        self._device = torch.device(device)
+
+    @staticmethod
+    def _stack_caches(adapters: Sequence[Any]):
+        """Concatenate the per-session DynamicCaches into one batched cache.
+
+        Every adapter must have an incremental ``_past`` (DynamicCache) at the
+        SAME sequence length (synchronized cohort). Returns a new batched
+        DynamicCache ``[K, heads, T, dim]`` per layer.
+        """
+        from transformers.cache_utils import DynamicCache
+
+        pasts = [a._past for a in adapters]
+        if any(p is None for p in pasts):
+            raise ValueError("all adapters must be prefilled (incremental _past)")
+        lengths = {int(a._past_len) for a in adapters}
+        if len(lengths) != 1:
+            raise ValueError(f"cohort must share one cache length; got {lengths}")
+        n_layers = len(pasts[0].layers)
+        batched = DynamicCache()
+        for li in range(n_layers):
+            k = torch.cat([p.layers[li].keys for p in pasts], dim=0)
+            v = torch.cat([p.layers[li].values for p in pasts], dim=0)
+            batched.update(k, v, li)
+        return batched
+
+    @torch.no_grad()
+    def run_cohort(
+        self,
+        adapters: List[Any],
+        *,
+        max_tokens: int,
+        eos_ids: Optional[set] = None,
+    ) -> Dict[str, Any]:
+        """Decode ``adapters`` in lockstep via batched forwards.
+
+        Each adapter is a prefilled restored verifier (own KV) at the same
+        cache length. Returns per-session generated tokens + timing.
+        """
+        eos_ids = eos_ids or set()
+        K = len(adapters)
+        if K == 0:
+            return {"tokens": [], "decode_s": 0.0, "decode_tokens_per_s": 0.0}
+        T = int(adapters[0]._past_len)
+        cache = self._stack_caches(adapters)
+        # batched next-token logits from each adapter's prefill
+        logits = torch.cat([a.next_token_logits.view(1, -1) for a in adapters], dim=0)
+        gen: List[List[int]] = [[] for _ in range(K)]
+        active = list(range(K))            # rows still generating
+        # map current batch row -> original session index
+        row_to_sess = list(range(K))
+        torch.cuda.synchronize(self._device) if self._device.type == "cuda" else None
+        t0 = time.perf_counter()
+        step = 0
+        while active and step < max_tokens:
+            nxt = logits.argmax(-1)                         # [B]
+            B = nxt.size(0)
+            keep_rows = []
+            for r in range(B):
+                sidx = row_to_sess[r]
+                tok = int(nxt[r].item())
+                gen[sidx].append(tok)
+                if tok not in eos_ids:
+                    keep_rows.append(r)
+            step += 1
+            if step >= max_tokens or not keep_rows:
+                break
+            cur = nxt.view(B, 1)
+            pos = torch.full((B, 1), T + step - 1, device=self._device, dtype=torch.long)
+            cpos = torch.tensor([T + step - 1], device=self._device)
+            out = self._model(input_ids=cur, position_ids=pos, cache_position=cpos,
+                              past_key_values=cache, use_cache=True)
+            cache = out.past_key_values
+            logits = out.logits[:, -1, :]
+            if len(keep_rows) != B:
+                # Drop finished rows from the batch (shrink) — keeps the
+                # forward dense over only-active sessions.
+                idx = torch.tensor(keep_rows, device=self._device)
+                logits = logits.index_select(0, idx)
+                for layer in cache.layers:
+                    layer.keys = layer.keys.index_select(0, idx).contiguous()
+                    layer.values = layer.values.index_select(0, idx).contiguous()
+                row_to_sess = [row_to_sess[r] for r in keep_rows]
+            active = keep_rows
+        if self._device.type == "cuda":
+            torch.cuda.synchronize(self._device)
+        dt = time.perf_counter() - t0
+        total = sum(len(g) for g in gen)
+        return {
+            "tokens": gen,
+            "decode_s": dt,
+            "decode_tokens_per_s": round(total / dt, 3) if dt > 0 else 0.0,
+            "sessions": K,
+            "total_tokens": total,
+        }

results/research/k3_served_batched_scheduler_gpu.json

@@ -0,0 +1,22 @@
+{
+  "kind": "k3_served_batched_scheduler",
+  "config": {
+    "sessions": 8,
+    "modal_prompt_len": 1238,
+    "max_new_tokens": 24,
+    "sink": 4,
+    "window": 64
+  },
+  "env": {
+    "gpu": "NVIDIA H200 NVL"
+  },
+  "serialized": {
+    "aggregate_tps": 26.615,
+    "recall": 1.0
+  },
+  "batched_scheduler": {
+    "aggregate_tps": 224.928,
+    "recall": 1.0
+  },
+  "batched_speedup_vs_serialized": 8.45
+}