Skip to content

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

Merged
valarLip merged 52 commits into
mainfrom
guanbao/k3-pp2tp4
Nov 5, 2026
Merged

valarLip merged 52 commits into
mainfrom
guanbao/k3-pp2tp4

Conversation

@gbyu-amd

@gbyu-amd gbyu-amd commented Sep 23, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

Kimi-K3 only ran at TP8. This PR makes it serve at --pipeline-parallel-size 2 --tensor-parallel-size 4 on MI355X ×8, covering plain decode, prefix caching (KDA state checkpoints), LMCache offload and DSpark speculative decoding. None of these paths had been exercised under PP for this model.

The first commit's message lists every original change; the review: commits on top each address one review point.

Technical Details

Model and transport

  • K3's block_residual (AttnRes) now crosses the stage boundary with the hidden states. It used to be dropped there, and stage 1 raised KeyError during warmup.
  • The DSpark draft loads its own checkpoint embedding on the last stage (587 MiB/rank at TP4), where the target's embed_tokens is a PPMissingLayer.
  • PPAuxRelay (atom/spec_decode/pp_aux_relay.py) relays DSpark's aux hidden states across stages. K3 taps layers 2 and 23, which live on stage 0, so without it the drafter read zero buffers. The embedding tap (-1) resolves as Drafter.arm_aux_capture resolves it.
  • MLA's query-head padding rounds up to a width aiter dispatches (24 → 32 at TP4). TP8 already padded 12 → 16, and head counts that already dispatch are unchanged.

Speculative decoding under PP

  • A PP engine now verifies its own drafts. PP output is undeferred, so:

    • the sampler reports this step's rejection and bonus counts on the non-deferred path;
    • every stage publishes those counts to its attention backends, so the stages that don't sample still roll back KDA state;
    • postprocess leaves exactly one run of mtp_k + 1 placeholders after every step, whatever was rejected.
  • The attention backends rebuild the verify window from tokenIDProcessor.verify_ctx_lens:

    • deferred output shifts by num_rejected (unchanged);
    • undeferred output shifts by 1.

    The scheduler stages the undeferred window from the recorded seq.num_placeholder_tokens.

  • The verify window is gated and sized on the run's speculative config, not on whether this stage drafts. Stage 0 used to allocate 3 conv rows per KDA slot while the kernel addressed 10.

  • A P/D producer still never speculates, with or without PP. Under PP a producer keeps any drafter; the other architectures and spec methods allowed under PP are declared in registry.pp_speculative_methods.

  • TBO is refused under PP (UBatchWrapper.forward takes no intermediate_tensors).

  • On the last stage the drafter reads the relayed aux block in place (Drafter.take_relayed_aux), with no per-tap split and concatenation.

KDA state and prefix caching under PP

  • Checkpoint hold. A chunk that ends on a checkpoint rung holds the next chunk until the rung's checkpoint is filed, because progress advances at schedule time. Otherwise the archived state ran ahead of its hash (filed as 768 tokens while holding 884).
  • State fork on decode. The checkpoint is filed after the prompt, so its state fork lands on the first decode. That decode now publishes the source slot and moves conv and recurrent state onto the write slot before the kernels run.
    • With DSpark, that first decode is a verify step, so the spec path does the same; every column of a forked row's spec read table holds the source slot. ReplaySSM seeds forked slots from host-known (write, source) pairs, with no host sync.
    • Only a backend whose model does that copy may take a fork on the spec path (GDNStateMixin.spec_honours_state_fork, which K3 sets). Every other GDN backend keeps the assert.

CI

  • A new accuracy case, Kimi-K3 DSpark PP2, runs pp2 × tp4 with DSpark, prefix caching and ptpc_fp8 online quant on. It is the only accuracy job that reaches the PP paths.

