feat(kimi-k3): PP8 prefill for P/D serving — KDA state under PP, PP→DCP transfer, MLA staging by wanzhenchn · Pull Request #2709 · ROCm/ATOM

wanzhenchn · 2026-10-26T05:30:02Z

Motivation

#2365, now on main, serves Kimi-K3 at PP > 1 in one engine. For AgentX 1M P/D at high concurrency (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%). Before this PR:

that topology was refused at startup;
once it was allowed, there was no prefix reuse under PP;
moving KV from the PP stages to the DCP decode took seconds per request.

This PR adds what that topology needs, on top of #2365. Everything in it is specific to PP or to the PP → DCP transfer.

Stack and split

#2365 has landed on main as 64412f1, so this PR now targets main. The original #2709 was split into three PRs to make review easier:

[feature] support pp>1 for Kimi-K3 #2365: Kimi-K3 at PP > 1, merged as 64412f1
This PR: PP-specific support and fixes
perf(pd): P/D first-token latency and API-server GC fixes (from Kimi-K3 AgentX 1M) #2786: P/D and API-server changes, on main, independent of PP
ci(k3-agentic): Kimi-K3 AgentX 1M P/D cells, one per concurrency on its best measured prefill #2787: CI cells, one per concurrency. Merges after this PR and perf(pd): P/D first-token latency and API-server GC fixes (from Kimi-K3 AgentX 1M) #2786.

While #2365 was open, this PR merged its head each time it moved, the last 9fb45ff in b82e2d0. #2365 landed as the squash 64412f1, whose tree is 9fb45ff merged with main d8525dd (#2913). 0ab7e65 merges 64412f1. The two files the squash repeats and this PR edits further, scheduler.py and tests/test_pp.py, keep this PR's versions, so the diff against main is the one this PR had against 9fb45ff. Merging #2365's heads changed two things from #2365. Its test_head_flushes_a_held_chunk_whether_or_not_it_is_acked bounds the head's in-flight window through pp_size, which this PR replaces with pp_max_inflight (item 14), so it now sets both. Its 5b0b85d drops the start + chunk < num_prompt_tokens guard on the PP checkpoint hold because PP refused PAGE-backed checkpoints. This PR lets them span PP stages (item 1), so a rung can sit at prompt end again, and the guard stays.

Changes

Feature support (each was a startup refusal or a failed first transfer)

PAGE-backed KDA state checkpoints under PP.
- Problem: PAGE checkpoints refused PP. Only the head schedules, so a store or restore op it builds must name the same unit ids on every stage. Each stage, however, holds its own layers' state in its own unit size.
- Change: at startup the stages agree on one joint spec over the PP group: every stage takes the largest units_per_checkpoint. Each op is checked against a joint layout id and the summed image bytes of every stage. A stage without a linear-attention layer holds no state for an image, so every stage refuses to start with one error naming it.
- Benefit: prefix caching of KDA state works under PP.
Kimi-K3 LMCache state tier under PP, with a producer fence on state stores.
- Problem: the state tier was refused under PP. Separately, a store could gather PAGE units before the copy into them had executed; on a PP stage that copy can be a pipeline depth behind.
- Change: each stage builds its own tier, keyed by its PP-aware worker id. State stores wait on an event recorded on the compute stream at dispatch. Stages hold images of different sizes, and each stage's tier evicts by its own, so the head sizes its index of what the tier holds by the largest image any stage holds, which the stages agree on with the joint spec (item 1).
- Only Kimi-K3's PAGE checkpoints span PP stages. Qwen's GDN layouts keep fork checkpoints under PP, so they keep refusing the state tier there, on the engine side as on the worker side.
- Benefit: joint (KV + state) LMCache hits under PP. On PP4×TP2, the hit rate went from 22.47% to 91.52%, mean TTFT from 777 s to 39 s, and throughput from 568 to 4,392 tok/s/GPU. The race reproduces in a unit test: fresh bytes read are 0.000 without the fence and 1.000 with it.
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.
- 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.
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.

Correctness under PP

Scale K3's default save bound with pipeline stages.
- Problem: under PP a save completes only after every stage stores, so the default bound of two outstanding saves held K3 PP8 to ~0.7 saves/s. The slot pool sat at 141–144/144, and admission stalled with 1–3 requests running.
- Change: the default is multiplied by pipeline_parallel_size; explicit settings are unchanged.
Put the prompt-end anchor on the joint chunk grid.
- Problem: with a hash block finer than the LMCache chunk (PP: 128 vs 1024), 7 of 8 anchors were lost once the prompt left HBM.
- Change: the anchor is floored to the chunk grid when the joint load is configured.
Floor demand checkpoints onto the LMCache chunk grid.
- Problem: an HBM hit that stops on an off-grid rung cannot continue from LMCache. In a C88 run, 559 requests recomputed 11.75M prompt tokens this way: 54% of the avoidable prefill recompute.
- Change: demand rungs are floored like the anchor.
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.
NUMA binding uses the PP-aware GPU index.
- Problem: the index omitted pp_rank, so all eight PP8×TP1 stages bound to GPU 0's node. With ATOM_NUMA_NODE=0,0,0,0,1,1,1,1, stages 4–7 ran on the wrong node.
- Change: the index is (dp_local_rank * pp_size + pp_rank) * tp_size * pcp_size + rank. The binding and ModelRunner's device choice now share one local_device_index in config.py, so they cannot drift apart again.

Performance

