[MiniMax-M3.1] P/D from a PP4 prefill into a TP4 decode (2/2: DSpark) by Jasen2201 · Pull Request #3038 · ROCm/ATOM

Jasen2201 · 2026-11-05T12:15:31Z

Summary

Part 2 of 2 of MiniMax-M3.1-Flash P/D from a PP4 prefill (TP=1 per stage) into a TP4 decode: the checkpoint's DSpark draft on both sides. #2721 (merged) moves the target model's KV; this PR lets the PP producer write the draft's context K/V, which the decode drafts from.

Producers never propose. A P/D producer's drafts are dropped (mtp_k 0). One whose drafter writes its context K/V apart from propose (DSpark: M3.1, Kimi-K3, DeepSeek-V4/V4.1) now skips propose(), its draft graph warmup and the DP draft lockstep, on both output paths. It still commits the target's tentative step (_commit_target_step, which propose also calls): DeepSeek-V4.1 refuses a step whose predecessor was left uncommitted. EAGLE still proposes, since its first draft step writes the final chunk's context K/V.
The M3.1 draft below TP4. Allowed on a producer, which never runs the block pass; the block pass refuses more than one KV head per rank before writing anything. Context-only producers skip sharing the target embedding and LM head, while retaining KV initialization. Drafter.context_kv_only drives sharing, propose, draft warmup and DP lockstep; PP engines that generate draft tokens still share their heads.
Page-16 draft runs. The DSpark binding reads and writes its cache through page-16 SHUFFLE views, so it names that page and DraftKvBuilder passes it to MhaKvPool.transfer_pages. A TP1 stage's draft block is then cut for each TP4 rank as 8 runs of 4 heads, as [MiniMax-M3.1] P/D from a PP4 prefill into a TP4 decode (1/2: KV transfer) #2721 cuts the sparse target layers.
PP warmup rows. Every stage warms up on the rows the last stage does: the block drafter's two-sequence floor comes from the config, not from the stage's drafter.
PP refusal. A drafter that taps target layers without a config-only aux_capture_spec (EAGLE3) is refused under PP: PPAuxRelay relays only a config-declared spec, so a producer would write wrong context K/V.

Results

One MI355X node, rocm/atom-dev:nightly_202611110332, commit 51c8e5c unless noted; prefill on GPUs 0-3, decode on GPUs 4-7, through the atomesh router. Chat GSM8K with the M3.1 recipe's settings (5-shot, thinking on), full set, DSpark with 7 speculative tokens.

Deployment	GSM8K flexible / strict	DSpark acceptance (tok/fwd)
P/D PP4 → TP4, DSpark	0.9553 / 0.9166	52.4% (4.66)
P/D PP4 → TP4, DSpark, draft cut per 128-token block (`1f4e679`, before the page-16 runs)	0.9530 / 0.9212	45.8% (4.20)
P/D TP4 → TP4, DSpark	0.9538 / 0.9166	53.9% (4.77)

The second row is the cut without the binding's page: each TP4 rank received another head's draft context, with no error, and acceptance fell 6.6 points.
Long context: retrieval prompts of 81.7K tokens through the router, five depths from 1% to 99%: 5/5.
DeepSeek-V4.1-Flash TP4 + DSpark as a producer, serving a request it decodes itself: without _commit_target_step (1f4e679) its second decode step raised Commit the accepted prefix before reusing or checkpointing state on every rank; with it, 16-, 64- and 4 × 128-token requests complete.

Prefill cost of DSpark, PP4 vs TP4

Prefill only, straight to the producer (max_tokens 1), 8K-token random prompts, concurrency 16, no prefix caching, GPUs 0-3. TP4 at --max-num-batched-tokens 32768, PP4 at 8192.

Prefill	DSpark	Throughput	Mean TTFT	KV blocks
TP4	off	66.8K tok/s	1913 ms	61,461
TP4	on	63.9K tok/s (−4.4%)	2004 ms	57,062
PP4	off	99.6K tok/s	1276 ms	86,322
PP4	on	88.2K tok/s (−11.5%)	1443 ms	65,527
PP4, `VLLM_PP_LAYER_PARTITION=16,16,16,12`	on	88.4K tok/s	1439 ms	79,874

