[gbyu-amd]

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 P/D prefill 1/4 — PAGE-backed state checkpoints under PP #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

• None outstanding. [Config] Tune a8w8 bpreshuffle for Kimi-K3's KDA gate at tp4 aiter#5783 (tuned a8w8_bpreshuffle rows for K3's KDA gate at TP4, N=3072/K=128) merged on 2026-09-30.

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

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[github-actions]

🏷️ 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]

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]

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)

ATOM/atom/spec_decode/pp_aux_relay.py

Lines 45 to 47 in 501364c

	# 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)

ATOM/atom/spec_decode/pp_aux_relay.py

Lines 82 to 86 in 501364c

	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

ATOM/atom/model_ops/attentions/gdn_attn.py

Lines 1204 to 1207 in 501364c

	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]

@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.

[valar-cop]

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

[valar-cop]

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

[ganyi1996ppo]

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.

[valar-cop]

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]

@valarLip would you review again at the latest head?