Phase 4: disaggregated prefill/decode

README.md

@@ -6,16 +6,19 @@ execution, multi-GPU expert-parallel serving, and disaggregated
 prefill/decode, benchmarked against production serving engines (vLLM,
 SGLang) on real, measured throughput/latency/cost numbers, not estimates.
 
-**Status: Phase 1 (custom Triton grouped-GEMM kernel) complete.**
-A naive and a persistent, cache-aware Triton grouped-GEMM kernel, swapped
-into `deepseek-ai/deepseek-moe-16b-base`'s real 27 MoE layers, measured
-**~65-67% faster decode throughput than DeepSeek's own stock `moe_infer`**
-(12.55 → 20.98 tokens/sec, bf16, single NVIDIA L40, unbatched eager-mode
-decode) at perfect mutual top-5 and top-1 logit agreement across every
-tested position. One honest null result: the persistent kernel's grouped
-launch ordering showed no measurable benefit over the naive kernel at
-this project's decode-shaped workload — recorded rather than buried. See
-[`docs/findings/2026-09-15-phase-1-grouped-gemm-run.md`](docs/findings/2026-09-15-phase-1-grouped-gemm-run.md)
+**Status: Phase 4 (disaggregated prefill/decode) complete, open as a PR.**
+Phases 0-3 are merged: a plain-HF baseline, a naive and a persistent
+Triton grouped-GEMM kernel (**~65-67% faster decode throughput** than
+DeepSeek's own stock `moe_infer`, perfect mutual top-5/top-1 logit
+agreement), a skewed-load benchmark contribution upstreamed to vLLM, and
+real 2-GPU expert-parallel serving via DeepSeek's DeepEP. Phase 4 adds a
+continuous-batching prefill/decode scheduler and tests it across two real
+4-GPU topologies — a co-located 4-rank EP pool and a disaggregated
+2+2-rank EP pool connected by a real cross-rank KV-cache handoff — both
+proven byte-exact correct against a single-GPU reference on real
+hardware. Three real bugs found and fixed along the way, all now covered
+by regression tests. See
+[`docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md`](docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md)
 for the full measured account, and
 [`docs/design/2026-09-14-dispatch-system-design.md`](docs/design/2026-09-14-dispatch-system-design.md)
 for the architecture — model choice, what's custom vs. reused, the 7-phase

docs/STATUS.md

@@ -284,9 +284,68 @@ Full account, including the exact bugs, fixes, and every number:
 **Total Phase 3 GPU cost: $13.03** of the $25 cap, 2x H200 SXM, 85
 minutes.
 