File a PAGE state checkpoint when its chunk is scheduled.
- Problem: the chunk after a checkpoint rung waited for the rung's batch to retire, a full pipeline traversal. That hold was ~3.3 s of a ~7 s fixed prefill latency on K3 PP8.
- Change: a PAGE image is a copy made in batch order on every stage, so it is filed at schedule time. Fork classes keep the hold.
Stage MLA KV for a DCP consumer below a page of interleave.
- Problem: at interleave 1 the producer wrote one 576-byte descriptor per token per MLA region, about 1M per 117K-token request. The decode waited ~11.5 s per request for its KV.
- Change: the consumer rank's tokens are gathered into staging pages on a side stream, and one descriptor is written per destination page. The slots are held back from the KV cache budget, like the plain-MLA staging pool: 144 MiB per producer GPU at 128 pages and 16 send workers. That costs KV cache whatever the decode runs, so it is opt-in: ATOM_PD_MLA_STAGING_CHUNK_PAGES defaults to 0, the old path, and ci(k3-agentic): Kimi-K3 AgentX 1M P/D cells, one per concurrency on its best measured prefill #2787's PP8 cells set 128. The slots are handed out through the connector's SlotPool, as for the other staging pools.
- Only the request's KV-ready event orders the gather behind the batch's KV writes, so a request without one keeps the per-token path.
- Benefit: together with 12, the decode's KV wait for 64–128K-token prompts is 165 ms at C88.
Give every MLA staging slot its own gather stream (ATOM_PD_MLA_STAGING_STREAM_PRIORITY).
- Problem: with one gather stream for all slots, as this PR first had it, each gather's synchronize waited for every send worker's queued gathers.
- Benefit at C88, vs one shared stream:
  - send queue / exec: 440 / 330 → 46 / 97 ms;
  - TTFT P50 / P90: 5.08 / 15.35 → 4.55 / 13.26 s;
  - throughput: 17,908 → 18,163 tok/s/GPU.
- Priority −1 vs 0: the decode's KV wait drops from 373 to 165 ms, and TTFT P90 from 15.13 to 13.26 s.
Warm the MLA staging slot streams at startup.
- Problem: under load, the first gather on a slot stream stalled each stage's first push for 20–100 s.
- Benefit: one single-row gather per stream at startup removes that stall.
Let the PP head keep more batches in flight (ATOM_PP_MAX_INFLIGHT, 0 = stage count).
- Problem: with variable prefill batch sizes, a cap equal to the stage count leaves the stages behind a long batch idle.
- Benefit: at C88, 16 in flight gives ~15% more throughput than 8. 12 was no better than 16.

Tests

New or extended unit tests:
- test_pp_paged_checkpoints.py, test_page_unit_checkpoint.py, test_kda_layout_id.py (1);
- test_lmcache_offload_connector.py, test_block_manager.py (2), including that a Qwen layout keeps no state tier under PP and that the index counts the largest stage image;
- test_pd_pp.py, test_gdn_state_relocation.py (3, 4, 8, 11, 14), entrypoints/test_pd_prompt_token_ids.py (3);
- test_lmcache_offload_v4_page_slot.py (5);
- test_state_checkpoint.py (6, 7, 10);
- test_pd_kv_staging.py, test_pd_producer_config.py (11). The first imports the connector through stubbed_aiter, so its CPU cases, staging at 1 and 4 pages per chunk, run without aiter. The second checks the staging reserve, and the first that the slots fit it.
In a CPU-only environment, the suite Pre Checkin runs (tests/ without tests/plugin) fails nowhere on 9e9f07e that it passes on the base, [feature] support pp>1 for Kimi-K3 #2365's head b62bf7a. The failures it has on both need what this environment lacks (aiter, host UVA). The two test_multiprocess_metrics.py cases that give a subprocess 30 s to import a server entry point time out here on both trees, one or the other from run to run. On 0ab7e65, merged with main 64412f1, the 18 test files this PR changes, the scheduler tests and feat(api-server): --api-server-count N, front-end worker processes for the OpenAI server #2913's front-end tests pass: 1,936 passed, 6 skipped. Black and ruff are clean on the changed files.
Pre Checkin passed on 0ab7e65, the head, with 11,524 passed, 1,176 skipped and 3 xfailed (42025058447), against 11,458, 1,174 and 3 on main 64412f1 (42024391441). The 66 more are this PR's tests. Its runner has no aiter, so test_gdn_state_relocation.py, test_kda_layout_id.py and test_pp_paged_checkpoints.py are skipped there as whole modules. test_pd_kv_staging.py now runs its CPU cases there, and test_pd_pp.py all 155 tests, each through an aiter stub.
On a GPU node, those three modules, test_pd_kv_staging.py, test_pd_pp.py, test_pp_draft_input_ids.py and test_pp_aux_relay.py passed on afcaebe in rocm/atom-dev:nightly_202610291449: 250 passed, none skipped. On its first rebase onto main, 252 passed. On bd26837, before the merges of [feature] support pp>1 for Kimi-K3 #2365's later heads, 258 passed, none skipped. They have not been run on a GPU since.
End-to-end results 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).

github-actions · 2026-10-26T05:30:18Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2709 --add-label <label>

Yuechguo

LGTM

