[wanzhenchn]

Stack

Part 3 of 4 of the Kimi-K3 PP8 prefill for P/D serving, split from #2709 to fit ATOM's limit of 1,000 changed lines and 15 files. Merge in this order:

1. feat(kimi-k3): PP8 P/D prefill 1/4 — PAGE-backed state checkpoints under PP #2709: PAGE-backed KDA state checkpoints under PP. Merged.
2. feat(kimi-k3): PP8 P/D prefill 2/4 — LMCache state tier under PP #3068: Kimi-K3 LMCache state tier under PP. Targets main.
3. This PR: KV and KDA state from a PP prefill to a DCP × TP decode. Targets feat(kimi-k3): PP8 P/D prefill 2/4 — LMCache state tier under PP #3068.
4. perf(kimi-k3): PP8 P/D prefill 4/4 — MLA staging for a DCP consumer, more PP batches in flight #3070: MLA staging for a DCP consumer, and more PP batches in flight. Targets this PR.
5. ci(k3-agentic): Kimi-K3 AgentX 1M P/D cells, one per concurrency on its best measured prefill #2787: the Kimi-K3 AgentX CI cells. Targets perf(kimi-k3): PP8 P/D prefill 4/4 — MLA staging for a DCP consumer, more PP batches in flight #3070.

The diff shown is against #3068. Once #3068 merges, this PR is rebased onto main.

Motivation

For AgentX 1M P/D at C32–C88, the fastest prefill measured was PP8×TP1 feeding a TP8×DCP8 decode: at C32, 10,409 vs 9,196 tok/s/GPU for a TP8EP8 prefill (+13%). That topology was refused at startup.

Changes

1. KV and KDA state from a PP prefill to a DCP × TP decode.
   • Problem: PP8×TP1 → TP8×DCP8 was refused at startup. Once allowed, four things stood in the way:
     • region placement: per-kind row numbering;
     • TP-sharded state: a TP1 stage holds every KDA head, while a TP8 rank holds one eighth;
     • asymmetric DCP;
     • consumer pairing: tp_rank % remote_tp_size is right only for replicated KV.
   • Change: the K3 builder places rows by module layer and publishes each slot region's TP segments. It refuses at startup a region it cannot split. The producer writes the consumer rank's share. A consumer outside the producer's head group is refused, and so is a decode TP that is not a multiple of the prefill TP, naming both instead of reporting a slot size mismatch.
   • main's MoRIIO connector also moves K3's state, but it reads none of this placement, so it keeps refusing PP.
   • A sparse MLA model that shares this builder (has_mla_indexer, e.g. glm5_next_text) has indexer regions that no stage places, so it keeps refusing PP too, now by name.
   • Chunked transfer still sends no per-request state to a DCP decode. Neither side can refuse that at startup: chunked transfer is the producer's setting, and the producer learns the decode's DCP only from a request. The decode now rejects such a request at admission with HTTP 400, before enqueue, as it does the other chunked topologies it cannot serve.
   • Benefit: enables the topology.
2. Split K3 ReplaySSM records from a TP1 PP stage to a TP8 decode.
   • Problem: with DSpark on both sides, the first transfer was refused ("cannot split 8 ways").
   • Change: records are published as a head-split segment. The per-slot cursor is a TP-replicated region, written whole.
   • Benefit: DSpark decode with ReplaySSM behind a PP8 prefill.
3. Let PP stages 1..7 drop a request's Mooncake producer entry.
   • Problem: only stage 0 receives the release, so the other stages kept one entry per request for the life of the worker.
   • Benefit: bounded memory. The longest gen-2 pause went from 329–363 ms to 245–293 ms; end-to-end throughput is flat.

Tests

• ruff and black are clean.
• The PR adds no test files. It changes five existing tests:
  • test_gdn_state_relocation.py: the builder fixture now describes which layers a PP stage holds, and the TP size the builder reads. test_pipeline_parallelism_is_refused becomes test_moriio_keeps_refusing_pp, since only MoRIIO still refuses PP.
  • test_deepseek_v4_transfer_regions.py, test_pd_chunked_transfer.py and test_v41_index_layout.py: their connector fakes set the three new slot-region attributes to None.
  • entrypoints/test_pd_prompt_token_ids.py: one admission test's decode config sets decode_context_parallel_size=1.
• Pre Checkin on 4725803 (feat(kimi-k3): PP8 P/D prefill 2/4 — LMCache state tier under PP #3068's head a54d273 merged in, no conflicts here): 13,203 passed, 1,299 skipped, 3 xfailed (43223683838). On 56bcaa3 (feat(kimi-k3): PP8 P/D prefill 2/4 — LMCache state tier under PP #3068 with main merged in) it failed three cases of test_v41_index_layout.py, which came in with main (feat(indexer): migrate V4 and V4.1 FP4 scoring to the OPUS config API #2809) and whose connector fake lacked those attributes. 3f2e304 adds them.
• The five changed files on 3f2e304, locally on CPU: 299 passed, 7 skipped. test_gdn_state_relocation.py needs aiter, so CPU runs skip it. On 56bcaa3, the offload, state, PP, DCP and P/D tests pass locally on CPU: 3,576 passed, 352 skipped. On a GPU (rocm/atom-dev:nightly_202611021458), test_kda_layout_id.py, test_gdn_state_relocation.py, test_pd_pp.py, test_lmcache_offload_connector.py, test_pd_mla_staging.py and test_pd_mla_landing.py give 676 passed and none skipped on ad9d6d2, before main was merged in.

Related

• End-to-end results for the whole stack are in ci(k3-agentic): Kimi-K3 AgentX 1M P/D cells, one per concurrency on its best measured prefill #2787 (PP8 cells C32–C88).

[valar-cop]

some changes needed..

gsm8k/lm_eval for the PP8xTP1 -> TP8xDCP8 cell?

ci green.

[valar-cop]

the sparse-mla pp refusal looks right now, one thing left..

ci Run non-GPU unit tests (Pre Checkin) failed — not sure yet (low): The captured log is empty (0 bytes), so the failed step and the failing test cannot be identified. The same job passed 8/8 recent runs on main and 5/7 on other recent runs, so neither the PR nor CI can be blamed from this evidence.

next: Re-run the job and share the failing step's log; if it reproduces on the head, check whether it also fails at the merge base. https://github.com/ROCm/ATOM/actions/runs/42758671954/job/129638942868

[Jasen2201]

@valar-cop pls review

[valar-cop]

@Jasen2201 already looked at 4e06431, nothing new since my last review. say @valar-cop re-review if you want a fresh one.

[Jasen2201]

@valar-cop re-review