Repository navigation
Phase 4: disaggregated prefill/decode - #4
Merged
Merged
Conversation
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.
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.
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.
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.
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.
…evice 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.
…l_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.
…rency 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.
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.
Was still describing Phase 1 as current despite Phase 3 already merged -- refreshed to describe what's actually built through Phase 4.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
src/dispatch/serving/), proven correct entirely on CPU (kv_cache.pygrounded in transformers v5's realDynamicCacheAPI,handoff.pyproven overgloobefore any GPU) before ever touching rented hardware.deepseek-ai/deepseek-moe-16b-base.DynamicCache;local_expert_contribution(Phase 3 code) crashes when a rank receives zero tokens for any local expert at all, a case 4-way EP's finer sharding made reachable; two device-placement bugs in the new handoff/padding code (NCCL needs CUDA tensors,gloo-only tests never caught it).Test plan
make checkgreen (lint,mypy --strict, 102 tests) throughout every commitdocs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md