Kimi-K3 under pipeline parallelism (verified as pp2 x tp4), including DSpark speculative decoding and prefix caching of the KDA state. PP still needs --enforce-eager: decode has no CUDA graphs yet. Model and transport - Carry K3's block_residual across the stage boundary with the hidden states, and load the DSpark draft's own embedding on the last stage, where the target's embedding is a PPMissingLayer. - Relay DSpark's aux hidden states across stages (PPAuxRelay). The drafter lives on the last stage, so a hook on a tap owned by an earlier stage lands on a PPMissingLayer and never fires. Each stage captures its own taps and forwards them, concatenated in layer_ids order, in the intermediate-tensor payload. The embedding tap (-1) is resolved as the drafter resolves it: the embedding module, identity extract. Speculative decoding under PP - Let a PP engine verify its own drafts. PP output is undeferred, so: the sampler reports this step's rejection and bonus counts on the non-deferred path; they are published to every stage's attention backends, so non-sampling stages roll back linear-attention state; postprocess keeps ONE run of exactly mtp_k + 1 placeholders after every step, whatever was rejected. - Rebuild the verify window from that run. All four backends read `verify_context_shift()`: `num_rejected` under deferred output, 1 under undeferred. The arithmetic lives in token_layout.decode so it is testable without aiter. The window's positions and KV slots are anchored where the scheduler staged them. - Gate and size the verify window on the run's speculative config, not on whether this stage drafts: only the last stage has a drafter, but every stage verifies. Per-stage buffers are sized the same way. - A P/D producer still never speculates, under PP or not. - Round MLA's query-head padding to a width aiter dispatches (24 -> 32 at TP4; TP8 already pads 12 -> 16). KDA state and prefix caching under PP - Progress advances at schedule time, so a chunk ending on a checkpoint rung would let the next chunk overwrite the state slot before the checkpoint is filed. Hold the next chunk (both prefill scans skip it) until the batch that carries the rung files its hashes; the hold is released with that batch even if the seq was preempted, aborted or retired meanwhile. - The checkpoint is filed after the prompt, so its state fork lands on the first decode. Honour it there: publish the source slot and move conv and recurrent state onto the write slot before the kernels run. With DSpark that decode is a verify step, so the spec path does the same. Only a backend whose model does that copy may take a fork on the spec path (`GDNStateMixin.spec_honours_state_fork`, set by K3); every other backend keeps the assert. K3's mono plan refuses a forked step, since a mono step replaces the branch that does the copy. CI - Add a Kimi-K3 pp2 x tp4 DSpark accuracy case with prefix caching and ptpc_fp8 online quant on, as the TP8 Kimi-K3 case runs it, the only accuracy job that reaches any of the paths above. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- tokenIDProcessor.verify_ctx_lens(context_lens) replaces token_layout.decode.verify_context_shift: the attention builders now take their verify-window context lengths straight from the runner. - Rename the runner's `spec_enabled` -> `is_spec_run` (speculative decoding configured, every PP stage runs verify windows) and `use_spec` -> `has_drafter` (this stage owns the drafter and rejection sampler; under PP the last stage only). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

A producer never speculates, so the Kimi-K3 + DSpark gate refused GLM-5.2's PP prefill + MTP Atomesh cells for nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The kernel resumes from column num_accepted - 1, so a forked verify with a bonus read a recycled write slot; only column 0 held the source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

GDNStateMixin.spec_honours_state_fork() defaults to False; only Kimi-K3's KDA forward (kimi_linear) reads spec_state_indices_in_tensor and opts in, so glm5_next, Qwen3-Next/3.5 and Qwen4Exp keep the assert. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

UBatchWrapper.forward takes no intermediate_tensors, so TBO never ran under PP; refuse it in config instead of keeping a grow path nothing reaches. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

seed_forked_slots indexed with a GPU bool mask, a host sync in every linear-attention layer of a fork step. The builder already knows the (write, source) pairs: ship them once per step in the pinned copy that parks write_pos, and index with them directly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The lowest own tap opens the forward's send buffer and every own tap writes its columns; the reset and the _filled check read state that the compiled forward bakes to the warmup trace, so they guarded nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…der PP pp_speculative_methods sits beside mono_decode_installers, so the next model wired to the relay adds a row there instead of editing Config. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

_pp_checkpoint_hold -> _checkpoint_holds (keyed on advance_on_schedule), holds_checkpoint_hostage -> carries_held_checkpoint, and _defer_prefill_hashes -> _queue_or_flush_prefill_hashes, since it flushes the whole queue; mark the hold as scaffolding #2709 replaces. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Read seq.num_placeholder_tokens, which postprocess refills to one full window, instead of re-deriving the run's width from the scheduled count. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Drafter.take_relayed_aux keeps the received [N, k*H] block and aux_for serves it ahead of the own-tap buffers, so the last stage no longer splits it per tap for compute_draft_kv to concatenate back (DSpark skips the cat when it is the only block). The relay takes the drafter class's spec on every stage instead of the drafter's private _aux_capture_spec, and the runner no longer hands it _aux_buffers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

_aux_capture_spec now parses dspark_target_layer_ids once, through aux_capture_spec, and drops its unused target_model parameter. PPAuxRelay keeps only spec (no layer_ids / hidden_size copies), drops the dead _armed guard, and raises the missing-payload error from one place. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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

Brings #2365's fifteen review commits since c4975cd (2edbab7 through 30829a3). One conflict, in `Scheduler`'s PP state-checkpoint hold. #2365 turned `_pp_checkpoint_hold` into `_checkpoint_holds`, keyed by the holding chunk's end so only that chunk releases it; this branch files the checkpoint at schedule time where the class allows it and skips a rung at prompt end. Keep #2365's map and release, and this branch's conditions for taking a hold. The note that this branch would file at schedule time is dropped, as the merge does so. Co-authored-by: Cursor <cursoragent@cursor.com>

valar-cop

LGTM

ci Run non-GPU unit tests (Pre Checkin) failed — not sure yet (low): The log artifact is 0 bytes, so there is no failure text to attribute. The same job fails 4 of 8 recent main runs in the same window (14:17-14:52), while every other job at this head passes and the head is merges-only over a tree that passed 10,737 before.

log file .valar/ci/job_126043848503.log is empty — 0 bytes
same job on main: ['2026-11-03T15:45 success', '2026-11-03T15:33 success', '2026-11-03T14:52 failure', '2026-11-03T14:40 failure', '2026-11-03T14:39 failure', '2026-11-03T14:17 failure']
next: Re-run the Pre Checkin job on 1a07eb2 and read the log; if it is green, compare the passed/skipped counts against the parent head (10,739 / 1,133 on 41422091906) before merge. https://github.com/ROCm/ATOM/actions/runs/41672378533/job/126043848503

