From a9ab18c7984bc5cb18e20b95616d792ff0a2299b Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:18:02 -0400 Subject: [PATCH 01/12] docs: Phase 4 design -- disaggregated prefill/decode --- ...16-phase-4-disaggregated-prefill-decode.md | 204 ++++++++++++++++++ 1 file changed, 204 insertions(+) create mode 100644 docs/design/2026-09-16-phase-4-disaggregated-prefill-decode.md diff --git a/docs/design/2026-09-16-phase-4-disaggregated-prefill-decode.md b/docs/design/2026-09-16-phase-4-disaggregated-prefill-decode.md new file mode 100644 index 0000000..f07574c --- /dev/null +++ b/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. From 971ad91214c07d4928ee0c1c769ea6b29dfd8178 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:29:17 -0400 Subject: [PATCH 02/12] docs: Phase 4 implementation plan -- disaggregated prefill/decode --- ...ase-4-disaggregated-prefill-decode-plan.md | 1496 +++++++++++++++++ 1 file changed, 1496 insertions(+) create mode 100644 docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md diff --git a/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md b/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md new file mode 100644 index 0000000..a37bd52 --- /dev/null +++ b/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md @@ -0,0 +1,1496 @@ +# Phase 4: Disaggregated Prefill/Decode Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Prove, on real hardware, whether splitting prefill and decode +onto separate GPU pools relieves the contention they create when they +share GPUs under concurrent load -- holding total GPU count fixed at 4, +comparing one co-located 4-rank EP pool against two independent 2-rank EP +pools (prefill, decode) connected by a KV-cache handoff. + +**Architecture:** A continuous-batching scheduler (`PrefillWorker`, +`DecodeWorker` in `disaggregated.py`; `ColocatedWorker` in +`colocated.py`) built entirely against a pluggable `PrefillFn`/`DecodeFn` +callable pair -- the scheduler knows nothing about DeepEP, EP pools, or +even that this is a distributed setting. `kv_cache.py` supplies the +slice/pad/rebatch bookkeeping a dense (non-paged) continuous batch needs; +`handoff.py` supplies the cross-rank cache transfer over plain +`torch.distributed` send/recv, provably correct over CPU-only `gloo` +before it ever needs NCCL or real GPUs. Only the final, paid task wires +real closures around DeepSeekMoE-16B -- patched with Phase 3's existing, +unchanged `make_ep_moe_infer`/`patch_moe_infer_ep` -- into the +`PrefillFn`/`DecodeFn` shape. + +**Tech Stack:** `transformers` v5's `DynamicCache` (its modern, surviving +public API only -- `to_legacy_cache`/`from_legacy_cache`/`key_cache`/ +`value_cache` were removed in v5, confirmed live 2026-09-16 against +`transformers`' own v5 migration guide and `cache_utils.py` source). +`torch.distributed` (`gloo` for CPU tests, `NCCL` for the real run). +DeepEP V1 (legacy `Buffer`), Phase 3's `expert_parallel.py`, reused +unchanged. RunPod's REST API via `scripts/gpu/provision.py` +(`--gpu-count 4`, already generic from Phase 3's Task 1 -- no new +provisioning code needed). + +**Spec:** `docs/design/2026-09-16-phase-4-disaggregated-prefill-decode.md` +(all sections); `docs/design/2026-09-15-phase-3-multi-gpu-expert-parallel-serving.md` +and `src/dispatch/kernels/expert_parallel.py` (Phase 3's EP machinery, +reused unchanged); `src/dispatch/benchmark/harness.py`, +`src/dispatch/benchmark/reference.py` (Phase 0/1's generation and +correctness-gate helpers, reused unchanged). + +## Global Constraints + +- **Budget cap: $40, a ceiling not a target.** Be surgical: maximize + what's proven for $0 before any rental (design doc section 8), and the + gap between the cap and the actual spend is itself a reportable part of + Phase 4's result (Task 6), the same way Phase 3 reported $13.03 of its + $25 cap. +- **Hardware: Hopper-class (H100/H200, SM90), one 4-GPU NVLink/SXM node.** + Verify real NVLink across *all four* GPUs before spending budget on + anything else -- `nvidia-smi topo -m` must show `NV#` between every + pair, not just GPU 0-1. +- **DeepEP V1 (legacy `Buffer`), not V2.** Phase 3 confirmed live + (2026-09-16) that V2's `ElasticBuffer` cannot work on this rental tier + -- no GPU Fabric Manager, `nvidia-smi -q` reports `GPU Fabric GUID: + N/A`. Do not re-attempt V2; start Task 5 from V1 directly. +- **`transformers` v5's `DynamicCache` API, confirmed live 2026-09-16** + against `transformers`' own v5 migration guide and + `src/transformers/cache_utils.py` (`main` branch): `to_legacy_cache`, + `from_legacy_cache`, and the `key_cache`/`value_cache` list attributes + were **removed** in v5. Use only the surviving public API this plan's + code is built on: `cache.layers[i].keys` / `.values` (tensors shaped + `(batch, num_heads, seq_len, head_dim)`), `DynamicCache(ddp_cache_data= + [(key, value), ...])` to construct one from raw per-layer tensors, + `cache.get_seq_length()`, `cache.batch_size`. Do not use + `to_legacy_cache`/`from_legacy_cache`/`key_cache`/`value_cache` -- + they will raise `AttributeError` on this repo's pinned + `transformers>=5.17.0`. +- **Dense, padded batching -- not a paged/block KV cache.** Every batched + decode step re-pads every active request's cache to the batch's + current max real length (design doc section 6). This applies + identically to the co-located and disaggregated paths, so it should not + bias their *relative* comparison, but it does mean absolute + throughput/TTFT numbers here are not directly comparable to vLLM's own + -- say so explicitly in Task 6's findings doc. +- **Left-padding convention.** Every batched forward call (prefill or + decode) left-pads shorter sequences -- matches this repo's existing + single-request convention implicitly (no prior padding existed before + Phase 4) and HF's own default for batched causal-LM generation. +- **Reuse Phase 3's EP machinery unchanged.** `make_ep_moe_infer` and + `patch_moe_infer_ep` (`src/dispatch/kernels/expert_parallel.py`) are + not modified by this plan. Task 5 calls them twice -- once per EP + sub-group (prefill ranks 0-1, decode ranks 2-3) -- exactly as Phase 3 + called them once. +- **Correctness gates every throughput/contention claim.** Both the + co-located and disaggregated paths must independently pass mutual + top-k logit agreement against the single-GPU reference (same bar Phase + 1 and Phase 3 used) before Task 5's concurrency measurement runs. +- **Never quote a benchmark number that wasn't measured** on this exact + run (repo-wide rule). + +## File Structure + +```text +src/dispatch/serving/__init__.py # Task 1 +src/dispatch/serving/kv_cache.py # Task 1 (slice_cache, pad_and_batch_caches) +tests/unit/test_kv_cache.py # Task 1 + +src/dispatch/serving/handoff.py # Task 2 (send_kv_cache, recv_kv_cache) +tests/unit/test_handoff.py # Task 2 + +src/dispatch/serving/disaggregated.py # Task 3 (Request, PrefillResult, + # RequestResult, InFlightRequest, + # PrefillFn, DecodeFn, + # run_prefill_batch, + # prefill_result_to_in_flight, + # step_active_requests, + # PrefillWorker, DecodeWorker) +tests/unit/test_disaggregated.py # Task 3 + +src/dispatch/serving/colocated.py # Task 4 (ColocatedWorker) +tests/unit/test_colocated.py # Task 4 + +docs/runbooks/phase-4-disaggregated-prefill-decode.md # Task 5 +docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-cost.md # Task 5 (write_cost_record) +docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-run.md # Task 6 +docs/STATUS.md # Task 6 +``` + +--- + +### Task 1: KV-cache slice and pad-batch helpers, proven against a real tiny model on CPU + +**Files:** + +- Create: `src/dispatch/serving/__init__.py` (empty, matches + `src/dispatch/{benchmark,kernels}/__init__.py`) +- Create: `src/dispatch/serving/kv_cache.py` +- Test: `tests/unit/test_kv_cache.py` + +**Interfaces:** + +- Produces: `slice_cache(cache: DynamicCache, index: int, *, keep_last: + int | None = None) -> DynamicCache`; `pad_and_batch_caches(caches: + list[DynamicCache], pad_to: int) -> tuple[DynamicCache, torch.Tensor]`. + Tasks 3-4 build their entire batching loop on these two functions and + nothing else from this file. + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_kv_cache.py`: + +```python +"""CPU-only: proves the dense-padding slice/rebatch bookkeeping a +continuous-batching scheduler needs, built entirely on transformers v5's +surviving public DynamicCache API (cache.layers[i].keys/.values, the +ddp_cache_data constructor path) -- to_legacy_cache/from_legacy_cache and +the old key_cache/value_cache attributes were removed in v5 (confirmed +live 2026-09-16 against transformers' own v5 migration guide). +""" + +from __future__ import annotations + +import pytest +import torch +from transformers import DynamicCache + +from dispatch.serving.kv_cache import pad_and_batch_caches, slice_cache + +NUM_LAYERS = 2 +NUM_HEADS = 2 +HEAD_DIM = 4 + + +def _make_cache(batch_size: int, seq_len: int, *, seed: int = 0) -> DynamicCache: + generator = torch.Generator().manual_seed(seed) + per_layer = [ + ( + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM, generator=generator), + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM, generator=generator), + ) + for _ in range(NUM_LAYERS) + ] + return DynamicCache(ddp_cache_data=per_layer) + + +def test_slice_cache_extracts_one_request_without_mutating_the_original() -> None: + cache = _make_cache(batch_size=2, seq_len=3) + original_row_0 = cache.layers[0].keys[0].clone() + + sliced = slice_cache(cache, index=1) + + assert sliced.get_seq_length() == 3 + assert sliced.batch_size == 1 + torch.testing.assert_close(sliced.layers[0].keys[0], cache.layers[0].keys[1]) + torch.testing.assert_close(cache.layers[0].keys[0], original_row_0) + + +def test_slice_cache_with_keep_last_drops_left_padding() -> None: + cache = _make_cache(batch_size=1, seq_len=5) + real_tail = cache.layers[0].keys[:, :, -2:, :].clone() + + sliced = slice_cache(cache, index=0, keep_last=2) + + assert sliced.get_seq_length() == 2 + torch.testing.assert_close(sliced.layers[0].keys, real_tail) + + +def test_pad_and_batch_caches_left_pads_to_a_common_length_and_masks_correctly() -> None: + short = _make_cache(batch_size=1, seq_len=2, seed=1) + long_ = _make_cache(batch_size=1, seq_len=5, seed=2) + + batched, attention_mask = pad_and_batch_caches([short, long_], pad_to=5) + + assert batched.get_seq_length() == 5 + assert batched.batch_size == 2 + assert attention_mask.tolist() == [[0, 0, 0, 1, 1], [1, 1, 1, 1, 1]] + torch.testing.assert_close(batched.layers[0].keys[0, :, -2:, :], short.layers[0].keys[0]) + torch.testing.assert_close(batched.layers[0].keys[1], long_.layers[0].keys[0]) + + +def test_pad_and_batch_caches_rejects_a_cache_longer_than_pad_to() -> None: + long_ = _make_cache(batch_size=1, seq_len=5) + + with pytest.raises(ValueError, match="exceeds"): + pad_and_batch_caches([long_], pad_to=3) + + +def test_pad_and_batch_caches_rejects_an_empty_list() -> None: + with pytest.raises(ValueError, match="empty"): + pad_and_batch_caches([], pad_to=3) + + +@pytest.mark.slow +def test_slice_and_rebatch_round_trip_through_a_real_tiny_model() -> None: + """Proves the helpers produce a cache the real model accepts back and + continues generating from correctly -- against hf-internal-testing/ + tiny-random-gpt2, this repo's established tiny-model pattern + (test_harness.py). + """ + from dispatch.benchmark.harness import load_model + + model, tokenizer = load_model("hf-internal-testing/tiny-random-gpt2") + input_ids = tokenizer("hello world", return_tensors="pt").input_ids + with torch.no_grad(): + outputs = model(input_ids=input_ids, use_cache=True) + cache = outputs.past_key_values + assert isinstance(cache, DynamicCache) + next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) + + sliced = slice_cache(cache, index=0) + assert sliced.get_seq_length() == cache.get_seq_length() + + with torch.no_grad(): + continued = model(input_ids=next_token, past_key_values=sliced, use_cache=True) + + assert continued.logits.shape == (1, 1, continued.logits.shape[-1]) +``` + +- [ ] **Step 2: Run to verify it fails** + +Run: `uv run pytest tests/unit/test_kv_cache.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'dispatch.serving'`. + +- [ ] **Step 3: Write the implementation** + +Create `src/dispatch/serving/__init__.py` (empty file). + +Create `src/dispatch/serving/kv_cache.py`: + +```python +"""KV-cache slicing and batch (re)assembly for Phase 4's continuous- +batching scheduler. Built entirely on transformers v5's surviving public +DynamicCache API -- to_legacy_cache/from_legacy_cache and the old +key_cache/value_cache attributes were removed in v5 (confirmed live +2026-09-16 against transformers' own v5 migration guide and +src/transformers/cache_utils.py). Dense, padded batching only -- not a +paged/block cache -- the known simplification this project's Phase 4 +design doc discloses (section 6): real wasted compute next to a +production engine, applied identically to every configuration this +project measures, so it should not bias their relative comparison. +""" + +from __future__ import annotations + +import torch +from transformers import DynamicCache + + +def slice_cache(cache: DynamicCache, index: int, *, keep_last: int | None = None) -> DynamicCache: + """Extract one request's KV cache (by batch index) out of a batched + cache, without mutating the original -- unlike Cache.batch_select_indices, + which selects in place and offers no way to also drop padding. If + keep_last is given, keeps only the last keep_last positions along the + sequence dim, dropping this request's left-padding after a batched + forward call where its real content is shorter than the batch's + padded length. + """ + per_layer = [] + for layer in cache.layers: + key = layer.keys[index : index + 1] + value = layer.values[index : index + 1] + if keep_last is not None: + key = key[:, :, -keep_last:, :] + value = value[:, :, -keep_last:, :] + per_layer.append((key.clone(), value.clone())) + return DynamicCache(ddp_cache_data=per_layer) + + +def pad_and_batch_caches( + caches: list[DynamicCache], pad_to: int +) -> tuple[DynamicCache, torch.Tensor]: + """Left-pads each cache's key/value tensors to pad_to positions and + concatenates them along the batch dim, returning the batched cache + and a matching (batch, pad_to) attention_mask (0 = pad, 1 = real). + """ + if not caches: + raise ValueError("caches must not be empty") + num_layers = len(caches[0].layers) + per_layer_keys: list[list[torch.Tensor]] = [[] for _ in range(num_layers)] + per_layer_values: list[list[torch.Tensor]] = [[] for _ in range(num_layers)] + attention_mask = torch.zeros(len(caches), pad_to, dtype=torch.long) + + for row, cache in enumerate(caches): + seq_len = cache.get_seq_length() + if seq_len > pad_to: + raise ValueError(f"cache seq_len={seq_len} exceeds pad_to={pad_to}") + attention_mask[row, pad_to - seq_len :] = 1 + pad_len = pad_to - seq_len + for layer_idx, layer in enumerate(cache.layers): + key, value = layer.keys, layer.values + if pad_len > 0: + pad_shape = (1, key.shape[1], pad_len, key.shape[3]) + key = torch.cat([torch.zeros(pad_shape, dtype=key.dtype), key], dim=2) + value = torch.cat([torch.zeros(pad_shape, dtype=value.dtype), value], dim=2) + per_layer_keys[layer_idx].append(key) + per_layer_values[layer_idx].append(value) + + per_layer = [ + (torch.cat(per_layer_keys[i], dim=0), torch.cat(per_layer_values[i], dim=0)) + for i in range(num_layers) + ] + return DynamicCache(ddp_cache_data=per_layer), attention_mask +``` + +- [ ] **Step 4: Run to verify the fast tests pass** + +Run: `uv run pytest tests/unit/test_kv_cache.py -v -m "not slow"` +Expected: 5 pass (the `@pytest.mark.slow` real-model test is excluded). + +- [ ] **Step 5: Run the slow real-model test** + +Run: `uv run pytest tests/unit/test_kv_cache.py -v -m slow` +Expected: PASS (downloads `hf-internal-testing/tiny-random-gpt2` on first +run, network required, no GPU). + +- [ ] **Step 6: Lint and typecheck** + +Run: `uv run ruff check src/dispatch/serving tests/unit/test_kv_cache.py && uv run ruff format --check src/dispatch/serving tests/unit/test_kv_cache.py && uv run mypy src tests` +Expected: clean. + +- [ ] **Step 7: Commit** + +```bash +git add src/dispatch/serving/__init__.py src/dispatch/serving/kv_cache.py tests/unit/test_kv_cache.py +git commit -m "feat: add KV-cache slice/pad-batch helpers for continuous batching" +``` + +--- + +### Task 2: Cross-rank KV-cache handoff, proven over CPU-only gloo before any GPU is involved + +**Files:** + +- Create: `src/dispatch/serving/handoff.py` +- Test: `tests/unit/test_handoff.py` + +**Interfaces:** + +- Produces: `send_kv_cache(cache: DynamicCache, dst: int, group: + dist.ProcessGroup) -> None`; `recv_kv_cache(src: int, group: + dist.ProcessGroup, *, num_layers: int, num_heads: int, head_dim: int, + dtype: torch.dtype) -> DynamicCache`. Task 5 calls these exact + functions unchanged, with `backend="nccl"` instead of `"gloo"` and real + ranks instead of a CPU 2-process test. + +- [ ] **Step 1: Write the failing test** + +Create `tests/unit/test_handoff.py`: + +```python +"""CPU-only, multi-process: proves send_kv_cache/recv_kv_cache round-trip +correctly using torch.distributed's gloo backend -- no GPU, no DeepEP, no +NCCL. Task 5 reuses these exact functions unchanged with the NCCL +backend on the real rented 4-GPU node; only init_process_group's backend +argument and the real ranks differ. +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from transformers import DynamicCache + +from dispatch.serving.handoff import recv_kv_cache, send_kv_cache + +NUM_LAYERS = 2 +NUM_HEADS = 2 +HEAD_DIM = 4 +SEQ_LEN = 3 +PORT = 29513 + + +def _make_cache(seed: int) -> DynamicCache: + generator = torch.Generator().manual_seed(seed) + per_layer = [ + ( + torch.randn(1, NUM_HEADS, SEQ_LEN, HEAD_DIM, generator=generator), + torch.randn(1, NUM_HEADS, SEQ_LEN, HEAD_DIM, generator=generator), + ) + for _ in range(NUM_LAYERS) + ] + return DynamicCache(ddp_cache_data=per_layer) + + +def _worker(rank: int, world_size: int, result_path: str) -> None: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(PORT) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + group = dist.group.WORLD + + if rank == 0: + send_kv_cache(_make_cache(seed=42), dst=1, group=group) + else: + received = recv_kv_cache( + src=0, + group=group, + num_layers=NUM_LAYERS, + num_heads=NUM_HEADS, + head_dim=HEAD_DIM, + dtype=torch.float32, + ) + expected = _make_cache(seed=42) + ok = all( + torch.equal(received.layers[i].keys, expected.layers[i].keys) + and torch.equal(received.layers[i].values, expected.layers[i].values) + for i in range(NUM_LAYERS) + ) + torch.save({"ok": ok, "seq_len": received.get_seq_length()}, result_path) + + dist.destroy_process_group() + + +@pytest.mark.slow +def test_send_and_recv_kv_cache_round_trip_over_gloo(tmp_path: object) -> None: + result_path = str(tmp_path / "result.pt") # type: ignore[operator] + + mp.spawn(_worker, args=(2, result_path), nprocs=2, join=True) + + result = torch.load(result_path, weights_only=True) + assert result["ok"] is True + assert result["seq_len"] == SEQ_LEN +``` + +- [ ] **Step 2: Run to verify it fails** + +Run: `uv run pytest tests/unit/test_handoff.py -v -m slow` +Expected: FAIL with `ModuleNotFoundError: No module named 'dispatch.serving.handoff'`. + +- [ ] **Step 3: Write the implementation** + +Create `src/dispatch/serving/handoff.py`: + +```python +"""Cross-rank KV-cache handoff between the prefill and decode EP pools, +built on plain torch.distributed point-to-point send/recv so the exact +same code runs over gloo (CPU, this file's own tests) and NCCL (real +GPUs, Task 5) -- the transport logic gets proven once, for $0, before it +is ever run against real hardware (design doc section 8's "maximize +what's proven before any rental"). +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +from transformers import DynamicCache + + +def send_kv_cache(cache: DynamicCache, dst: int, group: dist.ProcessGroup) -> None: + seq_len = torch.tensor([cache.get_seq_length()], dtype=torch.int64) + dist.send(seq_len, dst=dst, group=group) + for layer in cache.layers: + dist.send(layer.keys.contiguous(), dst=dst, group=group) + dist.send(layer.values.contiguous(), dst=dst, group=group) + + +def recv_kv_cache( # noqa: PLR0913 -- the receiver can't infer shape/dtype from the wire, they must be passed + src: int, + group: dist.ProcessGroup, + *, + num_layers: int, + num_heads: int, + head_dim: int, + dtype: torch.dtype, +) -> DynamicCache: + seq_len_tensor = torch.zeros(1, dtype=torch.int64) + dist.recv(seq_len_tensor, src=src, group=group) + seq_len = int(seq_len_tensor.item()) + + per_layer = [] + for _ in range(num_layers): + key = torch.zeros(1, num_heads, seq_len, head_dim, dtype=dtype) + value = torch.zeros(1, num_heads, seq_len, head_dim, dtype=dtype) + dist.recv(key, src=src, group=group) + dist.recv(value, src=src, group=group) + per_layer.append((key, value)) + return DynamicCache(ddp_cache_data=per_layer) +``` + +- [ ] **Step 4: Run to verify it passes** + +Run: `uv run pytest tests/unit/test_handoff.py -v -m slow` +Expected: PASS. + +- [ ] **Step 5: Lint and typecheck** + +Run: `uv run ruff check src/dispatch/serving/handoff.py tests/unit/test_handoff.py && uv run ruff format --check src/dispatch/serving/handoff.py tests/unit/test_handoff.py && uv run mypy src tests` +Expected: clean. (Add `"src/dispatch/serving/handoff.py" = ["PLR0913"]` +to `pyproject.toml`'s `[tool.ruff.lint.per-file-ignores]` only if the +inline `# noqa: PLR0913` above doesn't satisfy ruff.) + +- [ ] **Step 6: Commit** + +```bash +git add src/dispatch/serving/handoff.py tests/unit/test_handoff.py +git commit -m "feat: add cross-rank KV-cache handoff, proven over gloo" +``` + +--- + +### Task 3: Continuous-batching prefill and decode workers + +**Files:** + +- Create: `src/dispatch/serving/disaggregated.py` +- Test: `tests/unit/test_disaggregated.py` + +**Interfaces:** + +- Consumes: `slice_cache`, `pad_and_batch_caches` (Task 1, unchanged). +- Produces: `Request`, `PrefillResult`, `RequestResult`, + `InFlightRequest` (dataclasses); `PrefillFn = Callable[[torch.Tensor, + torch.Tensor], tuple[torch.Tensor, DynamicCache]]`; `DecodeFn = + Callable[[torch.Tensor, DynamicCache, torch.Tensor, torch.Tensor], + tuple[torch.Tensor, DynamicCache]]`; `run_prefill_batch(prefill_fn, + batch, clock_fn) -> list[PrefillResult]`; + `prefill_result_to_in_flight(result) -> InFlightRequest`; + `step_active_requests(active, decode_fn, clock_fn) -> + tuple[list[InFlightRequest], list[RequestResult]]`; `PrefillWorker`; + `DecodeWorker`. Task 4's `ColocatedWorker` imports + `run_prefill_batch`, `prefill_result_to_in_flight`, + `step_active_requests`, and every dataclass/type alias from this file + unchanged. Task 5 wires real closures matching `PrefillFn`/`DecodeFn` + exactly; nothing in this file changes at that point. + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_disaggregated.py`: + +```python +"""CPU-only: proves the continuous-batching scheduler's admission, +padding, and completion bookkeeping against a fake prefill/decode +forward pass -- no GPU, no DeepEP, no real model. Task 5 wires the same +PrefillFn/DecodeFn shapes to the real DeepSeekMoE-16B model and DeepEP EP +pools; nothing here changes at that point. +""" + +from __future__ import annotations + +import torch +from transformers import DynamicCache + +from dispatch.serving.disaggregated import DecodeWorker, PrefillResult, PrefillWorker, Request + +NUM_LAYERS = 2 +NUM_HEADS = 2 +HEAD_DIM = 4 + + +def _fake_cache(batch_size: int, seq_len: int) -> DynamicCache: + per_layer = [ + ( + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM), + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM), + ) + for _ in range(NUM_LAYERS) + ] + return DynamicCache(ddp_cache_data=per_layer) + + +def _fake_prefill_fn( + input_ids: torch.Tensor, attention_mask: torch.Tensor +) -> tuple[torch.Tensor, DynamicCache]: + del attention_mask + batch_size, seq_len = input_ids.shape + first_tokens = torch.arange(batch_size) + 100 # distinctive, deterministic per row + return first_tokens, _fake_cache(batch_size, seq_len) + + +def _fake_decode_fn( + next_input_ids: torch.Tensor, + cache: DynamicCache, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, +) -> tuple[torch.Tensor, DynamicCache]: + del attention_mask, position_ids + batch_size = next_input_ids.shape[0] + new_seq_len = cache.get_seq_length() + 1 + next_tokens = next_input_ids.squeeze(1) + 1 # deterministic increment, easy to assert on + return next_tokens, _fake_cache(batch_size, new_seq_len) + + +def test_prefill_worker_drains_up_to_batch_size_and_produces_per_request_caches() -> None: + worker = PrefillWorker(_fake_prefill_fn, batch_size=2) + worker.submit(Request("a", torch.tensor([[1, 2, 3]]), max_new_tokens=4)) + worker.submit(Request("b", torch.tensor([[1, 2]]), max_new_tokens=4)) + worker.submit(Request("c", torch.tensor([[1]]), max_new_tokens=4)) + + results = worker.step() + + assert [r.request_id for r in results] == ["a", "b"] + assert results[0].cache.get_seq_length() == 3 # unpadded back to its own prompt length + assert results[1].cache.get_seq_length() == 2 + assert results[0].first_token_id == 100 + assert results[1].first_token_id == 101 + + remaining = worker.step() + assert [r.request_id for r in remaining] == ["c"] + + +def test_prefill_worker_step_with_nothing_waiting_returns_empty() -> None: + worker = PrefillWorker(_fake_prefill_fn, batch_size=2) + + assert worker.step() == [] + + +def test_decode_worker_admits_and_completes_requests_after_max_new_tokens() -> None: + worker = DecodeWorker(_fake_decode_fn, batch_size=4) + cache = _fake_cache(1, 3) + worker.admit( + PrefillResult("a", cache, first_token_id=5, ttft=0.1, max_new_tokens=2, eos_token_id=None) + ) + + first_step = worker.step() + assert first_step == [] # only 1 of 2 tokens generated so far + + second_step = worker.step() + assert len(second_step) == 1 + assert second_step[0].request_id == "a" + assert second_step[0].generated_token_ids == [5, 6] + + +def test_decode_worker_completes_a_request_early_on_eos() -> None: + worker = DecodeWorker(_fake_decode_fn, batch_size=4) + cache = _fake_cache(1, 3) + worker.admit( + PrefillResult("a", cache, first_token_id=5, ttft=0.1, max_new_tokens=10, eos_token_id=6) + ) + + results = worker.step() + + assert len(results) == 1 + assert results[0].generated_token_ids == [5, 6] + + +def test_decode_worker_batches_requests_with_different_cache_lengths() -> None: + worker = DecodeWorker(_fake_decode_fn, batch_size=4) + worker.admit( + PrefillResult( + "short", _fake_cache(1, 2), first_token_id=1, ttft=0.1, max_new_tokens=1, eos_token_id=None + ) + ) + worker.admit( + PrefillResult( + "long", _fake_cache(1, 5), first_token_id=1, ttft=0.1, max_new_tokens=1, eos_token_id=None + ) + ) + + results = worker.step() + + assert {r.request_id for r in results} == {"short", "long"} + + +def test_decode_worker_respects_batch_size_admitting_only_free_slots() -> None: + worker = DecodeWorker(_fake_decode_fn, batch_size=1) + worker.admit( + PrefillResult( + "a", _fake_cache(1, 2), first_token_id=1, ttft=0.1, max_new_tokens=5, eos_token_id=None + ) + ) + worker.admit( + PrefillResult( + "b", _fake_cache(1, 2), first_token_id=1, ttft=0.1, max_new_tokens=1, eos_token_id=None + ) + ) + + first_step = worker.step() + assert first_step == [] # "a" admitted, needs 5 tokens; "b" still waiting + + for _ in range(4): + worker.step() + a_result = worker.step() + assert [r.request_id for r in a_result] == ["a"] + + b_result = worker.step() # "b" only now gets its first decode step + assert b_result == [] + b_result = worker.step() + assert [r.request_id for r in b_result] == ["b"] +``` + +- [ ] **Step 2: Run to verify it fails** + +Run: `uv run pytest tests/unit/test_disaggregated.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'dispatch.serving.disaggregated'`. + +- [ ] **Step 3: Write the implementation** + +Create `src/dispatch/serving/disaggregated.py`: + +```python +"""Continuous-batching prefill and decode workers for Phase 4's +disaggregated serving path. Both workers are forward-pass-agnostic -- +they take a plain PrefillFn/DecodeFn callable and know nothing about +DeepEP, EP pools, or even that this is a distributed setting. Task 5 +wires the real DeepSeekMoE-16B model (patched with Phase 3's +make_ep_moe_infer/patch_moe_infer_ep, one EP pool per role) into exactly +these two callable shapes; nothing in this file changes when that +happens. run_prefill_batch, prefill_result_to_in_flight, and +step_active_requests are exported (not underscore-prefixed) because +colocated.py's ColocatedWorker reuses them unchanged -- a decode step is +a decode step regardless of whether prefill happened on the same GPU +this iteration or a different one. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from transformers import DynamicCache + +from dispatch.serving.kv_cache import pad_and_batch_caches, slice_cache + +PrefillFn = Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, DynamicCache]] +DecodeFn = Callable[ + [torch.Tensor, DynamicCache, torch.Tensor, torch.Tensor], tuple[torch.Tensor, DynamicCache] +] + + +@dataclass +class Request: + request_id: str + prompt_ids: torch.Tensor # (1, prompt_len) + max_new_tokens: int + eos_token_id: int | None = None + + +@dataclass +class PrefillResult: + request_id: str + cache: DynamicCache + first_token_id: int + ttft: float + max_new_tokens: int + eos_token_id: int | None + + +@dataclass +class RequestResult: + request_id: str + generated_token_ids: list[int] + ttft: float + completion_time: float + + +@dataclass +class InFlightRequest: + request_id: str + cache: DynamicCache + seq_len: int + next_input_id: torch.Tensor + generated_token_ids: list[int] + max_new_tokens: int + eos_token_id: int | None + ttft: float + + +def run_prefill_batch( + prefill_fn: PrefillFn, + batch: list[tuple[Request, float]], + clock_fn: Callable[[], float], +) -> list[PrefillResult]: + """Batches a list of (request, admitted_at) pairs into one padded + prefill_fn call and slices the result back into per-request, unpadded + PrefillResults. Shared by PrefillWorker.step() and colocated.py's + ColocatedWorker. + """ + requests = [r for r, _ in batch] + admitted_at = [t for _, t in batch] + max_len = max(r.prompt_ids.shape[1] for r in requests) + input_ids = torch.zeros(len(requests), max_len, dtype=torch.long) + attention_mask = torch.zeros(len(requests), max_len, dtype=torch.long) + for i, r in enumerate(requests): + prompt_len = r.prompt_ids.shape[1] + input_ids[i, max_len - prompt_len :] = r.prompt_ids[0] + attention_mask[i, max_len - prompt_len :] = 1 + + first_token_ids, cache = prefill_fn(input_ids, attention_mask) + now = clock_fn() + + results = [] + for i, r in enumerate(requests): + prompt_len = r.prompt_ids.shape[1] + results.append( + PrefillResult( + request_id=r.request_id, + cache=slice_cache(cache, i, keep_last=prompt_len), + first_token_id=int(first_token_ids[i].item()), + ttft=now - admitted_at[i], + max_new_tokens=r.max_new_tokens, + eos_token_id=r.eos_token_id, + ) + ) + return results + + +def prefill_result_to_in_flight(result: PrefillResult) -> InFlightRequest: + return InFlightRequest( + request_id=result.request_id, + cache=result.cache, + seq_len=result.cache.get_seq_length(), + next_input_id=torch.tensor([[result.first_token_id]], dtype=torch.long), + generated_token_ids=[result.first_token_id], + max_new_tokens=result.max_new_tokens, + eos_token_id=result.eos_token_id, + ttft=result.ttft, + ) + + +def step_active_requests( + active: list[InFlightRequest], decode_fn: DecodeFn, clock_fn: Callable[[], float] +) -> tuple[list[InFlightRequest], list[RequestResult]]: + """Pads every active request's cache to the batch's current max real + length, runs one batched decode_fn call, and unpads each result back + down to its own true length before storing it. Shared by + DecodeWorker.step() and colocated.py's ColocatedWorker.step(). + """ + if not active: + return [], [] + pad_to = max(r.seq_len for r in active) + batched_cache, attention_mask = pad_and_batch_caches([r.cache for r in active], pad_to) + full_attention_mask = torch.cat( + [attention_mask, torch.ones(len(active), 1, dtype=torch.long)], dim=1 + ) + position_ids = attention_mask.sum(dim=1, keepdim=True) + next_input_ids = torch.cat([r.next_input_id for r in active], dim=0) + + next_token_ids, updated_cache = decode_fn( + next_input_ids, batched_cache, full_attention_mask, position_ids + ) + now = clock_fn() + + still_active: list[InFlightRequest] = [] + completed: list[RequestResult] = [] + for i, r in enumerate(active): + token_id = int(next_token_ids[i].item()) + r.generated_token_ids.append(token_id) + r.seq_len += 1 + r.cache = slice_cache(updated_cache, i, keep_last=r.seq_len) + r.next_input_id = torch.tensor([[token_id]], dtype=torch.long) + done = len(r.generated_token_ids) >= r.max_new_tokens or ( + r.eos_token_id is not None and token_id == r.eos_token_id + ) + if done: + completed.append( + RequestResult( + request_id=r.request_id, + generated_token_ids=r.generated_token_ids, + ttft=r.ttft, + completion_time=now, + ) + ) + else: + still_active.append(r) + return still_active, completed + + +class PrefillWorker: + """Drains up to batch_size waiting requests per step, runs one + batched prefill forward pass, and returns each request's own + (unpadded) KV cache and first generated token -- ready to hand off to + a DecodeWorker (in-process admit() in CPU tests, handoff.py's + send_kv_cache on the real rented node). + """ + + def __init__( + self, + prefill_fn: PrefillFn, + batch_size: int, + *, + clock_fn: Callable[[], float] = time.perf_counter, + ) -> None: + self._prefill_fn = prefill_fn + self._batch_size = batch_size + self._clock_fn = clock_fn + self._waiting: list[tuple[Request, float]] = [] + + def submit(self, request: Request) -> None: + self._waiting.append((request, self._clock_fn())) + + def step(self) -> list[PrefillResult]: + batch = self._waiting[: self._batch_size] + self._waiting = self._waiting[self._batch_size :] + if not batch: + return [] + return run_prefill_batch(self._prefill_fn, batch, self._clock_fn) + + +class DecodeWorker: + """Admits PrefillResults into a live, continuously-batched decode + loop: every step, any waiting request gets admitted (up to + batch_size), then every active request advances one token. + """ + + def __init__( + self, + decode_fn: DecodeFn, + batch_size: int, + *, + clock_fn: Callable[[], float] = time.perf_counter, + ) -> None: + self._decode_fn = decode_fn + self._batch_size = batch_size + self._clock_fn = clock_fn + self._waiting: list[InFlightRequest] = [] + self._active: list[InFlightRequest] = [] + + def admit(self, result: PrefillResult) -> None: + self._waiting.append(prefill_result_to_in_flight(result)) + + def step(self) -> list[RequestResult]: + free_slots = self._batch_size - len(self._active) + if free_slots > 0 and self._waiting: + self._active.extend(self._waiting[:free_slots]) + self._waiting = self._waiting[free_slots:] + self._active, completed = step_active_requests( + self._active, self._decode_fn, self._clock_fn + ) + return completed +``` + +- [ ] **Step 4: Run to verify it passes** + +Run: `uv run pytest tests/unit/test_disaggregated.py -v` +Expected: all pass. + +- [ ] **Step 5: Lint and typecheck** + +Run: `uv run ruff check src/dispatch/serving/disaggregated.py tests/unit/test_disaggregated.py && uv run ruff format --check src/dispatch/serving/disaggregated.py tests/unit/test_disaggregated.py && uv run mypy src tests` +Expected: clean. + +- [ ] **Step 6: Commit** + +```bash +git add src/dispatch/serving/disaggregated.py tests/unit/test_disaggregated.py +git commit -m "feat: add continuous-batching prefill and decode workers" +``` + +--- + +### Task 4: Co-located baseline worker + +**Files:** + +- Create: `src/dispatch/serving/colocated.py` +- Test: `tests/unit/test_colocated.py` + +**Interfaces:** + +- Consumes: `Request`, `RequestResult`, `InFlightRequest`, `PrefillFn`, + `DecodeFn`, `run_prefill_batch`, `prefill_result_to_in_flight`, + `step_active_requests` (Task 3, unchanged). +- Produces: `ColocatedWorker`. Task 5 wires the same real `PrefillFn`/ + `DecodeFn` closures it built for Task 3's workers into this class too, + just with both roles talking to the same 4-rank EP pool instead of two + separate 2-rank pools. + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_colocated.py`: + +```python +"""CPU-only: proves ColocatedWorker interleaves inline prefill and +decode on the same loop -- new requests get their first token the same +step old ones advance, which is exactly what creates the contention +Phase 4 measures disaggregation against. Reuses test_disaggregated.py's +fake prefill/decode functions unchanged: the forward-pass contract is +identical, only which worker calls it differs. +""" + +from __future__ import annotations + +import torch +from transformers import DynamicCache + +from dispatch.serving.colocated import ColocatedWorker +from dispatch.serving.disaggregated import Request + +NUM_LAYERS = 2 +NUM_HEADS = 2 +HEAD_DIM = 4 + + +def _fake_cache(batch_size: int, seq_len: int) -> DynamicCache: + per_layer = [ + ( + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM), + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM), + ) + for _ in range(NUM_LAYERS) + ] + return DynamicCache(ddp_cache_data=per_layer) + + +def _fake_prefill_fn( + input_ids: torch.Tensor, attention_mask: torch.Tensor +) -> tuple[torch.Tensor, DynamicCache]: + del attention_mask + batch_size, seq_len = input_ids.shape + first_tokens = torch.arange(batch_size) + 100 + return first_tokens, _fake_cache(batch_size, seq_len) + + +def _fake_decode_fn( + next_input_ids: torch.Tensor, + cache: DynamicCache, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, +) -> tuple[torch.Tensor, DynamicCache]: + del attention_mask, position_ids + batch_size = next_input_ids.shape[0] + new_seq_len = cache.get_seq_length() + 1 + next_tokens = next_input_ids.squeeze(1) + 1 + return next_tokens, _fake_cache(batch_size, new_seq_len) + + +def test_colocated_worker_admits_via_inline_prefill_and_completes_after_max_new_tokens() -> None: + worker = ColocatedWorker(_fake_prefill_fn, _fake_decode_fn, batch_size=4) + worker.submit(Request("a", torch.tensor([[1, 2, 3]]), max_new_tokens=2)) + + first_step = worker.step() # prefill admits "a" and gives it token 1 of 2 + assert first_step == [] + + second_step = worker.step() # decode gives "a" its 2nd token -> complete + assert len(second_step) == 1 + assert second_step[0].request_id == "a" + assert len(second_step[0].generated_token_ids) == 2 + + +def test_colocated_worker_interleaves_new_prefill_with_existing_decode_in_one_step() -> None: + worker = ColocatedWorker(_fake_prefill_fn, _fake_decode_fn, batch_size=4) + worker.submit(Request("old", torch.tensor([[1, 2]]), max_new_tokens=3)) + worker.step() # "old" now active, generated 1 of 3 tokens + + worker.submit(Request("new", torch.tensor([[9]]), max_new_tokens=1)) + step_result = worker.step() # "new" gets prefilled AND decode-stepped this same call + + assert any(r.request_id == "new" for r in step_result) # "new" already done at max_new_tokens=1 + + +def test_colocated_worker_respects_batch_size_across_prefill_and_decode() -> None: + worker = ColocatedWorker(_fake_prefill_fn, _fake_decode_fn, batch_size=1) + worker.submit(Request("a", torch.tensor([[1]]), max_new_tokens=5)) + worker.submit(Request("b", torch.tensor([[2]]), max_new_tokens=1)) + + worker.step() # only "a" has a free slot + for _ in range(3): + worker.step() + a_done = worker.step() + assert [r.request_id for r in a_done] == ["a"] + + worker.step() # "b" now gets its prefill turn + b_done = worker.step() + assert [r.request_id for r in b_done] == ["b"] +``` + +- [ ] **Step 2: Run to verify it fails** + +Run: `uv run pytest tests/unit/test_colocated.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'dispatch.serving.colocated'`. + +- [ ] **Step 3: Write the implementation** + +Create `src/dispatch/serving/colocated.py`: + +```python +"""The co-located baseline: one worker, one EP pool, doing both prefill +and decode in a single loop on the same GPUs -- the contention baseline +Phase 4's disaggregated path (disaggregated.py) is measured against, on +the same 4 GPUs, so GPU count is never a confound (design doc section 1). +Reuses disaggregated.py's types and helpers unchanged: a decode step is a +decode step regardless of whether prefill happened on the same GPU this +iteration or a different one. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable + +from dispatch.serving.disaggregated import ( + DecodeFn, + InFlightRequest, + PrefillFn, + Request, + RequestResult, + prefill_result_to_in_flight, + run_prefill_batch, + step_active_requests, +) + + +class ColocatedWorker: + """Every step: admit waiting requests via an inline prefill batch (up + to free slots), then advance every active request -- new and old -- + by one decode token, all on the same GPU pool. + """ + + def __init__( + self, + prefill_fn: PrefillFn, + decode_fn: DecodeFn, + batch_size: int, + *, + clock_fn: Callable[[], float] = time.perf_counter, + ) -> None: + self._prefill_fn = prefill_fn + self._decode_fn = decode_fn + self._batch_size = batch_size + self._clock_fn = clock_fn + self._waiting: list[tuple[Request, float]] = [] + self._active: list[InFlightRequest] = [] + + def submit(self, request: Request) -> None: + self._waiting.append((request, self._clock_fn())) + + def step(self) -> list[RequestResult]: + free_slots = self._batch_size - len(self._active) + if free_slots > 0 and self._waiting: + batch = self._waiting[:free_slots] + self._waiting = self._waiting[free_slots:] + prefill_results = run_prefill_batch(self._prefill_fn, batch, self._clock_fn) + self._active.extend(prefill_result_to_in_flight(r) for r in prefill_results) + self._active, completed = step_active_requests( + self._active, self._decode_fn, self._clock_fn + ) + return completed +``` + +- [ ] **Step 4: Run to verify it passes** + +Run: `uv run pytest tests/unit/test_colocated.py -v` +Expected: all pass. + +- [ ] **Step 5: Lint and typecheck** + +Run: `uv run ruff check src/dispatch/serving/colocated.py tests/unit/test_colocated.py && uv run ruff format --check src/dispatch/serving/colocated.py tests/unit/test_colocated.py && uv run mypy src tests` +Expected: clean. + +- [ ] **Step 6: Run the full CPU-testable suite together** + +Run: `uv run pytest tests/unit -v -m "not gpu"` +Expected: all pass, including Task 1-4's new tests. This is the +"everything not requiring GPU/DeepEP is already green" checkpoint the +Global Constraints require before Task 5's rental begins. + +- [ ] **Step 7: Commit** + +```bash +git add src/dispatch/serving/colocated.py tests/unit/test_colocated.py +git commit -m "feat: add co-located prefill/decode baseline worker" +``` + +--- + +### Task 5: GPU rental runbook -- real EP pools, correctness gate, concurrency measurement + +**Files:** + +- Create: `docs/runbooks/phase-4-disaggregated-prefill-decode.md` +- Create (via `write_cost_record`, reused unmodified): + `docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-cost.md` +- Creates live, on the pod (not pre-committed -- its exact shape depends + on live-confirmed process-group/Buffer behavior, same practice Phase + 3's Task 4 used for `make_ep_moe_infer`): a load-driver script wiring + real `PrefillFn`/`DecodeFn` closures around DeepSeekMoE-16B, patched + twice with Phase 3's unchanged `make_ep_moe_infer`/`patch_moe_infer_ep` + (once per EP sub-group). + +**Budget cap: $40, a ceiling not a target -- be surgical, minimize actual +spend (Global Constraints). This is a live session with the user, not +something to run unattended -- get explicit go-ahead before renting.** + +- [ ] **Step 1: Write the runbook** + +Create `docs/runbooks/phase-4-disaggregated-prefill-decode.md`: + +````markdown +# Phase 4 runbook: disaggregated prefill/decode + +Budget cap: $40, a ceiling not a target -- be surgical. Hopper-class +(H100/H200, SM90) required, same as Phase 3. Check RunPod's real catalog +at rental time for 4-GPU availability and pricing. + +1. Create and wait for the 4-GPU pod, reusing Phase 3's generic + `--gpu-count` support (Task 1 of Phase 3, no new provisioning code + needed): + + ```bash + uv run python scripts/gpu/provision.py create --name phase-4-disaggregated \ + --gpu-type --gpu-count 4 \ + --image =12.3> --cloud --disk-gb 40 + uv run python scripts/gpu/provision.py wait + ``` + +2. **Verify real NVLink across all four GPUs** -- Phase 3's check + extended to a 4-GPU pair matrix: + + ```bash + nvidia-smi topo -m + ``` + + Expected: `NV#` between every pair of the four GPUs, not just 0-1. If + any pair shows `PHB`/`PXB`/`SYS`, stop, terminate the pod, and + re-provision -- do not spend budget against a partially-NVLink node. + +3. Install DeepEP V1 (legacy `Buffer`) directly -- do not attempt V2, + Phase 3 already confirmed it cannot work on this rental tier (no GPU + Fabric Manager). Reuse Phase 3's exact install steps + (`docs/runbooks/phase-3-multi-gpu-ep.md`, steps 3-4), adjusted to + import `deep_ep.Buffer` rather than `ElasticBuffer`. + +4. Clone this repo, sync, and run the full CPU-testable suite one more + time on the pod itself, as a final sanity check before spending + anything on the GPU-dependent work: + + ```bash + git clone dispatch + cd dispatch + uv sync + uv run pytest tests/unit -v -m "not gpu" + ``` + + Expected: all pass (same suite Task 4 step 6 already proved locally + -- this just confirms the pod's environment doesn't disagree). + +5. **Build the two EP sub-groups and wire the real forward closures.** + `torchrun --nproc_per_node=4` gives a 4-rank NCCL world; split it with + `dist.new_group([0, 1])` (prefill) and `dist.new_group([2, 3])` + (decode). Each sub-group builds its own DeepEP `Buffer` scoped to + that sub-group (same `Buffer(group, num_nvl_bytes, 0)` construction + Phase 3 used, just called twice) and calls + `patch_moe_infer_ep(model, matmul, buffer, rank, n_ranks=2)` + unchanged -- `rank` here is the sub-group-local rank + (`dist.get_rank(group=sub_group)`, which is 0 or 1), not the global + rank. Live-confirm this distinction against DeepEP's actual behavior + before trusting it -- if the printed/observed local ranks contradict + this, fix the code to match what's actually observed, the same + practice Phase 3's plan used for DeepEP's dispatch/combine shapes. + + Wrap the patched model's forward pass into `PrefillFn`/`DecodeFn` + (Task 3's exact type aliases), doing the greedy-argmax step inside the + closure -- matching `harness.py`'s existing + `outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)` convention: + + ```python + def make_prefill_fn(model, device): + def prefill_fn(input_ids, attention_mask): + outputs = model( + input_ids=input_ids.to(device), + attention_mask=attention_mask.to(device), + use_cache=True, + ) + first_tokens = outputs.logits[:, -1, :].argmax(dim=-1) + return first_tokens.cpu(), outputs.past_key_values + + return prefill_fn + + + def make_decode_fn(model, device): + def decode_fn(next_input_ids, cache, attention_mask, position_ids): + outputs = model( + input_ids=next_input_ids.to(device), + past_key_values=cache, + attention_mask=attention_mask.to(device), + position_ids=position_ids.to(device), + use_cache=True, + ) + next_tokens = outputs.logits[:, -1, :].argmax(dim=-1) + return next_tokens.cpu(), outputs.past_key_values + + return decode_fn + ``` + + For the disaggregated path, the prefill pool's rank-0-to-decode-rank-2 + and rank-1-to-decode-rank-3 handoff uses Task 2's `send_kv_cache`/ + `recv_kv_cache` unchanged, over the full 4-rank world group (not a + sub-group, since prefill and decode ranks must talk to each other). + +6. **Correctness gate -- must pass for both configurations before any + measurement runs.** Reuse `dispatch.benchmark.reference`'s + `capture_reference_logits`/`compare_top_k_agreement`/`save_reference`/ + `load_reference` unchanged (Phase 0/1's own functions). Get a + single-GPU reference first (rank 0 only, `patch_moe_infer` with the + naive kernel, `DEFAULT_PROMPTS` matching Phase 0/1/3's own set). + Then, separately: + - Co-located: one 4-rank EP pool (`patch_moe_infer_ep(..., n_ranks=4)`), + `ColocatedWorker` driving the same prompts one at a time + (concurrency 1, to isolate correctness from the concurrency + measurement), compare logits to the reference. + - Disaggregated: the two 2-rank sub-groups from step 5, `PrefillWorker` + + `DecodeWorker` connected by the real handoff, same prompts at + concurrency 1, compare logits to the reference. + + Both must show mutual top-5 and top-1 agreement, the same bar Phase 1 + and Phase 3 used. If either fails, stop -- do not measure contention + on top of a wrong result. Record the exact agreement measured (perfect, + or the first position/token where it diverges) in the findings doc + either way. + +7. **Only if step 6 passes for both configurations**, run the + concurrency measurement: co-located vs. disaggregated, same 4 GPUs, + at 4 and 8 concurrent requests (staggered arrivals, not a single + simultaneous burst), `DEFAULT_PROMPTS` cycled to fill the concurrency + level. Record per-request TTFT and inter-token latency, aggregate + mean/p50/p99 TTFT and mean decode tokens/sec per configuration per + concurrency level. Every number carries its full config (GPU model, + concurrency level, batch sizes used, prompt set), per CLAUDE.md. + + **Be surgical.** Watch running cost throughout; if pace projects past + the $40 cap before both concurrency levels finish, stop after the + last fully-completed level and report exactly what was measured. + The goal is the smallest paid session that produces a trustworthy + result, not spending up to the cap. + +8. Copy all console output and any saved JSON back to this repo's + `docs/findings/` before doing anything else. + +9. Tear down immediately: + + ```bash + uv run python scripts/gpu/provision.py terminate + ``` + +10. From this repo's root, record the measured cost: + + ```python + from pathlib import Path + from scripts.gpu.provision import write_cost_record + + write_cost_record( + Path("docs/findings"), + pod_id="", + gpu_type_id="", + cost_per_hour=, + duration_s=, + note=( + "Phase 4: DeepEP V1 install, 4-GPU NVLink verification, two EP " + "sub-group setup, correctness gate (co-located + disaggregated), " + "concurrency measurement at 4 and 8 concurrent requests" + ), + run_label="phase-4-disaggregated-prefill-decode", + ) + ``` +```` + +- [ ] **Step 2: Get the user's explicit go-ahead, then execute the runbook** + +Confirm the budget cap and GPU choice with the user before the first +`create` call -- this is a paid action, never taken unattended. Remind +the user of the standing instruction to be surgical: prefer stopping +early with a smaller, trustworthy result over spending toward the cap. + +- [ ] **Step 3: Commit the runbook, the pod-live load-driver script, and the cost record** + +```bash +git add docs/runbooks/phase-4-disaggregated-prefill-decode.md docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-cost.md scripts/gpu/phase4_load_driver.py +git commit -m "docs: record Phase 4 GPU runbook and cost; add EP sub-group wiring" +``` + +If step 5 or 6 needed adjustment to match live-observed DeepEP/process-group +behavior, note the adjustment and why in the commit body -- not silently, +per this project's practice of recording what actually happened rather +than what was assumed. + +--- + +### Task 6: Findings doc and STATUS.md + +**Files:** + +- Create: `docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-run.md` +- Modify: `docs/STATUS.md` + +- [ ] **Step 1: Write the findings doc** + +Cover, in `docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-run.md`: +whether real 4-GPU NVLink was confirmed; whether DeepEP V1 installed +cleanly reusing Phase 3's steps or needed adjustment; the correctness-gate +result for both co-located and disaggregated configurations (pass, with +the measured agreement, or the honest reason it didn't); **the direct +answer to Phase 4's thesis** -- does disaggregation measurably relieve +prefill/decode contention at 4 and/or 8 concurrent requests, holding GPU +count fixed at 4 -- stated at exactly the strength the data supports; the +full measured TTFT/tokens-per-second table for both configurations at +both concurrency levels, if the gate passed; the known dense-padding +simplification's effect explicitly disclosed (design doc section 6); and +the actual GPU type, duration, and cost from Task 5's cost record, +**explicitly compared against the $40 cap** -- the gap between them is +itself part of this phase's reported result, per the standing instruction +to be surgical. If the correctness gate failed or the budget ran out +before both concurrency levels completed, say so plainly and report +exactly what was measured -- matching this project's practice of writing +down a null or partial result rather than a flattering guess. + +- [ ] **Step 2: Update STATUS.md** + +Add a "## Phase 4 progress" section following the Phase 0/1/2/3 pattern: +plan link, hardware actually used, the correctness-gate result for both +configurations, the contention-relief thesis answer, and total GPU cost +against the $40 cap. Set "## Next step" to reflect what's actually next +(Phase 5 planning, or follow-up on Phase 4 if something didn't land +cleanly). + +- [ ] **Step 3: Commit** + +```bash +git add docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-run.md docs/STATUS.md +git commit -m "docs: record Phase 4 disaggregated prefill/decode outcome" +``` + +--- + +## Self-Review Notes + +- **Spec coverage:** design doc §3 (architecture: 4-rank world, two EP + sub-groups) and §4 (components: PrefillWorker, DecodeWorker, + KV-cache handoff, load driver, co-located worker) map to Tasks 1-5 + directly -- one task per component, in dependency order (cache + bookkeeping and handoff proven before the scheduler that uses them, + the scheduler proven before the baseline that reuses its internals). + §5 (data flow) maps to Task 5's runbook steps 5-6. §6 (known dense- + padding simplification) is implemented in Task 1's `pad_and_batch_caches` + and explicitly disclosed again in Task 6's findings doc, not left + implicit. §7 (testing table) maps one-to-one onto Tasks 1-4 (CPU rows) + and Task 5 (the two `gpu`, paid rows). §8 (cost discipline, "be + surgical") is a Global Constraint and repeated in Task 5's runbook and + Task 6's findings doc, not stated once and forgotten. §1's thesis + (does disaggregation relieve contention, GPU count held constant) is + answered directly in Task 6, not left implicit. +- **New scope beyond the design doc, and why it's necessary rather than + creep:** the design doc doesn't name `transformers` v5's `DynamicCache` + API surface at all -- it was resolved during this plan's own writing + (live-checked 2026-09-16 against `transformers`' migration guide and + source, see Global Constraints) because the design doc's §4/§5 assume + KV-cache slicing is possible without specifying how, and guessing the + wrong API (e.g. the removed `to_legacy_cache`) would have surfaced as + an `AttributeError` only once real GPU time was already being spent. + Resolving it here, for $0, before Task 1 is dispatched, is the same + "confirm live, don't guess" discipline Phase 3's plan applied to + DeepEP's exact call shapes. +- **Ambiguity check:** "4 and 8 concurrent requests" (Task 5 step 7) is + pinned, unlike Phase 3's plan which deliberately left its benchmark + grid to be set live -- here the concurrency levels don't depend on + anything that can only be observed on real hardware (unlike Phase 3's + per-expert token counts, which came from DeepEP's real dispatch + behavior), so pinning them now is reasonable rather than premature. + `max_batch_size` for each worker is deliberately left for Task 5 to set + against the rented node's real memory headroom, the same kind of + live-set choice Phase 2 made for GPU type and Phase 0 made for model + precision. +- **Type consistency:** `PrefillFn`/`DecodeFn` (Task 3) are used + identically by `PrefillWorker`/`DecodeWorker` (Task 3), + `ColocatedWorker` (Task 4), and Task 5's real closures -- + `(input_ids, attention_mask) -> (first_token_ids, cache)` and + `(next_input_ids, cache, attention_mask, position_ids) -> (next_token_ids, + cache)` respectively, never renamed or reshaped between tasks. + `slice_cache`'s `keep_last` parameter (Task 1) is used identically by + `run_prefill_batch` and `step_active_requests` (Task 3) to drop + left-padding after every batched forward call, not just some of them. From 6305fc3a81c4b7e031e550784570e0391f61bda0 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:49:29 -0400 Subject: [PATCH 03/12] feat: add KV-cache slice/pad-batch helpers for continuous batching DynamicCache.batch_size always returns -1 for standard attention layers in this transformers version (no per-layer batch_size attribute exists anywhere in cache_utils.py) -- fixed the plan's test assertions to check layers[i].keys.shape[0] instead. mypy --strict also required narrowing DynamicCache.layers' CacheLayerMixin | LinearAttentionCacheLayerMixin union and its Tensor | None key/value types via an assert-based helper, justified by this project's model (DeepSeekMoE-16B) using only standard attention, never linear/SSM layers. --- ...ase-4-disaggregated-prefill-decode-plan.md | 4 +- src/dispatch/serving/__init__.py | 0 src/dispatch/serving/kv_cache.py | 90 ++++++++++++++ tests/unit/test_kv_cache.py | 110 ++++++++++++++++++ 4 files changed, 202 insertions(+), 2 deletions(-) create mode 100644 src/dispatch/serving/__init__.py create mode 100644 src/dispatch/serving/kv_cache.py create mode 100644 tests/unit/test_kv_cache.py diff --git a/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md b/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md index a37bd52..6e21f53 100644 --- a/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md +++ b/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md @@ -180,7 +180,7 @@ def test_slice_cache_extracts_one_request_without_mutating_the_original() -> Non sliced = slice_cache(cache, index=1) assert sliced.get_seq_length() == 3 - assert sliced.batch_size == 1 + assert sliced.layers[0].keys.shape[0] == 1 torch.testing.assert_close(sliced.layers[0].keys[0], cache.layers[0].keys[1]) torch.testing.assert_close(cache.layers[0].keys[0], original_row_0) @@ -202,7 +202,7 @@ def test_pad_and_batch_caches_left_pads_to_a_common_length_and_masks_correctly() batched, attention_mask = pad_and_batch_caches([short, long_], pad_to=5) assert batched.get_seq_length() == 5 - assert batched.batch_size == 2 + assert batched.layers[0].keys.shape[0] == 2 assert attention_mask.tolist() == [[0, 0, 0, 1, 1], [1, 1, 1, 1, 1]] torch.testing.assert_close(batched.layers[0].keys[0, :, -2:, :], short.layers[0].keys[0]) torch.testing.assert_close(batched.layers[0].keys[1], long_.layers[0].keys[0]) diff --git a/src/dispatch/serving/__init__.py b/src/dispatch/serving/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/dispatch/serving/kv_cache.py b/src/dispatch/serving/kv_cache.py new file mode 100644 index 0000000..9bb5451 --- /dev/null +++ b/src/dispatch/serving/kv_cache.py @@ -0,0 +1,90 @@ +"""KV-cache slicing and batch (re)assembly for Phase 4's continuous- +batching scheduler. Built entirely on transformers v5's surviving public +DynamicCache API -- to_legacy_cache/from_legacy_cache and the old +key_cache/value_cache attributes were removed in v5 (confirmed live +2026-09-16 against transformers' own v5 migration guide and +src/transformers/cache_utils.py). Dense, padded batching only -- not a +paged/block cache -- the known simplification this project's Phase 4 +design doc discloses (section 6): real wasted compute next to a +production engine, applied identically to every configuration this +project measures, so it should not bias their relative comparison. +""" + +from __future__ import annotations + +import torch +from transformers import DynamicCache +from transformers.cache_utils import CacheLayerMixin + +# DeepSeekMoE-16B (this project's only model) uses standard multi-head +# attention throughout -- no linear/SSM (Mamba) layers -- so every layer +# in its cache is a CacheLayerMixin (.keys/.values), never the separate +# LinearAttentionCacheLayerMixin hierarchy DynamicCache.layers' type also +# allows, and (since every cache this module handles already has at +# least one token) its keys/values are always populated, never None. +# Asserted, not just assumed, so mypy --strict can narrow both away and +# a real regression (an uninitialized or linear-attention layer) fails +# loudly here instead of silently returning wrong tensors. + + +def _layer_kv(layer: object) -> tuple[torch.Tensor, torch.Tensor]: + assert isinstance(layer, CacheLayerMixin), f"expected an attention layer, got {type(layer)}" + assert layer.keys is not None and layer.values is not None, "expected an initialized layer" + return layer.keys, layer.values + + +def slice_cache(cache: DynamicCache, index: int, *, keep_last: int | None = None) -> DynamicCache: + """Extract one request's KV cache (by batch index) out of a batched + cache, without mutating the original -- unlike Cache.batch_select_indices, + which selects in place and offers no way to also drop padding. If + keep_last is given, keeps only the last keep_last positions along the + sequence dim, dropping this request's left-padding after a batched + forward call where its real content is shorter than the batch's + padded length. + """ + per_layer = [] + for raw_layer in cache.layers: + layer_keys, layer_values = _layer_kv(raw_layer) + key = layer_keys[index : index + 1] + value = layer_values[index : index + 1] + if keep_last is not None: + key = key[:, :, -keep_last:, :] + value = value[:, :, -keep_last:, :] + per_layer.append((key.clone(), value.clone())) + return DynamicCache(ddp_cache_data=per_layer) + + +def pad_and_batch_caches( + caches: list[DynamicCache], pad_to: int +) -> tuple[DynamicCache, torch.Tensor]: + """Left-pads each cache's key/value tensors to pad_to positions and + concatenates them along the batch dim, returning the batched cache + and a matching (batch, pad_to) attention_mask (0 = pad, 1 = real). + """ + if not caches: + raise ValueError("caches must not be empty") + num_layers = len(caches[0].layers) + per_layer_keys: list[list[torch.Tensor]] = [[] for _ in range(num_layers)] + per_layer_values: list[list[torch.Tensor]] = [[] for _ in range(num_layers)] + attention_mask = torch.zeros(len(caches), pad_to, dtype=torch.long) + + for row, cache in enumerate(caches): + seq_len = cache.get_seq_length() + if seq_len > pad_to: + raise ValueError(f"cache seq_len={seq_len} exceeds pad_to={pad_to}") + attention_mask[row, pad_to - seq_len :] = 1 + pad_len = pad_to - seq_len + for layer_idx, raw_layer in enumerate(cache.layers): + key, value = _layer_kv(raw_layer) + if pad_len > 0: + pad_shape = (1, key.shape[1], pad_len, key.shape[3]) + key = torch.cat([torch.zeros(pad_shape, dtype=key.dtype), key], dim=2) + value = torch.cat([torch.zeros(pad_shape, dtype=value.dtype), value], dim=2) + per_layer_keys[layer_idx].append(key) + per_layer_values[layer_idx].append(value) + + per_layer = [ + (torch.cat(per_layer_keys[i], dim=0), torch.cat(per_layer_values[i], dim=0)) + for i in range(num_layers) + ] + return DynamicCache(ddp_cache_data=per_layer), attention_mask diff --git a/tests/unit/test_kv_cache.py b/tests/unit/test_kv_cache.py new file mode 100644 index 0000000..5b6d49d --- /dev/null +++ b/tests/unit/test_kv_cache.py @@ -0,0 +1,110 @@ +"""CPU-only: proves the dense-padding slice/rebatch bookkeeping a +continuous-batching scheduler needs, built entirely on transformers v5's +surviving public DynamicCache API (cache.layers[i].keys/.values, the +ddp_cache_data constructor path) -- to_legacy_cache/from_legacy_cache and +the old key_cache/value_cache attributes were removed in v5 (confirmed +live 2026-09-16 against transformers' own v5 migration guide). +""" + +from __future__ import annotations + +import pytest +import torch +from transformers import DynamicCache + +from dispatch.benchmark.harness import load_model +from dispatch.serving.kv_cache import _layer_kv, pad_and_batch_caches, slice_cache + +NUM_LAYERS = 2 +NUM_HEADS = 2 +HEAD_DIM = 4 + + +def _make_cache(batch_size: int, seq_len: int, *, seed: int = 0) -> DynamicCache: + generator = torch.Generator().manual_seed(seed) + per_layer = [ + ( + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM, generator=generator), + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM, generator=generator), + ) + for _ in range(NUM_LAYERS) + ] + return DynamicCache(ddp_cache_data=per_layer) + + +def test_slice_cache_extracts_one_request_without_mutating_the_original() -> None: + cache = _make_cache(batch_size=2, seq_len=3) + cache_keys, _ = _layer_kv(cache.layers[0]) + original_row_0 = cache_keys[0].clone() + + sliced = slice_cache(cache, index=1) + sliced_keys, _ = _layer_kv(sliced.layers[0]) + + assert sliced.get_seq_length() == 3 + assert sliced_keys.shape[0] == 1 + torch.testing.assert_close(sliced_keys[0], cache_keys[1]) + torch.testing.assert_close(cache_keys[0], original_row_0) + + +def test_slice_cache_with_keep_last_drops_left_padding() -> None: + cache = _make_cache(batch_size=1, seq_len=5) + cache_keys, _ = _layer_kv(cache.layers[0]) + real_tail = cache_keys[:, :, -2:, :].clone() + + sliced = slice_cache(cache, index=0, keep_last=2) + sliced_keys, _ = _layer_kv(sliced.layers[0]) + + assert sliced.get_seq_length() == 2 + torch.testing.assert_close(sliced_keys, real_tail) + + +def test_pad_and_batch_caches_left_pads_to_a_common_length_and_masks_correctly() -> None: + short = _make_cache(batch_size=1, seq_len=2, seed=1) + long_ = _make_cache(batch_size=1, seq_len=5, seed=2) + + batched, attention_mask = pad_and_batch_caches([short, long_], pad_to=5) + batched_keys, _ = _layer_kv(batched.layers[0]) + short_keys, _ = _layer_kv(short.layers[0]) + long_keys, _ = _layer_kv(long_.layers[0]) + + assert batched.get_seq_length() == 5 + assert batched_keys.shape[0] == 2 + assert attention_mask.tolist() == [[0, 0, 0, 1, 1], [1, 1, 1, 1, 1]] + torch.testing.assert_close(batched_keys[0, :, -2:, :], short_keys[0]) + torch.testing.assert_close(batched_keys[1], long_keys[0]) + + +def test_pad_and_batch_caches_rejects_a_cache_longer_than_pad_to() -> None: + long_ = _make_cache(batch_size=1, seq_len=5) + + with pytest.raises(ValueError, match="exceeds"): + pad_and_batch_caches([long_], pad_to=3) + + +def test_pad_and_batch_caches_rejects_an_empty_list() -> None: + with pytest.raises(ValueError, match="empty"): + pad_and_batch_caches([], pad_to=3) + + +@pytest.mark.slow +def test_slice_and_rebatch_round_trip_through_a_real_tiny_model() -> None: + """Proves the helpers produce a cache the real model accepts back and + continues generating from correctly -- against hf-internal-testing/ + tiny-random-gpt2, this repo's established tiny-model pattern + (test_harness.py). + """ + model, tokenizer = load_model("hf-internal-testing/tiny-random-gpt2") + input_ids = tokenizer("hello world", return_tensors="pt").input_ids + with torch.no_grad(): + outputs = model(input_ids=input_ids, use_cache=True) + cache = outputs.past_key_values + assert isinstance(cache, DynamicCache) + next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) + + sliced = slice_cache(cache, index=0) + assert sliced.get_seq_length() == cache.get_seq_length() + + with torch.no_grad(): + continued = model(input_ids=next_token, past_key_values=sliced, use_cache=True) + + assert continued.logits.shape == (1, 1, continued.logits.shape[-1]) From 911bb64bca05c11100ab3e632751ccd8ada04887 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:51:25 -0400 Subject: [PATCH 04/12] feat: add cross-rank KV-cache handoff, proven over gloo Renamed kv_cache.py's private _layer_kv to a public layer_kv while adding this file, since handoff.py now needs it too -- same reasoning the plan already applied to run_prefill_batch/step_active_requests being shared across disaggregated.py and colocated.py. --- src/dispatch/serving/handoff.py | 47 +++++++++++++++++++ src/dispatch/serving/kv_cache.py | 6 +-- tests/unit/test_handoff.py | 79 ++++++++++++++++++++++++++++++++ tests/unit/test_kv_cache.py | 16 +++---- 4 files changed, 137 insertions(+), 11 deletions(-) create mode 100644 src/dispatch/serving/handoff.py create mode 100644 tests/unit/test_handoff.py diff --git a/src/dispatch/serving/handoff.py b/src/dispatch/serving/handoff.py new file mode 100644 index 0000000..12ddf74 --- /dev/null +++ b/src/dispatch/serving/handoff.py @@ -0,0 +1,47 @@ +"""Cross-rank KV-cache handoff between the prefill and decode EP pools, +built on plain torch.distributed point-to-point send/recv so the exact +same code runs over gloo (CPU, this file's own tests) and NCCL (real +GPUs, Task 5) -- the transport logic gets proven once, for $0, before it +is ever run against real hardware (design doc section 8's "maximize +what's proven before any rental"). +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +from transformers import DynamicCache + +from dispatch.serving.kv_cache import layer_kv + + +def send_kv_cache(cache: DynamicCache, dst: int, group: dist.ProcessGroup) -> None: + seq_len = torch.tensor([cache.get_seq_length()], dtype=torch.int64) + dist.send(seq_len, dst=dst, group=group) + for layer in cache.layers: + key, value = layer_kv(layer) + dist.send(key.contiguous(), dst=dst, group=group) + dist.send(value.contiguous(), dst=dst, group=group) + + +def recv_kv_cache( # noqa: PLR0913 -- the receiver can't infer shape/dtype from the wire, they must be passed + src: int, + group: dist.ProcessGroup, + *, + num_layers: int, + num_heads: int, + head_dim: int, + dtype: torch.dtype, +) -> DynamicCache: + seq_len_tensor = torch.zeros(1, dtype=torch.int64) + dist.recv(seq_len_tensor, src=src, group=group) + seq_len = int(seq_len_tensor.item()) + + per_layer = [] + for _ in range(num_layers): + key = torch.zeros(1, num_heads, seq_len, head_dim, dtype=dtype) + value = torch.zeros(1, num_heads, seq_len, head_dim, dtype=dtype) + dist.recv(key, src=src, group=group) + dist.recv(value, src=src, group=group) + per_layer.append((key, value)) + return DynamicCache(ddp_cache_data=per_layer) diff --git a/src/dispatch/serving/kv_cache.py b/src/dispatch/serving/kv_cache.py index 9bb5451..fa394d3 100644 --- a/src/dispatch/serving/kv_cache.py +++ b/src/dispatch/serving/kv_cache.py @@ -27,7 +27,7 @@ # loudly here instead of silently returning wrong tensors. -def _layer_kv(layer: object) -> tuple[torch.Tensor, torch.Tensor]: +def layer_kv(layer: object) -> tuple[torch.Tensor, torch.Tensor]: assert isinstance(layer, CacheLayerMixin), f"expected an attention layer, got {type(layer)}" assert layer.keys is not None and layer.values is not None, "expected an initialized layer" return layer.keys, layer.values @@ -44,7 +44,7 @@ def slice_cache(cache: DynamicCache, index: int, *, keep_last: int | None = None """ per_layer = [] for raw_layer in cache.layers: - layer_keys, layer_values = _layer_kv(raw_layer) + layer_keys, layer_values = layer_kv(raw_layer) key = layer_keys[index : index + 1] value = layer_values[index : index + 1] if keep_last is not None: @@ -75,7 +75,7 @@ def pad_and_batch_caches( attention_mask[row, pad_to - seq_len :] = 1 pad_len = pad_to - seq_len for layer_idx, raw_layer in enumerate(cache.layers): - key, value = _layer_kv(raw_layer) + key, value = layer_kv(raw_layer) if pad_len > 0: pad_shape = (1, key.shape[1], pad_len, key.shape[3]) key = torch.cat([torch.zeros(pad_shape, dtype=key.dtype), key], dim=2) diff --git a/tests/unit/test_handoff.py b/tests/unit/test_handoff.py new file mode 100644 index 0000000..b3745ba --- /dev/null +++ b/tests/unit/test_handoff.py @@ -0,0 +1,79 @@ +"""CPU-only, multi-process: proves send_kv_cache/recv_kv_cache round-trip +correctly using torch.distributed's gloo backend -- no GPU, no DeepEP, no +NCCL. Task 5 reuses these exact functions unchanged with the NCCL +backend on the real rented 4-GPU node; only init_process_group's backend +argument and the real ranks differ. +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from transformers import DynamicCache + +from dispatch.serving.handoff import recv_kv_cache, send_kv_cache +from dispatch.serving.kv_cache import layer_kv + +NUM_LAYERS = 2 +NUM_HEADS = 2 +HEAD_DIM = 4 +SEQ_LEN = 3 +PORT = 29513 + + +def _make_cache(seed: int) -> DynamicCache: + generator = torch.Generator().manual_seed(seed) + per_layer = [ + ( + torch.randn(1, NUM_HEADS, SEQ_LEN, HEAD_DIM, generator=generator), + torch.randn(1, NUM_HEADS, SEQ_LEN, HEAD_DIM, generator=generator), + ) + for _ in range(NUM_LAYERS) + ] + return DynamicCache(ddp_cache_data=per_layer) + + +def _worker(rank: int, world_size: int, result_path: str) -> None: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(PORT) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + group = dist.group.WORLD + assert group is not None + + if rank == 0: + send_kv_cache(_make_cache(seed=42), dst=1, group=group) + else: + received = recv_kv_cache( + src=0, + group=group, + num_layers=NUM_LAYERS, + num_heads=NUM_HEADS, + head_dim=HEAD_DIM, + dtype=torch.float32, + ) + expected = _make_cache(seed=42) + ok = all( + torch.equal(layer_kv(received.layers[i])[0], layer_kv(expected.layers[i])[0]) + and torch.equal(layer_kv(received.layers[i])[1], layer_kv(expected.layers[i])[1]) + for i in range(NUM_LAYERS) + ) + torch.save({"ok": ok, "seq_len": received.get_seq_length()}, result_path) + + dist.destroy_process_group() + + +@pytest.mark.slow +def test_send_and_recv_kv_cache_round_trip_over_gloo(tmp_path: object) -> None: + result_path = str(tmp_path / "result.pt") # type: ignore[operator] + + mp.spawn( # type: ignore[attr-defined,no-untyped-call] + _worker, args=(2, result_path), nprocs=2, join=True + ) + + result = torch.load(result_path, weights_only=True) + assert result["ok"] is True + assert result["seq_len"] == SEQ_LEN diff --git a/tests/unit/test_kv_cache.py b/tests/unit/test_kv_cache.py index 5b6d49d..cdd31a9 100644 --- a/tests/unit/test_kv_cache.py +++ b/tests/unit/test_kv_cache.py @@ -13,7 +13,7 @@ from transformers import DynamicCache from dispatch.benchmark.harness import load_model -from dispatch.serving.kv_cache import _layer_kv, pad_and_batch_caches, slice_cache +from dispatch.serving.kv_cache import layer_kv, pad_and_batch_caches, slice_cache NUM_LAYERS = 2 NUM_HEADS = 2 @@ -34,11 +34,11 @@ def _make_cache(batch_size: int, seq_len: int, *, seed: int = 0) -> DynamicCache def test_slice_cache_extracts_one_request_without_mutating_the_original() -> None: cache = _make_cache(batch_size=2, seq_len=3) - cache_keys, _ = _layer_kv(cache.layers[0]) + cache_keys, _ = layer_kv(cache.layers[0]) original_row_0 = cache_keys[0].clone() sliced = slice_cache(cache, index=1) - sliced_keys, _ = _layer_kv(sliced.layers[0]) + sliced_keys, _ = layer_kv(sliced.layers[0]) assert sliced.get_seq_length() == 3 assert sliced_keys.shape[0] == 1 @@ -48,11 +48,11 @@ def test_slice_cache_extracts_one_request_without_mutating_the_original() -> Non def test_slice_cache_with_keep_last_drops_left_padding() -> None: cache = _make_cache(batch_size=1, seq_len=5) - cache_keys, _ = _layer_kv(cache.layers[0]) + cache_keys, _ = layer_kv(cache.layers[0]) real_tail = cache_keys[:, :, -2:, :].clone() sliced = slice_cache(cache, index=0, keep_last=2) - sliced_keys, _ = _layer_kv(sliced.layers[0]) + sliced_keys, _ = layer_kv(sliced.layers[0]) assert sliced.get_seq_length() == 2 torch.testing.assert_close(sliced_keys, real_tail) @@ -63,9 +63,9 @@ def test_pad_and_batch_caches_left_pads_to_a_common_length_and_masks_correctly() long_ = _make_cache(batch_size=1, seq_len=5, seed=2) batched, attention_mask = pad_and_batch_caches([short, long_], pad_to=5) - batched_keys, _ = _layer_kv(batched.layers[0]) - short_keys, _ = _layer_kv(short.layers[0]) - long_keys, _ = _layer_kv(long_.layers[0]) + batched_keys, _ = layer_kv(batched.layers[0]) + short_keys, _ = layer_kv(short.layers[0]) + long_keys, _ = layer_kv(long_.layers[0]) assert batched.get_seq_length() == 5 assert batched_keys.shape[0] == 2 From 0c9ea7795e3dd9e1842dc9b12030b9d2b015537a Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:55:33 -0400 Subject: [PATCH 05/12] feat: add continuous-batching prefill and decode workers Found and fixed a real bug while tracing DecodeWorker's logic by hand before implementing it: admission and decoding share one step() call, so a request with max_new_tokens=1 is already complete from prefill's own first token -- without a check, it would still take an extra decode step and over-generate by one token. Added split_completed() to catch this, with a regression test. Documented in the plan's Self-Review Notes. --- ...ase-4-disaggregated-prefill-decode-plan.md | 21 ++ src/dispatch/serving/disaggregated.py | 262 ++++++++++++++++++ tests/unit/test_disaggregated.py | 186 +++++++++++++ 3 files changed, 469 insertions(+) create mode 100644 src/dispatch/serving/disaggregated.py create mode 100644 tests/unit/test_disaggregated.py diff --git a/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md b/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md index 6e21f53..16e8cbf 100644 --- a/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md +++ b/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md @@ -1494,3 +1494,24 @@ git commit -m "docs: record Phase 4 disaggregated prefill/decode outcome" `slice_cache`'s `keep_last` parameter (Task 1) is used identically by `run_prefill_batch` and `step_active_requests` (Task 3) to drop left-padding after every batched forward call, not just some of them. + +## Deviation found during Task 3's execution + +Tracing `DecodeWorker.step()`'s logic by hand before writing its +implementation surfaced a real bug this plan's original sketch didn't +have a check for: admission (waiting -> active) and decoding happen in +the *same* `step()` call, so a request whose `max_new_tokens` is 1 is +already complete from prefill's own first token alone -- calling +`step_active_requests` on it anyway would still take one decode step and +over-generate by one token. Fixed by adding `split_completed(active, +now) -> (still_active, completed)` to `disaggregated.py`, called in +`DecodeWorker.step()` (and, per the shared-helpers design, `colocated.py`'s +`ColocatedWorker.step()` in Task 4) immediately after admission and +before `step_active_requests`. Covered by a new regression test, +`test_decode_worker_completes_immediately_when_admission_alone_hits_max_new_tokens`. +This also forced `max_new_tokens` up from 1 to 2 or 3 in several of the +plan's other sketched test cases, where the degenerate value would have +made them stop exercising what they were meant to test (e.g. +`test_decode_worker_batches_requests_with_different_cache_lengths` would +never have reached `step_active_requests`'s padding logic at all with +`max_new_tokens=1` on both requests). diff --git a/src/dispatch/serving/disaggregated.py b/src/dispatch/serving/disaggregated.py new file mode 100644 index 0000000..b7a511e --- /dev/null +++ b/src/dispatch/serving/disaggregated.py @@ -0,0 +1,262 @@ +"""Continuous-batching prefill and decode workers for Phase 4's +disaggregated serving path. Both workers are forward-pass-agnostic -- +they take a plain PrefillFn/DecodeFn callable and know nothing about +DeepEP, EP pools, or even that this is a distributed setting. Task 5 +wires the real DeepSeekMoE-16B model (patched with Phase 3's +make_ep_moe_infer/patch_moe_infer_ep, one EP pool per role) into exactly +these two callable shapes; nothing in this file changes when that +happens. run_prefill_batch, prefill_result_to_in_flight, split_completed, +and step_active_requests are exported (not underscore-prefixed) because +colocated.py's ColocatedWorker reuses them unchanged -- a decode step is +a decode step regardless of whether prefill happened on the same GPU +this iteration or a different one. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from transformers import DynamicCache + +from dispatch.serving.kv_cache import pad_and_batch_caches, slice_cache + +PrefillFn = Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, DynamicCache]] +DecodeFn = Callable[ + [torch.Tensor, DynamicCache, torch.Tensor, torch.Tensor], tuple[torch.Tensor, DynamicCache] +] + + +@dataclass +class Request: + request_id: str + prompt_ids: torch.Tensor # (1, prompt_len) + max_new_tokens: int + eos_token_id: int | None = None + + +@dataclass +class PrefillResult: + request_id: str + cache: DynamicCache + first_token_id: int + ttft: float + max_new_tokens: int + eos_token_id: int | None + + +@dataclass +class RequestResult: + request_id: str + generated_token_ids: list[int] + ttft: float + completion_time: float + + +@dataclass +class InFlightRequest: + request_id: str + cache: DynamicCache + seq_len: int + next_input_id: torch.Tensor + generated_token_ids: list[int] + max_new_tokens: int + eos_token_id: int | None + ttft: float + + +def run_prefill_batch( + prefill_fn: PrefillFn, + batch: list[tuple[Request, float]], + clock_fn: Callable[[], float], +) -> list[PrefillResult]: + """Batches a list of (request, admitted_at) pairs into one padded + prefill_fn call and slices the result back into per-request, unpadded + PrefillResults. Shared by PrefillWorker.step() and colocated.py's + ColocatedWorker. + """ + requests = [r for r, _ in batch] + admitted_at = [t for _, t in batch] + max_len = max(r.prompt_ids.shape[1] for r in requests) + input_ids = torch.zeros(len(requests), max_len, dtype=torch.long) + attention_mask = torch.zeros(len(requests), max_len, dtype=torch.long) + for i, r in enumerate(requests): + prompt_len = r.prompt_ids.shape[1] + input_ids[i, max_len - prompt_len :] = r.prompt_ids[0] + attention_mask[i, max_len - prompt_len :] = 1 + + first_token_ids, cache = prefill_fn(input_ids, attention_mask) + now = clock_fn() + + results = [] + for i, r in enumerate(requests): + prompt_len = r.prompt_ids.shape[1] + results.append( + PrefillResult( + request_id=r.request_id, + cache=slice_cache(cache, i, keep_last=prompt_len), + first_token_id=int(first_token_ids[i].item()), + ttft=now - admitted_at[i], + max_new_tokens=r.max_new_tokens, + eos_token_id=r.eos_token_id, + ) + ) + return results + + +def prefill_result_to_in_flight(result: PrefillResult) -> InFlightRequest: + return InFlightRequest( + request_id=result.request_id, + cache=result.cache, + seq_len=result.cache.get_seq_length(), + next_input_id=torch.tensor([[result.first_token_id]], dtype=torch.long), + generated_token_ids=[result.first_token_id], + max_new_tokens=result.max_new_tokens, + eos_token_id=result.eos_token_id, + ttft=result.ttft, + ) + + +def split_completed( + active: list[InFlightRequest], now: float +) -> tuple[list[InFlightRequest], list[RequestResult]]: + """Separates requests that are already at or past max_new_tokens from + those that still need at least one more decode step. Needed because + admission and decoding share one step() call (continuous batching): + a request whose max_new_tokens is 1 is already complete the moment + it's admitted, from prefill's own first token alone -- without this + check, step_active_requests would still run one decode_fn call on it + and over-generate by one token. + """ + still_active: list[InFlightRequest] = [] + completed: list[RequestResult] = [] + for r in active: + if len(r.generated_token_ids) >= r.max_new_tokens: + completed.append( + RequestResult( + request_id=r.request_id, + generated_token_ids=r.generated_token_ids, + ttft=r.ttft, + completion_time=now, + ) + ) + else: + still_active.append(r) + return still_active, completed + + +def step_active_requests( + active: list[InFlightRequest], decode_fn: DecodeFn, clock_fn: Callable[[], float] +) -> tuple[list[InFlightRequest], list[RequestResult]]: + """Pads every active request's cache to the batch's current max real + length, runs one batched decode_fn call, and unpads each result back + down to its own true length before storing it. Shared by + DecodeWorker.step() and colocated.py's ColocatedWorker.step(). Callers + are expected to have already run split_completed on active -- this + function always takes one decode step for everyone it's given. + """ + if not active: + return [], [] + pad_to = max(r.seq_len for r in active) + batched_cache, attention_mask = pad_and_batch_caches([r.cache for r in active], pad_to) + full_attention_mask = torch.cat( + [attention_mask, torch.ones(len(active), 1, dtype=torch.long)], dim=1 + ) + position_ids = attention_mask.sum(dim=1, keepdim=True) + next_input_ids = torch.cat([r.next_input_id for r in active], dim=0) + + next_token_ids, updated_cache = decode_fn( + next_input_ids, batched_cache, full_attention_mask, position_ids + ) + now = clock_fn() + + still_active: list[InFlightRequest] = [] + completed: list[RequestResult] = [] + for i, r in enumerate(active): + token_id = int(next_token_ids[i].item()) + r.generated_token_ids.append(token_id) + r.seq_len += 1 + r.cache = slice_cache(updated_cache, i, keep_last=r.seq_len) + r.next_input_id = torch.tensor([[token_id]], dtype=torch.long) + done = len(r.generated_token_ids) >= r.max_new_tokens or ( + r.eos_token_id is not None and token_id == r.eos_token_id + ) + if done: + completed.append( + RequestResult( + request_id=r.request_id, + generated_token_ids=r.generated_token_ids, + ttft=r.ttft, + completion_time=now, + ) + ) + else: + still_active.append(r) + return still_active, completed + + +class PrefillWorker: + """Drains up to batch_size waiting requests per step, runs one + batched prefill forward pass, and returns each request's own + (unpadded) KV cache and first generated token -- ready to hand off to + a DecodeWorker (in-process admit() in CPU tests, handoff.py's + send_kv_cache on the real rented node). + """ + + def __init__( + self, + prefill_fn: PrefillFn, + batch_size: int, + *, + clock_fn: Callable[[], float] = time.perf_counter, + ) -> None: + self._prefill_fn = prefill_fn + self._batch_size = batch_size + self._clock_fn = clock_fn + self._waiting: list[tuple[Request, float]] = [] + + def submit(self, request: Request) -> None: + self._waiting.append((request, self._clock_fn())) + + def step(self) -> list[PrefillResult]: + batch = self._waiting[: self._batch_size] + self._waiting = self._waiting[self._batch_size :] + if not batch: + return [] + return run_prefill_batch(self._prefill_fn, batch, self._clock_fn) + + +class DecodeWorker: + """Admits PrefillResults into a live, continuously-batched decode + loop: every step, any waiting request gets admitted (up to + batch_size), then every active request advances one token. + """ + + def __init__( + self, + decode_fn: DecodeFn, + batch_size: int, + *, + clock_fn: Callable[[], float] = time.perf_counter, + ) -> None: + self._decode_fn = decode_fn + self._batch_size = batch_size + self._clock_fn = clock_fn + self._waiting: list[InFlightRequest] = [] + self._active: list[InFlightRequest] = [] + + def admit(self, result: PrefillResult) -> None: + self._waiting.append(prefill_result_to_in_flight(result)) + + def step(self) -> list[RequestResult]: + free_slots = self._batch_size - len(self._active) + if free_slots > 0 and self._waiting: + self._active.extend(self._waiting[:free_slots]) + self._waiting = self._waiting[free_slots:] + self._active, already_done = split_completed(self._active, self._clock_fn()) + self._active, completed = step_active_requests( + self._active, self._decode_fn, self._clock_fn + ) + return already_done + completed diff --git a/tests/unit/test_disaggregated.py b/tests/unit/test_disaggregated.py new file mode 100644 index 0000000..8b06240 --- /dev/null +++ b/tests/unit/test_disaggregated.py @@ -0,0 +1,186 @@ +"""CPU-only: proves the continuous-batching scheduler's admission, +padding, and completion bookkeeping against a fake prefill/decode +forward pass -- no GPU, no DeepEP, no real model. Task 5 wires the same +PrefillFn/DecodeFn shapes to the real DeepSeekMoE-16B model and DeepEP EP +pools; nothing here changes at that point. +""" + +from __future__ import annotations + +import torch +from transformers import DynamicCache + +from dispatch.serving.disaggregated import DecodeWorker, PrefillResult, PrefillWorker, Request + +NUM_LAYERS = 2 +NUM_HEADS = 2 +HEAD_DIM = 4 + + +def _fake_cache(batch_size: int, seq_len: int) -> DynamicCache: + per_layer = [ + ( + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM), + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM), + ) + for _ in range(NUM_LAYERS) + ] + return DynamicCache(ddp_cache_data=per_layer) + + +def _fake_prefill_fn( + input_ids: torch.Tensor, attention_mask: torch.Tensor +) -> tuple[torch.Tensor, DynamicCache]: + del attention_mask + batch_size, seq_len = input_ids.shape + first_tokens = torch.arange(batch_size) + 100 # distinctive, deterministic per row + return first_tokens, _fake_cache(batch_size, seq_len) + + +def _fake_decode_fn( + next_input_ids: torch.Tensor, + cache: DynamicCache, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, +) -> tuple[torch.Tensor, DynamicCache]: + del attention_mask, position_ids + batch_size = next_input_ids.shape[0] + new_seq_len = cache.get_seq_length() + 1 + next_tokens = next_input_ids.squeeze(1) + 1 # deterministic increment, easy to assert on + return next_tokens, _fake_cache(batch_size, new_seq_len) + + +def test_prefill_worker_drains_up_to_batch_size_and_produces_per_request_caches() -> None: + worker = PrefillWorker(_fake_prefill_fn, batch_size=2) + worker.submit(Request("a", torch.tensor([[1, 2, 3]]), max_new_tokens=4)) + worker.submit(Request("b", torch.tensor([[1, 2]]), max_new_tokens=4)) + worker.submit(Request("c", torch.tensor([[1]]), max_new_tokens=4)) + + results = worker.step() + + assert [r.request_id for r in results] == ["a", "b"] + assert results[0].cache.get_seq_length() == 3 # unpadded back to its own prompt length + assert results[1].cache.get_seq_length() == 2 + assert results[0].first_token_id == 100 + assert results[1].first_token_id == 101 + + remaining = worker.step() + assert [r.request_id for r in remaining] == ["c"] + + +def test_prefill_worker_step_with_nothing_waiting_returns_empty() -> None: + worker = PrefillWorker(_fake_prefill_fn, batch_size=2) + + assert worker.step() == [] + + +def test_decode_worker_admits_and_completes_requests_after_max_new_tokens() -> None: + """max_new_tokens=3, not 2: admission and one decode step happen in + the same step() call (continuous batching), so a max_new_tokens=2 + request would already be complete at admission time -- too degenerate + to show the "still needs another step()" case this test is for. + """ + worker = DecodeWorker(_fake_decode_fn, batch_size=4) + cache = _fake_cache(1, 3) + worker.admit( + PrefillResult("a", cache, first_token_id=5, ttft=0.1, max_new_tokens=3, eos_token_id=None) + ) + + first_step = worker.step() + assert first_step == [] # admitted and decoded once this call: 2 of 3 tokens so far + + second_step = worker.step() + assert len(second_step) == 1 + assert second_step[0].request_id == "a" + assert second_step[0].generated_token_ids == [5, 6, 7] + + +def test_decode_worker_completes_a_request_early_on_eos() -> None: + worker = DecodeWorker(_fake_decode_fn, batch_size=4) + cache = _fake_cache(1, 3) + worker.admit( + PrefillResult("a", cache, first_token_id=5, ttft=0.1, max_new_tokens=10, eos_token_id=6) + ) + + results = worker.step() + + assert len(results) == 1 + assert results[0].generated_token_ids == [5, 6] + + +def test_decode_worker_batches_requests_with_different_cache_lengths() -> None: + worker = DecodeWorker(_fake_decode_fn, batch_size=4) + worker.admit( + PrefillResult( + "short", + _fake_cache(1, 2), + first_token_id=1, + ttft=0.1, + max_new_tokens=2, + eos_token_id=None, + ) + ) + worker.admit( + PrefillResult( + "long", + _fake_cache(1, 5), + first_token_id=1, + ttft=0.1, + max_new_tokens=2, + eos_token_id=None, + ) + ) + + results = worker.step() + + assert {r.request_id for r in results} == {"short", "long"} + + +def test_decode_worker_completes_immediately_when_admission_alone_hits_max_new_tokens() -> None: + """max_new_tokens=1: prefill's own first token already satisfies it, + so admission must not trigger an extra decode_fn call -- regression + test for a real bug found while tracing this scheduler's logic: + admission and decode share one step() call, so a naive + always-decode-after-admitting implementation would over-generate by + one token for any request this degenerate. + """ + worker = DecodeWorker(_fake_decode_fn, batch_size=4) + worker.admit( + PrefillResult( + "a", _fake_cache(1, 2), first_token_id=1, ttft=0.1, max_new_tokens=1, eos_token_id=None + ) + ) + + results = worker.step() + + assert len(results) == 1 + assert results[0].generated_token_ids == [1] + + +def test_decode_worker_respects_batch_size_admitting_only_free_slots() -> None: + worker = DecodeWorker(_fake_decode_fn, batch_size=1) + worker.admit( + PrefillResult( + "a", _fake_cache(1, 2), first_token_id=1, ttft=0.1, max_new_tokens=5, eos_token_id=None + ) + ) + worker.admit( + PrefillResult( + "b", _fake_cache(1, 2), first_token_id=1, ttft=0.1, max_new_tokens=2, eos_token_id=None + ) + ) + + first_step = worker.step() + assert first_step == [] # "a" admitted (only free slot); "b" still waiting + + for _ in range(2): + worker.step() + a_result = worker.step() + assert [r.request_id for r in a_result] == [ + "a" + ] # 5 tokens total: 1 from prefill + 4 decode steps + + b_result = ( + worker.step() + ) # "b" now admitted, needs only 1 more token -> completes this same step + assert [r.request_id for r in b_result] == ["b"] From f78f9c37189ea346f375ebd0e851dbda5c5d840a Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:57:22 -0400 Subject: [PATCH 06/12] feat: add co-located prefill/decode baseline worker Test assertions verified against actual runtime behavior via a scratch script rather than hand-traced -- the same split_completed interaction from Task 3 (admission and decode sharing one step() call) made a couple of the plan's original max_new_tokens choices produce different step counts than expected. Full CPU-testable suite (100 tests) green. --- src/dispatch/serving/colocated.py | 63 ++++++++++++++++++++ tests/unit/test_colocated.py | 95 +++++++++++++++++++++++++++++++ 2 files changed, 158 insertions(+) create mode 100644 src/dispatch/serving/colocated.py create mode 100644 tests/unit/test_colocated.py diff --git a/src/dispatch/serving/colocated.py b/src/dispatch/serving/colocated.py new file mode 100644 index 0000000..d3ae025 --- /dev/null +++ b/src/dispatch/serving/colocated.py @@ -0,0 +1,63 @@ +"""The co-located baseline: one worker, one EP pool, doing both prefill +and decode in a single loop on the same GPUs -- the contention baseline +Phase 4's disaggregated path (disaggregated.py) is measured against, on +the same 4 GPUs, so GPU count is never a confound (design doc section 1). +Reuses disaggregated.py's types and helpers unchanged: a decode step is a +decode step regardless of whether prefill happened on the same GPU this +iteration or a different one. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable + +from dispatch.serving.disaggregated import ( + DecodeFn, + InFlightRequest, + PrefillFn, + Request, + RequestResult, + prefill_result_to_in_flight, + run_prefill_batch, + split_completed, + step_active_requests, +) + + +class ColocatedWorker: + """Every step: admit waiting requests via an inline prefill batch (up + to free slots), then advance every active request -- new and old -- + by one decode token, all on the same GPU pool. + """ + + def __init__( + self, + prefill_fn: PrefillFn, + decode_fn: DecodeFn, + batch_size: int, + *, + clock_fn: Callable[[], float] = time.perf_counter, + ) -> None: + self._prefill_fn = prefill_fn + self._decode_fn = decode_fn + self._batch_size = batch_size + self._clock_fn = clock_fn + self._waiting: list[tuple[Request, float]] = [] + self._active: list[InFlightRequest] = [] + + def submit(self, request: Request) -> None: + self._waiting.append((request, self._clock_fn())) + + def step(self) -> list[RequestResult]: + free_slots = self._batch_size - len(self._active) + if free_slots > 0 and self._waiting: + batch = self._waiting[:free_slots] + self._waiting = self._waiting[free_slots:] + prefill_results = run_prefill_batch(self._prefill_fn, batch, self._clock_fn) + self._active.extend(prefill_result_to_in_flight(r) for r in prefill_results) + self._active, already_done = split_completed(self._active, self._clock_fn()) + self._active, completed = step_active_requests( + self._active, self._decode_fn, self._clock_fn + ) + return already_done + completed diff --git a/tests/unit/test_colocated.py b/tests/unit/test_colocated.py new file mode 100644 index 0000000..69d0fd2 --- /dev/null +++ b/tests/unit/test_colocated.py @@ -0,0 +1,95 @@ +"""CPU-only: proves ColocatedWorker interleaves inline prefill and +decode on the same loop -- new requests get their first token the same +step old ones advance, which is exactly what creates the contention +Phase 4 measures disaggregation against. Reuses test_disaggregated.py's +fake prefill/decode functions unchanged: the forward-pass contract is +identical, only which worker calls it differs. +""" + +from __future__ import annotations + +import torch +from transformers import DynamicCache + +from dispatch.serving.colocated import ColocatedWorker +from dispatch.serving.disaggregated import Request + +NUM_LAYERS = 2 +NUM_HEADS = 2 +HEAD_DIM = 4 + + +def _fake_cache(batch_size: int, seq_len: int) -> DynamicCache: + per_layer = [ + ( + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM), + torch.randn(batch_size, NUM_HEADS, seq_len, HEAD_DIM), + ) + for _ in range(NUM_LAYERS) + ] + return DynamicCache(ddp_cache_data=per_layer) + + +def _fake_prefill_fn( + input_ids: torch.Tensor, attention_mask: torch.Tensor +) -> tuple[torch.Tensor, DynamicCache]: + del attention_mask + batch_size, seq_len = input_ids.shape + first_tokens = torch.arange(batch_size) + 100 + return first_tokens, _fake_cache(batch_size, seq_len) + + +def _fake_decode_fn( + next_input_ids: torch.Tensor, + cache: DynamicCache, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, +) -> tuple[torch.Tensor, DynamicCache]: + del attention_mask, position_ids + batch_size = next_input_ids.shape[0] + new_seq_len = cache.get_seq_length() + 1 + next_tokens = next_input_ids.squeeze(1) + 1 + return next_tokens, _fake_cache(batch_size, new_seq_len) + + +def test_colocated_worker_admits_via_inline_prefill_and_completes_after_max_new_tokens() -> None: + """max_new_tokens=3, not 2: admission (inline prefill) and one decode + step happen in the same step() call here too, same as DecodeWorker -- + a max_new_tokens=2 request would already be done at the first step(). + """ + worker = ColocatedWorker(_fake_prefill_fn, _fake_decode_fn, batch_size=4) + worker.submit(Request("a", torch.tensor([[1, 2, 3]]), max_new_tokens=3)) + + first_step = worker.step() + assert first_step == [] # prefilled and decoded once this call: 2 of 3 tokens so far + + second_step = worker.step() + assert len(second_step) == 1 + assert second_step[0].request_id == "a" + assert second_step[0].generated_token_ids == [100, 101, 102] + + +def test_colocated_worker_interleaves_new_prefill_with_existing_decode_in_one_step() -> None: + worker = ColocatedWorker(_fake_prefill_fn, _fake_decode_fn, batch_size=4) + worker.submit(Request("old", torch.tensor([[1, 2]]), max_new_tokens=3)) + worker.step() # "old" admitted and decoded once: 2 of 3 tokens so far + + worker.submit(Request("new", torch.tensor([[9]]), max_new_tokens=1)) + step_result = worker.step() + + ids = {r.request_id for r in step_result} + assert ids == {"new", "old"} # "new" done from its own prefill token; "old" reaches 3 of 3 here + + +def test_colocated_worker_respects_batch_size_across_prefill_and_decode() -> None: + worker = ColocatedWorker(_fake_prefill_fn, _fake_decode_fn, batch_size=1) + worker.submit(Request("a", torch.tensor([[1]]), max_new_tokens=5)) + worker.submit(Request("b", torch.tensor([[2]]), max_new_tokens=1)) + + for _ in range(3): + worker.step() # "a" holds the only slot; "b" stays waiting + a_done = worker.step() + assert [r.request_id for r in a_done] == ["a"] # 5 tokens: 1 from prefill + 4 decode steps + + b_done = worker.step() # "b" now gets its prefill turn, already done at max_new_tokens=1 + assert [r.request_id for r in b_done] == ["b"] From e864aa6d9bbc9f929c560d6e5993063165ed3255 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:43:35 -0400 Subject: [PATCH 07/12] fix: make recv_kv_cache device-aware for NCCL NCCL (unlike gloo, which Task 2's own test uses) requires every send/recv tensor to be on the correct CUDA device. Found while wiring Task 5's real cross-rank handoff on the rented pod, before spending GPU time debugging it live -- recv_kv_cache's buffers defaulted to CPU, which would hang or error under the real 4-rank NCCL world. send_kv_cache now infers its device from the cache being sent; recv_kv_cache takes an explicit device parameter, defaulting to "cpu" so the existing gloo test is unaffected. --- src/dispatch/serving/handoff.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/dispatch/serving/handoff.py b/src/dispatch/serving/handoff.py index 12ddf74..bcf0cfc 100644 --- a/src/dispatch/serving/handoff.py +++ b/src/dispatch/serving/handoff.py @@ -16,7 +16,11 @@ def send_kv_cache(cache: DynamicCache, dst: int, group: dist.ProcessGroup) -> None: - seq_len = torch.tensor([cache.get_seq_length()], dtype=torch.int64) + # NCCL (unlike gloo) requires every send/recv tensor to be on the + # correct CUDA device -- matching the cache's own device keeps this + # gloo-safe too (the existing CPU test's caches default to "cpu"). + device = layer_kv(cache.layers[0])[0].device + seq_len = torch.tensor([cache.get_seq_length()], dtype=torch.int64, device=device) dist.send(seq_len, dst=dst, group=group) for layer in cache.layers: key, value = layer_kv(layer) @@ -24,7 +28,7 @@ def send_kv_cache(cache: DynamicCache, dst: int, group: dist.ProcessGroup) -> No dist.send(value.contiguous(), dst=dst, group=group) -def recv_kv_cache( # noqa: PLR0913 -- the receiver can't infer shape/dtype from the wire, they must be passed +def recv_kv_cache( # noqa: PLR0913 -- the receiver can't infer shape/dtype/device from the wire, they must be passed src: int, group: dist.ProcessGroup, *, @@ -32,15 +36,16 @@ def recv_kv_cache( # noqa: PLR0913 -- the receiver can't infer shape/dtype from num_heads: int, head_dim: int, dtype: torch.dtype, + device: torch.device | str = "cpu", ) -> DynamicCache: - seq_len_tensor = torch.zeros(1, dtype=torch.int64) + seq_len_tensor = torch.zeros(1, dtype=torch.int64, device=device) dist.recv(seq_len_tensor, src=src, group=group) seq_len = int(seq_len_tensor.item()) per_layer = [] for _ in range(num_layers): - key = torch.zeros(1, num_heads, seq_len, head_dim, dtype=dtype) - value = torch.zeros(1, num_heads, seq_len, head_dim, dtype=dtype) + key = torch.zeros(1, num_heads, seq_len, head_dim, dtype=dtype, device=device) + value = torch.zeros(1, num_heads, seq_len, head_dim, dtype=dtype, device=device) dist.recv(key, src=src, group=group) dist.recv(value, src=src, group=group) per_layer.append((key, value)) From 91a9f4e0961365f5feb61750df4495c9f8d4e6a8 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:44:52 -0400 Subject: [PATCH 08/12] fix: allocate pad_and_batch_caches' zero-padding on the cache's own device Found the same way as handoff.py's device bug (reading the real-GPU code path before running it, not by a live crash): torch.zeros(pad_shape, dtype=...) with no device arg always allocates on CPU, so concatenating it against a CUDA cache's key/value tensors would raise a device mismatch. CPU-only tests never caught this since every tensor involved was already on CPU. Same class of bug as Phase 3's local_expert_contribution device-placement fix. --- src/dispatch/serving/kv_cache.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/dispatch/serving/kv_cache.py b/src/dispatch/serving/kv_cache.py index fa394d3..e2b6c1b 100644 --- a/src/dispatch/serving/kv_cache.py +++ b/src/dispatch/serving/kv_cache.py @@ -78,8 +78,10 @@ def pad_and_batch_caches( key, value = layer_kv(raw_layer) if pad_len > 0: pad_shape = (1, key.shape[1], pad_len, key.shape[3]) - key = torch.cat([torch.zeros(pad_shape, dtype=key.dtype), key], dim=2) - value = torch.cat([torch.zeros(pad_shape, dtype=value.dtype), value], dim=2) + key_pad = torch.zeros(pad_shape, dtype=key.dtype, device=key.device) + value_pad = torch.zeros(pad_shape, dtype=value.dtype, device=value.device) + key = torch.cat([key_pad, key], dim=2) + value = torch.cat([value_pad, value], dim=2) per_layer_keys[layer_idx].append(key) per_layer_values[layer_idx].append(value) From f68a8e88cf8785c23f2623fd20b766a634058b98 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:51:42 -0400 Subject: [PATCH 09/12] fix: handle a rank receiving zero tokens for any local expert in local_expert_contribution Found running Phase 4's 4-way EP correctness gate: finer sharding than Phase 3's 2-way EP (16 experts/rank vs 32) plus short prompts makes it possible for a rank's topk_idx to come back completely empty after DeepEP dispatch -- topk_idx.max() then raises on the empty tensor. This is a different case from the earlier "missing the highest local expert" fix, where topk_idx was still non-empty. Regression test added. --- src/dispatch/kernels/expert_parallel.py | 11 ++++++++++- tests/unit/test_expert_parallel.py | 24 ++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/dispatch/kernels/expert_parallel.py b/src/dispatch/kernels/expert_parallel.py index 00b298a..77662a0 100644 --- a/src/dispatch/kernels/expert_parallel.py +++ b/src/dispatch/kernels/expert_parallel.py @@ -65,7 +65,16 @@ def local_expert_contribution( # routing inputs plus a pluggable GEMM and tile # 64-expert model: local_expert_ids can reach higher than whatever this # particular batch's topk_idx happens to touch), so sizing on topk_idx # alone left local_index_of[local_expert_ids] indexing out of bounds. - index_span = max(int(topk_idx.max().item()), int(local_expert_ids.max().item())) + 1 + # topk_idx can also be entirely empty -- confirmed live 2026-09-17 with + # 4-way EP (finer sharding than Phase 3's 2-way) on short prompts, where + # a rank can receive zero tokens for any of its local experts at all -- + # and .max() on an empty tensor raises, so that case sizes off + # local_expert_ids alone. + index_span = ( + int(local_expert_ids.max().item()) + 1 + if topk_idx.numel() == 0 + else max(int(topk_idx.max().item()), int(local_expert_ids.max().item())) + 1 + ) local_index_of = torch.zeros(index_span, dtype=torch.int64, device=topk_idx.device) local_index_of[local_expert_ids] = torch.arange( local_expert_ids.numel(), device=topk_idx.device diff --git a/tests/unit/test_expert_parallel.py b/tests/unit/test_expert_parallel.py index 4bf5197..faefb62 100644 --- a/tests/unit/test_expert_parallel.py +++ b/tests/unit/test_expert_parallel.py @@ -96,6 +96,30 @@ def test_local_expert_contribution_handles_batch_missing_the_highest_local_exper assert torch.equal(actual, torch.zeros_like(actual)) +def test_local_expert_contribution_handles_a_batch_with_no_tokens_for_this_rank_at_all() -> None: + """Regression test for a real bug (found 2026-09-17 running Phase 4's + 4-way EP on short prompts): with finer sharding than Phase 3's 2-way + EP, a rank can receive zero tokens for any of its local experts -- + topk_idx.max() then raises on the empty tensor, unlike the earlier + "missing the highest local expert" case where topk_idx was still + non-empty.""" + moe = ReferenceMoE(TOY_CONFIG) + weights = stack_expert_weights(moe.experts) + local_expert_ids = torch.tensor([2, 3]) + local_weights = StackedExpertWeights( + gate=weights.gate[2:4], up=weights.up[2:4], down=weights.down[2:4] + ) + x = torch.randn(0, TOY_CONFIG.hidden_size) + topk_idx = torch.zeros(0, 2, dtype=torch.int64) + topk_weight = torch.zeros(0, 2, dtype=torch.float32) + + actual = local_expert_contribution( + x, topk_idx, topk_weight, local_weights, torch_grouped_matmul, local_expert_ids + ) + + assert actual.shape == (0, TOY_CONFIG.hidden_size) + + @pytest.mark.gpu def test_local_expert_contribution_matches_cpu_when_inputs_are_on_cuda() -> None: """Regression test for a real bug (found 2026-09-16 while wiring Task 4's From a0106c71aee6773b42b6ef5f6b5850d502484e18 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:23:10 -0400 Subject: [PATCH 10/12] docs: record Phase 4 GPU runbook, cost, and correctness gate + concurrency results 4x H100 SXM, real NVLink confirmed across all pairs, DeepEP V1 install reused Phase 3's exact CUDA-13-toolkit fix cleanly. Both correctness gates (co-located 4-rank EP, disaggregated 2+2-rank EP) passed with exact greedy-token-sequence match against a single-GPU reference. Concurrency measurement at 4 and 8 concurrent requests collected for both topologies, though the comparison itself is confounded by kernel-warmup state between independent process launches -- reported as genuinely inconclusive rather than forced, per this project's practice of writing down what the data actually supports. $10.94 of the $40 cap, 47 minutes total. --- .../2026-09-16-phase-4-colocated-c4.json | 10 + .../2026-09-16-phase-4-colocated-c8.json | 10 + ...-09-16-phase-4-colocated-gate-results.json | 59 +++ ...09-16-phase-4-disaggregated-decode-c4.json | 7 + ...09-16-phase-4-disaggregated-decode-c8.json | 7 + ...16-phase-4-disaggregated-gate-results.json | 59 +++ ...9-16-phase-4-disaggregated-prefill-c4.json | 5 + ...9-16-phase-4-disaggregated-prefill-c8.json | 5 + ...26-09-16-phase-4-single-gpu-reference.json | 56 +++ ...ase-4-disaggregated-prefill-decode-cost.md | 7 + .../phase-4-disaggregated-prefill-decode.md | 176 ++++++++ pyproject.toml | 9 + scripts/gpu/phase4_concurrency.py | 376 ++++++++++++++++++ scripts/gpu/phase4_correctness_gate.py | 333 ++++++++++++++++ 14 files changed, 1119 insertions(+) create mode 100644 docs/findings/2026-09-16-phase-4-colocated-c4.json create mode 100644 docs/findings/2026-09-16-phase-4-colocated-c8.json create mode 100644 docs/findings/2026-09-16-phase-4-colocated-gate-results.json create mode 100644 docs/findings/2026-09-16-phase-4-disaggregated-decode-c4.json create mode 100644 docs/findings/2026-09-16-phase-4-disaggregated-decode-c8.json create mode 100644 docs/findings/2026-09-16-phase-4-disaggregated-gate-results.json create mode 100644 docs/findings/2026-09-16-phase-4-disaggregated-prefill-c4.json create mode 100644 docs/findings/2026-09-16-phase-4-disaggregated-prefill-c8.json create mode 100644 docs/findings/2026-09-16-phase-4-single-gpu-reference.json create mode 100644 docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-cost.md create mode 100644 docs/runbooks/phase-4-disaggregated-prefill-decode.md create mode 100644 scripts/gpu/phase4_concurrency.py create mode 100644 scripts/gpu/phase4_correctness_gate.py diff --git a/docs/findings/2026-09-16-phase-4-colocated-c4.json b/docs/findings/2026-09-16-phase-4-colocated-c4.json new file mode 100644 index 0000000..a2092f8 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/docs/findings/2026-09-16-phase-4-colocated-c8.json b/docs/findings/2026-09-16-phase-4-colocated-c8.json new file mode 100644 index 0000000..30a1423 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/docs/findings/2026-09-16-phase-4-colocated-gate-results.json b/docs/findings/2026-09-16-phase-4-colocated-gate-results.json new file mode 100644 index 0000000..099f2b0 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/docs/findings/2026-09-16-phase-4-disaggregated-decode-c4.json b/docs/findings/2026-09-16-phase-4-disaggregated-decode-c4.json new file mode 100644 index 0000000..d7ab9c8 --- /dev/null +++ b/docs/findings/2026-09-16-phase-4-disaggregated-decode-c4.json @@ -0,0 +1,7 @@ +{ + "topology": "disaggregated", + "concurrency": 4, + "max_new_tokens": 8, + "total_wall_s": 4.255717797204852, + "decode_tokens_per_s": 6.579383627925317 +} \ No newline at end of file diff --git a/docs/findings/2026-09-16-phase-4-disaggregated-decode-c8.json b/docs/findings/2026-09-16-phase-4-disaggregated-decode-c8.json new file mode 100644 index 0000000..7955867 --- /dev/null +++ b/docs/findings/2026-09-16-phase-4-disaggregated-decode-c8.json @@ -0,0 +1,7 @@ +{ + "topology": "disaggregated", + "concurrency": 8, + "max_new_tokens": 8, + "total_wall_s": 4.7249182015657425, + "decode_tokens_per_s": 11.852057032742437 +} \ No newline at end of file diff --git a/docs/findings/2026-09-16-phase-4-disaggregated-gate-results.json b/docs/findings/2026-09-16-phase-4-disaggregated-gate-results.json new file mode 100644 index 0000000..099f2b0 --- /dev/null +++ b/docs/findings/2026-09-16-phase-4-disaggregated-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 +} \ No newline at end of file diff --git a/docs/findings/2026-09-16-phase-4-disaggregated-prefill-c4.json b/docs/findings/2026-09-16-phase-4-disaggregated-prefill-c4.json new file mode 100644 index 0000000..8a910cb --- /dev/null +++ b/docs/findings/2026-09-16-phase-4-disaggregated-prefill-c4.json @@ -0,0 +1,5 @@ +{ + "mean_ttft_s": 0.45823238184675574, + "p50_ttft_s": 0.15552851185202599, + "p99_ttft_s": 1.330020813848823 +} \ No newline at end of file diff --git a/docs/findings/2026-09-16-phase-4-disaggregated-prefill-c8.json b/docs/findings/2026-09-16-phase-4-disaggregated-prefill-c8.json new file mode 100644 index 0000000..c1fe5f5 --- /dev/null +++ b/docs/findings/2026-09-16-phase-4-disaggregated-prefill-c8.json @@ -0,0 +1,5 @@ +{ + "mean_ttft_s": 0.983093936694786, + "p50_ttft_s": 0.7487213974818587, + "p99_ttft_s": 1.986582173947245 +} \ No newline at end of file diff --git a/docs/findings/2026-09-16-phase-4-single-gpu-reference.json b/docs/findings/2026-09-16-phase-4-single-gpu-reference.json new file mode 100644 index 0000000..eeafd3b --- /dev/null +++ b/docs/findings/2026-09-16-phase-4-single-gpu-reference.json @@ -0,0 +1,56 @@ +{ + "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 + ] +} \ No newline at end of file diff --git a/docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-cost.md b/docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-cost.md new file mode 100644 index 0000000..07ba175 --- /dev/null +++ b/docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-cost.md @@ -0,0 +1,7 @@ +# phase-4-disaggregated-prefill-decode -- GPU rental cost + +- Pod: `ds338r4byldpqu` (NVIDIA H100 80GB HBM3) +- Rate: $13.9600/hr +- Duration: 2820s (0.783hr) +- Cost: $10.9353 +- Note: Phase 4: CUDA 13.0 toolkit + DeepEP V1 install, 4-GPU NVLink verification, co-located (4-rank) and disaggregated (2+2-rank) EP correctness gates (both exact-token-match pass), concurrency measurement at 4 and 8 concurrent requests for both topologies diff --git a/docs/runbooks/phase-4-disaggregated-prefill-decode.md b/docs/runbooks/phase-4-disaggregated-prefill-decode.md new file mode 100644 index 0000000..a0d4cc1 --- /dev/null +++ b/docs/runbooks/phase-4-disaggregated-prefill-decode.md @@ -0,0 +1,176 @@ +# Phase 4 runbook: disaggregated prefill/decode + +Budget cap: $40, a ceiling not a target. Hopper-class (H100/H200, SM90) +required, same as Phase 3. + +**What actually happened (2026-09-17), read this before following the +steps below as written:** everything in this runbook ran successfully, +but real hardware and a real live `get-capacity` price quote both +corrected the plan before any measurement happened. Total cost: +**$10.94** of the $40 cap (docs/findings/ +2026-09-17-phase-4-disaggregated-prefill-decode-cost.md). Full account, +including the correctness gate results and the measured contention +comparison: `docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md`. + +1. Create and wait for the 4-GPU pod, reusing Phase 3's generic + `--gpu-count` support: + + ```bash + uv run python scripts/gpu/provision.py create --name phase-4-disaggregated \ + --gpu-type --gpu-count 4 \ + --image runpod/pytorch:1.0.2-cu1281-torch280-ubuntu2404 --cloud secure --disk-gb 60 + uv run python scripts/gpu/provision.py wait + ``` + + **What actually happened:** `get-capacity` quoted 4x H100 SXM at + $10.76/hr (CUDA 13.0, Secure cloud) -- the pod actually billed at + **$13.96/hr**, ~30% higher than the catalog check. Not a planning + error to avoid, just a real gap between the catalog-lookup price and + the actual create-pod price worth expecting; confirm the real rate + from the pod object itself (its own `cost` field) immediately after + creation, not just from the pre-rental catalog check. Disk bumped to + 60GB (from Phase 3's 40GB) specifically to avoid Phase 3's + disk-space-vs-32.8GB-model-download issue; it was enough (47GB free + after all installs, before the download). + +2. **Verify real NVLink across all four GPUs**: + + ```bash + nvidia-smi topo -m + ``` + + Expected: `NV#` between every pair of the four GPUs. **Confirmed**: + every pair showed `NV18` (18 bonded links). + +3. Install DeepEP V1 (legacy `Buffer`) directly -- do not attempt V2, + Phase 3 already confirmed it cannot work on this rental tier (no GPU + Fabric Manager). Install the matching CUDA toolkit first (this + image's nvcc was missing on PATH; NVIDIA's apt repo was already + configured): + + ```bash + apt-get update -qq + DEBIAN_FRONTEND=noninteractive apt-get install -y -qq \ + cuda-toolkit-13-0 libibverbs-dev rdma-core ninja-build + export PATH=/usr/local/cuda-13.0/bin:$PATH + + uv sync # pulls torch==2.14.0+cu130, matching the 13.0 toolkit + + git clone https://github.com/deepseek-ai/DeepEP.git + cd DeepEP + uv pip install --python ../.venv/bin/python "nvidia-nccl-cu13>=2.30.4" + PATH=../.venv/bin:$PATH ../.venv/bin/python setup.py install + cd .. + .venv/bin/python -c "from deep_ep import Buffer; print('V1 Buffer OK')" + ``` + + **What actually happened:** this worked cleanly on the first attempt + -- DeepEP V1 built and imported with no further issues, reusing + exactly the CUDA-13-toolkit fix Phase 3 had to discover live. + +4. Clone this repo, sync, and run the full CPU-testable suite one more + time on the pod itself: + + ```bash + git clone https://github.com/bsreecharanreddy/dispatch.git dispatch + cd dispatch + git checkout phase-4-disaggregated-prefill-decode + uv sync + uv run pytest tests/unit -v -m "not gpu" + ``` + + Expected and confirmed: all 100 pass (private repo -- cloned via an + `x-access-token` HTTPS URL piped through stdin, never as a bare + command-line argument). + +5. **Build the two EP sub-groups and wire the real forward closures.** + `torchrun --nproc_per_node=4` gives a 4-rank NCCL world; every rank + calls `dist.new_group([0, 1])` then `dist.new_group([2, 3])` + identically (a collective call every rank must join regardless of + membership), then branches on `rank in (0, 1)` to pick its own + sub-group. Each sub-group builds its own DeepEP `Buffer` scoped to + that sub-group and calls `patch_moe_infer_ep(model, matmul, buffer, + pool_rank, n_ranks=2)` unchanged, where `pool_rank = + dist.get_rank(group=pool_group)` -- the **sub-group-local** rank (0 or + 1), not the global rank. + + **What actually happened, two real fixes needed before this worked:** + - DeepSeek's remote-code modeling file (unchanged since Phase 0) + predates the `transformers` Cache-class refactor entirely: it + returns `past_key_values` as the **legacy tuple-of-(key, + value)-per-layer** format, not a `DynamicCache` object, confirmed + live by inspecting a real forward call's output type. Every + `PrefillFn`/`DecodeFn` closure converts at this boundary only, via + `DynamicCache.from_legacy_cache(...)` / `cache.to_legacy_cache()` + (both still present in `transformers==4.57.6`, though removed in + v5) -- `kv_cache.py`/`disaggregated.py`'s own committed + `DynamicCache`-based contract, proven by Tasks 1-4's CPU tests, + needed no changes at all. + - `local_expert_contribution` (`expert_parallel.py`, committed in + Phase 3) crashed with `RuntimeError: max(): Expected reduction dim + to be specified for input.numel() == 0` -- a rank's `topk_idx` came + back **completely empty** after DeepEP dispatch. Finer sharding + (4-way EP, 16 experts/rank) than Phase 3's 2-way (32 experts/rank) + plus short prompts made this a real, reachable case Phase 3 never + hit. Fixed in `src/dispatch/kernels/expert_parallel.py` + (`index_span` now sizes off `local_expert_ids` alone when + `topk_idx` is empty), with a CPU regression test, found and fixed + for $0 before spending more GPU time debugging it live. + + Rank-to-rank KV-cache handoff for the disaggregated path uses + `handoff.py`'s `send_kv_cache`/`recv_kv_cache` over the full 4-rank + `dist.group.WORLD`, not a sub-group (prefill and decode ranks must + talk to each other). **A third real fix was needed here too, found + by re-reading the committed code before running it, not by a live + crash**: NCCL (unlike the `gloo` backend Task 2's own CPU test uses) + requires every send/recv tensor to be on the correct CUDA device; + `recv_kv_cache`'s buffers defaulted to CPU. Fixed by making + `send_kv_cache` infer its device from the cache being sent and adding + an explicit `device` parameter to `recv_kv_cache`. A companion device + bug in `pad_and_batch_caches` (its zero-padding tensors also defaulted + to CPU) was caught the same way, before it could crash a real batched + decode step. + +6. **Correctness gate -- must pass for both configurations before any + measurement runs.** `scripts/gpu/phase4_correctness_gate.py --mode + single` first (also warms the model cache), then `--mode colocated` + and `--mode disaggregated` via `torchrun --nproc_per_node=4`. Every + rank in a pool runs the identical scheduler logic (`PrefillWorker`/ + `DecodeWorker`/`ColocatedWorker`) against the identical prompt list -- + since that logic is pure and deterministic, every rank's `model(...)` + calls stay naturally synchronized without a driver/follower split, + the same pattern Phase 3's own correctness gate used. Correctness is + judged by exact greedy-token-sequence match against the single-GPU + reference (not top-k logit agreement -- the scheduler's `PrefillFn`/ + `DecodeFn` contract only returns argmax token ids, dropping raw + logits by design), which is at least as strict a bar for a + greedy-decode pipeline. + + **Confirmed: both passed, byte-exact token match**, all 3 prompts, 16 + tokens each. See the findings doc for the full token sequences. + +7. **Only after both gates passed**, ran the concurrency measurement: + co-located vs. disaggregated, same 4 GPUs, at 4 and 8 concurrent + requests (`scripts/gpu/phase4_concurrency.py`, staggered arrivals 50ms + apart, 8 max new tokens per request -- short by design, to keep this + step cheap once the hard correctness work was already done). The + disaggregated decode side uses `dist.irecv` + polling (`is_completed()`) + to interleave receiving new prefill handoffs with stepping its own + already-active decode requests, rather than blocking on `dist.recv`. + + Full measured numbers: findings doc. **Be surgical** held throughout + -- total session was 47 minutes end to end. + +8. Copied all console output and result JSON off the pod via `scp` + before doing anything else. + +9. Tore down immediately: + + ``` + pod-action terminate ds338r4byldpqu + ``` + + Confirmed via a follow-up `list-pods` returning an empty list. + +10. Recorded the measured cost via `write_cost_record`, reused + unmodified: `docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-cost.md`. diff --git a/pyproject.toml b/pyproject.toml index 278bdf2..b28fa60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,15 @@ max-complexity = 10 # Expert-parallel sharding/dispatch functions carry routing inputs plus a # pluggable GEMM and tile size, same shape as moe_forward.py's own functions. "src/dispatch/kernels/expert_parallel.py" = ["PLR0913", "PLR0917"] +# Phase 4's pod-live rank-orchestration scripts (torchrun-only, not +# imported by anything else): the prefill/decode role branch inside one +# function is what keeps every rank's collective calls in the same +# deterministic order (see their own docstrings), so splitting it up +# would risk the exact synchronization these scripts were debugged +# against. Rank ids (0/1 prefill, 2/3 decode) are the real protocol, not +# magic numbers to extract into constants. +"scripts/gpu/phase4_correctness_gate.py" = ["PLR0915", "PLR2004"] +"scripts/gpu/phase4_concurrency.py" = ["PLR0915", "PLR2004", "C901"] [tool.mypy] python_version = "3.12" diff --git a/scripts/gpu/phase4_concurrency.py b/scripts/gpu/phase4_concurrency.py new file mode 100644 index 0000000..485cdca --- /dev/null +++ b/scripts/gpu/phase4_concurrency.py @@ -0,0 +1,376 @@ +"""Phase 4 concurrency/contention measurement: co-located (one 4-rank EP +pool doing both prefill and decode) vs disaggregated (2-rank prefill pool ++ 2-rank decode pool, connected by handoff.py's real cross-rank cache +transfer), same 4 GPUs, at a few concurrency levels, staggered arrivals. +Only run after both correctness gates in phase4_correctness_gate.py have +passed -- this script does not re-verify correctness, per this project's +"correctness gates the benchmark" rule. + +Co-located: + NCCL_NVLS_ENABLE=0 torchrun --nproc_per_node=4 phase4_concurrency.py \ + --mode colocated --concurrency 4 + NCCL_NVLS_ENABLE=0 torchrun --nproc_per_node=4 phase4_concurrency.py \ + --mode colocated --concurrency 8 + +Disaggregated: + NCCL_NVLS_ENABLE=0 torchrun --nproc_per_node=4 phase4_concurrency.py \ + --mode disaggregated --concurrency 4 + NCCL_NVLS_ENABLE=0 torchrun --nproc_per_node=4 phase4_concurrency.py \ + --mode disaggregated --concurrency 8 +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from deep_ep import Buffer +from transformers import DynamicCache, PreTrainedTokenizerBase + +from dispatch.benchmark.harness import load_model +from dispatch.kernels.backends import resolve_backend +from dispatch.kernels.expert_parallel import patch_moe_infer_ep +from dispatch.serving.colocated import ColocatedWorker +from dispatch.serving.disaggregated import ( + DecodeFn, + DecodeWorker, + PrefillFn, + PrefillResult, + PrefillWorker, + Request, + RequestResult, +) +from dispatch.serving.kv_cache import layer_kv + +MODEL_NAME = "deepseek-ai/deepseek-moe-16b-base" +DEFAULT_PROMPTS = [ + "The quick brown fox jumps over the lazy dog.", + "In a distant galaxy, a small crew of explorers", + "def fibonacci(n):", +] +MAX_NEW_TOKENS = 8 +ARRIVAL_STAGGER_S = 0.05 +OUTPUT_DIR = Path("docs/findings") + + +# See phase4_correctness_gate.py's own module docstring/comments for why +# this patch and the from_legacy_cache/to_legacy_cache conversions below +# are needed -- same DeepSeek remote-code quirk, same transformers 4.57.6 +# pin, not repeated here. +def _get_usable_length( + self: DynamicCache, new_seq_length: int | None = None, layer_idx: int = 0 +) -> int: + return self.get_seq_length(layer_idx) + + +DynamicCache.get_usable_length = _get_usable_length # type: ignore[attr-defined] + + +def make_prefill_fn(model: torch.nn.Module, device: str) -> PrefillFn: + def prefill_fn(input_ids: torch.Tensor, attention_mask: torch.Tensor) -> Any: + with torch.no_grad(): + outputs = model( + input_ids=input_ids.to(device), + attention_mask=attention_mask.to(device), + use_cache=True, + ) + first_tokens = outputs.logits[:, -1, :].argmax(dim=-1) + cache = DynamicCache.from_legacy_cache(outputs.past_key_values) # type: ignore[attr-defined] + return first_tokens.cpu(), cache + + return prefill_fn + + +def make_decode_fn(model: torch.nn.Module, device: str) -> DecodeFn: + def decode_fn( + next_input_ids: torch.Tensor, + cache: DynamicCache, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + ) -> Any: + with torch.no_grad(): + outputs = model( + input_ids=next_input_ids.to(device), + past_key_values=cache.to_legacy_cache(), # type: ignore[attr-defined] + attention_mask=attention_mask.to(device), + position_ids=position_ids.to(device), + use_cache=True, + ) + next_tokens = outputs.logits[:, -1, :].argmax(dim=-1) + updated_cache = DynamicCache.from_legacy_cache(outputs.past_key_values) # type: ignore[attr-defined] + return next_tokens.cpu(), updated_cache + + return decode_fn + + +def _make_requests(concurrency: int, tokenizer: PreTrainedTokenizerBase) -> list[Request]: + requests = [] + for i in range(concurrency): + prompt = DEFAULT_PROMPTS[i % len(DEFAULT_PROMPTS)] + input_ids = tokenizer(prompt, return_tensors="pt").input_ids + requests.append(Request(str(i), input_ids, max_new_tokens=MAX_NEW_TOKENS)) + return requests + + +def _percentile(values: list[float], fraction: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + if len(ordered) == 1: + return ordered[0] + index = fraction * (len(ordered) - 1) + lower = int(index) + upper = min(lower + 1, len(ordered) - 1) + weight = index - lower + return ordered[lower] * (1 - weight) + ordered[upper] * weight + + +def run_colocated(concurrency: int) -> None: + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank) + dist.barrier() + world_group = dist.group.WORLD + assert world_group is not None + + model, tokenizer = load_model( + MODEL_NAME, device="cuda", dtype=torch.bfloat16, trust_remote_code=True + ) + hidden_bytes = model.config.hidden_size * 2 + Buffer.set_num_sms(24) + dispatch_config = Buffer.get_dispatch_config(world_size) + combine_config = Buffer.get_combine_config(world_size) + num_nvl_bytes = max( + dispatch_config.get_nvl_buffer_size_hint(hidden_bytes, world_size), + combine_config.get_nvl_buffer_size_hint(hidden_bytes, world_size), + ) + buffer = Buffer(world_group, num_nvl_bytes, 0) + patch_moe_infer_ep(model, resolve_backend("naive"), buffer, rank, world_size) + + requests = _make_requests(concurrency, tokenizer) + arrival_times = [i * ARRIVAL_STAGGER_S for i in range(concurrency)] + colocated_worker = ColocatedWorker( + make_prefill_fn(model, "cuda"), make_decode_fn(model, "cuda"), batch_size=4 + ) + + dist.barrier() + t0 = time.perf_counter() + pending = list(zip(arrival_times, requests, strict=True)) + completed: list[RequestResult] = [] + while len(completed) < concurrency: + now = time.perf_counter() - t0 + while pending and pending[0][0] <= now: + _, req = pending.pop(0) + colocated_worker.submit(req) + completed.extend(colocated_worker.step()) + total_wall_s = time.perf_counter() - t0 + + if rank == 0: + ttfts = [r.ttft for r in completed] + total_tokens = sum( + len(r.generated_token_ids) - 1 for r in completed + ) # decode-generated only + result = { + "topology": "colocated", + "concurrency": concurrency, + "max_new_tokens": MAX_NEW_TOKENS, + "total_wall_s": total_wall_s, + "mean_ttft_s": statistics.mean(ttfts), + "p50_ttft_s": _percentile(ttfts, 0.5), + "p99_ttft_s": _percentile(ttfts, 0.99), + "decode_tokens_per_s": total_tokens / total_wall_s, + } + print(f"[rank 0] colocated concurrency={concurrency}: {result}") + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + (OUTPUT_DIR / f"2026-09-16-phase-4-colocated-c{concurrency}.json").write_text( + json.dumps(result, indent=2) + ) + + dist.destroy_process_group() + + +def _run_prefill_side( # noqa: PLR0913, PLR0917 -- one call site (run_disaggregated), each arg independently needed + concurrency: int, + rank: int, + world_group: dist.ProcessGroup, + model: torch.nn.Module, + tokenizer: PreTrainedTokenizerBase, + t0: float, +) -> None: + requests = _make_requests(concurrency, tokenizer) + arrival_times = [i * ARRIVAL_STAGGER_S for i in range(concurrency)] + prefill_worker = PrefillWorker(make_prefill_fn(model, "cuda"), batch_size=4) + dst = rank + 2 + pending = list(zip(arrival_times, requests, strict=True)) + sent = 0 + ttfts: list[float] = [] + while sent < concurrency: + now = time.perf_counter() - t0 + while pending and pending[0][0] <= now: + _, req = pending.pop(0) + prefill_worker.submit(req) + for result in prefill_worker.step(): + send_kv_cache_inline(result.cache, dst, world_group) + meta = torch.tensor( + [int(result.request_id), result.first_token_id], dtype=torch.int64, device="cuda" + ) + dist.send(meta, dst=dst, group=world_group) + ttfts.append(result.ttft) + sent += 1 + if rank == 0: + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + (OUTPUT_DIR / f"2026-09-16-phase-4-disaggregated-prefill-c{concurrency}.json").write_text( + json.dumps( + { + "mean_ttft_s": statistics.mean(ttfts), + "p50_ttft_s": _percentile(ttfts, 0.5), + "p99_ttft_s": _percentile(ttfts, 0.99), + }, + indent=2, + ) + ) + print(f"[rank 0] disaggregated prefill side done, mean_ttft={statistics.mean(ttfts):.4f}s") + + +def _run_decode_side( # noqa: PLR0913, PLR0917 -- one call site (run_disaggregated), each arg independently needed + concurrency: int, + rank: int, + world_group: dist.ProcessGroup, + model: torch.nn.Module, + num_layers: int, + num_heads: int, + head_dim: int, + t0: float, +) -> None: + decode_worker = DecodeWorker(make_decode_fn(model, "cuda"), batch_size=8) + src = rank - 2 + received = 0 + completed: list[RequestResult] = [] + seq_len_buf = torch.zeros(1, dtype=torch.int64, device="cuda") + pending_recv: Any = dist.irecv(seq_len_buf, src=src, group=world_group) + while received < concurrency or len(completed) < concurrency: + if pending_recv is not None and pending_recv.is_completed(): + seq_len = int(seq_len_buf.item()) + per_layer = [] + for _ in range(num_layers): + key = torch.zeros( + 1, num_heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda" + ) + value = torch.zeros( + 1, num_heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda" + ) + dist.recv(key, src=src, group=world_group) + dist.recv(value, src=src, group=world_group) + per_layer.append((key, value)) + cache = DynamicCache(ddp_cache_data=per_layer) + meta = torch.zeros(2, dtype=torch.int64, device="cuda") + dist.recv(meta, src=src, group=world_group) + request_idx, first_token_id = int(meta[0].item()), int(meta[1].item()) + decode_worker.admit( + PrefillResult( + request_id=str(request_idx), + cache=cache, + first_token_id=first_token_id, + ttft=0.0, + max_new_tokens=MAX_NEW_TOKENS, + eos_token_id=None, + ) + ) + received += 1 + if received < concurrency: + seq_len_buf = torch.zeros(1, dtype=torch.int64, device="cuda") + pending_recv = dist.irecv(seq_len_buf, src=src, group=world_group) + else: + pending_recv = None + completed.extend(decode_worker.step()) + total_wall_s = time.perf_counter() - t0 + decode_pool_local_rank_0 = 2 # rank 2 is the decode pool's local rank 0, the protocol + if rank == decode_pool_local_rank_0: + total_tokens = sum(len(r.generated_token_ids) - 1 for r in completed) + result = { + "topology": "disaggregated", + "concurrency": concurrency, + "max_new_tokens": MAX_NEW_TOKENS, + "total_wall_s": total_wall_s, + "decode_tokens_per_s": total_tokens / total_wall_s, + } + print(f"[rank 2] disaggregated decode side concurrency={concurrency}: {result}") + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + (OUTPUT_DIR / f"2026-09-16-phase-4-disaggregated-decode-c{concurrency}.json").write_text( + json.dumps(result, indent=2) + ) + + +def run_disaggregated(concurrency: int) -> None: + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + torch.cuda.set_device(rank) + world_group = dist.group.WORLD + prefill_group = dist.new_group([0, 1]) + decode_group = dist.new_group([2, 3]) + assert world_group is not None + dist.barrier() + + is_prefill = rank in (0, 1) + pool_group = prefill_group if is_prefill else decode_group + assert isinstance(pool_group, dist.ProcessGroup) + pool_rank = dist.get_rank(group=pool_group) + + model, tokenizer = load_model( + MODEL_NAME, device="cuda", dtype=torch.bfloat16, trust_remote_code=True + ) + hidden_bytes = model.config.hidden_size * 2 + Buffer.set_num_sms(24) + dispatch_config = Buffer.get_dispatch_config(2) + combine_config = Buffer.get_combine_config(2) + num_nvl_bytes = max( + dispatch_config.get_nvl_buffer_size_hint(hidden_bytes, 2), + combine_config.get_nvl_buffer_size_hint(hidden_bytes, 2), + ) + buffer = Buffer(pool_group, num_nvl_bytes, 0) + patch_moe_infer_ep(model, resolve_backend("naive"), buffer, pool_rank, 2) + + num_heads = model.config.num_attention_heads + head_dim = model.config.hidden_size // num_heads + num_layers = model.config.num_hidden_layers + + dist.barrier() + t0 = time.perf_counter() + + if is_prefill: + _run_prefill_side(concurrency, rank, world_group, model, tokenizer, t0) + else: + _run_decode_side(concurrency, rank, world_group, model, num_layers, num_heads, head_dim, t0) + + dist.destroy_process_group() + + +def send_kv_cache_inline(cache: DynamicCache, dst: int, group: dist.ProcessGroup) -> None: + seq_len = torch.tensor([cache.get_seq_length()], dtype=torch.int64, device="cuda") + dist.send(seq_len, dst=dst, group=group) + for layer in cache.layers: + key, value = layer_kv(layer) + dist.send(key.contiguous(), dst=dst, group=group) + dist.send(value.contiguous(), dst=dst, group=group) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--mode", required=True, choices=["colocated", "disaggregated"]) + parser.add_argument("--concurrency", type=int, required=True) + args = parser.parse_args() + if args.mode == "colocated": + run_colocated(args.concurrency) + else: + run_disaggregated(args.concurrency) + + +if __name__ == "__main__": + main() diff --git a/scripts/gpu/phase4_correctness_gate.py b/scripts/gpu/phase4_correctness_gate.py new file mode 100644 index 0000000..b1c44ef --- /dev/null +++ b/scripts/gpu/phase4_correctness_gate.py @@ -0,0 +1,333 @@ +"""Phase 4 correctness gate: does the real DeepSeekMoE-16B model produce +the exact same greedily-generated tokens under (a) a 4-rank co-located +EP pool and (b) two 2-rank disaggregated EP pools (prefill + decode, +connected by handoff.py's real cross-rank cache transfer) as a plain +single-GPU reference? Every rank in a pool runs the identical scheduler +logic against the identical prompt list -- PrefillWorker/DecodeWorker/ +ColocatedWorker are pure and deterministic, so this keeps every rank's +model(...) calls naturally synchronized without a driver/follower split, +the same pattern Phase 3's own correctness gate already used (every EP +rank calling capture_reference_logits on the identical prompt). + +Single-GPU reference: + uv run python phase4_correctness_gate.py --mode single + +Co-located (4 ranks, one EP pool): + NCCL_NVLS_ENABLE=0 torchrun --nproc_per_node=4 phase4_correctness_gate.py --mode colocated + +Disaggregated (4 ranks, prefill=[0,1] decode=[2,3]): + NCCL_NVLS_ENABLE=0 torchrun --nproc_per_node=4 phase4_correctness_gate.py --mode disaggregated +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from deep_ep import Buffer +from transformers import DynamicCache + +from dispatch.benchmark.harness import load_model +from dispatch.kernels.backends import resolve_backend +from dispatch.kernels.expert_parallel import patch_moe_infer_ep +from dispatch.kernels.integration import patch_moe_infer +from dispatch.serving.colocated import ColocatedWorker +from dispatch.serving.disaggregated import ( + DecodeFn, + DecodeWorker, + PrefillFn, + PrefillResult, + PrefillWorker, + Request, +) +from dispatch.serving.handoff import recv_kv_cache, send_kv_cache + + +# Pod-local, never committed to the repo -- DeepSeek's own remote-code +# modeling file (unchanged since Phase 0) still calls the pre-v5 +# DynamicCache.get_usable_length, renamed to get_seq_length even in +# transformers 4.57.6 (the pin Phase 0/1/3 all needed for the *other* +# direction: transformers 5.17.0 removed is_torch_fx_available, which +# the same modeling file also imports). Confirmed equivalent for this +# model (no sliding-window attention) by Phase 0's own investigation. +# transformers' own type stubs don't declare this pre-v5 method (or +# from_legacy_cache/to_legacy_cache below), hence the type: ignore +# comments throughout this file -- real, version-specific runtime +# behavior the stubs don't (and can't) capture. +def _get_usable_length( + self: DynamicCache, new_seq_length: int | None = None, layer_idx: int = 0 +) -> int: + return self.get_seq_length(layer_idx) + + +DynamicCache.get_usable_length = _get_usable_length # type: ignore[attr-defined] + +MODEL_NAME = "deepseek-ai/deepseek-moe-16b-base" +DEFAULT_PROMPTS = [ + "The quick brown fox jumps over the lazy dog.", + "In a distant galaxy, a small crew of explorers", + "def fibonacci(n):", +] +MAX_NEW_TOKENS = 16 +OUTPUT_DIR = Path("docs/findings") +REFERENCE_PATH = OUTPUT_DIR / "2026-09-16-phase-4-single-gpu-reference.json" + + +# DeepSeek's remote-code modeling file (unchanged since Phase 0) predates +# the Cache-class refactor: it accepts and returns past_key_values as the +# legacy tuple-of-(key, value)-per-layer format, not a DynamicCache +# object, confirmed live by inspecting a real forward call's output type +# (plain tuple, len == num_layers). Converting at this boundary only -- +# via from_legacy_cache/to_legacy_cache, still present in transformers +# 4.57.6 though removed in v5 -- keeps kv_cache.py/disaggregated.py's own +# DynamicCache-based contract, proven by Tasks 1-4's CPU tests, unchanged. +def make_prefill_fn(model: torch.nn.Module, device: str) -> PrefillFn: + def prefill_fn(input_ids: torch.Tensor, attention_mask: torch.Tensor) -> Any: + with torch.no_grad(): + outputs = model( + input_ids=input_ids.to(device), + attention_mask=attention_mask.to(device), + use_cache=True, + ) + first_tokens = outputs.logits[:, -1, :].argmax(dim=-1) + cache = DynamicCache.from_legacy_cache(outputs.past_key_values) # type: ignore[attr-defined] + return first_tokens.cpu(), cache + + return prefill_fn + + +def make_decode_fn(model: torch.nn.Module, device: str) -> DecodeFn: + def decode_fn( + next_input_ids: torch.Tensor, + cache: DynamicCache, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + ) -> Any: + with torch.no_grad(): + outputs = model( + input_ids=next_input_ids.to(device), + past_key_values=cache.to_legacy_cache(), # type: ignore[attr-defined] + attention_mask=attention_mask.to(device), + position_ids=position_ids.to(device), + use_cache=True, + ) + next_tokens = outputs.logits[:, -1, :].argmax(dim=-1) + updated_cache = DynamicCache.from_legacy_cache(outputs.past_key_values) # type: ignore[attr-defined] + return next_tokens.cpu(), updated_cache + + return decode_fn + + +def run_single_gpu_reference() -> None: + model, tokenizer = load_model( + MODEL_NAME, device="cuda", dtype=torch.bfloat16, trust_remote_code=True + ) + patched = patch_moe_infer(model, resolve_backend("naive")) + if patched == 0: + raise RuntimeError("patched no MoE layers -- naive kernel never ran") + print(f"single-GPU reference: patched {patched} MoE layers") + + results = {} + for prompt in DEFAULT_PROMPTS: + input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda") + generated: list[int] = [] + past = None + next_input = input_ids + with torch.no_grad(): + for _ in range(MAX_NEW_TOKENS): + outputs = model(input_ids=next_input, past_key_values=past, use_cache=True) + past = outputs.past_key_values + next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) + generated.append(int(next_token.item())) + next_input = next_token + results[prompt] = generated + print(f" {prompt!r} -> {generated}") + + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + REFERENCE_PATH.write_text(json.dumps(results, indent=2)) + print(f"wrote {REFERENCE_PATH}") + + +def _run_worker_to_completion(worker: ColocatedWorker | DecodeWorker, request_id: str) -> list[int]: + while True: + for result in worker.step(): + if result.request_id == request_id: + return result.generated_token_ids + + +def run_colocated() -> None: + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank) + dist.barrier() + world_group = dist.group.WORLD + assert world_group is not None + + model, tokenizer = load_model( + MODEL_NAME, device="cuda", dtype=torch.bfloat16, trust_remote_code=True + ) + hidden_bytes = model.config.hidden_size * 2 + Buffer.set_num_sms(24) + dispatch_config = Buffer.get_dispatch_config(world_size) + combine_config = Buffer.get_combine_config(world_size) + num_nvl_bytes = max( + dispatch_config.get_nvl_buffer_size_hint(hidden_bytes, world_size), + combine_config.get_nvl_buffer_size_hint(hidden_bytes, world_size), + ) + buffer = Buffer(world_group, num_nvl_bytes, 0) + patched = patch_moe_infer_ep(model, resolve_backend("naive"), buffer, rank, world_size) + if patched == 0: + raise RuntimeError(f"[rank {rank}] patched no MoE layers") + print(f"[rank {rank}] colocated: patched {patched} MoE layers, world_size={world_size}") + + colocated_worker = ColocatedWorker( + make_prefill_fn(model, "cuda"), make_decode_fn(model, "cuda"), batch_size=1 + ) + results: dict[str, list[int]] = {} + for prompt in DEFAULT_PROMPTS: + input_ids = tokenizer(prompt, return_tensors="pt").input_ids + colocated_worker.submit(Request(prompt, input_ids, max_new_tokens=MAX_NEW_TOKENS)) + results[prompt] = _run_worker_to_completion(colocated_worker, prompt) + + if rank == 0: + print("[rank 0] colocated results:") + for prompt, tokens in results.items(): + print(f" {prompt!r} -> {tokens}") + reference = json.loads(REFERENCE_PATH.read_text()) + match = all(results[p] == reference[p] for p in DEFAULT_PROMPTS) + print(f"[rank 0] MATCH: {match}") + (OUTPUT_DIR / "2026-09-16-phase-4-colocated-gate-results.json").write_text( + json.dumps({"results": results, "match": match}, indent=2) + ) + if not match: + dist.destroy_process_group() + raise SystemExit("colocated EP disagrees with the single-GPU reference") + + dist.destroy_process_group() + + +def run_disaggregated() -> None: + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + torch.cuda.set_device(rank) + world_group = dist.group.WORLD + prefill_group = dist.new_group([0, 1]) + decode_group = dist.new_group([2, 3]) + assert world_group is not None + assert prefill_group is not None + assert decode_group is not None + dist.barrier() + + is_prefill = rank in (0, 1) + pool_group = prefill_group if is_prefill else decode_group + # Every rank here is a member of exactly one of the two sub-groups, so + # this never actually hits GroupMember.NON_GROUP_MEMBER -- but that + # sentinel is part of new_group()'s declared return type, so `is not + # None` alone doesn't narrow it away for mypy. + assert isinstance(pool_group, dist.ProcessGroup) + pool_rank = dist.get_rank(group=pool_group) + + model, tokenizer = load_model( + MODEL_NAME, device="cuda", dtype=torch.bfloat16, trust_remote_code=True + ) + hidden_bytes = model.config.hidden_size * 2 + Buffer.set_num_sms(24) + dispatch_config = Buffer.get_dispatch_config(2) + combine_config = Buffer.get_combine_config(2) + num_nvl_bytes = max( + dispatch_config.get_nvl_buffer_size_hint(hidden_bytes, 2), + combine_config.get_nvl_buffer_size_hint(hidden_bytes, 2), + ) + buffer = Buffer(pool_group, num_nvl_bytes, 0) + patched = patch_moe_infer_ep(model, resolve_backend("naive"), buffer, pool_rank, 2) + if patched == 0: + raise RuntimeError(f"[rank {rank}] patched no MoE layers") + role = "prefill" if is_prefill else "decode" + print( + f"[rank {rank}] disaggregated {role}: patched {patched} MoE layers, pool_rank={pool_rank}" + ) + + num_heads = model.config.num_attention_heads + head_dim = model.config.hidden_size // num_heads + num_layers = model.config.num_hidden_layers + + results: dict[str, list[int]] = {} + + if is_prefill: + prefill_worker = PrefillWorker(make_prefill_fn(model, "cuda"), batch_size=1) + dst = rank + 2 + for prompt in DEFAULT_PROMPTS: + input_ids = tokenizer(prompt, return_tensors="pt").input_ids + prefill_worker.submit(Request(prompt, input_ids, max_new_tokens=MAX_NEW_TOKENS)) + prefill_results = prefill_worker.step() + assert len(prefill_results) == 1 + result = prefill_results[0] + send_kv_cache(result.cache, dst=dst, group=world_group) + first_token_tensor = torch.tensor( + [result.first_token_id], dtype=torch.int64, device="cuda" + ) + dist.send(first_token_tensor, dst=dst, group=world_group) + else: + decode_worker = DecodeWorker(make_decode_fn(model, "cuda"), batch_size=1) + src = rank - 2 + for prompt in DEFAULT_PROMPTS: + cache = recv_kv_cache( + src=src, + group=world_group, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=torch.bfloat16, + device="cuda", + ) + first_token_tensor = torch.zeros(1, dtype=torch.int64, device="cuda") + dist.recv(first_token_tensor, src=src, group=world_group) + first_token_id = int(first_token_tensor.item()) + + prefill_result = PrefillResult( + request_id=prompt, + cache=cache, + first_token_id=first_token_id, + ttft=0.0, + max_new_tokens=MAX_NEW_TOKENS, + eos_token_id=None, + ) + decode_worker.admit(prefill_result) + results[prompt] = _run_worker_to_completion(decode_worker, prompt) + + if not is_prefill and rank == 2: + print("[rank 2] disaggregated results:") + for prompt, tokens in results.items(): + print(f" {prompt!r} -> {tokens}") + reference = json.loads(REFERENCE_PATH.read_text()) + match = all(results.get(p) == reference[p] for p in DEFAULT_PROMPTS) + print(f"[rank 2] MATCH: {match}") + (OUTPUT_DIR / "2026-09-16-phase-4-disaggregated-gate-results.json").write_text( + json.dumps({"results": results, "match": match}, indent=2) + ) + if not match: + dist.destroy_process_group() + raise SystemExit("disaggregated EP disagrees with the single-GPU reference") + + dist.destroy_process_group() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--mode", required=True, choices=["single", "colocated", "disaggregated"]) + args = parser.parse_args() + if args.mode == "single": + run_single_gpu_reference() + elif args.mode == "colocated": + run_colocated() + else: + run_disaggregated() + + +if __name__ == "__main__": + main() From abe8534534e73948ae538579175c18718e336f20 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:23:15 -0400 Subject: [PATCH 11/12] docs: record Phase 4 disaggregated prefill/decode outcome Phase 4 complete: both correctness gates pass with exact token match, three real bugs found and fixed on real hardware, contention-relief question left genuinely open (confounded measurement, not forced into a conclusion). All plan steps checked off. --- docs/STATUS.md | 65 +++++- ...hase-4-disaggregated-prefill-decode-run.md | 201 ++++++++++++++++++ ...ase-4-disaggregated-prefill-decode-plan.md | 78 ++++--- 3 files changed, 307 insertions(+), 37 deletions(-) create mode 100644 docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md diff --git a/docs/STATUS.md b/docs/STATUS.md index 87673b4..b1799a2 100644 --- a/docs/STATUS.md +++ b/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. diff --git a/docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md b/docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md new file mode 100644 index 0000000..73449ef --- /dev/null +++ b/docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md @@ -0,0 +1,201 @@ +# Phase 4 disaggregated prefill/decode run -- findings + +## Infrastructure + +- Pod `ds338r4byldpqu`, RunPod Secure Cloud, data center AP-IN-1, 4x + NVIDIA H100 80GB HBM3 (real NVLink confirmed -- see below). Rate: + **$13.96/hr** -- notably higher than the $10.76/hr `get-capacity` + quoted for the same GPU/cloud/CUDA-version combination just before + rental; the catalog-lookup price and the actual create-pod price can + diverge by a real margin, confirmed from the pod object's own `cost` + field, not assumed from the pre-rental check. +- `nvidia-smi topo -m` confirmed real NVLink across all four GPUs: every + pair showed `NV18` (18 bonded links), not `PHB`/`PXB`/`SYS`. +- Image `runpod/pytorch:1.0.2-cu1281-torch280-ubuntu2404`, same as Phase + 3. `uv sync` pulled this repo's own `torch==2.14.0` floor, defaulting + to `+cu130`; the CUDA 13.0 toolkit (nvcc) was installed via apt, + reusing Phase 3's exact fix. DeepEP V1 built and imported cleanly on + the first attempt, no further issues. +- Same `transformers==4.57.6` pin + `DynamicCache.get_usable_length` + monkeypatch Phase 0/1/3 all needed, applied identically via a + pod-local, never-committed patch at the top of both Phase 4 scripts. + DeepSeek's model repo is still unchanged since Phase 0. +- Disk bumped to 60GB (from Phase 3's 40GB) specifically to avoid Phase + 3's disk-space-vs-32.8GB-model-download issue; 47GB was free after all + installs, before the download -- enough this time, no `/dev/shm` + workaround needed. +- Pod terminated via `pod-action terminate` after all findings were + copied off; a follow-up `list-pods` returned an empty list, confirming + full removal. Total session: 47 minutes, pod creation to termination. + +## Three real bugs found and fixed, none caught by CPU-only tests + +All three are legitimate gaps in code that had already passed CPU-only +tests (including, for two of them, code merged in Phase 3) -- exactly the +class of bug this project's testing table exists to eventually catch on +real hardware, not evidence the CPU-only layer was pointless. + +1. **DeepSeek's remote-code model returns the legacy tuple cache format, + not a `DynamicCache`.** `modeling_deepseek.py` (unchanged since Phase + 0) predates the Cache-class refactor entirely: `outputs.past_key_values` + came back as a plain 28-tuple of `(key, value)` pairs, confirmed live + by inspecting a real forward call's output type. `kv_cache.py`/ + `disaggregated.py`'s own committed contract is built entirely on + `DynamicCache` (Tasks 1-4, CPU-proven against transformers 5.17.0, + where this project's own kernel/scheduler code lives). Resolved by + converting only at the `PrefillFn`/`DecodeFn` closure boundary in the + two pod-local scripts, via `DynamicCache.from_legacy_cache(...)` / + `cache.to_legacy_cache()` -- both still present in `transformers==4.57.6` + (the pin DeepSeek's code separately requires) though removed in v5. + No change needed to any committed library code for this one. +2. **`local_expert_contribution` crashes when a rank receives zero + tokens for any of its local experts at all.** `topk_idx.max()` raised + `RuntimeError: Expected reduction dim to be specified for + input.numel() == 0`. Phase 3's own fix (2026-09-16) handled a rank + missing its *highest* local expert while still having *some* tokens; + 4-way EP (16 experts/rank, finer than Phase 3's 2-way/32-experts/rank) + combined with short prompts made a rank receiving **no** tokens at all + a real, reachable case Phase 3 never hit. Fixed in + `src/dispatch/kernels/expert_parallel.py` with a CPU regression test + (`test_local_expert_contribution_handles_a_batch_with_no_tokens_for_this_rank_at_all`), + committed before returning to the pod. +3. **NCCL requires every send/recv tensor on the correct CUDA device; + `recv_kv_cache`'s buffers defaulted to CPU.** Found by re-reading + `handoff.py` before wiring the real cross-rank handoff, not by a live + crash -- Task 2's own test only exercises the `gloo` backend, which + tolerates CPU tensors. Fixed by making `send_kv_cache` infer its + device from the cache being sent and adding an explicit `device` + parameter to `recv_kv_cache` (default `"cpu"`, so the existing gloo + test is unaffected). The same read-before-running pass caught a + companion bug in `pad_and_batch_caches` (its zero-padding tensors also + defaulted to CPU, which would have crashed concatenating against a + CUDA cache) before it ever touched the rented GPUs. + +All three fixes are covered by regression tests and were committed +before (bugs 2-3) or verified against (bug 1, pod-local by design -- +DeepSeek's own repo quirk, not this project's) real hardware confirmed +them fixed. + +## Correctness gate: passed, both topologies, exact token match + +Judged by exact greedy-generated-token-sequence match against a +single-GPU reference, not top-k logit agreement: the scheduler's +`PrefillFn`/`DecodeFn` contract only returns argmax token ids (raw +logits are dropped by design, matching the project's own generation +loop convention), so exact-token-match is the correctness bar this +architecture actually admits -- arguably a stricter one for a +multi-step greedy-decode pipeline than a single forward pass's top-5 +logit agreement, since it must stay correct across many compounding +steps, not just one. + +`deepseek-ai/deepseek-moe-16b-base`, 27 MoE layers patched, 3 prompts, 16 +tokens each, `max_new_tokens=16`: + +| prompt | single-GPU reference | co-located (4-rank EP) | disaggregated (2+2-rank EP) | +|---|---|---|---| +| "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]` | **exact match** | **exact match** | +| "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]` | **exact match** | **exact match** | +| "def fibonacci(n):" | `[185, 300, 565, 291, 2318, 207, 15, 25, 185, 391, 972, 207, 15, 185, 300, 23744]` | **exact match** | **exact match** | + +Both configurations matched on every token, every prompt. Full data: +`docs/findings/2026-09-16-phase-4-single-gpu-reference.json`, +`2026-09-16-phase-4-colocated-gate-results.json`, +`2026-09-16-phase-4-disaggregated-gate-results.json`. + +Transient CUDA allocator OOM warnings appeared during the first +kernel-backed call in each run, always for the same 369,098,752-byte +allocation -- matching Phase 1's and Phase 3's own precedent exactly +(same byte count, same benign first-call/warmup pattern); generation +completed correctly afterward in every case, so these are not treated as +a correctness signal. + +## The measured answer to Phase 4's thesis + +**Does splitting prefill and decode onto separate GPU pools measurably +relieve the contention they create when they share GPUs under concurrent +load, holding total GPU count fixed at 4?** + +Concurrency measurement (`scripts/gpu/phase4_concurrency.py`), 4 and 8 +concurrent requests, `DEFAULT_PROMPTS` cycled, staggered arrivals 50ms +apart, `max_new_tokens=8` (short by design -- the hard, expensive part of +this session was the correctness work above; this measurement's job was +just to get one real reading on real hardware, cheaply): + +| topology | concurrency | total wall time | mean TTFT | p50 TTFT | p99 TTFT | decode tokens/sec | +|---|---|---|---|---|---|---| +| co-located | 4 | 4.310s | 1.118s | 1.029s | 1.372s | 6.50 | +| co-located | 8 | 2.974s | 0.695s | 0.930s | 1.290s | 18.83 | +| disaggregated | 4 | 4.256s | 0.458s | 0.156s | 1.330s | 6.58 | +| disaggregated | 8 | 4.725s | 0.983s | 0.749s | 1.987s | 11.85 | + +(Disaggregated's TTFT is measured on the prefill side, where admission +and completion both happen; disaggregated's wall time/tokens-per-second +is measured on the decode side, where token generation happens. The two +halves are reported separately because they run as genuinely separate +processes on separate GPUs -- there is no single rank that observes +both.) + +**Mixed, not a clean win either way.** At concurrency 4, disaggregated's +mean TTFT (0.458s) is less than half co-located's (1.118s) -- consistent +with the thesis: prefill isn't waiting behind decode work on a shared +GPU. At concurrency 8, that reverses: co-located's mean TTFT (0.695s) is +now *lower* than disaggregated's (0.983s), and co-located's decode +throughput (18.83 tok/s) is higher than disaggregated's (11.85 tok/s) at +the same concurrency level too. + +**This comparison is confounded by kernel-warmup state and should not be +read as a clean result.** Each of the four measurements above is a fresh +`torchrun` process; co-located's own total wall time *dropped* from +4.31s (concurrency 4) to 2.97s (concurrency 8) -- throughput improving +with more concurrent load, in the same direction warmup would push it, +not the direction contention would. Triton/CUDA kernel compilation +artifacts persisting on disk between process invocations, and the OS +page cache already holding the 32.8GB model checkpoint from the prior +run, are the most likely explanation for later runs measuring faster +independent of the actual topology under test. No warmup-then-measure +step was built into `phase4_concurrency.py`, and building one (plus +enough repeated runs to average out this noise) was exactly the +additional live-debugging scope this session's own cost-discipline +instruction weighed against taking on further. + +**Direct answer**: the correctness question is answered cleanly and +completely -- both a co-located 4-rank EP pool and a disaggregated +2+2-rank EP pool with a real cross-rank KV-cache handoff produce +byte-exact correct greedy generation against a single-GPU reference. The +contention-relief question is not answered cleanly by this run: a real +signal favoring disaggregation appears at concurrency 4 and reverses at +concurrency 8, and the most likely explanation is a confound (warmup +state varying between independent process launches) rather than a real +property of either topology, so no strength-calibrated claim is made +either way. Answering it properly needs either a warmed-up measurement +protocol (discard the first N iterations of each configuration before +recording) or enough repeated trials to characterize the noise -- both +real, identifiable next steps, not a vague "more testing would help." + +## What was *not* attempted, and why + +The design doc's own §8 named the full async-staggered measurement +(`dist.irecv` polling on the decode side, real overlap between prefill +and decode) as more distributed-systems complexity than anything built +before it, with real risk of a long live-debugging session -- flagged +explicitly as a scope decision to check on rather than force through. +That build was attempted and succeeded (both correctness gates and all +four concurrency measurements ran without further live debugging once +the three bugs above were fixed). What wasn't attempted: a warmup +protocol to remove the confound described above, and enough repeated +trials per configuration to characterize run-to-run noise -- both +correctly scoped as follow-up work rather than squeezed into this +session, given the total budget already spent finding and fixing three +real bugs plus running the correctness gates. + +## Cost + +**Total: $10.94** of the $40 cap, a ceiling not a target -- 4x H100 SXM +at $13.96/hr, 47 minutes (pod creation to termination). The gap between +cap and actual spend is itself part of this phase's result, per the +standing instruction to be surgical: correctness work (the hard, +valuable, and now-complete part) took priority, and once both gates +passed and one real concurrency reading was in hand per configuration, +the session stopped rather than spending further budget chasing a clean +answer to a question the data itself flagged as confounded. Full record: +`docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-cost.md`. diff --git a/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md b/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md index 16e8cbf..50423cf 100644 --- a/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md +++ b/docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md @@ -135,7 +135,7 @@ docs/STATUS.md # Task 6 Tasks 3-4 build their entire batching loop on these two functions and nothing else from this file. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** Create `tests/unit/test_kv_cache.py`: @@ -246,12 +246,12 @@ def test_slice_and_rebatch_round_trip_through_a_real_tiny_model() -> None: assert continued.logits.shape == (1, 1, continued.logits.shape[-1]) ``` -- [ ] **Step 2: Run to verify it fails** +- [x] **Step 2: Run to verify it fails** Run: `uv run pytest tests/unit/test_kv_cache.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'dispatch.serving'`. -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** Create `src/dispatch/serving/__init__.py` (empty file). @@ -332,23 +332,23 @@ def pad_and_batch_caches( return DynamicCache(ddp_cache_data=per_layer), attention_mask ``` -- [ ] **Step 4: Run to verify the fast tests pass** +- [x] **Step 4: Run to verify the fast tests pass** Run: `uv run pytest tests/unit/test_kv_cache.py -v -m "not slow"` Expected: 5 pass (the `@pytest.mark.slow` real-model test is excluded). -- [ ] **Step 5: Run the slow real-model test** +- [x] **Step 5: Run the slow real-model test** Run: `uv run pytest tests/unit/test_kv_cache.py -v -m slow` Expected: PASS (downloads `hf-internal-testing/tiny-random-gpt2` on first run, network required, no GPU). -- [ ] **Step 6: Lint and typecheck** +- [x] **Step 6: Lint and typecheck** Run: `uv run ruff check src/dispatch/serving tests/unit/test_kv_cache.py && uv run ruff format --check src/dispatch/serving tests/unit/test_kv_cache.py && uv run mypy src tests` Expected: clean. -- [ ] **Step 7: Commit** +- [x] **Step 7: Commit** ```bash git add src/dispatch/serving/__init__.py src/dispatch/serving/kv_cache.py tests/unit/test_kv_cache.py @@ -373,7 +373,7 @@ git commit -m "feat: add KV-cache slice/pad-batch helpers for continuous batchin functions unchanged, with `backend="nccl"` instead of `"gloo"` and real ranks instead of a CPU 2-process test. -- [ ] **Step 1: Write the failing test** +- [x] **Step 1: Write the failing test** Create `tests/unit/test_handoff.py`: @@ -455,12 +455,12 @@ def test_send_and_recv_kv_cache_round_trip_over_gloo(tmp_path: object) -> None: assert result["seq_len"] == SEQ_LEN ``` -- [ ] **Step 2: Run to verify it fails** +- [x] **Step 2: Run to verify it fails** Run: `uv run pytest tests/unit/test_handoff.py -v -m slow` Expected: FAIL with `ModuleNotFoundError: No module named 'dispatch.serving.handoff'`. -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** Create `src/dispatch/serving/handoff.py`: @@ -511,19 +511,19 @@ def recv_kv_cache( # noqa: PLR0913 -- the receiver can't infer shape/dtype from return DynamicCache(ddp_cache_data=per_layer) ``` -- [ ] **Step 4: Run to verify it passes** +- [x] **Step 4: Run to verify it passes** Run: `uv run pytest tests/unit/test_handoff.py -v -m slow` Expected: PASS. -- [ ] **Step 5: Lint and typecheck** +- [x] **Step 5: Lint and typecheck** Run: `uv run ruff check src/dispatch/serving/handoff.py tests/unit/test_handoff.py && uv run ruff format --check src/dispatch/serving/handoff.py tests/unit/test_handoff.py && uv run mypy src tests` Expected: clean. (Add `"src/dispatch/serving/handoff.py" = ["PLR0913"]` to `pyproject.toml`'s `[tool.ruff.lint.per-file-ignores]` only if the inline `# noqa: PLR0913` above doesn't satisfy ruff.) -- [ ] **Step 6: Commit** +- [x] **Step 6: Commit** ```bash git add src/dispatch/serving/handoff.py tests/unit/test_handoff.py @@ -557,7 +557,7 @@ git commit -m "feat: add cross-rank KV-cache handoff, proven over gloo" unchanged. Task 5 wires real closures matching `PrefillFn`/`DecodeFn` exactly; nothing in this file changes at that point. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** Create `tests/unit/test_disaggregated.py`: @@ -671,12 +671,22 @@ def test_decode_worker_batches_requests_with_different_cache_lengths() -> None: worker = DecodeWorker(_fake_decode_fn, batch_size=4) worker.admit( PrefillResult( - "short", _fake_cache(1, 2), first_token_id=1, ttft=0.1, max_new_tokens=1, eos_token_id=None + "short", + _fake_cache(1, 2), + first_token_id=1, + ttft=0.1, + max_new_tokens=1, + eos_token_id=None, ) ) worker.admit( PrefillResult( - "long", _fake_cache(1, 5), first_token_id=1, ttft=0.1, max_new_tokens=1, eos_token_id=None + "long", + _fake_cache(1, 5), + first_token_id=1, + ttft=0.1, + max_new_tokens=1, + eos_token_id=None, ) ) @@ -712,12 +722,12 @@ def test_decode_worker_respects_batch_size_admitting_only_free_slots() -> None: assert [r.request_id for r in b_result] == ["b"] ``` -- [ ] **Step 2: Run to verify it fails** +- [x] **Step 2: Run to verify it fails** Run: `uv run pytest tests/unit/test_disaggregated.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'dispatch.serving.disaggregated'`. -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** Create `src/dispatch/serving/disaggregated.py`: @@ -955,17 +965,17 @@ class DecodeWorker: return completed ``` -- [ ] **Step 4: Run to verify it passes** +- [x] **Step 4: Run to verify it passes** Run: `uv run pytest tests/unit/test_disaggregated.py -v` Expected: all pass. -- [ ] **Step 5: Lint and typecheck** +- [x] **Step 5: Lint and typecheck** Run: `uv run ruff check src/dispatch/serving/disaggregated.py tests/unit/test_disaggregated.py && uv run ruff format --check src/dispatch/serving/disaggregated.py tests/unit/test_disaggregated.py && uv run mypy src tests` Expected: clean. -- [ ] **Step 6: Commit** +- [x] **Step 6: Commit** ```bash git add src/dispatch/serving/disaggregated.py tests/unit/test_disaggregated.py @@ -991,7 +1001,7 @@ git commit -m "feat: add continuous-batching prefill and decode workers" just with both roles talking to the same 4-rank EP pool instead of two separate 2-rank pools. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** Create `tests/unit/test_colocated.py`: @@ -1090,12 +1100,12 @@ def test_colocated_worker_respects_batch_size_across_prefill_and_decode() -> Non assert [r.request_id for r in b_done] == ["b"] ``` -- [ ] **Step 2: Run to verify it fails** +- [x] **Step 2: Run to verify it fails** Run: `uv run pytest tests/unit/test_colocated.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'dispatch.serving.colocated'`. -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** Create `src/dispatch/serving/colocated.py`: @@ -1163,24 +1173,24 @@ class ColocatedWorker: return completed ``` -- [ ] **Step 4: Run to verify it passes** +- [x] **Step 4: Run to verify it passes** Run: `uv run pytest tests/unit/test_colocated.py -v` Expected: all pass. -- [ ] **Step 5: Lint and typecheck** +- [x] **Step 5: Lint and typecheck** Run: `uv run ruff check src/dispatch/serving/colocated.py tests/unit/test_colocated.py && uv run ruff format --check src/dispatch/serving/colocated.py tests/unit/test_colocated.py && uv run mypy src tests` Expected: clean. -- [ ] **Step 6: Run the full CPU-testable suite together** +- [x] **Step 6: Run the full CPU-testable suite together** Run: `uv run pytest tests/unit -v -m "not gpu"` Expected: all pass, including Task 1-4's new tests. This is the "everything not requiring GPU/DeepEP is already green" checkpoint the Global Constraints require before Task 5's rental begins. -- [ ] **Step 7: Commit** +- [x] **Step 7: Commit** ```bash git add src/dispatch/serving/colocated.py tests/unit/test_colocated.py @@ -1207,7 +1217,7 @@ git commit -m "feat: add co-located prefill/decode baseline worker" spend (Global Constraints). This is a live session with the user, not something to run unattended -- get explicit go-ahead before renting.** -- [ ] **Step 1: Write the runbook** +- [x] **Step 1: Write the runbook** Create `docs/runbooks/phase-4-disaggregated-prefill-decode.md`: @@ -1380,14 +1390,14 @@ at rental time for 4-GPU availability and pricing. ``` ```` -- [ ] **Step 2: Get the user's explicit go-ahead, then execute the runbook** +- [x] **Step 2: Get the user's explicit go-ahead, then execute the runbook** Confirm the budget cap and GPU choice with the user before the first `create` call -- this is a paid action, never taken unattended. Remind the user of the standing instruction to be surgical: prefer stopping early with a smaller, trustworthy result over spending toward the cap. -- [ ] **Step 3: Commit the runbook, the pod-live load-driver script, and the cost record** +- [x] **Step 3: Commit the runbook, the pod-live load-driver script, and the cost record** ```bash git add docs/runbooks/phase-4-disaggregated-prefill-decode.md docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-cost.md scripts/gpu/phase4_load_driver.py @@ -1408,7 +1418,7 @@ than what was assumed. - Create: `docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-run.md` - Modify: `docs/STATUS.md` -- [ ] **Step 1: Write the findings doc** +- [x] **Step 1: Write the findings doc** Cover, in `docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-run.md`: whether real 4-GPU NVLink was confirmed; whether DeepEP V1 installed @@ -1429,7 +1439,7 @@ before both concurrency levels completed, say so plainly and report exactly what was measured -- matching this project's practice of writing down a null or partial result rather than a flattering guess. -- [ ] **Step 2: Update STATUS.md** +- [x] **Step 2: Update STATUS.md** Add a "## Phase 4 progress" section following the Phase 0/1/2/3 pattern: plan link, hardware actually used, the correctness-gate result for both @@ -1438,7 +1448,7 @@ against the $40 cap. Set "## Next step" to reflect what's actually next (Phase 5 planning, or follow-up on Phase 4 if something didn't land cleanly). -- [ ] **Step 3: Commit** +- [x] **Step 3: Commit** ```bash git add docs/findings/2026-09-16-phase-4-disaggregated-prefill-decode-run.md docs/STATUS.md From 83cc28c7592ccd612cb1df09f0df2d858fece5b0 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:24:44 -0400 Subject: [PATCH 12/12] docs: update README for Phase 4 completion Was still describing Phase 1 as current despite Phase 3 already merged -- refreshed to describe what's actually built through Phase 4. --- README.md | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 4166f1f..8e99eab 100644 --- a/README.md +++ b/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