The draft costs PP4 more than TP4: its 5 layers and 4 KV heads all sit on the last stage at TP1, so that stage's block grows from 2.27 to 2.95 MB and caps every stage's block count. A layer split that gives the last stage fewer target layers wins back most of the blocks, not the throughput.
On the AgentX traces (AIPerf's inferencex-agentx-mvp replay, prefill only, HBM prefix caching on, 300 s per point), PP4 + DSpark served 11.58 req/s at concurrency 16 against 9.60 for TP4 + DSpark, both at 97% prefix hits. At concurrency 32, with no CPU tier, both pools thrash (about 10% hits); PP4 without the draft (79.5K blocks) held 96.5%.

Test plan

test_drafter_weight_sharing.py: 10 CPU cases cover PP1/PP4, DSpark and serial producers, retained KV initialization, and the real M3 sharing hook with a missing PP embedding. External AITER/Triton/einops/flydsl kernels are stubbed; the Python loading and sharing code is real.
test_draft_kv_builder.py: a TP1 draft filled per head through its binding's pages, cut for four TP4 ranks and read back the same way, holds only that rank's head, for a page-128 and a page-16 binding (CPU, no aiter), and for M3.1's own KVCacheBinding through to_page16_shuffle. Both page-16 cases fail without the fix.
test_postprocess_width.py: which producers skip propose, that they leave the DP draft lockstep, and that on both output paths they skip it, keep their bookkeeping and commit the target step once. The commit check fails without the fix.
test_warmup_batch_phase.py: every PP stage warms up on the same rows.
Related suites (72 files: P/D, draft KV, DSpark, MTP, GDN, LMCache MP, PP relay) under an aiter/einops/flydsl stub, against main f7b8962: no new failures, 12 new passes. black 26.10.1 and ruff are clean.
GPU: the tables above, the needle check, the V4.1 producer, and P/D with the 16,16,16,12 split (GSM8K 0.9568 / 0.9128, acceptance 52.3%, needle 5/5).

Limits and follow-ups

DSpark costs the PP4 prefill a quarter of its KV blocks and 11.5% of its throughput (above). An uneven VLLM_PP_LAYER_PARTITION wins back most of the blocks.
Not from this PR: PP4 at --gpu-memory-utilization 0.9, with prefix caching and 300K-740K-token contexts, aborted in an RCCL kernel (HSA_STATUS_ERROR_OUT_OF_RESOURCES) at concurrency 16, with and without DSpark; at 0.85 both ran 16 and 32 cleanly. It will be tracked separately.
The draft's context K/V write runs the fused qkv_proj and drops q, 64 of its 72 heads' columns; a K/V-only GEMM would cut most of that on the producer.
Image input and dynamic PP chunking under PP are not covered.

🤖 Generated with Claude Code

valar-cop

some changes needed..

gpu run of this head for the P/D PP4->TP4 gsm8k chat + dspark acceptance, and TP4->TP4 with producer skipping propose? also pp4 prefill throughput with dspark taps on?

ci green.

Copilot

🟡 Changes recommended

The standard CPU unit-test job skips the entire new test module because of GPU-dependent imports, leaving relay, missing-tap, payload, and warmup regressions untested in that job.

1 open issue

CPU unit-test job skips the relay regression tests · New

What changed in this PR

This PR adds DSpark draft support to MiniMax-M3.1’s PP4 prefill → TP4 decode deployment, building on #2721’s target KV transfer.

Changes:

Relays tapped target-layer states across pipeline stages and aligns warmup shapes.
Adds head-aware draft KV transfer and consumer-relative region mapping.
Skips unnecessary producer drafting, adds regression tests, and documents the deployment recipe.

File	Description
`tests/test_pp_aux_relay.py`	Tests tap relay, missing-tap rejection, payload retention, and warmup alignment.
`tests/test_postprocess_width.py`	Tests producer proposal skipping and deferred-output bookkeeping.
`tests/test_pd_pp.py`	Updates mapping validation cases for negative indices.
`tests/test_page_consumer_mapping.py`	Tests appended-region mapping and contributor restrictions.
`tests/test_mha_pp_pd_transfer.py`	Extends transfer tests to draft pools and incompatible consumers.
`tests/test_draft_kv_builder.py`	Supplies draft row bindings for transfer export tests.
`recipes/mesh/MiniMax-M3.md`	Documents PP4 → TP4 DSpark setup and prefill limits.
`atom/spec_decode/pp_aux_relay.py`	Introduces cross-stage DSpark tap relay.
`atom/spec_decode/drafter.py`	Exposes reusable capture hooks and accepts relayed taps.
`atom/spec_decode/draft_kv.py`	Declares draft cache head runs for transfer.
`atom/models/minimax_m3/dspark.py`	Allows lower-TP producers and handles pipeline-stage embeddings.
`atom/model_engine/model_runner.py`	Integrates relay, aligned warmup, and producer proposal skipping.
`atom/kv_transfer/disaggregation/types.py`	Maps appended regions relative to the consumer’s end.
`atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py`	Resolves and validates negative consumer-region indices.
`atom/distributed/pp_comm.py`	Preserves draft taps in pipeline payloads.

_{Review effort: Balanced}

A PP4 prefill runs TP1 stages: each holds 15 of MiniMax-M3.1's 60 layers and all four KV heads of each, while a TP4 decode rank holds every layer and one head. So a stage cannot find its consumer regions by position, and its units are four heads wide where the consumer's are one. - The MHA backend's export says how each unit splits by KV head (`PageRegion.head_major_runs`): head-major per page, and a layer that re-views its cache in smaller SHUFFLE pages (M3's sparse layers, page-16: `SparseMHAPagedAttentionImpl.shuffle_page_tokens`) keeps a 128-token block as eight runs. The indexer keys are the same on every rank. - Under PP, when every layer of a stage holds one KV row of the one geometry, the export numbers each unit as an unsplit model does, in its name (`mha.<geometry>.<field>.layer_<row>`) and as its consumer region: the layer for KV, the sparse-layer slot (`sparse_layer_ordinal`) for the indexer key. Without PP nothing is renamed, so TP4 -> TP4 peers on either side of this change and LMCache MP tensor names are unchanged. - For a consumer at another TP size, the Mooncake producer writes each consumer rank only its heads of every run, batched by the same loop as the whole-block and DCP paths (`queue_descriptors`). Only a TP1 producer (every consumer's routing target, holding every head) may cut; another producer TP or DCP is refused before anything is written, whether or not the consumer describes its runs. A head cut matches region names up to the per-rank head count in the geometry. - An explicit consumer map that names a region the consumer lacks, or sends two regions to one, is refused. The head cut follows Mengqing Cao's PP4 prefill commit 0021a4a (mengqing/minimax-m3-gather-concurrency-sweep). Co-authored-by: Mengqing Cao <cmq0113@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Review on 84d110f: - A GDN hybrid transfers through the block+slot path, which copied whole units and never asked `_consumer_head_cut`. Its per-request state is TP-sharded and never cut, so a consumer at another TP size is now refused there before any write, from a TP1 producer as from any other. - A head cut matched region names up to the whole geometry token, so a bf16 d64 producer unit could pair with an fp8 d128 consumer unit of equal bytes per head. The head dim must agree too. - The PP numbering assumes the consumer's target pool holds a row per layer: an MHA draft caches into a pool of its own (`DRAFT_OWNS_KV_POOL`). The docstring now says so. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

A P/D producer never drafts, but it must write the draft's context K/V for the prompt, which the decode drafts from. Under PP4 the drafter is built on the last stage only: - Its taps (layers 3/17/31/45/59) mostly run on earlier stages; main's `PPAuxRelay` (#2365) carries DSpark's to the last stage. A PP drafter that taps target layers without a config-only aux spec (EAGLE3) is refused at startup: nothing relays its taps, and a producer would write wrong context K/V. - A block drafter warms up on at least two sequences. That floor is read off the config, not the stage's drafter, so every stage warms up on the rows the last stage does. - The draft pool's units say how they split by KV head (DSpark's binding reads and writes page-16: eight runs a block), so a TP1 stage's draft K/V is cut per decode rank like the target's. `merge_pages` numbers regions appended after a mapped contributor from the consumer's end, where the consumer appends them too; the producer resolves them against the consumer's region count at handshake. - Any P/D producer whose drafter writes its context K/V apart from propose (DSpark: M3.1, K3, V4.1) skips `propose()` and its draft graph warmup, on the deferred output path as on the direct one, PP or not, and leaves the DP draft lockstep (every rank skips alike). Its drafts were discarded anyway (`mtp_k` 0). EAGLE still proposes, since it leaves a final chunk's context K/V to propose's first step. - So the M3.1 draft may run below TP4 on a producer. Its one-KV-head check runs at the top of the block pass, before anything is written. The last stage's embedding is a placeholder, so the vocab check reads the LM head. - Adds a recipe section with the prefill at `--max-num-batched-tokens 8192`: the sparse prefill launches one `pa_fwd` row per (token, KV head), HIP caps a launch at 65535 rows, and a TP1 stage holds all four KV heads. Co-authored-by: Mengqing Cao <cmq0113@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Copilot

🟡 Changes recommended

V4.1 DSpark producer requests with multiple output tokens can leave speculative state uncommitted when drafting is skipped, causing the next forward pass or checkpoint to fail.

2 open issues

_{Review effort: Balanced}

Resolve the MHA attention import conflict by keeping get_pp_group for PP region mapping and main's local_num_kv_heads for Ulysses-aware head sizing.

A TP1 producer cutting heads for a consumer at another TP size now refuses a request without `consumer_head_major_runs`, instead of falling back to whole-unit copies that the optional role and byte checks may not catch. Also give main's new test_v41_index_layout connector fake the `_consumer_head_cut` it now needs, and correct the `_pp_num_global_layers` docstring for single-geometry stages of multi-geometry models. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Export unit identity, head dimension, dtype and layout from the attention backend. Validate every mapped region before queuing writes and require structured metadata for cross-TP head cuts. Validation: 786 CPU tests passed, 3 GPU tests skipped using the existing AITER/Triton stubs. Black, Ruff and git diff --check passed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> # Conflicts: # atom/model_engine/model_runner.py # tests/test_mha_pp_pd_transfer.py

#2721 now validates MHA units by `PageUnitSpec` and refuses a head cut on any per-head unit without one, so a TP1 PP prefill's draft rows could no longer reach a TP4 decode. Describe them as the backend describes its own. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Copilot

🟡 Changes recommended

The new tests fail before reaching their assertions because _served is unset, the event mock rejects the stream argument, and non-producer cases lack the hidden-state rows needed to avoid requiring a forward context.

3 open issues

_{Review effort: Balanced}

valar-cop

some changes needed..

ci Run non-GPU unit tests failed — not sure yet (low): The job log under .valar/ci is empty (0 bytes), so the failure cannot be attributed; the same job is green on main (7 recent runs) and 5/7 on other recent PR runs. The new postprocess tests would not explain it if their module self-skips for lack of aiter.

next: Re-run the job / fetch its log and JUnit artifact before attributing the failure to this PR. https://github.com/ROCm/ATOM/actions/runs/43020655723/job/130522744337

The merged main brought tests whose fakes predate this PR: the DeepSeek-V4.1 mono runner's `tokenID_processor` has no `is_pd_producer`, the warmup config no `speculative_config`. The producer-skip test now patches what main's postprocess reads (PCP/SP helpers, `_served`, the forward mode). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Copilot

🔵 Needs a closer look

The head-verdict regression test still uses an object() drafter double without context_kv_only, causing AttributeError before it can reach the verdict and broadcast assertions.

2 open issues

_{Review effort: Balanced}

Copilot

🔵 Needs a closer look

Cross-TP cache layouts and shared DSpark/DP lifecycle changes require GPU validation across affected drafter families, as source review alone cannot establish distributed runtime safety.

1 open issue

Commit V4.1 producer state when skipping draft-token proposals

1 resolved since last review

Update drafter mock with context_kv_only attribute

_{Review effort: Balanced}

valarLip

⚠️ Action required before merge (5)

Drop the getattr default on kv_transfer_config. atom/models/minimax_m3/dspark.py:220-222: Config always has kv_transfer_config (atom/config.py:2178), so the default is only there because tests/test_dspark.py:2660 passes atom_config=None. Read atom_config.kv_transfer_config directly and give that test SimpleNamespace(kv_transfer_config={}). While you're in that test, add the producer case, which no test covers yet: TP1 with a kv_producer config builds num_kv_heads == 4, and the block pass raises RuntimeError (dspark.py:333).
Make page_tokens required instead of defaulting it. atom/spec_decode/draft_kv.py:203 uses getattr(m, "page_tokens", pool.block_size). If a binding forgets to declare its page, the draft is silently cut with the wrong page, which is the 45.8% to 52.4% acceptance bug this PR fixes. Declare page_tokens on every draft-KV binding (or set it in DraftKvBuilder when it binds the module) and read it directly.
Add a test for the PP refusal of EAGLE3, and fetch drafter once. model_runner.py:1013-1025: the guard can fire, because EagleProposer.arm_aux_capture sets _captures_aux without a config spec (eagle_proposer.py:308), but no test raises it. In the same block, drafter = getattr(self, "drafter", None) repeats the lookup on line 1010. Hoist it above line 1010 and use the one value for both.
Merge the duplicated propose-or-commit branch into one helper. model_runner.py:3804-3814 and 3835-3847 are the same if self.drafter.context_kv_only: _commit_target_step(...) else: propose_draft_token_ids(...). Move it into one helper that returns None on the producer path. Also import kv_config_has_producer from atom.kv_transfer.disaggregation, as the other callers do, rather than from .types (dspark.py:63).
Settle the ordering with #3023. Both PRs edit the same share_with_target block in Drafter.load_model (drafter.py:701-706). Whichever lands second must keep #3023's check_shared_embedding inside the not self.context_kv_only branch. Also, #3023's pp_holds_embedding builds a full-vocab embedding on a Kimi-K3 PP last stage even when that stage is a producer. After this PR such a producer never shares the embedding, so make that check false on a producer.

_{Reviewed at 58d2384. Base #2721 (1/2) is merged; CI is green (9 pass, 11 skipped). Already fixed: valar-cop and Copilot asked for the target-step commit when propose is skipped (_commit_target_step, model_runner.py:4101), the producer-gate test setup, and the V4.1 test double; the requested GPU numbers are in the PR body. Rejected: an agent claimed the PP guard can never fire; it can (see item 3).}

valar-cop · 2026-11-12T13:19:08Z

ci triage for 58d2384:

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/43225116704

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/43225117064

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/43225116852

Running Copilot Code Review passed. https://github.com/ROCm/ATOM/actions/runs/43216758070

Pre Checkin passed. https://github.com/ROCm/ATOM/actions/runs/43216739305

- dspark: read `kv_transfer_config` directly (Config always has it) and import `kv_config_has_producer` from the package; test_dspark gives the TP-floor test a config and adds the producer case (TP1 holds all four KV heads, the block pass refuses). - Every draft-KV binding declares `page_tokens` (`Attention` None for its own block, the M3.1 DSpark binding page-16); DraftKvBuilder reads it directly, so a binding that forgets fails at startup instead of being cut with the wrong page. - ModelRunner: one drafter lookup in `load_model`; the PP refusal of an unrelayable tap spec is `_refuse_unrelayed_pp_taps`, tested for EAGLE3 and for a config-declared spec; both output paths call one `_draft_or_commit` helper. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Jasen2201 · 2026-11-12T13:54:48Z

Thanks. Items 1-4 are in c79ccc5, on top of 58d2384:

dspark.py reads atom_config.kv_transfer_config directly. test_dspark.py gives the TP-floor test SimpleNamespace(kv_transfer_config={}) and adds the producer case: TP1 with a kv_producer config builds num_kv_heads == 4, and the block pass raises RuntimeError. The test fails if the block-pass check is removed.
Both draft-KV bindings declare page_tokens: Attention sets None (its own block) and M3.1's KVCacheBinding sets page-16. DraftKvBuilder reads it directly, so a binding without it fails at startup.
load_model looks the drafter up once. The PP refusal is now _refuse_unrelayed_pp_taps, tested for EAGLE3 (refused under PP, allowed at PP1) and for a config-declared spec. The EAGLE3 case fails if the refusal is removed.
Both output paths call _draft_or_commit, which returns None on a context-only producer. kv_config_has_producer now comes from atom.kv_transfer.disaggregation.
[fix] Kimi-K3.1 DSpark under PP: share the target's embedding on the last stage #3023: see the note on that PR. Whichever lands second keeps check_shared_embedding inside the not self.context_kv_only branch and makes pp_holds_embedding false on a producer.

With real aiter on MI355X: the M3.1 cases in test_dspark.py (3) and the relay, draft KV, postprocess, warmup, weight-sharing and V4.1 runner suites (110) pass. black and ruff are clean.

Copilot

🔵 Needs a closer look

Cross-TP cache layout and speculative-decoding collective changes need human assessment because supplied GPU coverage excludes multi-DP producers and distributed behavior cannot be independently validated here.

No open issues 🎉

1 resolved since last review

Commit V4.1 producer state when skipping draft-token proposals

_{Review effort: Balanced}

valarLip

⚠️ Action required before merge (2)

Pass the drafter you already fetched to the mono-decode installer. atom/model_engine/model_runner.py:1020 still calls getattr(self, "drafter", None), ten lines after drafter is fetched at line 1010. Pass drafter there so the block looks it up once, as item 3 of the last review asked.
Wait for the GPU accuracy jobs to finish green before merging. At c79ccc5, the Pre Checkin, offline smoke and unit jobs pass, but the accuracy matrix (MiniMax-M3, Kimi-K3/K3.1 DSpark, DeepSeek-V4-Pro/V4.1-Flash DSpark, and others) is still pending. If DeepSeek-V4-Pro DSpark fails with "VllmBackend can only be called once", that failure also happens on main and is not caused by this PR.

Still open, not blocking this PR: #3023 is open (CHANGES_REQUESTED). Whichever PR lands second must keep `check_shared_embedding` inside the `not self.context_kv_only` branch (`atom/spec_decode/drafter.py:704`) and must make `pp_holds_embedding` false on a producer. Your note on #3023 says you will do this.

_{Re-reviewed at c79ccc5. Already fixed: dspark.py:220 reads atom_config.kv_transfer_config directly, and the test passes SimpleNamespace(kv_transfer_config={}). A new TP1 kv_producer test checks num_kv_heads == 4 and the block-pass RuntimeError (test_dspark.py:2664). draft_kv.py:203 reads m.page_tokens directly: Attention declares None, which means its own block, and M3.1's KVCacheBinding declares 16. _refuse_unrelayed_pp_taps has an EAGLE3 PP refusal test (test_pp_aux_relay.py:275). The two propose/commit branches are now one helper, _draft_or_commit. kv_config_has_producer is imported from atom.kv_transfer.disaggregation. Copilot's issues are all resolved. Rejected: a request for a Protocol type and a page_tokens assert (over-design; no current binding is affected), and a note that these suites skip on the aiter-less CPU runner (true before this PR; the author ran them on MI355X).}

Jasen2201 · 2026-11-12T14:22:45Z

@valarLip Fixed the remaining lookup in 17d2e69: the mono-decode installer now receives the drafter fetched above.

Black, Ruff, and git diff --check pass for this change. Understood on waiting for the GPU accuracy jobs to finish green before merging. Please re-review the latest commit when you have a chance.

Copilot

🔵 Needs a closer look

Shared speculative-decode lifecycle and cross-TP KV layout changes require maintainer GPU integration review, as this source-only review could not independently validate distributed execution, graph replay, or DP collective behavior.

No open issues 🎉

_{Review effort: Balanced}

valarLip

⚠️ Action required before merge (1)

Wait for the GPU accuracy matrix to finish green before merging. At 17d2e69, Pre Checkin (Black, Ruff, non-GPU unit tests, catalog schema), the offline smoke test and model-config load pass. All 14 ATOM Test accuracy jobs are still queued: MiniMax-M3, Kimi-K3/K3 DSpark/K3 DSpark PP2, Kimi-K3.1-Pro (+DSpark), DeepSeek-V4-Pro (+DSpark), V4.1-Flash DSpark, R1 MTP, GLM-5.3-Flash, gpt-oss, Qwen3.5 and Qwen3-Next. The c79ccc5 run was cancelled by the new push, so no accuracy result exists for this PR yet. If DeepSeek-V4-Pro DSpark fails with "VllmBackend can only be called once", that failure also happens on main and is not caused by this PR.

Still open, not blocking this PR: #3023 is open (CHANGES_REQUESTED). Whichever PR lands second must keep `check_shared_embedding` inside the `not self.context_kv_only` branch (`atom/spec_decode/drafter.py`) and must make `pp_holds_embedding` false on a producer.

_{Re-reviewed at 17d2e69. Already fixed: load_model looks up drafter once (model_runner.py:1010) and passes it to _refuse_unrelayed_pp_taps and to the mono-decode installer (model_runner.py:1020). The earlier asks still hold: dspark.py reads atom_config.kv_transfer_config directly, with a TP1 kv_producer test; draft_kv.py reads m.page_tokens directly; the EAGLE3 PP refusal is tested; there is one _draft_or_commit helper; and kv_config_has_producer comes from atom.kv_transfer.disaggregation. Copilot has no open issues. Rejected: requests for a Protocol type and a construction-time head assert (over-design; a TP4 producer legitimately holds one head), a getattr default on context_kv_only (a defensive default, which this repo does not use), and type-annotation nits.}

Jasen2201 · 2026-11-13T01:45:55Z

@valarLip only Accuracy (GLM-5.3-Flash-kpool-16shot Looks like this is a resource issue. Please approve.

Copilot

🔵 Needs a closer look

Shared runner changes affect collective participation and state-commit sequencing across DSpark producers, while reported validation covers only M3.1 and V4.1, leaving DP-attention lockstep and Kimi-K3/DeepSeek-V4 paths needing human runtime review.

No open issues 🎉

_{Review effort: Balanced}

Jasen2201 · 2026-11-13T02:07:21Z

Resolved the conflicts with main in 9784dff and checked the resolution against both parents. The M3 DSpark binding still declares 16-token pages for KV transfer, while the shared binding's default preserves the existing pool block size for other drafts. The context-only producer's weight-sharing guard is also preserved.

Validation: Black and Ruff passed. Targeted CPU regression tests with external-kernel stubs reported 286 passed and 12 skipped; one test requiring the real AITER private API was excluded.

Based on the code review and these checks, the conflict resolution is not expected to affect model correctness. This is a preliminary assessment: GPU accuracy has not yet been verified for this commit, and the vLLM/SGLang accuracy jobs were skipped by the changes-requested CI gate.

Jasen2201 mentioned this pull request Nov 5, 2026

[MiniMax-M3.1] P/D from a PP4 prefill into a TP4 decode (1/2: KV transfer) #2721

Merged

4 of 5 tasks

Jasen2201 added this pull request to stack #3039 November 5, 2026 12:15

valar-cop reviewed Nov 5, 2026

View reviewed changes

Jasen2201 force-pushed the Jasen/m31-pd-pp4-dspark branch from e052e08 to 60263a0 Compare November 5, 2026 14:00

Copilot AI balanced review requested due to automatic review settings November 5, 2026 14:00

Copilot started reviewing on behalf of Jasen2201 November 5, 2026 14:00

View session

Copilot AI reviewed Nov 5, 2026

View reviewed changes

Comment thread tests/test_pp_aux_relay.py Resolved

Jasen2201 force-pushed the Jasen/m31-pd-pp4-dspark branch from 60263a0 to aab63f3 Compare November 5, 2026 14:12

Jasen2201 and others added 3 commits November 6, 2026 03:05

Jasen2201 force-pushed the Jasen/m31-pd-pp4-dspark branch from aab63f3 to 4ea4cc2 Compare November 6, 2026 03:17

Copilot AI balanced review requested due to automatic review settings November 6, 2026 03:17

Copilot started reviewing on behalf of Jasen2201 November 6, 2026 03:18

View session

Copilot AI reviewed Nov 6, 2026

View reviewed changes

Comment thread atom/model_engine/model_runner.py Resolved

Jasen2201 and others added 4 commits November 9, 2026 18:26

Merge main into Jasen/m31-pd-pp4-tp4

5905e7f

Resolve the MHA attention import conflict by keeping get_pp_group for PP region mapping and main's local_num_kv_heads for Ulysses-aware head sizing.

refactor: separate Mooncake block transfer preparation

ff732fa

Jasen2201 mentioned this pull request Nov 11, 2026

[fix] Kimi-K3.1 DSpark under PP: share the target's embedding on the last stage #3023

Open

Jasen2201 and others added 3 commits November 11, 2026 11:40

Merge main into Jasen/m31-pd-pp4-tp4

6a58ef9

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge Jasen/m31-pd-pp4-tp4 into Jasen/m31-pd-pp4-dspark

dc634ac

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> # Conflicts: # atom/model_engine/model_runner.py # tests/test_mha_pp_pd_transfer.py

Copilot AI balanced review requested due to automatic review settings November 11, 2026 12:28

Copilot started reviewing on behalf of Jasen2201 November 11, 2026 12:29

View session

Copilot AI reviewed Nov 11, 2026

View reviewed changes

Comment thread tests/test_postprocess_width.py Resolved

valar-cop reviewed Nov 11, 2026

View reviewed changes

Comment thread atom/model_engine/model_runner.py Resolved

Comment thread tests/test_postprocess_width.py Resolved

Copilot AI balanced review requested due to automatic review settings November 12, 2026 11:47

Copilot started reviewing on behalf of Jasen2201 November 12, 2026 11:48

View session

Copilot AI reviewed Nov 12, 2026

View reviewed changes

test: update V4.1 drafter double for context-only API

58d2384

Copilot AI balanced review requested due to automatic review settings November 12, 2026 12:18

Copilot started reviewing on behalf of Jasen2201 November 12, 2026 12:19

View session

Copilot AI reviewed Nov 12, 2026

View reviewed changes

valarLip requested changes Nov 12, 2026

View reviewed changes

Copilot AI balanced review requested due to automatic review settings November 12, 2026 13:54

Copilot started reviewing on behalf of Jasen2201 November 12, 2026 13:55

View session

Jasen2201 added the ci:atom label Nov 12, 2026

Copilot AI reviewed Nov 12, 2026

View reviewed changes

valarLip reviewed Nov 12, 2026

View reviewed changes

refactor(runner): reuse drafter for mono installation

17d2e69

Copilot AI balanced review requested due to automatic review settings November 12, 2026 14:20

Copilot started reviewing on behalf of Jasen2201 November 12, 2026 14:21

View session

Copilot AI reviewed Nov 12, 2026

View reviewed changes

valarLip reviewed Nov 12, 2026

View reviewed changes

zufayu requested a review from valarLip November 13, 2026 01:28

zufayu assigned valarLip Nov 13, 2026

Merge main into PR 3038 and preserve draft KV page layouts

9784dff

Copilot AI balanced review requested due to automatic review settings November 13, 2026 01:59

Copilot started reviewing on behalf of Jasen2201 November 13, 2026 02:00

View session

Copilot AI reviewed Nov 13, 2026

View reviewed changes

Conversation

Jasen2201 commented Nov 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Results

Prefill cost of DSpark, PP4 vs TP4

Test plan

Limits and follow-ups

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

valarLip left a comment

Choose a reason for hiding this comment

⚠️ Action required before merge (5)

Uh oh!

valar-cop commented Nov 12, 2026

Uh oh!

Jasen2201 commented Nov 12, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

valarLip left a comment

Choose a reason for hiding this comment

⚠️ Action required before merge (2)

Uh oh!

Jasen2201 commented Nov 12, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

valarLip left a comment

Choose a reason for hiding this comment

⚠️ Action required before merge (1)

Still open, not blocking this PR: #3023 is open (CHANGES_REQUESTED). Whichever PR lands second must keep check_shared_embedding inside the not self.context_kv_only branch (atom/spec_decode/drafter.py) and must make pp_holds_embedding false on a producer.

Uh oh!

Jasen2201 commented Nov 13, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Jasen2201 commented Nov 13, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Jasen2201 commented Nov 5, 2026 •

edited

Loading

Still open, not blocking this PR: #3023 is open (CHANGES_REQUESTED). Whichever PR lands second must keep `check_shared_embedding` inside the `not self.context_kv_only` branch (`atom/spec_decode/drafter.py`) and must make `pp_holds_embedding` false on a producer.