wanzhenchn · 2026-11-03T16:37:09Z

Thanks. The log of job 126043848803 is readable through the API, and the failure is not from this head:

The job checks out 645ae1ed8, which merges 1a07eb21a onto 9269fc82b, i.e. [feature] support pp>1 for Kimi-K3 #2365 on main af97c262d.
At af97c262d, main fails the same job. tests/model_ops/test_dp_gather_routed_rows.py (fix(moe): restore TP MoE num_local_tokens #2905) stubs triton and triton.language but not triton.language.extra, which perf(moe): Triton softmax top-k routing, bit-identical to AITER #2665's atom/model_ops/triton_topk_softmax.py imports, so collection stops with ModuleNotFoundError: No module named 'triton.language.extra'.
[feature] support pp>1 for Kimi-K3 #2365's own run on 9269fc82b fails identically (job 126011160276).
fix(ci): stub triton.language.extra in the DP gather CPU test #2954 (a8946ff73) adds the missing stub, and main's Pre Checkin passes from there.

A re-run reuses the same merge commit, so it would fail the same way. The next trigger against current main picks up the fix; nothing to change in this PR.

Taps owned by an earlier PP stage arrive through the relay, so the last stage gives them a 0-row buffer and no hook (~470 MB/rank for K3 at pp2). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Brings #2365's merge of main a8946ff, which carries #2954's stub of `triton.language.extra` in the DP gather CPU test. Without it the non-GPU unit test job cannot collect `test_dp_gather_routed_rows.py`. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

Both now ask whether the tap's module is a PPMissingLayer, the embedding included, instead of the relay keying the -1 tap on is_first_rank. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Brings #2365's review commit a1436cd, which allocates the drafter's aux buffers only for this stage's own taps. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

Brings #2365's review commit b62bf7a, which shares one tap residency rule between the drafter and the PP aux relay. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

- State tier: refuse pipeline parallelism again for every layout except kimi_k3. Under PP, Qwen keeps fork checkpoints and its worker turns the tier off, so the engine side must not install an index or floor rungs to the LMCache chunk for it. - Under PP the head bounds its state index by the largest image any stage holds (`pipeline_image_bytes` in the agreed spec), not its own: each stage's codec evicts by its own image. - MLA shard staging (ATOM_PD_MLA_STAGING_CHUNK_PAGES) is now held back from the KV cache budget (`mla_shard_staging_reserve_bytes`), and its slots fit that reserve. They are handed out by the connector's `SlotPool` like the other staging pools, and a gather waits on its slot's stream. - The scheduled-checkpoint hash reuses `_chain_to`; one helper drops a completed prefill with its ready event; ModelRunner and the NUMA binding share `local_device_index`; both engine cores share the idle state publish; always-set connector fields are read directly; stale comments fixed. Co-authored-by: Cursor <cursoragent@cursor.com>

wanzhenchn · 2026-11-04T09:31:25Z

Updated in 9e9f07e after another review pass over the PR:

The state tier under PP is Kimi-K3 only again. Dropping the PP refusal from state_tier_capability had also turned the engine-side tier on for Qwen's GDN layouts under PP. Their worker still turns the tier off there, because they keep fork checkpoints under PP and publish no layout id. The engine then floored demand and anchor rungs to the LMCache chunk and installed an index that nothing feeds. The capability now refuses PP for every layout except kimi_k3. test_a_qwen_layout_under_pipeline_parallelism_hosts_no_state_tier fails without the fix.
The state index is bounded by the largest stage image. Each stage's tier evicts by its own image size, but the head sized its index by its own stage's image only. If a later stage's image were larger, that stage would evict entries the head still offered for loading. The agreed spec now carries the largest image (pipeline_image_bytes), and the index divides by it. test_under_pipeline_parallelism_the_index_counts_the_largest_stage_image fails without the fix.
MLA shard staging is held back from the KV cache budget. The ATOM_PD_MLA_STAGING_CHUNK_PAGES slots were allocated after the KV cache was sized, on top of --gpu-memory-utilization: 144 MiB per producer GPU at the CI cells' 128 pages and 16 send workers. They are now reserved (mla_shard_staging_reserve_bytes) and fitted into that reserve. Their slots come from the connector's SlotPool, like the other staging pools, and a gather waits on its slot's stream directly instead of on a new event per chunk.
Smaller fixes:
- The hash a checkpoint is filed under at schedule time now comes from the existing _chain_to.
- The unused _discard_kv_cache_ready_event became the one helper that drops a completed prefill together with its ready event.
- ModelRunner and the NUMA binding share one local_device_index, so they cannot pick different GPUs.
- Both engine cores share the idle state publish.
- Connector fields that are always set are read directly instead of through getattr.
- Stale comments are fixed.

GSM8K, which the review asked for before merge, both on the PP8×TP1 → TP8×DCP8 DSpark4 agentic cells with 128 staging pages per slot:

Cell	GSM8K (flexible-extract)	Run
C88, `ATOM_PP_MAX_INFLIGHT=16`	0.9591 ± 0.0055	41165801494
C72	0.9560 ± 0.0056	41268221658

Both runs used this PR at bd26837, on #2787's verification branch. Until this update, the PR's own code changed after that only in merge resolutions. This update's changes have not been through GSM8K yet; #2787's cells will run them.

Pre Checkin passed on 9e9f07e with 11,210 passed and 1,161 skipped (41805335038), 66 more passed than on #2365's head b62bf7a, both merged with the same main. The description is updated to match.

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

Brings #2365's merge of main ab4ada4, which adds Kimi-K3.1 day-0 support (#3003) and resolves its overlap with #2365 in models_accuracy.json and kimi_k3_dspark.py. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

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

Brings #2365's merge of main 058c219, which adds Chat token logprobs (#2856) and resolves its overlap with #2365 in tests/test_scheduler_mtp_max_tokens.py. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

* [feature] Support pp>1 for Kimi-K3, with DSpark speculative decoding Kimi-K3 under pipeline parallelism (verified as pp2 x tp4), including DSpark speculative decoding and prefix caching of the KDA state. PP still needs --enforce-eager: decode has no CUDA graphs yet. Model and transport - Carry K3's block_residual across the stage boundary with the hidden states, and load the DSpark draft's own embedding on the last stage, where the target's embedding is a PPMissingLayer. - Relay DSpark's aux hidden states across stages (PPAuxRelay). The drafter lives on the last stage, so a hook on a tap owned by an earlier stage lands on a PPMissingLayer and never fires. Each stage captures its own taps and forwards them, concatenated in layer_ids order, in the intermediate-tensor payload. The embedding tap (-1) is resolved as the drafter resolves it: the embedding module, identity extract. Speculative decoding under PP - Let a PP engine verify its own drafts. PP output is undeferred, so: the sampler reports this step's rejection and bonus counts on the non-deferred path; they are published to every stage's attention backends, so non-sampling stages roll back linear-attention state; postprocess keeps ONE run of exactly mtp_k + 1 placeholders after every step, whatever was rejected. - Rebuild the verify window from that run. All four backends read `verify_context_shift()`: `num_rejected` under deferred output, 1 under undeferred. The arithmetic lives in token_layout.decode so it is testable without aiter. The window's positions and KV slots are anchored where the scheduler staged them. - Gate and size the verify window on the run's speculative config, not on whether this stage drafts: only the last stage has a drafter, but every stage verifies. Per-stage buffers are sized the same way. - A P/D producer still never speculates, under PP or not. - Round MLA's query-head padding to a width aiter dispatches (24 -> 32 at TP4; TP8 already pads 12 -> 16). KDA state and prefix caching under PP - Progress advances at schedule time, so a chunk ending on a checkpoint rung would let the next chunk overwrite the state slot before the checkpoint is filed. Hold the next chunk (both prefill scans skip it) until the batch that carries the rung files its hashes; the hold is released with that batch even if the seq was preempted, aborted or retired meanwhile. - The checkpoint is filed after the prompt, so its state fork lands on the first decode. Honour it there: publish the source slot and move conv and recurrent state onto the write slot before the kernels run. With DSpark that decode is a verify step, so the spec path does the same. Only a backend whose model does that copy may take a fork on the spec path (`GDNStateMixin.spec_honours_state_fork`, set by K3); every other backend keeps the assert. K3's mono plan refuses a forked step, since a mono step replaces the branch that does the copy. CI - Add a Kimi-K3 pp2 x tp4 DSpark accuracy case with prefix caching and ptpc_fp8 online quant on, as the TP8 Kimi-K3 case runs it, the only accuracy job that reaches any of the paths above. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: move the verify-window shift into the runner, clarify spec flags - tokenIDProcessor.verify_ctx_lens(context_lens) replaces token_layout.decode.verify_context_shift: the attention builders now take their verify-window context lengths straight from the runner. - Rename the runner's `spec_enabled` -> `is_spec_run` (speculative decoding configured, every PP stage runs verify windows) and `use_spec` -> `has_drafter` (this stage owns the drafter and rejection sampler; under PP the last stage only). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test: skip test_verify_ctx_lens without AITER, like the other model_runner tests The non-GPU runner has no aiter, so importing model_runner aborted collection. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: derive the PP aux relay layout from ascending tap order and check it Every stage now orders its columns by sorted(layer_ids), and a stage's incoming taps are those below its first layer, so the layout no longer depends on the order the drafter lists its taps in. pack/absorb reject a payload whose width does not match the taps they expect. Round-trip test over pp 2/4/8 with unsorted taps and stages that own none. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: release a batch's checkpoint holds before any per-seq work A filing that raised mid-batch (logged and swallowed by the PP head loop) left the holds of every later seq in the batch set, keeping them out of the prefill scan forever. Drop the whole batch's holds up front in both register_prefill_hashes and postprocess. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: make the K3 DSpark skip prefixes a tuple the PP stage rebinds The non-first PP stage overrode skip_weight_prefixes on the instance while the class declared it a ClassVar list. It already rebound rather than mutated, but the annotation said the opposite and a list default invites in-place edits. Declare an immutable tuple default instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: fill the PP verify draft columns in one vectorized scatter The PP decode path staged each row's drafts with a per-row Python loop. Build the destination indices with a mask over the draft columns and write them in one assign, as the scheduler's draft marshal does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: read undeferred spec counts back in one D2H Both counts are rows of the rejection sampler's int32 verdict, already contiguous as spec_counts, so copy that once instead of syncing per row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: pack the PP aux relay into one preallocated send buffer pack() concatenated the received block and each own capture into a new tensor every step. Preallocate the payload instead: own taps' capture buffers are column views of it and received rows are copied into its leading columns, so packing allocates nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: size the PP aux relay send buffer to the forward's rows The preallocated send payload was max_num_batched_tokens rows wide and resident on every non-last stage (~470 MB per rank at pp2). Allocate it per forward instead: the first own tap's hook allocates the rows it captures, every own tap writes its columns, and pack() copies the received taps into the leading ones and hands the buffer to the send path, which clones it. A stage that owns no tap forwards the received rows as-is. PP stages run the target eagerly, so no recorded graph bakes the per-step address. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: drop the stale read-side reason from _checkpoint_room This PR adds the spec path's read-side state index, so the token count is the only reason a forking class skips checkpoints on spec decode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: flush a held chunk's hashes on the PP ack path too Under dynamic chunking a middle chunk is acked, and that branch only queued its prefill hashes. A seq held for the chunk's checkpoint is out of the prefill scan, so nothing flushed the queue and it stalled. Both no-output branches now go through one helper that queues and flushes when the batch holds a seq. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: let any PP aux relay tap allocate, and check every own tap fired Only the stage's first own tap allocated the send buffer, so a forward where it captured nothing left the later taps writing into the previous forward's buffer (or into None). Whichever tap fires first now allocates for the forward's rows, and pack() asserts every own tap was captured. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: build the PP aux relay in spec_decode, not in the runner ModelRunner reached into the drafter's private _aux_capture_spec and dspark_proposer's build_aux_capture_spec to assemble the relay. Move that into PPAuxRelay.build so the runner makes one call and keeps only the pp_aux_relay handle it uses at absorb/pack. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: size a P/D first decode's neighbours from the shifted context lens verify_ctx_lens returns a new array, so the is_first_decode_without_local_prefill branch read batch.context_lens unshifted: a row with num_rejected > 0 in the same batch got r extra pages in kv_indptr while block_tables and positions used the shifted lengths. Read the local context_lens there, as the other two branches do. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: read a forked state through the kernels' read-side index Under pp>1 a state fork can reach decode and spec verify. Instead of the model copying the source state onto the write slot, the builder now ships `spec_state_indices_in_tensor` (the write table with the fork source in the resume column), `fused_sigmoid_gating_delta_rule_update` takes `ssm_state_indices_in` like fused_recurrent already does, and the spec conv update gets `conv_state_indices_in`. This drops `spec_honours_state_fork`, its assert, and both copies in KimiKDAAttention, and covers Qwen GDN too. ReplaySSM kernels have no read index and fold records onto the write slot's checkpoint, so its wrappers seed that checkpoint from the source (`seed_forked_slots`) and the builder parks the forked write slot's cursor at -1. Before, the cursor was left at the previous tenant's value and its records were folded into the forked state. The mono K1 refusal stays: K1 reads `st_idx` only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: refuse unwired PP combinations at config time Under pipeline_parallel_size > 1, only Kimi-K3 with DSpark relays the target's aux taps across stages. Every other (model, spec method) pair now fails at config time instead of mid-serving: EAGLE3's relay zip, a Qwen GDN prompt-end fork, MiniMax-M3's embed on a PPMissingLayer. The DSpark confidence scheduler and a P/D kv_consumer, which the PP loop has no remote-KV path for, are refused as well. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: plan the vLLM plugin's dense MLA work at the layer's padded width `MLAAttention.padded_num_heads` now rounds up to a width aiter dispatches, and the plugin's MLA layer inherits it, so K3 at TP4 pads q from 24 to 32 heads. The plugin builder still planned the work descriptors for 24. Round it the same way, as aiter_mla.py's builder does. The sparse builder and layer keep max(heads, 16): they pad by repeat_interleave, which only divides evenly at that width, and they already agree with each other. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: drop the PP checkpoint hold's dead prompt-end case PP refuses PAGE-backed checkpoints, so every class that applies forks and keeps a positive successor_room: a rung (or the prompt-end anchor) sits strictly inside the prompt, and `start + chunk < num_prompt_tokens` always holds. Filing checkpoints by rung hash at schedule time, which would remove the hold altogether, is left as a follow-up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: install the PP aux relay's taps through the drafter's own helper `PPAuxRelay.arm` repeated `Drafter.arm_aux_capture`'s hook mechanism and skipped its checks. Both now go through drafter.py: - `install_aux_taps` validates the spec's ids (each in {-1} U [0, n), strictly increasing) and hooks a subset of them. The last stage hooks the drafter's buffers, other stages the relay's send buffer. - `_aux_hook` skips draft forwards and hands its writer `ctx.ubatch_token_offset`. The relay writes at that offset and grows its buffer for a later ubatch instead of reallocating over the first. - `Drafter.aux_capture_spec(config)` is each drafter class's config-only spec. A drafterless stage reads it off `factory.drafter_class(config)`, so the relay no longer imports DSpark's builder. `PP_AUX_KEY` is now the model-neutral "pp_aux". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: send the PP aux relay's payload without a second clone The relay allocates a fresh buffer per step and hands it to the send path, so cloning it again only doubles the copy. async_send_intermediate_tensors takes the keys the caller hands over and sends those as they are. The received aux block is now a run_model local instead of runner state, and PPAuxRelay is imported at the top of the runner. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: one use_spec on every stage, runner and scheduler alike Every PP stage runs the verify window, so the runner's is_spec_run / has_drafter pair collapses back to the single use_spec it had, now set on every stage (has_drafter's readers sit on the deferred path, which PP never takes). The scheduler's spec_decode_local was the same flag under a second name; its readers take use_spec directly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: one source each for the deferred-output mode and the verify width The scheduler's is_deferred_out (pp_size == 1) and tokenIDProcessor's (not PP, not a recurrent-state P/D producer) agreed only because a producer never speculates. Both now come from defers_output(config): the runner passes the attention family it already resolved, the scheduler reads it off the config. GDN sizes its verify window from model_runner.num_spec_tokens, which every PP stage carries, instead of the drafter's mtp_k with a num_speculative_tokens fallback; config already refuses a speculative config without num_speculative_tokens, so the fallback's ValueError was unreachable. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: top-level imports, and document PP + DSpark for Kimi-K3 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: keep the DSpark draft out of online quant in the PP2 accuracy case Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: keep the PP aux relay's hook free of in-place set mutation under Dynamo The stage-0 relay hook runs inside the compiled backbone. `_filled.add` mutated a set Dynamo did not create, so Dynamo replayed it with clear()+update(); with cudagraphs on, an `update` in the rewritten bytecode triggers the decorator's depyf import, which crashed PP2 warmup on images without depyf. Rebind a frozenset instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: read undeferred spec counts off the sampled ids, without a D2H The undeferred step already brings the sampled ids to the host. Each row keeps its accepted drafts plus one token before the -1 tail, so both counts are its length; drop the verdict's own .cpu().numpy(). The ragged sampler test pins that layout against the verdict. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: stage the forked write slots through pinned memory `torch.tensor(list, device=cuda)` is a pageable H2D, which PyTorch follows with a stream sync. Pin the list and copy it non_blocking, as the other host-staged indices do. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: the relay's docstring follows the spec's layer_ids order Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: pass the runner's attn_family to defers_output again The merge dropped it, so defers_output re-derived the family from config.hf_config and test_spec_deferred_status_queue's producer cases failed. The runner stub in test_h2d_runner_publication gains attn_family. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: gate PP placeholders on the batch's frozen chunk state Under schedule-time advancement a middle chunk's live is_partial_prefill flips once its final chunk is scheduled, so the middle chunk's postprocess appended mtp_k + 1 placeholders to a seq still in prefill. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: release a checkpoint hold only from the chunk that holds it Under PP an earlier chunk of the same seq can still be in flight when the chunk ending on a rung is scheduled; its completion released the hold, so the next chunk could run before the rung's checkpoint was filed. The hold now records the held chunk's end, and only a batch reaching it releases. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: let a P/D producer under PP keep any drafter A producer never speculates, so the Kimi-K3 + DSpark gate refused GLM-5.2's PP prefill + MTP Atomesh cells for nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: read a forked spec row's source in every column The kernel resumes from column num_accepted - 1, so a forked verify with a bonus read a recycled write slot; only column 0 held the source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: keep the spec-path fork assert unless the backend opts in GDNStateMixin.spec_honours_state_fork() defaults to False; only Kimi-K3's KDA forward (kimi_linear) reads spec_state_indices_in_tensor and opts in, so glm5_next, Qwen3-Next/3.5 and Qwen4Exp keep the assert. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: refuse TBO under PP and drop the relay's TBO grow branch UBatchWrapper.forward takes no intermediate_tensors, so TBO never ran under PP; refuse it in config instead of keeping a grow path nothing reaches. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: seed ReplaySSM forks from host-known slot pairs, without a sync seed_forked_slots indexed with a GPU bool mask, a host sync in every linear-attention layer of a fork step. The builder already knows the (write, source) pairs: ship them once per step in the pinned copy that parks write_pos, and index with them directly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: write relay captures without bookkeeping Dynamo freezes The lowest own tap opens the forward's send buffer and every own tap writes its columns; the reset and the _filled check read state that the compiled forward bakes to the warmup trace, so they guarded nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: let the registry declare which spec methods an arch serves under PP pp_speculative_methods sits beside mono_decode_installers, so the next model wired to the relay adds a row there instead of editing Config. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: name the checkpoint hold by what it does _pp_checkpoint_hold -> _checkpoint_holds (keyed on advance_on_schedule), holds_checkpoint_hostage -> carries_held_checkpoint, and _defer_prefill_hashes -> _queue_or_flush_prefill_hashes, since it flushes the whole queue; mark the hold as scaffolding #2709 replaces. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: stage the undeferred window from the recorded placeholder width Read seq.num_placeholder_tokens, which postprocess refills to one full window, instead of re-deriving the run's width from the scheduled count. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: let the drafter read the relayed aux block in place Drafter.take_relayed_aux keeps the received [N, k*H] block and aux_for serves it ahead of the own-tap buffers, so the last stage no longer splits it per tap for compute_draft_kv to concatenate back (DSpark skips the cat when it is the only block). The relay takes the drafter class's spec on every stage instead of the drafter's private _aux_capture_spec, and the runner no longer hands it _aux_buffers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: derive DSpark's instance aux spec from its classmethod _aux_capture_spec now parses dspark_target_layer_ids once, through aux_capture_spec, and drops its unused target_model parameter. PPAuxRelay keeps only spec (no layer_ids / hidden_size copies), drops the dead _armed guard, and raises the missing-payload error from one place. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: drop the CLAUDE.md key list and getattrs on always-set fields Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: allocate aux buffers only for this stage's own taps Taps owned by an earlier PP stage arrive through the relay, so the last stage gives them a 0-row buffer and no hook (~470 MB/rank for K3 at pp2). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: share one tap residency rule between the drafter and the relay Both now ask whether the tap's module is a PPMissingLayer, the embedding included, instead of the relay keying the -1 tap on is_first_rank. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: ganyi <ygan@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: wanzhenchn <wanzhenchn@gmail.com>