Known limits (follow-ups)

  • PP still requires --enforce-eager: PP decode has no CUDA graphs yet. This is the existing forward_vars ring assert, not something this PR adds. CUDA-graph support for PP decode is follow-up work.
  • The checkpoint hold costs prefill latency: in K3 PP8 prefill it accounts for about 3.3 s of a roughly 7 s fixed latency. feat(kimi-k3): PP8 prefill for P/D serving — KDA state under PP, PP→DCP transfer, MLA staging #2709, stacked on this PR, files a PAGE-backed checkpoint as soon as its chunk is scheduled and keeps the hold only for fork-class requests.
  • Aux-relay traffic grows with PP depth. At PP8 the 6 → 7 boundary carries 4 taps. Relaying a running context_proj partial sum would keep it at one fp32 column.
  • LMCache offload under PP is validated on the store path only: GSM8K has no repeated prefix, so loads from the CPU tier are not exercised yet. The KDA state tier stays off under PP (state_offload.py, CacheEngineKey has no PP component).

Dependency

Test Plan

Unit tests

  • Command: python -m pytest <file> for every file under tests/, one process per file. Compared with main at the same commit, on the same host.
  • New tests:
    • test_verify_ctx_lens.py: the verify window's positions, deferred and undeferred. It imports model_runner (where tokenIDProcessor.verify_ctx_lens lives, per @valarLip's review), so it needs aiter and skips on the non-GPU runner.
    • test_pp_draft_input_ids.py: the vectorized PP verify-draft scatter
    • test_pp_aux_relay.py: tap split at pp = 2/4/8, the -1 tap, the round trip through every stage, the send buffer sized per forward, and the last stage arming only its own taps
    • test_gdn_state_fork.py: a fork on the decode and spec paths
    • undeferred placeholder algebra in test_scheduler_mtp_max_tokens.py
    • checkpoint-hold lifecycle in test_pp.py, including the flush on the ack path
    • K3 mono refuses a forked step, spec and plain, in models/test_kimi_k3_mono_plan.py

Accuracy

  • GSM8K, full 1319 questions, lm_eval as in recipes/Kimi-K3.md, num_concurrent=32.
  • Server flags are the new CI case's: -tp 4 --pipeline-parallel-size 2 --enforce-eager --block-size 128 --enable_prefix_caching --method dspark --num-speculative-tokens 7 --max-num-seqs 64, plus the recipe's ptpc_fp8 online quant.
  • Image: rocm/atom-dev:nightly_202610251529, with its own aiter.

Test Result

Unit tests

  • At 5ab3d68f7 (this branch merged with main 99ebf633a), every test file was run, one process per file, against main itself in the same container. No file does worse. At head b62bf7ac6, every file the PR affects was re-run against the merge base a8946ff73: no regressions (the one failure, test_model_runner_exit::test_a_failed_megamoe_close_still_releases_everything, fails at the merge base too). Every file that differs does so by gaining tests:
    • new files: test_verify_ctx_lens (26), test_pp_aux_relay (24), test_pp_draft_input_ids (1), test_gdn_state_fork (5);
    • added tests: test_pp +7, test_scheduler_mtp_max_tokens +6, models/test_kimi_k3_mono_plan +2.
  • test_verify_ctx_lens, test_pp_aux_relay and test_gdn_state_fork import aiter, so the non-GPU pre-checkin runner skips them; they pass in the rocm/atom-dev nightly image.

GSM8K, pp2 × tp4, DSpark, prefix caching, ptpc_fp8, 1319 questions

Run Build flexible / strict DSpark acceptance
this PR f34cb8a8e, nightly_202610310756 0.9704 / 0.9704 (±0.0047) 53.7%, 4.76 tok/fwd
this PR this commit on main 713eb0614 (70dc6d37c), nightly_202610251529 0.9553 / 0.9560 (±0.0057) 53.4%, 4.74 tok/fwd
main CI, TP8, Kimi-K3 DSpark 2026-10-24 / 2026-10-26 0.9568 / 0.9545 (flexible) —

pp2 × tp4 scores the same as main's TP8 DSpark case from the same week, and its acceptance matches TP8 (53.7–55.2% on earlier same-host runs). Earlier numbers for this PR (0.9666 on nightly_202609220547) were measured on an older image, so they are not comparable.

GSM8K, pp2 × tp4, no DSpark, prefix caching + LMCache offload, ptpc_fp8, 1319 questions

flexible 0.9674 ± 0.0049 / strict 0.9682 ± 0.0048 (e1c80ceae, nightly_202610310756)

Submission Checklist