+## Phase 4 progress
+
+Design: `docs/design/2026-09-16-phase-4-disaggregated-prefill-decode.md`.
+Plan: `docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md`.
+Branch `phase-4-disaggregated-prefill-decode`.
+
+- [x] Task 1: KV-cache slice/pad-batch helpers (`src/dispatch/serving/kv_cache.py`),
+      grounded in transformers v5's real `DynamicCache` API (confirmed live
+      against transformers' own migration guide and source -- `to_legacy_cache`
+      etc. were removed in v5)
+- [x] Task 2: Cross-rank KV-cache handoff (`src/dispatch/serving/handoff.py`),
+      proven over CPU-only `gloo` before any GPU
+- [x] Task 3: Continuous-batching prefill/decode workers
+      (`src/dispatch/serving/disaggregated.py`) -- found and fixed a real
+      scheduler bug (admission and decode sharing one `step()` call could
+      over-generate by one token) before ever running on GPU
+- [x] Task 4: Co-located baseline worker (`src/dispatch/serving/colocated.py`)
+- [x] Task 5: 4x H100 SXM rental, real EP wiring, correctness gate, concurrency measurement
+- [x] Task 6: Findings doc and this update
+
+**Phase 4 is complete.** All 6 tasks done, `make check` green throughout
+(102 tests total, including the new `test_kv_cache.py`/`test_handoff.py`/
+`test_disaggregated.py`/`test_colocated.py`).
+
+**Three real bugs found and fixed on real hardware**, none caught by
+CPU-only tests: DeepSeek's remote-code model returns the legacy
+tuple-of-tensors cache format (not `DynamicCache`), handled at the
+pod-local script boundary; `local_expert_contribution` (Phase 3 code)
+crashed when a rank received zero tokens for any local expert at all --
+a case 4-way EP's finer sharding made reachable where Phase 3's 2-way EP
+never hit it; and `handoff.py`/`kv_cache.py` had two device-placement
+bugs (NCCL needs CUDA tensors; `gloo`-only tests never caught it) found
+by re-reading the code before running it, not by a live crash. All three
+fixed with regression tests.
+
+**Correctness gate passed for both topologies**: exact greedy-token-sequence
+match (the strict bar this scheduler's token-only interface actually
+admits) against a single-GPU reference, all 3 prompts, for both a
+co-located 4-rank EP pool and a disaggregated 2+2-rank EP pool connected
+by a real cross-rank KV-cache handoff.
+
+**The measured answer to Phase 4's thesis** (does disaggregation relieve
+prefill/decode contention, holding GPU count fixed at 4): mixed, not a
+clean result. Disaggregated TTFT beats co-located's at concurrency 4
+(0.46s vs 1.12s mean) but loses at concurrency 8 (0.98s vs 0.70s) --
+most likely a kernel-warmup confound between independent process
+launches rather than a real topology effect (co-located's own wall time
+*dropped* going from concurrency 4 to 8, the signature of warmup, not
+contention). Reported as genuinely inconclusive rather than forced into
+either direction; a warmed-up measurement protocol is the identified
+follow-up, not attempted this session per the standing cost-discipline
+instruction once both correctness gates had passed. Full account:
+`docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md`.
+
+**Total Phase 4 GPU cost: $10.94** of the $40 cap, 4x H100 SXM, 47
+minutes.
+
 ## Next step
 
-Phase 3's PR (#3) is ready to merge -- `make check` green, correctness
-gate passed, findings recorded. Phase 2's PR is still open and awaiting
+Phase 4 is complete and ready to push as a PR (`make check` green,
+correctness gates passed, findings recorded, cost well under cap).
+Phase 3's PR (#3) already merged. Phase 2's PR is still open and awaiting
 maintainer review; no further work planned on it beyond responding to
-review feedback. Phase 4 is not yet planned.
+review feedback. Phase 5 (quantization + speculative decoding) is not
+yet planned.

docs/design/2026-09-16-phase-4-disaggregated-prefill-decode.md

@@ -0,0 +1,204 @@
+# Phase 4 design: disaggregated prefill/decode
+
+Per `docs/design/2026-09-14-dispatch-system-design.md` §7 (phase plan,
+row 4: "Disaggregated prefill/decode across separate GPU pools").
+
+## 1. Purpose and thesis
+
+Every phase through Phase 3 stayed deliberately single-request and
+unbatched (Phase 0's own methodology: "unbatched eager-mode token-by-token
+decode"). At that scale there is no batching contention between prefill
+and decode -- there is only ever one request in flight, so co-locating
+both phases on the same GPU costs nothing extra. Disaggregated
+prefill/decode only earns its keep once concurrent requests are actually
+competing for the same GPU: prefill's large, batched, compute-bound
+forward pass and decode's small, sequential, memory-bound steps interfere
+with each other's ideal shape when several requests share one engine.
+
+**Phase 4's thesis: does splitting prefill and decode onto separate GPU
+pools measurably relieve that contention, for this project's own model
+and kernels, holding total GPU count fixed?** That means Phase 4
+introduces concurrent, continuous-batched serving for the first time in
+this project -- new scope beyond anything Phase 0-3 built -- specifically
+because the contention this phase measures does not exist without it.
+
+The comparison is controlled, not a bigger-hardware-wins comparison: the
+same 4 GPUs, once as one 4-rank EP pool doing both prefill and decode
+for concurrent requests, once as two independent 2-rank EP pools (one
+prefill, one decode). GPU count held constant; the only variable is
+whether the two phases share GPUs.
+
+## 2. Scope, decided during brainstorming
+
+- **Batching: a real continuous-batching scheduler**, not static
+  fixed-window batching. New requests are admitted into a live decode
+  batch as slots free up, the pattern that actually produces prefill/decode
+  contention on a shared engine -- a static batch-then-drain model would
+  under-state the contention this phase exists to measure.
+- **Combined with Phase 3's EP**, not isolated from it. Each pool (prefill,
+  decode) is itself a 2-GPU DeepEP expert-parallel pool, reusing
+  `patch_moe_infer_ep` and `expert_parallel.py` directly rather than
+  falling back to single-GPU dense MoE. This is closer to what a real
+  deployment would actually run, at the cost of stacking two previously
+  separate subsystems (EP, disaggregation) together for the first time.
+- **Topology: one 4-GPU NVLink/SXM node**, not two separate pods. GPUs 0-1
+  = prefill EP pool, GPUs 2-3 = decode EP pool, all four ranks joined in
+  one `torch.distributed` world so KV-cache handoff is a same-node NCCL
+  `send`/`recv`, not a real network transfer. Logically disaggregated
+  (separate process groups, separate schedulers, an explicit handoff
+  step) without taking on cross-node networking as new scope.
+- **Baseline: the same 4 GPUs as one co-located EP pool** (`n_ranks=4`)
+  running both phases through a single continuous-batching loop. Isolates
+  disaggregation's effect by holding GPU count constant between baseline
+  and treatment.
+- **Budget cap: $40**, set before rental, deliberately treated as a
+  ceiling, not a target -- see §7.
+
+## 3. Architecture
+
+Four ranks, one `torch.distributed` NCCL world (`torchrun
+--nproc_per_node=4`), split into two disjoint DeepEP EP sub-groups:
+
+- **Prefill pool** -- ranks 0-1, `patch_moe_infer_ep(..., n_ranks=2)`,
+  its own `Buffer`.
+- **Decode pool** -- ranks 2-3, a second, independent
+  `patch_moe_infer_ep(..., n_ranks=2)` and its own `Buffer`.
+
+Each pool loads a full copy of DeepSeekMoE-16B (unavoidable -- they are
+separate GPU groups with separate EP shards of the routed experts).
+Attention, the gate, embeddings, and shared experts stay replicated on
+every rank, exactly as in Phase 3 -- only the routed-expert compute is
+EP-sharded, independently within each pool.
+
+The co-located baseline reuses the same four ranks as a single EP pool
+(`n_ranks=4`), no prefill/decode split, one scheduler loop doing both.
+
+## 4. Components
+
+- **`PrefillWorker`** (ranks 0-1) -- loop: drain up to `prefill_batch_size`
+  waiting requests, batch and pad their prompts, one
+  `model(input_ids=..., past_key_values=None, use_cache=True)` call across
+  both ranks (EP dispatch/combine inside, unchanged from Phase 3), slice
+  the batched `past_key_values` back into per-request caches, compute
+  each request's first token, enqueue `(request_id, kv_cache, first_token)`
+  onto the handoff queue.
+- **`DecodeWorker`** (ranks 2-3) -- loop: admit anything waiting on the
+  handoff queue into its live batch (up to `decode_batch_size`), run one
+  batched decode step, drop any request that hit EOS or `max_new_tokens`,
+  record its completion timing.
+- **KV-cache handoff** -- EP shards only the MoE FFN, not attention, so
+  both ranks in the prefill pool hold identical KV cache after prefill.
+  Rank 0 -> rank 2 and rank 1 -> rank 3 each do a direct NCCL `send`/`recv`
+  of that request's cache tensors; no broadcast step needed.
+- **Load driver** -- issues a configurable number of concurrent requests
+  with staggered arrivals (not a single simultaneous burst) against
+  whichever mode (co-located or disaggregated) is under test, and records
+  per-request TTFT, inter-token latency, and completion time.
+- **Co-located baseline worker** -- one loop, four ranks, doing what
+  `PrefillWorker` and `DecodeWorker` do separately: admit new requests'
+  prefill inline, then step every in-flight request's decode by one
+  token, every iteration, same 4 GPUs.
+
+New file: `src/dispatch/serving/disaggregated.py` (workers, handoff
+queue, scheduler). New file: `src/dispatch/serving/colocated.py` (the
+baseline loop). Neither changes `grouped_moe_routed`, `expert_parallel.py`,
+or Phase 3's EP integration -- both are new orchestration around the
+existing, already-proven compute path.
+
+## 5. Data flow
+
+Per request, disaggregated mode:
+
+1. Request arrives at the load driver, is pushed to the prefill pool's
+   incoming queue. TTFT clock starts.
+2. `PrefillWorker` admits it into the next prefill batch (padded,
+   `attention_mask`-aware), runs one forward pass across ranks 0-1
+   (identical EP dispatch/combine path to Phase 3), produces the first
+   token and a per-request KV cache.
+3. KV cache moves rank 0 -> rank 2 and rank 1 -> rank 3 via NCCL
+   `send`/`recv`. First token and cache land on the decode pool.
+4. `DecodeWorker` admits the request into its live batch on its next
+   iteration. Every iteration thereafter, it steps one token for every
+   active request (including this one) until EOS or `max_new_tokens`.
+5. Completion timing recorded by the load driver.
+
+Co-located mode: the same five steps happen on the same 4 ranks inside
+one loop, with prefill and decode work interleaved on the same GPUs
+instead of split across pools -- no handoff step, no separate queue.
+
+## 6. A known simplification, disclosed up front
+
+Continuous batching here uses plain dense, padded tensors, not a real
+paged/block KV cache (PagedAttention). Every decode step re-pads all
+active requests' caches to the batch's current max length -- real wasted
+compute next to a production engine, and deliberately out of scope to
+build (a paged cache is most of what makes vLLM's own engine hard).
+It applies identically to the co-located baseline and the disaggregated
+treatment, so it should not bias the *relative* comparison this phase
+measures -- but absolute throughput/TTFT numbers here are not directly
+comparable to vLLM's own, and the findings doc must say so explicitly,
+same as Phase 3's H200-substitution and V1-API disclosures.
+
+## 7. Testing
+
+Extends CLAUDE.md's existing testing table:
+
+| Layer | What's covered |
+|---|---|
+| Scheduler logic (CPU, no GPU) | Admission/eviction against a fake model call: batch-size limits, request lifecycle, queue draining. No real tensors needed. |
+| KV-cache slice/reassemble (CPU, no GPU) | Splitting a batched `past_key_values` into per-request caches and back, against small CPU tensors and a tiny real HF model -- proves the bookkeeping before any GPU is involved. |
+| Handoff mechanism (CPU, no GPU, no DeepEP) | `torch.distributed`'s `gloo` backend supports `send`/`recv` -- 4 CPU processes prove the right cache reaches the right rank before DeepEP or real GPUs enter the picture at all. |
+| Batched correctness (`gpu`) | Batched prefill/decode output (padding, attention masking) matches Phase 0's single-request, unbatched reference for the same prompt. |
+| Full pipeline correctness (`gpu`, paid) | Real 4-rank run, both EP pools live: mutual top-k logit agreement between the disaggregated path and the single-GPU reference, same bar Phase 1 and Phase 3 used. **Gates the benchmark** -- a run that fails correctness is not benchmarked. |
+| Contention measurement (`gpu`, paid) | Co-located vs. disaggregated, same 4 GPUs, a couple of concurrency levels (e.g. 4 and 8 concurrent requests). Every number carries its full config. |
+
+## 8. Risk, cost, and rollout
+
+- **Cost discipline is itself part of this phase's result, not just a
+  constraint on it.** $40 is a ceiling, not a target -- the goal is the
+  smallest, most surgical paid session that still produces a trustworthy
+  measurement, and the gap between the cap and the actual spend is worth
+  reporting alongside the throughput numbers, the way Phase 3 reported
+  $13.03 of its $25 cap.
+- **Maximize what's proven before any rental.** Every layer in §7 above
+  the `gpu` line is real, CPU-only, and can be fully debugged for $0 --
+  including the NCCL handoff mechanism itself via `gloo`. The paid
+  session should start from scheduler + handoff + batching logic already
+  green, and spend rented time only on what genuinely requires DeepEP and
+  real multi-GPU hardware: EP correctness, then the measurement.
+- **Correctness gate before any throughput claim**, same discipline as
+  every prior phase -- a disaggregated run that disagrees with the
+  single-GPU reference does not get benchmarked, it gets debugged or the
+  phase stops and reports why.
+- **Session order**: confirm the CPU-testable suite (scheduler, cache
+  slicing, `gloo` handoff) is green -> rent the 4-GPU node -> verify real
+  NVLink across all 4 GPUs (`nvidia-smi topo -m`, same check Phase 3 used)
+  -> correctness gate on both the co-located and disaggregated paths ->
+  only if both pass, the concurrency-level measurement -> tear down
+  immediately.
+- **Named risk carried over from Phase 3**: DeepEP V1 (not V2) is the
+  confirmed-working API on this project's rental tier (no GPU Fabric
+  Manager available). Phase 4 should assume V1 from the start rather than
+  re-discovering this -- two independent `Buffer` instances, one per EP
+  pool, same as Phase 3's single pool.
+- **Named risk, new to this phase**: running two independent 2-rank EP
+  pools plus cross-pool NCCL `send`/`recv` in the same 4-rank world is
+  more process-group plumbing than Phase 3 needed. If DeepEP's `Buffer`
+  or process-group setup conflicts across sub-groups in a way that isn't
+  resolvable within budget, that is itself a real, reportable finding --
+  not a reason to force a workaround that muddies the correctness gate.
+- **GPU class**: Hopper-class (H100/H200, SM90), same DeepEP V1
+  requirement as Phase 3. Exact type/cloud chosen live at rental time
+  against real-time availability, same practice as Phases 2-3.
+
+## 9. Non-goals
+
+- Cross-node networking, real distributed KV-cache transfer -- one
+  NVLink node only (§2).
+- Paged/block KV cache, priority scheduling, preemption, dynamic
+  rebalancing between pools -- fixed 2+2 split, dense padded batching
+  (§6).
+- Beating vLLM/SGLang's real numbers -- Phase 6 owns that comparison.
+  Phase 4's claim is limited to "disaggregation helps this project's own
+  co-located baseline," an internal, controlled comparison.
+- Quantization and speculative decoding -- Phase 5.