#2365 has landed on main as the squash 64412f1, so this PR now targets main. 64412f1's tree is #2365's head 9fb45ff, which this branch already merges, merged with main d8525dd (#2913). The only new change this merge brings is #2913. scheduler.py and tests/test_pp.py conflicted only because the squash repeats #2365's changes that this branch edits further. They keep this branch's versions: the result is this branch merged with d8525dd, which merges cleanly, and its diff against main is this PR's diff against 9fb45ff unchanged. Co-authored-by: Cursor <cursoragent@cursor.com>

This was referenced Oct 26, 2026

ci(benchmark): run AgentX benchmarks with agentx-harness 1.0.9 and ship it in the image #2576

Merged

[feature] support pp>1 for Kimi-K3 #2365

Merged

wanzhenchn force-pushed the zwan/k3-pp branch from 7832617 to 7e1bf6c Compare October 26, 2026 05:53

wanzhenchn changed the base branch from main to guanbao/k3-pp2tp4 October 26, 2026 05:53

wanzhenchn added this pull request to stack #2711 October 26, 2026 05:53

wanzhenchn changed the title ~~feat(kimi-k3): PP8-prefill P/D serving for AgentX 1M, with one CI cell per concurrency~~ feat(kimi-k3): PP8-prefill P/D serving for AgentX 1M, with one CI cell per concurrency Oct 26, 2026

wanzhenchn requested review from Jasen2201, Yuechguo, gbyu-amd, valarLip and yhl-amd October 26, 2026 05:57

wanzhenchn added the ci:atom label Oct 26, 2026

Yuechguo previously approved these changes Oct 26, 2026

View reviewed changes

wanzhenchn dismissed Yuechguo’s stale review via d5e502a October 27, 2026 05:46

wanzhenchn force-pushed the zwan/k3-pp branch 2 times, most recently from d5e502a to 1c6fbeb Compare October 27, 2026 05:52

wanzhenchn force-pushed the zwan/k3-pp branch 2 times, most recently from 3dae59f to b5acae7 Compare October 28, 2026 10:30

This was referenced Oct 28, 2026

perf(pd): P/D first-token latency and API-server GC fixes (from Kimi-K3 AgentX 1M) #2786

Open

ci(k3-agentic): Kimi-K3 AgentX 1M P/D cells, one per concurrency on its best measured prefill #2787

Open

wanzhenchn changed the title ~~feat(kimi-k3): PP8-prefill P/D serving for AgentX 1M, with one CI cell per concurrency~~ feat(kimi-k3): PP8 prefill for P/D serving — KDA state under PP, PP→DCP transfer, MLA staging Oct 28, 2026

ganyi1996ppo and others added 2 commits October 30, 2026 02:19