@gbyu-amd
gbyu-amd marked this pull request as draft September 23, 2026 02:35
@github-actions

Copy link
Copy Markdown
Contributor

🏷️ 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 2365 --add-label <label>

@wanzhenchn

Copy link
Copy Markdown
Contributor

Update: I went ahead and stacked the three PRs (main ← #2576 ← #2365 ← #2709). guanbao/k3-pp2tp4 is now at 501364cf4, the 15 commits rebased onto #2576 and current main, and this PR's base is Jasen/atomesh-agentx-harness. If you have local work on the old head 5c811df1d, git fetch && git rebase --onto origin/guanbao/k3-pp2tp4 5c811df1d moves it over. If you'd rather keep this PR out of the stack, say so and I'll restore 5c811df1d and unstack it.

@wanzhenchn wanzhenchn left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review summary

Thanks for this PR. It adds pp>1 support for Kimi-K3, including DSpark speculative decoding under PP, the aux hidden-state relay, and KDA state and checkpoint handling. The approach looks right and the main fixes look sound.

  • CI is green: black, ruff, and all 20 accuracy jobs.
  • Unit tests pass locally at 501364cf4: 285 passed, 60 skipped.

Before merge, I'd like to see four things addressed:

  • one latent bug,
  • one assert path that needs confirmation,
  • test coverage,
  • the commit history and the PR description.

The performance items at the end are fine as follow-ups.

1. Correctness

1.1 [Must fix] PPAuxRelay mishandles the embedding tap (-1)

# The embedding tap (-1) is resident on the first stage.
if -1 in spec.layer_ids and pp.is_first_rank:
resident.insert(0, -1)

for lid in self.own_ids:
module = layers[lid]
hook = self._make_hook(
self.buffers[lid], self.spec.extract, is_draft_forward
)

  • What's wrong:
    • On the first stage, -1 is added to own_ids. arm() then hooks layers[-1], which is the last decoder layer and a PPMissingLayer on that stage.
    • It also uses spec.extract, whereas the base drafter resolves -1 with _resolve_embedding and _identity_extract.
    • The hook never fires, so downstream stages receive an all-zero buffer and nothing reports an error.
  • Impact: K3's taps are [2, 23, 47, 71, 89], so this doesn't trigger today. But the relay is the generic DSpark path.
  • Ask: Please mirror the base drafter's handling of -1, or reject -1 explicitly under PP.

1.2 [Please confirm] A fork landing on a spec-verify decode step hits an assert

assert not (with_spec and any(s >= 0 for s in fork_srcs)), (
"state fork on the spec-decode path: spec_state_indices_tensor has "
"no read-side counterpart (BlockManager only forks onto prefill)"
)

  • Concern:
    • 838ac51 notes that under PP the checkpoint can land after the prompt, so a fork can reach decode. It only handles the non-spec KDA decode branch.
    • With DSpark enabled, the first decode after the prompt is a spec-verify step. If that step carries a FORK checkpoint, this assert fires.
  • Ask: Can this ordering happen under PP + DSpark + prefix caching? If it can, the spec branch needs handling. If it can't, please add a comment or a test that explains why.

1.3 [Minor] The PP path drops the producer's num_spec_step == 0 guard

  • The description says the P/D producer keeps this assert, but the code only checks it on the non-PP path (model_runner.py L558-L567).
  • It's safe today, because the producer's scheduler sets use_spec = has_spec and not is_pd_producer.
  • Please either keep the assert for producers on the PP path or adjust the description.

2. Tests

2.1 [Must fix] The new test doesn't run in CI

tests/test_verify_context_shift.py imports model_runner, which needs AITER, so pytest.importorskip skips it on the non-GPU runner. Consider moving verify_context_shift into a pure helper (e.g. in token_layout/decode.py) so it can be tested on CPU.

2.2 [Suggested] The new logic has no unit coverage

The existing is_deferred_out=False cases are all non-spec. CPU tests worth adding:

  • Placeholder algebra for undeferred spec output: rejected counts from 0 to k, steps with no drafts, and the max_tokens boundary.
  • The _pp_checkpoint_hold lifecycle: a held sequence is skipped by both scan loops and is released on retire, preempt and abort.
  • The PPAuxRelay split into incoming, own and outgoing taps for pp=2/4/8, using fake layers and including -1.
  • has_state_fork on the decode path.

2.3 [Suggested] Add a PP accuracy CI case

Every fix here is PP-only, but all 20 accuracy jobs run TP8 or TP4. A Kimi-K3 pp2tp4 GSM8K job with DSpark and prefix caching would cover it.

3. Commit history and description

3.1 [Suggested] Bisectability

  • 43f5aa7 removes the NotImplementedError and enables spec decode under PP, but the following 10 commits are needed for correct output.
  • So any intermediate commit silently produces garbage with PP + DSpark. The commit messages themselves say "decodes noise" and "still broken".
  • Please move the enabling change to the end of the series, or squash the chain into a few commits.

3.2 [Suggested] Stale references in commit messages

Commit Stale reference
43f5aa7, d7fdb2e, 21034f8 k3_pp2tp4/K3_PP2TP4.md (not in the repo)
4547d54 9ba0d6a6 (pre-rebase SHA of 54b724f)
6fc1fde 0ac00173 (pre-rebase SHA of 21034f8)
924147e 3d7ff8a5 (not found)

Please also:

  • drop the debugging narrative ("not sufficient", "still broken", "decodes noise") from 43f5aa7 and d7fdb2e;
  • fold "Trim the comments this PR adds" back into the commits it touches.

3.3 [Suggested] Update the PR description

  • The GSM8K table predates the rebases, and its last three rows predate the verify-window fixes. Please re-run at the current head.
  • ROCm/aiter#5783 merged on 2026-09-30, so the dependency note is out of date.
  • There is no performance data. A pp2tp4 vs tp8 comparison of throughput and TPOT would help reviewers judge where PP decode is worth using.

4. Performance (fine as follow-ups)

4.1 The checkpoint hold is the largest cost.

  • In our Kimi-K3 PP8 prefill runs, _pp_checkpoint_hold accounts for about 3.3 s of a roughly 7 s fixed prefill latency.
  • In #2709, which is stacked on this PR, b81cf24 and 7d55402 file a PAGE-backed checkpoint as soon as its chunk is scheduled, and keep the hold only for fork-class requests.
  • Either note this as a known cost that #2709 addresses, or pull those two commits into this PR.

4.2 PP decode still requires --enforce-eager, so it has no CUDA graphs. This decides whether PP decode is practical. Please list it as follow-up work in the description.

4.3 Aux relay traffic grows with PP depth.

  • Current cost: At PP8, the stage 6 → 7 boundary carries 4 taps. At hidden size 7168 in bf16, that is about 470 MB per 8192-token chunk. The torch.cat in pack plus the clone before the send copy the payload twice.
  • Alternative: context_proj (weight [7168, 35840]) is linear over the concatenated taps. Each stage could apply its own weight slice and relay a single running sum in fp32.
  • Effect: Each boundary carries one fp32 column regardless of PP depth.
  • Trade-offs: Each stage stores its own weight slice (about 103 MB per tap), and the summation order changes.

4.4 Small items:

  • The per-row Python loop that stages draft columns in prepare_input_ids (model_runner.py L578) can be vectorized.
  • _build_output does two .cpu() syncs (model_runner.py L3438-L3439); stacking the tensors makes it one.
  • The KDA fork state copy (kimi_k3/model.py L1352) copies every decode row; it could copy only the fork rows. Forks are rare, so the gain is small.

5. Checked, no issues

  • The relay's async send is safe: async_send_intermediate_tensors clones before isend, so there is no race.
  • Enabling spec decode under PP doesn't affect P/D producers: see 1.3.
  • K3 tap split at PP8 (12×7+9) is correct: each stage receives and forwards the expected taps.
  • The MLA head padding in e39bcb6 (24 → 32 at TP4) is not a new cost: it is aiter's dispatch constraint, and TP8 already pads 12 → 16.

@ganyi1996ppo

ganyi1996ppo commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

@wanzhenchn Thanks for the review. All points are addressed, except the performance items you marked as follow-ups. The PR is now one commit, d064068f5, on main (rebased after #2576 merged), so the SHAs in your review no longer resolve. The pre-squash history is kept locally (tip 513b13346); I can push it if you want to diff against it.

1. Correctness

1.1 PPAuxRelay and the -1 tap: fixed. arm() now resolves (module, extract) the same way Drafter.arm_aux_capture does: -1 maps to _resolve_embedding + _identity_extract, and any other id to layers[lid] + spec.extract. -1 also stays on the post-hook when the spec captures inputs, as on the base path. tests/test_pp_aux_relay.py covers it. Two of its tests fail with the fix reverted: the embedding is not hooked, and the buffer stays zero.

1.2 A fork on a spec-verify decode: it can happen, and it is now handled. Under PP + DSpark + prefix caching, three things combine:

  • a decode batch carries forks (_consume_state_forks also runs on the decode return path);
  • with spec decode on, every decode takes the spec branch, and nothing exempts the first one;
  • the checkpoint is filed after the prompt, so the fork lands on that first decode.

The fix uses the same approach as the plain decode path:

  • prepare_state_indices publishes the read slot on the spec path as well. prepare_gdn_metadata ships it to the device only when the step carries a fork; every other spec step still ships None.
  • K3's KDA spec branch copies conv state and recurrent state from that slot onto spec_state_indices[:, 0] before the kernels run. Neither spec kernel takes a read-side index.
  • The assert stays for every other GDN backend. Only a builder that sets GDNStateMixin.spec_honours_state_fork may take a fork on the spec path, and only K3's does, so Qwen and GLM keep the guard.
  • K3's mono plan refuses a forked step, because a mono step replaces the branch that does the copy. Mono already refuses PP > 1, so this guard only matters if pp=1 ever puts a fork on a decode.

tests/test_gdn_state_fork.py covers both paths and the refusal. Three of its tests fail against the old prepare_state_indices.

1.3 Producer guard on the PP path: fixed in code. num_spec_step == 0 is now asserted for every P/D producer, and for every non-PP engine as before. The description no longer claims more than the code does.

2. Tests

2.1 The new test now runs in CI. verify_context_shift is a pure function in token_layout/decode.py. tokenIDProcessor.verify_context_shift() forwards to it, so the four backend call sites are unchanged. The test passes with aiter blocked from import.

2.2 Added:

  • Undeferred placeholder run (test_scheduler_mtp_max_tokens.py): exactly one run of mtp_k + 1 placeholders after every reject count from 0 to k, across mixed steps, from a prefill step that scores no drafts, and at a max_tokens cap crossed by an all-accepted step.
  • _pp_checkpoint_hold (test_pp.py): a held sequence is skipped by both prefill scans and is released when its batch files its hashes. It is also released, through either filer, when the sequence left running first (preempted, aborted or retired).
  • PPAuxRelay tap split at pp=2/4/8, including -1, plus has_state_fork on decode (see 1.1 and 1.2).

2.3 Added the CI case Kimi-K3 DSpark PP2: pp2 × tp4, DSpark, prefix caching and ptpc_fp8, with test_level: pr. If maintainers want to spare every PR an 8-GPU eager job, it can move to main.

3. History and description

3.1, 3.2 Squashed into a single commit. This removes the bisect hazard, the stale references, the debugging narrative and the separate "Trim the comments" commit in one step.

3.3 Description updated:

  • GSM8K re-run on nightly_202610251529 with this same commit applied on main 713eb0614 (70dc6d37c), full 1319 questions with the CI case's flags: 0.9553 / 0.9560, DSpark acceptance 53.4%. main's TP8 Kimi-K3 DSpark CI case scored 0.9568 and 0.9545 this week.
  • aiter#5783 is now noted as merged.
  • Performance comparison (pp2 × tp4 vs tp8): not added; that is a follow-up together with 4.x.

4. Performance

Left as follow-ups and listed under "Known limits" in the description: the hold's prefill cost and #2709's fix for it (4.1), --enforce-eager (4.2), and relay traffic (4.3). 4.4 is untouched.

@ganyi1996ppo
ganyi1996ppo force-pushed the guanbao/k3-pp2tp4 branch 3 times, most recently from 65ebed6 to d102bb5 Compare October 27, 2026 03:32
Base automatically changed from Jasen/atomesh-agentx-harness to main October 27, 2026 06:48
wanzhenchn
wanzhenchn previously approved these changes Oct 28, 2026
ganyi1996ppo and others added 3 commits November 3, 2026 09:05
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>
@ganyi1996ppo

Copy link
Copy Markdown
Contributor

@valarLip The review fixes are pushed in 30829a333, and each inline comment has a reply naming its commit.

Run non-GPU unit tests failing at this head (and therefore Check Pre Checkin Signal) is not caused by this PR:

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

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

ci Run non-GPU unit tests (Pre Checkin) failed — main is already broken the same way (high): The attached log is empty (0 bytes), but the same job fails 4x on main in the same window, and the author identifies the abort as collection of tests/model_ops/test_dp_gather_routed_rows.py (added by #2776, untouched here) on a missing triton.language.extra. Not caused by this PR.

same job on main: ['2026-11-03T14:52 failure', '2026-11-03T14:40 failure', '2026-11-03T14:39 failure', '2026-11-03T14:17 failure', ...]
Collection stops at tests/model_ops/test_dp_gather_routed_rows.py with ModuleNotFoundError: No module named 'triton.language.extra', before any test runs
next: Wait for the #2954 stub fix on main, merge main into the branch and re-run; this also unblocks the gated GPU accuracy jobs. https://github.com/ROCm/ATOM/actions/runs/41663095846/job/126011160276

ci Atomesh Accuracy Validation failed — ci script/config bug, not the pr (high): The workflow triggers on .github/benchmark/models_accuracy.json (touched by this PR) but on a pull_request its whitelist selects zero models, so models_json is [] and the matrix include expands to nothing; unlike atom-test.yaml there is no placeholder fallback. Kimi-K3 is not in atomesh_levels, so no atomesh job can exist for this change.

include: ${{ fromJson(needs.load-test-models.outputs.models_json) }}
include: ${{ fromJson(needs.load-test-models.outputs.models_json || '[{"runner":"ubuntu-latest","model_name":"skip"}]') }}
next: Add atom-test's placeholder-entry fallback to the atomesh matrix (or skip the workflow on catalog-only PRs), then rerun. https://github.com/ROCm/ATOM/actions/runs/41663095656

wanzhenchn added a commit that referenced this pull request Nov 3, 2026
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>

@ganyi1996ppo ganyi1996ppo left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the round of fixes. I re-read at 337bcd5 (merged with main a8946ff). All 15 earlier points are addressed. The affected unit tests show no regressions against the merge base: the one failure, test_model_runner_exit.py::test_a_failed_megamoe_close_still_releases_everything, fails at the merge base too.

Worth addressing

  1. The last stage still allocates an aux capture buffer for every tap, including the ones that now arrive through the relay (inline on atom/spec_decode/drafter.py:455).

spec = self._aux_capture_spec()
if spec is None:
return
self._aux_buffers = [

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

arm_aux_capture allocates one [max_num_batched_tokens, hidden_size] buffer for each id in spec.layer_ids. Under PP, the taps owned by earlier stages hook a PPMissingLayer that never fires, and aux_for substitutes _relayed_aux for them, so those buffers are never written or read.

For Kimi-K3 at pp2, the taps are [2, 23, 47, 71, 89] and stage 0 holds layers 0-46. That leaves 2 idle buffers on the last stage: 2 x 16384 x 7168 x bf16, about 470 MB per rank taken from its KV-cache budget. The waste grows with PP depth: 4 idle taps, about 940 MB, at PP8.

Could arm_aux_capture allocate (and hook) only the taps resident on this stage? PPAuxRelay already computes incoming_ids.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in a1436cd: arm_aux_capture now gives taps that sit on a PPMissingLayer (relayed from an earlier stage) a 0-row buffer and no hook, so only this stage's own taps take memory. Covered by test_last_stage_arms_only_its_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>
wanzhenchn added a commit that referenced this pull request Nov 4, 2026
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>

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

a few small things..

ci Atomesh Accuracy Validation failed — ci script/config bug, not the pr (high): The workflow triggers on .github/benchmark/models_accuracy.json (touched by this PR), but on a pull_request its whitelist selects zero models, so models_json is [] and the atomesh-test matrix include expands to nothing; unlike atom-test.yaml there is no placeholder fallback and GitHub fails an empty matrix. Kimi-K3 is not in atomesh_levels.

include: ${{ fromJson(needs.load-test-models.outputs.models_json) }} # atomesh-accuracy-validation.yaml:138, no fallback
include: ${{ fromJson(needs.load-test-models.outputs.models_json || '[{"runner":"ubuntu-latest","model_name":"skip"}]') }} # atom-test.yaml:265
atomesh_levels = {"DeepSeek-R1-0528": "main", "gpt-oss-120b": "nightly",}
next: Not a PR defect: add atom-test's placeholder-entry fallback to the atomesh matrix (or skip the workflow on catalog-only PRs), then rerun; no job log was attached. https://github.com/ROCm/ATOM/actions/runs/41762102572

Comment thread atom/spec_decode/drafter.py Outdated
def relayed(lid):
# On an earlier PP stage: `take_relayed_aux` serves it instead.
# Out-of-range ids are left for install_aux_taps to reject.
if lid == -1:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

relayed() uses its own residency rule while pp_aux_relay uses is_first_rank, can we share one helper?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in b62bf7a: drafter.tap_is_resident is the one rule now (is the tap's module a PPMissingLayer, the embedding included), and both arm_aux_capture and PPAuxRelay use it.

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>
wanzhenchn added a commit that referenced this pull request Nov 4, 2026
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>
wanzhenchn added a commit that referenced this pull request Nov 4, 2026
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>

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

ci Atomesh Accuracy Validation failed — ci script/config bug, not the pr (high): The workflow triggers on .github/benchmark/models_accuracy.json, which this PR touches, but on a pull_request its whitelist selects zero models, so models_json is [] and the atomesh matrix include expands to nothing. Unlike atom-test.yaml this workflow has no placeholder fallback, and GitHub fails a job whose matrix resolves empty. Kimi-K3 is not in atomesh_levels, so no atomesh job can exist for this change.

    include: ${{ fromJson(needs.load-test-models.outputs.models_json) }}   # atomesh-accuracy-validation.yaml:138, no fallback
    include: ${{ fromJson(needs.load-test-models.outputs.models_json || '[{"runner":"ubuntu-latest","model_name":"skip"}]') }}   # atom-test.yaml:265, with fallback
      filtered = [m for m in models if atomesh_levels.get(m["model_name"], "skip") in allowed]   # atomesh_levels has no Kimi-K3 entry

next: Not a PR defect: add atom-test's placeholder-entry fallback to the atomesh matrix (or skip the workflow on catalog-only PRs), then rerun. No job log was attached to this run. https://github.com/ROCm/ATOM/actions/runs/41766392543

@gbyu-amd

gbyu-amd commented Nov 4, 2026

Copy link
Copy Markdown
Contributor Author

@valarLip would you review again at the latest head?

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
wanzhenchn added a commit that referenced this pull request Nov 5, 2026
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>
wanzhenchn added a commit that referenced this pull request Nov 5, 2026
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>
@valarLip
valarLip merged commit 64412f1 into main Nov 5, 2026
37 checks passed
@valarLip
valarLip deleted the guanbao/k3-pp2tp4 branch November 5, 2026 12:06
wanzhenchn added a commit that referenced this pull request Nov 5, 2026
#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>
rbrugaro-amd added a commit to rbrugaro-amd/ATOM that referenced this pull request Nov 5, 2026
main's ROCm#2365 brings PP > 1 to Kimi-K3, so 'such an engine never verifies'
no longer holds. The refusal stands: PP rings forward_vars per in-flight
step while bind_target binds the draft to one slot's filters. Message,
comment and guide row now say that.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
rbrugaro-amd added a commit to rbrugaro-amd/ATOM that referenced this pull request Nov 5, 2026
…tn now reads

main's ROCm#2365 has GDNStateMixin._init_gdn_state read
model_runner.num_spec_tokens (every PP stage runs the verify window), but
this GPU-only test's ModelRunner stub never set it, so all 19 cases fail
with AttributeError before reaching what they test. The stub sets it to
num_spec, the value its drafter stub already carries as mtp_k.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Jasen2201 added a commit that referenced this pull request Nov 6, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants