From 58a46280f18225b3a48fb91b67ea5cf6cc7460c8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 07:56:24 +0000 Subject: [PATCH 1/2] =?UTF-8?q?feat(PR-A3c):=20BatchedDecodeScheduler=20?= =?UTF-8?q?=E2=80=94=20fuse=20per-session=20adapters'=20decode=20into=20on?= =?UTF-8?q?e=20batched=20forward=20(served-path=20throughput)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixed-cohort batched decoder over PerSessionVerifierRegistry adapters: stacks the per-session restored caches along batch dim, one verifier forward per step, drops finished rows. Served-path realization of the §3.5 parallel throughput. + bench (batched vs serialized, recall-checked) + guard unit tests. Co-authored-by: FluffyAIcode --- inference_engine/session/batch_scheduler.py | 136 ++++++++++++++ .../k3_served_batched_scheduler_bench.py | 166 ++++++++++++++++++ .../session/test_batch_scheduler.py | 61 +++++++ 3 files changed, 363 insertions(+) create mode 100644 inference_engine/session/batch_scheduler.py create mode 100644 scripts/research/k3_served_batched_scheduler_bench.py create mode 100644 tests/inference_engine/session/test_batch_scheduler.py diff --git a/inference_engine/session/batch_scheduler.py b/inference_engine/session/batch_scheduler.py new file mode 100644 index 00000000..79bab116 --- /dev/null +++ b/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, + } diff --git a/scripts/research/k3_served_batched_scheduler_bench.py b/scripts/research/k3_served_batched_scheduler_bench.py new file mode 100644 index 00000000..ab51a6ab --- /dev/null +++ b/scripts/research/k3_served_batched_scheduler_bench.py @@ -0,0 +1,166 @@ +"""PR-A3c throughput: batched scheduler vs serialized, on the served per-session +adapters (recall-preserving S5). + +§3.6 made the served path correct multi-tenant but RPC-serialized. This bench +takes N per-session restored adapters (from PerSessionVerifierRegistry), prefills +each, and decodes the cohort two ways on CUDA: + * serialized — each session's decode forward run alone, summed (the §3.6 path) + * batched — BatchedDecodeScheduler fuses all N into one forward per step + +and reports aggregate decode tok/s for each, the speedup, and per-session recall +(must be 1.0 — recall is the bottom line). +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from collections import Counter +from pathlib import Path +from typing import Any, Dict, List + +import torch + + +@torch.no_grad() +def _serial_decode(model, adapter, max_tokens, eos_ids, device): + """Single-session greedy decode against the adapter's own _past.""" + cache = adapter._past + T = int(adapter._past_len) + logits = adapter.next_token_logits + gen: List[int] = [] + for step in range(max_tokens): + tok = int(logits.argmax(-1).item()) + gen.append(tok) + if tok in eos_ids: + break + cur = torch.tensor([[tok]], device=device) + pos = torch.tensor([[T + step]], device=device) + out = model(input_ids=cur, position_ids=pos, + cache_position=torch.tensor([T + step], device=device), + past_key_values=cache, use_cache=True) + cache = out.past_key_values + logits = out.logits[0, -1, :] + return gen + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--verifier-id", default="google/gemma-4-26B-A4B-it") + ap.add_argument("--drafter-id", default="z-lab/gemma-4-26B-A4B-it-DFlash") + ap.add_argument("--f-theta-dir", default="results/research/f_theta_v5_s5_sliding") + ap.add_argument("--sessions", type=int, default=8) + ap.add_argument("--haystack-lines", type=int, default=60) + ap.add_argument("--max-new-tokens", type=int, default=24) + ap.add_argument("--sink", type=int, default=4) + ap.add_argument("--window", type=int, default=64) + ap.add_argument("--output", default=None) + args = ap.parse_args() + + if not torch.cuda.is_available(): + print("[sched] CUDA required.", file=sys.stderr) + return 2 + device = torch.device("cuda") + from transformers import AutoTokenizer + from inference_engine.v04 import make_niah_dataset + from inference_engine.v04.build_restored import load_restored_verifier + from inference_engine.session.verifier_registry import PerSessionVerifierRegistry + from inference_engine.session.batch_scheduler import BatchedDecodeScheduler + + tok = AutoTokenizer.from_pretrained(args.verifier_id) + eos_ids = {tok.eos_token_id} if tok.eos_token_id is not None else set() + + print("[sched] loading restored verifier ...", file=sys.stderr, flush=True) + base = load_restored_verifier( + verifier_id=args.verifier_id, drafter_id=args.drafter_id, + f_theta_dir=args.f_theta_dir, sink_size=args.sink, window_size=args.window, + s5_exact_full_attn=True, device="cuda", incremental=True) + model = base.model + registry = PerSessionVerifierRegistry(factory=base.spawn) + scheduler = BatchedDecodeScheduler(model, device) + + N = args.sessions + pool = make_niah_dataset(n_samples=N * 3, haystack_min_lines=args.haystack_lines, + haystack_max_lines=args.haystack_lines, seed=0) + + def encode(text): + ids = tok.apply_chat_template([{"role": "user", "content": text}], + add_generation_prompt=True, tokenize=True, + return_tensors="pt") + if hasattr(ids, "keys"): + ids = ids["input_ids"] + return ids[0].tolist() + + enc = [(encode(s.prompt_text), s.answer_text) for s in pool] + modal = Counter(len(e[0]) for e in enc).most_common(1)[0][0] + bucket = [(i, a) for i, a in enc if len(i) == modal][:N] + while len(bucket) < N: + bucket += bucket[: N - len(bucket)] + prompts = [b[0] for b in bucket] + answers = [b[1] for b in bucket] + print(f"[sched] {N} sessions, modal prompt len={modal}", file=sys.stderr, flush=True) + + def recall(toks, ans): + return ans in tok.decode(toks, skip_special_tokens=True) + + def fresh_adapters(): + ads = [] + for i in range(N): + registry.remove(f"s{i}") + a = registry.get(f"s{i}") + a.prefill(prompts[i]) + ads.append(a) + return ads + + # warmup + try: + wa = fresh_adapters()[:2] + scheduler.run_cohort(wa, max_tokens=4, eos_ids=eos_ids) + except Exception as e: # noqa: BLE001 + print(f"[sched] warmup note: {e}", file=sys.stderr) + + # --- batched --- + ads = fresh_adapters() + bres = scheduler.run_cohort(ads, max_tokens=args.max_new_tokens, eos_ids=eos_ids) + batched_tps = bres["decode_tokens_per_s"] + batched_recall = sum(recall(bres["tokens"][i], answers[i]) for i in range(N)) / N + + # --- serialized (§3.6 path) --- + ads = fresh_adapters() + torch.cuda.synchronize(device) + t0 = time.perf_counter() + ser_tokens = [] + for i in range(N): + ser_tokens.append(_serial_decode(model, ads[i], args.max_new_tokens, eos_ids, device)) + torch.cuda.synchronize(device) + ser_dt = time.perf_counter() - t0 + ser_total = sum(len(t) for t in ser_tokens) + serial_tps = round(ser_total / ser_dt, 3) + serial_recall = sum(recall(ser_tokens[i], answers[i]) for i in range(N)) / N + + speedup = round(batched_tps / serial_tps, 2) if serial_tps else None + report = { + "kind": "k3_served_batched_scheduler", + "config": {"sessions": N, "modal_prompt_len": modal, + "max_new_tokens": args.max_new_tokens, + "sink": args.sink, "window": args.window}, + "env": {"gpu": torch.cuda.get_device_name(0)}, + "serialized": {"aggregate_tps": serial_tps, "recall": round(serial_recall, 3)}, + "batched_scheduler": {"aggregate_tps": batched_tps, + "recall": round(batched_recall, 3)}, + "batched_speedup_vs_serialized": speedup, + } + print(f"[sched] N={N}: serialized {serial_tps} tok/s (recall {serial_recall}) | " + f"batched {batched_tps} tok/s (recall {batched_recall}) | " + f"speedup {speedup}x", file=sys.stderr, flush=True) + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + Path(args.output).write_text(json.dumps(report, indent=2)) + print(f"[sched] wrote {args.output}", file=sys.stderr) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/inference_engine/session/test_batch_scheduler.py b/tests/inference_engine/session/test_batch_scheduler.py new file mode 100644 index 00000000..1d76cd8b --- /dev/null +++ b/tests/inference_engine/session/test_batch_scheduler.py @@ -0,0 +1,61 @@ +"""Guard-logic tests for BatchedDecodeScheduler (Linux, no model forward). + +The batched-forward correctness is validated end-to-end on H200 (per-session +recall 1.0 in k3_served_batched_scheduler_bench.py). Here we lock the cheap +invariants: empty cohort, and the same-length cohort requirement. +""" + +from __future__ import annotations + +import pytest + +from inference_engine.session.batch_scheduler import BatchedDecodeScheduler + + +class _FakeLayer: + def __init__(self, t): + import torch + self.keys = torch.zeros(1, 2, t, 4) + self.values = torch.zeros(1, 2, t, 4) + + +class _FakePast: + def __init__(self, t): + self.layers = [_FakeLayer(t), _FakeLayer(t)] + + +class _FakeAdapter: + def __init__(self, t): + self._past = _FakePast(t) + self._past_len = t + import torch + self.next_token_logits = torch.zeros(10) + + +def test_empty_cohort_is_noop(): + sched = BatchedDecodeScheduler(verifier_model=None, device="cpu") + out = sched.run_cohort([], max_tokens=4) + assert out["tokens"] == [] and out["decode_tokens_per_s"] == 0.0 + + +def test_cohort_must_share_one_cache_length(): + sched = BatchedDecodeScheduler(verifier_model=None, device="cpu") + with pytest.raises(ValueError, match="one cache length"): + sched._stack_caches([_FakeAdapter(5), _FakeAdapter(7)]) + + +def test_stack_caches_batches_equal_length(): + sched = BatchedDecodeScheduler(verifier_model=None, device="cpu") + batched = sched._stack_caches([_FakeAdapter(5), _FakeAdapter(5), _FakeAdapter(5)]) + # 2 layers, each stacked to batch dim 3 + assert len(batched.layers) == 2 + assert batched.layers[0].keys.shape[0] == 3 + assert batched.layers[0].keys.shape[2] == 5 + + +def test_prefill_required(): + sched = BatchedDecodeScheduler(verifier_model=None, device="cpu") + a = _FakeAdapter(5) + a._past = None + with pytest.raises(ValueError, match="prefilled"): + sched._stack_caches([a]) From b228665dc57736d18621d28001ee299037b68e39 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 08:02:44 +0000 Subject: [PATCH 2/2] =?UTF-8?q?docs(adr0014=20=C2=A73.7=20+=20report):=20b?= =?UTF-8?q?atched=20scheduler=208.45x=20served=20throughput=20(recall=201.?= =?UTF-8?q?0)=20+=20detailed=20PR-A3c=20multi-tenant=20test=20report?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: FluffyAIcode --- ...-capacity-and-cross-host-topology-tests.md | 34 ++++- .../pr-a3c-multitenant-serving-test-report.md | 130 ++++++++++++++++++ .../k3_served_batched_scheduler_gpu.json | 22 +++ 3 files changed, 182 insertions(+), 4 deletions(-) create mode 100644 docs/reports/pr-a3c-multitenant-serving-test-report.md create mode 100644 results/research/k3_served_batched_scheduler_gpu.json diff --git a/docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md b/docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md index 6e8917f3..db35d512 100644 --- a/docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md +++ b/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):** diff --git a/docs/reports/pr-a3c-multitenant-serving-test-report.md b/docs/reports/pr-a3c-multitenant-serving-test-report.md new file mode 100644 index 00000000..ad769f02 --- /dev/null +++ b/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/`.) diff --git a/results/research/k3_served_batched_scheduler_gpu.json b/results/research/k3_served_batched_scheduler_gpu.json new file mode 100644 index 00000000..df442fd9 --- /dev/null +++ b/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 +} \ No newline at end of file