ganyi1996ppo and others added 14 commits November 3, 2026 08:11

review: let a P/D producer under PP keep any drafter

1a570a2

A producer never speculates, so the Kimi-K3 + DSpark gate refused GLM-5.2's PP prefill + MTP Atomesh cells for nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

review: read a forked spec row's source in every column

d9f26d5

The kernel resumes from column num_accepted - 1, so a forked verify with a bonus read a recycled write slot; only column 0 held the source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

review: refuse TBO under PP and drop the relay's TBO grow branch

7e45c5b

UBatchWrapper.forward takes no intermediate_tensors, so TBO never ran under PP; refuse it in config instead of keeping a grow path nothing reaches. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

review: let the registry declare which spec methods an arch serves un…

cd0d315

…der PP pp_speculative_methods sits beside mono_decode_installers, so the next model wired to the relay adds a row there instead of editing Config. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

review: stage the undeferred window from the recorded placeholder width

3a2258b

Read seq.num_placeholder_tokens, which postprocess refills to one full window, instead of re-deriving the run's width from the scheduled count. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

review: drop the CLAUDE.md key list and getattrs on always-set fields

30829a3

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

Merge main (a8946ff) into guanbao/k3-pp2tp4

337bcd5

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

valar-cop reviewed Nov 3, 2026

View reviewed changes

ganyi1996ppo and others added 6 commits November 3, 2026 20:38

review: allocate aux buffers only for this stage's own taps

a1436cd

Taps owned by an earlier PP stage arrive through the relay, so the last stage gives them a 0-row buffer and no hook (~470 MB/rank for K3 at pp2). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

review: share one tap residency rule between the drafter and the relay

b62bf7a

Both now ask whether the tap's module is a PPMissingLayer, the embedding included, instead of the relay keying the -1 tap on is_first_rank. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge guanbao/k3-pp2tp4 (a1436cd) into zwan/k3-pp

8c69b26

Brings #2365's review commit a1436cd, which allocates the drafter's aux buffers only for this stage's own taps. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

Merge guanbao/k3-pp2tp4 (b62bf7a) into zwan/k3-pp

99d1ff3

Brings #2365's review commit b62bf7a, which shares one tap residency rule between the drafter and the PP aux relay. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

ganyi1996ppo and others added 4 commits November 4, 2026 20:52

Merge main (ab4ada4) into guanbao/k3-pp2tp4

c10069c

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

Merge guanbao/k3-pp2tp4 (c10069c) into zwan/k3-pp

d718458

Brings #2365's merge of main ab4ada4, which adds Kimi-K3.1 day-0 support (#3003) and resolves its overlap with #2365 in models_accuracy.json and kimi_k3_dspark.py. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

Merge main (058c219) into guanbao/k3-pp2tp4

9fb45ff

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

Merge guanbao/k3-pp2tp4 (9fb45ff) into zwan/k3-pp

b82e2d0

Brings #2365's merge of main 058c219, which adds Chat token logprobs (#2856) and resolves its overlap with #2365 in tests/test_scheduler_mtp_max_tokens.py. No conflicts. Co-authored-by: Cursor <cursoragent@cursor.com>

Base automatically changed from guanbao/k3-pp2tp4 to main November 5, 2026 12:06

Navigation Menu

Repository navigation

feat(kimi-k3): PP8 prefill for P/D serving — KDA state under PP, PP→DCP transfer, MLA staging - #2709

feat(kimi-k3): PP8 prefill for P/D serving — KDA state under PP, PP→DCP transfer, MLA staging#2709
wanzhenchn wants to merge 90 commits into
mainfrom
zwan/k3-pp

wanzhenchn commented Oct 26, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 26, 2026

Uh oh!

Yuechguo left a comment

Uh oh!

valar-cop left a comment

Uh oh!

wanzhenchn commented Nov 3, 2026

Uh oh!

wanzhenchn commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

wanzhenchn commented Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Stack and split

Changes

Feature support (each was a startup refusal or a failed first transfer)

Correctness under PP

Performance

Tests

Uh oh!

github-actions Bot commented Oct 26, 2026

🏷️ CI Guide

Uh oh!

Yuechguo left a comment

Choose a reason for hiding this comment

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

wanzhenchn commented Nov 3, 2026

Uh oh!

wanzhenchn commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wanzhenchn commented Oct 26, 2026 •

edited

Loading