Conversation
Signed-off-by: zhuyuhua-v <yuhzhu@amd.com>
Signed-off-by: zhuyuhua-v <yuhzhu@amd.com>
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
There was a problem hiding this comment.
Pull request overview
Adds a CPU-testable transfer-planning layer for mapping contiguous CPP-prefill KV blocks into DCP-decode layouts, and updates KV-transfer completion handling so offload-style SaveOperationId(req_id, generation) completions correctly pair with per-request state and PP-wide save quorum logic.
Changes:
- Introduce
atom/kv_transfer/disaggregation/dcp_plan.pyto compute sharded/replicated DCP block/token copy plans and coalesce contiguous transfer runs. - Normalize completion identity pairing via
completion_req_key()and propagate it through PP KV aggregation and multi-connector send/save pairing. - Fix PD incremental offset calculation on DCP decode nodes by counting virtual blocks (
block_size * dcp_size) instead of plain blocks, with dedicated tests.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_pp_kv_status.py | Adds regression coverage for pairing finished_sending with SaveOperationId-shaped finished_saving via request-id normalization. |
| tests/test_multi_connector.py | Adds coverage for scheduler-side process_completions fan-out and for worker-side pairing when saves report SaveOperationId. |
| tests/test_dcp_transfer_plan.py | New comprehensive unit tests validating DCP sharded/replicated plans, partitioning, coalescing, and incremental slicing behavior. |
| tests/test_dcp_connector_offset.py | New tests ensuring DCP-aware PD incremental offset counts virtual blocks and handles partial/per-request cache cases. |
| tests/aiter_stub.py | Extends the AITER stub module set to include aiter.dist.parallel_state for import compatibility in tests. |
| atom/model_engine/pp_engine_core.py | Uses completion_req_key() so PP-head send-holding correctly pairs against save completions regardless of completion id shape. |
| atom/kv_transfer/disaggregation/types.py | Adds completion_req_key() helper to collapse completion variants (raw req id vs operation id) onto request identity. |
| atom/kv_transfer/disaggregation/multi/multi_connector.py | Uses completion_req_key() for save pairing keys and adds scheduler-side process_completions() fan-out. |
| atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py | Makes PD incremental num_computed_blocks DCP-aware by dividing by block_size * dcp_size. |
| atom/kv_transfer/disaggregation/dcp_plan.py | New DCP transfer planner implementing sharded/replicated mapping plus contiguous-run coalescing. |
| atom/kv_transfer/disaggregation/init.py | Exports completion_req_key. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Signed-off-by: zhuyuhua-v <yuhzhu@amd.com>
c7d3018 to
c3c4577
Compare
Integrate main's compact PP-aware index-cache mapping while retaining DCP-expanded replicated pages, and fix the PR's formatting checks.
f806e57 to
44116b2
Compare
Squashed import of xingjche/test_pd_dcp, taken as its diff against the
main it branches from, so only that branch's own work lands here -- the
main commits it had merged are already on this line.
- ATOM_DCP_REPLICATE_INDEX_CACHE (default off): under native GLM-5.2
DCP, keep MLA KV sharded but replicate the full IndexShare index
cache on every DCP rank, which drops the indexer candidate
all-gather and global merge. Gated to glm_moe_dsa, dcp>1, no spec
decode. Index pages, the persistent-path plan, and the offload PAGE
identity (layout version 3) all expand by dcp_world_size with it.
- PD prefix-cache offsets and received-prompt hashing count in
hash_block_size (block_size * dcp_size) instead of block_size: under
DCP one block-table entry is a virtual block, so hashing at the
physical size bound several token ranges to the same block.
- Dense offload loads the largest chunk-aligned prefix instead of the
unaligned full-prompt hit, which the worker was reading as a failed
load and recomputing whole.
- atomesh CI: DCP PD accuracy and benchmark cases.
One conflict resolved: aiter_mla's index_page_factor locals, where the
incoming definitions were needed by uses that had already merged.
A non-DCP prefill node holds a sequence's KV contiguously; a DCP decode
node holds it round-robin across dcp_size ranks at interleave_size
granularity, and one of its block-table entries is a virtual block
covering block_size * dcp_size global tokens. The producer's whole-block
descriptors no longer line up with what the rank will read back, so plan
the transfer in token runs instead, and tell the producer which shard to
write by shipping the consumer's DCP topology in the write_request.
Every block region shares the KV cache's slot mapping (the indexer cache
included) and DCP only runs on MLA, whose layout stores a token
contiguously, so one plan covers all of them, scaled by the region's own
per-token width. It takes one planner per region class, each given the
source and destination block id lists and returning coalesced
(src_offset, dst_offset, length) triples:
- plan_sharded (MLA latent): rank r takes the runs of interleave_size
tokens it owns. block_size / interleave_size descriptors per
destination block, degenerating to one whole-block descriptor at
interleave_size == block_size. Both sides derive their (block, offset)
pair by splitting a slot index on block_size -- the local slot for the
destination, the global position for the source -- so the two halves
read the same way.
- plan_replicated (index cache under ATOM_DCP_REPLICATE_INDEX_CACHE=1):
every rank holds the full rows, so a destination virtual block is
filled by dcp_size consecutive source blocks, independent of rank and
interleave.
The choice is per region, not per request: aiter_mla fills the
semantic_role the region type already carries, the consumer ships whether
it replicates, and the indexer regions take plan_replicated while the MLA
latent stays interleave-sharded. Relayouting the dcp_size-times-wider
index page with the 1/W formula writes real KV to the wrong slots, which
decode reads back as plausible text rather than a fault, so
_uses_standalone_lmcache_offload becomes _uses_replication_safe_kv_transfer
and admits mooncake instead of blocking the combination. The multi fan-out
stays out: that is the prefill node's shape, and prefill runs
dcp_world_size == 1, where this layout is rejected anyway.
No new geometry crosses the wire. The plan already counts in the
destination's token space, and token_bytes = src_unit_bytes // block_size
holds for both ends because only the page width differs, in whole tokens.
Deriving the block counts from the array sizes removes the count mismatch a
separate-arguments signature would make expressible. Layout validation is
left to config construction, which already asserts every one of these
invariants (interleave_size >= 1, 1 <= interleave_size <= kv_cache_block_size
and block_size % interleave_size == 0 at config.py:1444, 1708, 1713;
tensor_parallel_size % dcp_size == 0 at config.py:1685) -- a server that
could violate them would not have started.
What config cannot see, the producer checks per transfer: that its
per-region plan lands on the same kind of region on the consumer, which
matching region counts do not establish, and that a region's block bytes
divide into whole tokens, without which the per-token size truncates.
Both otherwise corrupt KV silently.
One destination block spanning dcp_size source blocks is also what the
incremental-prefix offset and the producer's src/dst count check have to
account for: the offset divided num_cached_tokens by block_size, which
counts virtual blocks dcp_size times over and would skip live prefix
blocks. Without DCP nothing changes: same whole-block descriptors, same
arithmetic.
scripts/start_glm52_pp4pd_dcp.sh brings the pair up -- PP4xTP1 prefill on
GPUs 0-3, TP4xDCP4 decode on 4-7 -- and sets the replicated index cache on
the decode node alone, since the check reads the bare environment variable
and a prefill rank that saw it would abort on "decode context parallel size
must be > 1".
The planners are numpy-only and reachable through the aiter stub, so the
mapping is tested without a GPU.
GSM8K over that pair, all 1319 5-shot samples: 0.9295 with the index cache
sharded, 0.9356 with it replicated. Producer bytes confirm the split plan --
a 20-MLA/7-index stage sending 48 source tokens writes
20*576*12 + 7*144*48 = 186624, the latent still one quarter and the indexer
whole. Costs HBM: the block pool drops 144111 -> 121736 blocks per rank, the
~19% the wider index page predicts. docs/dcp_decode_topk_bottleneck.md
records where the decode time goes once it runs.
The replicated index cache refused speculative decode outright, which cost the DCP decode node its tokens-per-step: the DPA baseline reached 797 tok/s against 431 here almost entirely because it ran MTP. Nothing about the replicated path needs qlen=1. Three things had to change. The compacting DCP top-k filter was written per request -- `qo_indptr` picked one query token per row -- so widen its row unit to the query token and let each draft position carry its own compacted region. `token_to_seq_idxs` supplies the owning request for the block-table lookup, which is the identity at qlen=1 and not at qlen>1; both metadata builders now populate it for non-MTP sparse DCP too, since the filter is shared. `index_slot_mapping` grows to one slot per query token for the same reason, over the same `[seq_len - max_seqlen_q, seq_len)` range the main KV slot mapping walks. The filter writes compacted regions and records their lengths in `dcp_sparse_kv_indptr_buffer`, but only the qlen=1 branch of `_forward_decode` read that buffer -- the MTP branch still used `sparse_kv_indptr`, which describes uniform top-k-wide regions, so attention walked past each written region into stale indices. `B` is `q.shape[0]`, a token count on both branches (sparse layers in MTP verify run at max_q_len=1), so one substitution covers them. Third, MTP faulted in aiter's `_fwd_kernel_stage2_asm` -- 382 wave dumps, one kernel in all of them. That kernel is the intra-rank split-KV merge, and the validated non-MTP config never reaches it: `speculative_config` being set cleared `sparse_dcp_metadata_rebuild`, which cleared `use_persistent_mode`, which dropped aiter onto the non-persistent path. So MTP was not breaking the merge, it was switching to a different one -- and the split-KV merge cannot describe a DCP row. It gives every row the same number of split slots and then caps the count at `cdiv(cur_kv_seq_len, mgc)`, where the length comes from the compacted indptr, but a DCP row only holds the top-k slots this rank owns, so the lengths vary per row. The persistent work plan encodes exactly that, which is why the non-MTP path has always been fine. Nothing here is specific to speculative decode: DPA runs MTP on the persistent kernel already. So stop clearing the flag, and point the work-plan rebuild at the `sparse_mtp_*` buffers when verifying -- those were already built and already selected by `_forward_decode`, only the rebuild call still wrote the unprefixed set, whose empty-prefix branch asserts qlen==1. The assert in `_dcp_decode_candidate_exchange` stays: it guards the sharded branch, which a replicated index cache never reaches. The PAGE namespace carried its own copy of the same predicate and so needed the same edit; it now calls `dcp_replicated_index_cache_enabled` instead. A desync between the two would silently change an offload cache key. GSM8K 1319 = 0.9303 +/- 0.007 with zero faults, between the 0.9295 (replicate=0) and 0.9356 (replicate=1) no-MTP baselines on the same harness. Throughput is still unmeasured -- tokens per step was the point, accuracy only the gate. The two filter tests track the new signature but skip without aiter installed, so they did not run here.
44116b2 to
445fa83
Compare
Signed-off-by: Phi-C <chenxjhit@163.com>
A preshuffled index page is not token-addressable. indexer_k_quant_and_cache writes it as an MFMA-16x16 tiled fp8 key run, then one fp32 scale plane for the whole page, then padding, so no 144-byte window of it is a token. The mooncake block transfer addressed every block region as token_index * (per_block_bytes / block_size), which is only true of the MLA latent, and so handed the decode node a page whose keys were shuffled against their scales. The damage showed up twice and looked like two bugs. The plane the indexer built from those bytes was 5-7 % NaN, which is a wrong answer on any context past the W * index_topk = 8192 short-circuit; and an all-NaN region is a single radix bucket the one-block top-k cannot split, which is where the 1.6 ms radix_topk_one_block_kernel came from. Moving pages page-aware answers both: needle retrieval goes 0/4 -> 4/4 at 7 k and 6/6 at 23 k and 46 k, and the same kernel drops to 0.40 ms with every other GPU bucket unchanged. The replicated plan now emits a descriptor per plane instead of per page, and the index region carries its two plane widths in KVTransferRegion so the plan knows where the split is. A sub-page interleave on an unreplicated index region has no valid plan at all and is now a hard error rather than silent corruption.
Decode runs under a whole-forward graph, so the Python body of sparse_attn_indexer only executes at capture time and the plane the top-k kernel actually consumed is unreachable from the eager side. record() captures a copy_ into a buffer allocated during the eager warmup forward that precedes each capture, so the copy replays with every decode step and flush(), called from the run_model tail, reads back the plane of the step that just ran. Off unless ATOM_DSA_LOGITS_DUMP names a directory. This is what showed the index-cache relayout was corrupting the plane rather than the kernel being slow on honest data.
The script posted to /v1/completions with no chat template and the default generation length, which lands around 0.930 and is not comparable to the 0.9742 / 0.9727 row recipes/GLM-5.md publishes for TP4 MXFP4. It now uses the recipe's arguments: the chat endpoint, --apply_chat_template, and max_tokens=16384,temperature=0,top_p=1 at 64 concurrent. The model class is local-chat-completions, not the recipe's local-completions: that class posts `prompt`, so against a chat endpoint every request comes back 400 missing field `messages`.
Against the recipe's published TP4 row the disaggregated path looks 1.6 pp short, but that row was not measured on this machine. A single-node TP4 server started from the recipe's own command, same box and image and model copy, scores 0.9644 / 0.9651, which puts PD+DCP 0.6 pp behind it -- 0.8 sigma, no measurable cost to disaggregation. The remaining 1.0 pp is this machine against the published reference.
The GSM8K number in this doc predates the DSA index-page relayout fix, and nothing here said what the MTP head's real acceptance rate is. Both are now measured on the fixed stack: MTP costs no accuracy and runs the same eval 1.56x faster, and acceptance is 0.60 eval-wide but 0.62-0.65 while the batch is full. Single-request probes at 222, 11k and 40k prompt tokens pin the variation on generated content rather than on context length or batch size, which matters because 40k is past the point where the DSA top-k stops short-circuiting.
The script serves the throughput runs, where MTP=3 at a fixed acceptance rate is part of the configuration being measured, so every perf launch was prefixing ENABLE_MTP=1. Make that the default and keep ENABLE_MTP=0 for isolating a DCP change from the draft KV layer and the q>1 verify path. Note that DCP_INTERLEAVE > 1 now needs an explicit ENABLE_MTP=0: the existing guard rejects the combination rather than silently dropping the speculative flags.
The .github benchmark/workflow edits and the DCP design notes under docs/ are not part of the KV-transfer change itself. Restore them to main so the PR carries only the engine, kernel, script, and test changes.
Signed-off-by: Mengqing Cao <cmq0113@gmail.com>
Signed-off-by: Mengqing Cao <cmq0113@gmail.com>
Signed-off-by: Mengqing Cao <cmq0113@gmail.com>
Signed-off-by: Mengqing Cao <cmq0113@gmail.com>
The guards protect a (num_tokens,) Triton grid launched over raw pointers, so the failure they stand in front of is an out-of-bounds read and write, not a wrong answer. `python -O` strips a bare assert, which is exactly the configuration a production run is most likely to use. _assert_dcp_filter_rows is renamed _check_dcp_filter_rows to match, and the argument checks it makes raise ValueError/TypeError.
sharded_plan addresses a block in MLA token units, so it is meaningful for the latent cache and, through the replicated planes, for the index cache. A region registered with any other role -- an EAGLE3 draft KV region is the one that exists today -- was relaid out under that rule anyway, silently writing the consumer's blocks at the wrong offsets. Read the role once and raise on anything else. The descriptor loop also boxed one Python int per token on the replicated index plan; scale the offsets in the numpy arrays and tolist() them instead, and skip the debug format when DEBUG is off.
The key/scale plane split the replicated relayout addresses through assumes one fp32 scale per token, which holds only while the indexer quantizes a whole key row as a single block. The old check compared the row against the page bytes, which a wider row can still satisfy while carrying index_head_dim // quant_block_size scales -- so the plane boundary would be short and each destination sub-page would overwrite the previous one's scales. Compare against Indexer.quant_block_size instead.
The replicated index path set it on the metadata object and both readers reached it through getattr(..., None), so the field existed only where it happened to be assigned. Declare it beside sparse_kv_last_page_lens, which is the same shape of optional sparse buffer.
Both the comment and the guide still described the persistent path as non-speculative, from before the MTP verify step rebuilt into its own sparse_mtp_ work buffers. A verify step is per-token q_len=1 rows and takes the same rebuild, so what remains non-persistent is a missing per-layer rebuild: gfx942, or a page size above 1.
The dashboard derived prefill PP from the topology label, so `cpp4-dcp4`
worked and a run whose label does not spell it out did not; the decode side
had no PP term at all, so a decode PP run would under-report its GPUs. Read
--pipeline-parallel-size out of each role's extra server args and let it
outrank the label, and multiply the decode GPU count by it. DECODE_EXTRA_
SERVER_ARGS is now carried into the payload so the decode side has a source.
The label regex also matched the `pp` inside a longer word; anchor it to a
separator.
apply_role_env substituted the global HANDSHAKE_PORT into a ${HANDSHAKE_PORT}
placeholder, not the port of the launch it is building.
| plan_src, plan_dst, plan_len = plan | ||
| src_addrs.extend((src_base + plan_src * unit).tolist()) | ||
| dst_addrs.extend((dst_base + plan_dst * unit).tolist()) | ||
| sizes.extend((plan_len * unit).tolist()) |
Three files conflicted, all where main's #2102 (published DCP local context lens) landed next to this branch's replicated index cache: - forward_context.py / aiter_mla.py: both sides declare a new optional field and a new __init__ line in the same spot. Kept both; they are independent. - test_mla_index_cache.py: both sides append tests. Kept both, and dropped the function-local `import torch` now that main imports it at module scope. The two features do not overlap at runtime: dcp_local_context_lens localizes context for scoring a *sharded* index cache, and the replicated path skips that consumer entirely via _dcp_index_comm_required().
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
atom/model_ops/dcp_ops.py:1346
- This function still relies on multiple
assertchecks (dtype/shape,dcp_rankbounds, andout is not None). Since asserts are stripped underpython -O, the Triton launch can proceed with invalid inputs and potentially read/write out of bounds despite the new_check_dcp_filter_rowsguard. Convert these asserts into explicit exceptions (and validatetoken_indicesshape/dtype) so the safety checks are always enforced.
assert 0 <= dcp_rank < dcp_world_size
assert out is not None, "sparse_kv_indices_buffer (out) is required"
| assert 0 <= dcp_rank < dcp_world_size | ||
| assert out is not None, "sparse_kv_indices_buffer (out) is required" |
There was a problem hiding this comment.
🔵 Needs a closer look
It changes core sparse-attention/DCP indexing semantics and production RDMA KV-transfer behavior across multiple subsystems, so it needs careful human validation on real deployments beyond unit tests.
Review details
Suppressed comments (1)
atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py:1377
- When
has_slot_regionsis true andconsumer_dcp_size > 1, this path logs and returns without sending any notification to the consumer. The consumer has no timeout/cleanup for_pending_recv_*, so the receive can hang indefinitely. If this configuration is truly unsupported, it should fail in a way that unblocks the consumer (e.g., an explicit error notification + consumer-side cleanup, or a consumer-side preflight rejection before enqueueing the request).
if has_slot_data and consumer_dcp_size > 1:
# _execute_block_slot_transfer pairs whole blocks, so it has no
# way to express a consumer block that holds every dcp_size-th
# token. The block-count check below now admits that ratio, so
# refuse here rather than write another rank's KV.
logger.error(
"[PRODUCER] req %s carries per-request state regions and "
"the consumer runs dcp_size=%d; the slot-transfer path has "
"no DCP relayout. Aborting instead of writing misaligned "
"KV.",
req_id,
consumer_dcp_size,
)
return
- Files reviewed: 23/23 changed files
- Comments generated: 2
- Review effort level: Lite
| if getattr(attn_metadata, "index_slot_mapping", None) is not None: | ||
| # Absolute cache addresses, so a token slice needs no rebase. Its | ||
| # presence is also what selects the replicated index layout in the | ||
| # indexer -- dropping it would silently fall back to the sharded one. | ||
| ub_attn.index_slot_mapping = attn_metadata.index_slot_mapping[ts] |
| # Absolute addresses into the replicated DSA index cache, one per query | ||
| # token. Its presence is the switch: not None makes the indexer use it in | ||
| # place of slot_mapping, None leaves the dcp-sharded index layout. | ||
| index_slot_mapping: torch.Tensor | None = None |
Two DCP<->MTP adaptation defects found in PR review.
1. Draft steps reused the target step's token_to_seq_idxs.
The DCP top-k filter's row unit is the query token; it reads
token_to_seq_idxs to pick the owning request's block_tables row. The
target step publishes one entry per verify token
(arange(bs).repeat_interleave(max_seqlen_q)), but a draft step runs one
row per sequence, so the map must be the identity. Nobody rebuilt it,
and draft steps indexed the block table through the verify-step map.
Add AiterMLAMetadataBuilder.rebuild_draft_token_to_seq_idxs() and call
it from the proposer's per-step metadata rewrite. Padded rows carry no
top-k but still index the table, so the rebuild covers running_bs.
2. index_share_for_mtp_iteration never worked and is removed.
compact_topk_indices resolved self_attn.mla_attn -- the
paged_attention.Attention wrapper -- while sparse_kv_indices_buffer
lives one level down on mla_attn.impl (MLAAttention). The getattr
returned None, so the buffer-copy branch has never executed. Even
reaching the buffer, the copy indexed a flat 1-D allocation as if it
were [bs, index_topk], moving elements rather than rows, and under DCP
the layout is ragged: correct compaction has to repack
sparse_kv_indices_buffer, dcp_sparse_kv_indptr_buffer and
dcp_owned_counts_buffer together.
None of that surfaced because set_skip_topk had no effect at the
serving compilation level. use_custom_dispatcher (level >= 2) replays
compiled code without Dynamo guards, so the
if self.is_v32 and self.indexer is not None and not self.skip_topk
branch in the draft's attention is baked at first trace with
skip_topk=False. The indexer therefore ran on every draft step, which
is also why the flag showed no measurable effect. One frozen branch
cannot serve step 0 (must compute) and steps 1+ (must skip), so no
warmup ordering recovers it; the indexer call would have to move behind
a Dynamo-opaque custom op the way the PCP path does.
Forcing eager to make the sharing real measures the trade at 8k prompts,
concurrency 8: acceptance 0.5254 vs 0.6399 and 443 vs 490 requests, for
0.39 ms/forward (~1.6%) saved. At bs=2 a draft step reads rows 0 and 1,
so sequence 1 attends through sequence 0's top-k. Not worth repacking
three GPU buffers per step to keep.
Serving behaviour is unchanged -- level 3 already recomputed every step.
Removing the flag drops a real correctness bug at level 0/1.
tests/test_mtp_index_share.py goes with it: it re-implemented both
helpers inside the test file against a fake whose buffer was 2-D and
hung off mla_attn directly, so it asserted against the two properties
production does not have.
There was a problem hiding this comment.
🔵 Needs a closer look
It makes broad, high-impact changes across distributed RDMA transfer, DCP sparse attention kernels, and speculative decode metadata paths that require careful end-to-end validation on real multi-GPU topologies.
Review details
Suppressed comments (1)
atom/utils/forward_context.py:592
- The comment describes
index_slot_mappingas “absolute addresses”, but this field is populated with cache slot indices (e.g.,block_id * page_width + offset, likeslot_mapping), not byte addresses/pointers. Clarifying this avoids confusion for future callers and keeps terminology consistent with the rest of AttentionMetaData.
# Absolute addresses into the replicated DSA index cache, one per query
# token. Its presence is the switch: not None makes the indexer use it in
# place of slot_mapping, None leaves the dcp-sharded index layout.
- Files reviewed: 25/25 changed files
- Comments generated: 0 new
- Review effort level: Lite
Thanks — went through these one at a time. MTP findings are in 6337d9b. Fixed MTP 1 — draft token_to_seq_idxs. The target step writes one entry per verify token; a draft step is one row per sequence, so the map must be the identity, and nobody rebuilt it. Added AiterMLAMetadataBuilder.rebuild_draft_token_to_seq_idxs() and call it in the proposer's per-step rewrite. Covers running_bs — padded rows carry no top-k but still index the block table. MTP 2 — index_share_for_mtp_iteration removed, not fixed. compact_topk_indices read self_attn.mla_attn (the wrapper); the buffer is on .impl, so the copy branch has never executed. It also indexed a flat 1-D buffer as [bs, index_topk], and under DCP the layout is ragged — correct compaction has to repack sparse_kv_indices_buffer, dcp_sparse_kv_indptr_buffer and dcp_owned_counts_buffer together. None of that showed up because set_skip_topk is inert at level ≥ 2: use_custom_dispatcher (decorators.py:336) replays without Dynamo guards, so the not self.skip_topk branch (deepseek_v2.py:2682) is baked at first trace. The indexer ran every draft step. One frozen branch can't serve step 0 and steps 1+, so the call would have to move behind an opaque op like the PCP path does. Forcing eager to make sharing real (8k prompts, conc 8): acceptance 0.5254 vs 0.6399, 443 vs 490 requests, for 0.39 ms/forward — 1.6%, and 1.2–2.3% at level 3. At bs=2 a draft step reads rows 0 and 1, so sequence 1 attends through sequence 0's top-k. Not worth repacking three buffers per step. Serving behaviour is unchanged (level 3 already recomputed); the removal drops a real level-0/1 bug. Finding 11 (filter row guard). _assert_dcp_filter_rows → _check_dcp_filter_rows, bare asserts → ValueError/TypeError. python -O strips asserts and these kernels take raw pointers, so the stripped build puts the OOB launch back, silently. Test added in tests/test_dcp_sparse_filter.py. pd_server_atom.sh handshake port. ${HANDSHAKE_PORT} now substitutes the launch's port, not the global base — decode reads 6302, not 6301. No number moves (only the producer binds _side_channel_port), fixed because the next role that binds would get it wrong with no error. process_result.py. CPP_PP_RE anchored like the other two regexes in the file, so app4/dpp2/1p1d_dcp4 stop matching. --pipeline-parallel-size now outranks the topology label. Added symmetric decode_pp into num_decode_gpu (1 today, so no row changes). Dropped the unread "prefill_pp" key. Not doing the pd_submit.sh route: PP lives inside extra_args, and emitting PREFILL_PP/DECODE_PP without teaching the launcher to consume them creates a second source of truth. index_slot_mapping undeclared. Declared on AttentionMetaData, next to sparse_kv_last_page_lens (same precedent). Call sites keep their getattr checks; the comment now says its presence is what selects the replicated layout. plan_replicated_index docstring. Dropped the one clause that contradicted the rest of the paragraph. hash_block_size vs block_size. Behaviour kept, invariant written down. num_computed_blocks counts consumer blocks and the producer converts back with * consumer_dcp_size, so the check has to be physical — the old comparison is what disabled incremental transfer on CPP→DCP. The wire value equals physical only while the producer is dcp=1; the comment says so. No new wire key for a topology that doesn't exist yet. Not changing
Tests The transfer tests aren't skipped — tests/test_pd_pp.py runs on CPU (61 cases over plan_sharded, plan_replicated_index, _coalesce, region maps, write_done nonce/dedup), as does tests/test_lmcache_offload_config.py (34). Only the kernel files skip. replicated_index_cache* and rebuild_draft_token_to_seq_idxs are uncovered for a structural reason, not oversight: they live in aiter_mla.py, which does import triton at module scope, so any test importing them skips on CPU. Covering them means relocating them — the move argued against in finding 12. tests/conftest.py:89-94 documents this trap. tests/test_mtp_index_share.py deleted with the feature. It was green because it never imported deepseek_mtp — it re-implemented both helpers against a fake whose buffer was 2-D and hung off mla_attn directly, i.e. the two properties production doesn't have. Full suite (minus tests/plugin): 37 failed / 4883 passed. All pre-existing — clean HEAD in the same image gives 8, and the extra 29 are tests/test_sparsekv_coordinator.py, untracked from another branch. The 3 tests/plugin errors are the image lacking sglang/vllm. |
Preserve replicated index-cache decode while adopting the fused sharded DCP top-k path. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
🟡 Changes recommended
The new triton_filter_and_convert_dcp_index still relies on assert for required safety invariants, which can be stripped under python -O, reintroducing unsafe kernel launches.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
atom/model_ops/dcp_ops.py:1165
- This function intentionally adds runtime size checks because Triton kernels take raw pointers and out-of-bounds accesses can be silent, but it still uses
assertfor several required invariants (dtype/shape, divisibility, rank bounds, andoutnot None). Those asserts are stripped underpython -O, reintroducing the unsafe behavior this helper is trying to prevent.
assert token_indices.dtype == torch.int32
assert token_indices.shape[1] == NUM_TOPK_TOKENS
assert NUM_TOPK_TOKENS % BLOCK_N == 0, (
f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by"
f"BLOCK_N ({BLOCK_N})"
- Files reviewed: 25/25 changed files
- Comments generated: 1
- Review effort level: Lite
| rank keeps a position ``g`` only if ``g % W == dcp_rank`` and maps it to its | ||
| physical slot via the round-robin (virtual-block) layout, computed directly | ||
| from ``block_table`` (like vLLM): | ||
| vbs = block_size * W | ||
| slot = block_table[req, g // vbs] * block_size + (g % vbs) // W |
What this PR does
Makes a CPP (PP4) prefill node feed a TP4 x DCP4 decode node over Mooncake RDMA,
for GLM-5.2 (DSA / sparse MLA). Three pieces:
1. CPP prefill -> DCP decode KV transfer
The producer is not DCP, so its KV blocks do not match the consumer's layout: a
DCP rank owns global token
giff(g // S) % W == r. The connector now plansthe push per region and relays the blocks into the consumer's interleaved
layout, instead of assuming both sides share an addressing scheme.
mooncake_connector.py: sharded transfer plan, per-region dispatch, virtualblock accounting for prefix-cache offsets.
types.py: region roles the relayout dispatches on. A DSA index page ispreshuffled (MFMA-tiled fp8 keys + a plane of fp32 scales) and is not
token-addressable, so it moves plane by plane rather than by token range.
2.
ATOM_DCP_REPLICATE_INDEX_CACHEOpt-in (default off, decode-side only). The sharded index cache reaches its
global top-k through a cross-rank candidate all-gather. Replicating the index
cache across the DCP ranks removes that exchange, at the cost of
W xindexcache memory. Enabling it on prefill raises at startup: prefill runs pp4/dcp1
and the layout rejects both.
envs.py,dcp_utils.py: the flag and its model/topology gate (glm_moe_dsadcp > 1).aiter_mla.py:index_slot_mappingbuffers, startup validation of thesupported transfer topology.
deepseek_v2.py: index pages addressed atblock_size * dcp_world, and thecandidate exchange skipped when the cache is replicated.
offload/config.py: the PAGE namespace now carriesreplicated_index_cacheand
indexer_types, so a replicated layout can never reuse a sharded one'scached objects.
3. MTP on the DCP decode node
Previously
sparse_dcp_metadata_rebuildexcluded speculative decode, so MTPfell off the persistent kernel. The blocker was the DCP top-k filter keying on
requests (
qlen == 1); MTP verify forwardsmax_seqlen_qdraft positions perrequest, each with its own top-k.
dcp_ops.py:triton_filter_and_convert_dcp_indexnow keys on the querytoken (
token_to_seq_idxs+num_tokens) instead ofqo_indptr/global_kv_indptr. The block table is still per request, so a draft positionlooks up its owner. At
qlen == 1the two are identical.attention_mla.py: drop thespeculative_config is Nonegate; extend theDCP
paged_kv_indptroverride to the MTP verify branch (the filter's per-tokencumsum, not
sparse_kv_indptr's uniform stride).deepseek_v2.py: DCP takes priority over the dense MTP gather path.