docs/findings/2026-09-16-phase-4-colocated-c4.json

@@ -0,0 +1,10 @@
+{
+  "topology": "colocated",
+  "concurrency": 4,
+  "max_new_tokens": 8,
+  "total_wall_s": 4.310341304168105,
+  "mean_ttft_s": 1.117574252653867,
+  "p50_ttft_s": 1.0292676258832216,
+  "p99_ttft_s": 1.371898397002369,
+  "decode_tokens_per_s": 6.496005310049106
+}

docs/findings/2026-09-16-phase-4-colocated-c8.json

@@ -0,0 +1,10 @@
+{
+  "topology": "colocated",
+  "concurrency": 8,
+  "max_new_tokens": 8,
+  "total_wall_s": 2.974014701321721,
+  "mean_ttft_s": 0.6953176653478295,
+  "p50_ttft_s": 0.9301474057137966,
+  "p99_ttft_s": 1.289631819501519,
+  "decode_tokens_per_s": 18.829765695210686
+}

docs/findings/2026-09-16-phase-4-colocated-gate-results.json

@@ -0,0 +1,59 @@
+{
+  "results": {
+    "The quick brown fox jumps over the lazy dog.": [
+      185,
+      549,
+      3399,
+      10176,
+      32431,
+      33747,
+      855,
+      254,
+      24547,
+      5025,
+      13,
+      185,
+      549,
+      3399,
+      10176,
+      32431
+    ],
+    "In a distant galaxy, a small crew of explorers": [
+      317,
+      331,
+      245,
+      8723,
+      276,
+      1275,
+      245,
+      761,
+      1719,
+      327,
+      704,
+      1245,
+      13,
+      1955,
+      463,
+      32339
+    ],
+    "def fibonacci(n):": [
+      185,
+      300,
+      565,
+      291,
+      2318,
+      207,
+      15,
+      25,
+      185,
+      391,
+      972,
+      207,
+      15,
+      185,
+      300,
+      23744
+    ]
+  },
+  "match": true
+}