[davzhuAMD]

Motivation

Main records prefill token graphs (atom/utils/region_graph.py, #2688) for an attention backend with supports_token_graph. Until now only DeepSeek-V4.1's had it. Kimi-K3's model was already a @token_graph unit, and its KDA mixer already an @eager_region(), but its builder could not take a padded prefill, so every Kimi-K3 prefill ran eagerly. Short prefills are host-bound: #2688's docs put the saving at 30–40 ms of TTFT per prefill below 3072 tokens on DeepSeek-V4.1.

This PR lets Kimi-K3 take them. A prefill of at most 3072 tokens pads to the smallest recorded bucket (ATOM_TOKEN_GRAPH_SIZES, default 256,512,1024,2048,3072). The dense layers between the KDA and MLA regions then replay as graphs, and the regions run the step's real rows.

It rewrites an earlier internal patch of ours ("0016": piecewise prefill graphs on CUDAGraphWrapper, on ATOM 6e9b3e1c5) onto #2688's region graphs. #2688 removed that wrapper.

Technical Details

Changes

• kimi_mla_gdn_attn.py (the Kimi-K3 builder):
  • supports_token_graph = True.
  • prepare_prefill: the base builder already builds prefill metadata for the scheduled rows only (cu_seqlens_q, slots, positions, context lengths). On a step padded to a bucket, Kimi-K3 records the real row count (AttentionMetaData.prefill_real_rows) and sets the pad rows' slots to -1, so no recorded write lands anywhere for them.
  • step_needs_eager: a prefill carrying images runs eagerly, as on DeepSeek-V4.1.
  • build_for_token_graph_capture(tokens): the request-less prefill ModelRunner.dummy_execution(prefill=True) runs, on placeholder page 0 with no state slot. KDA gets no metadata, so it writes no state. MLA writes to the placeholder page only.
• Eager regions (only on a marked step; unmarked steps, decode and other models are unchanged):
  • MLA (base_attention.unified_attention_with_output_base): cuts q, k, v, positions and qkv to the real rows, rebuilds the group-major FP8 q_scale (transpose_scale=True) at the real height, and hands the output back at bucket height.
  • KDA (kimi_k3.model.kda_attention_with_output): _forward_impl already runs the real rows; the output is handed back at bucket height.
• atom/utils/token_graph_rows.py (new, no AITER): padded_prefill_rows(), real_rows() and to_rows(). atom/utils/forward_context.py gets the AttentionMetaData.prefill_real_rows field (default None).
• ATOM_TOKEN_GRAPH_EAGER_PAD (default 0; region_graph.py, envs.py, docs row): a padded step runs eagerly instead of replaying, for A/B checks. It is generic and about 3 lines in region_graph.py, so it could be split into its own PR if preferred.
• docs/environment_variables.md: ATOM_TOKEN_GRAPH_SIZES now names Kimi-K3, and the new row.

Checked

• DP = 1: _pad_to_token_graph takes the group's largest row count, which is the rank's own, so the bucket is the one the rank would choose. Above the largest bucket the step runs eagerly. Unit-tested.
• Draft:
  • The DSpark proposer reads the target's rows through last_token_indices, taken from the real rows' cu_seqlens_q, and aux_for slices to the same row count, so pad rows are not read.
  • A pure middle chunk hands the drafter the whole model_output (model_runner.py:3203)
  • The DFlash draft (feat(kimi-k3): DFlash GQA speculative draft model #2581, and perf(kimi-k3): one-launch q/k norm + rope + paged K/V write for the DFlash draft (ATOM_K3_DRAFT_QKROPE) #2953 on top of it) is not on main yet.
• Sparse indexer (MLA layers with has_mla_indexer): the shared atom/model_ops/indexer is already region-graph aware. Its band logits are an eager region reading index_rows live, and its K write skips absent slots (nope.py:67), which is why pad slots are -1.
• Shared-expert side stream: it forks and joins inside the MoE forward, within one recorded segment, so it needs no eager_region(stream=...).
• Embedding: the runner pads ids and positions (_padded_inputs), TokenGraph copies every tensor argument into its own buffers, and the text embedding runs inside the recorded unit.

Relation to the internal patch ("0016")

• Carried over: cutting to the real rows, including the group-major FP8 q_scale rebuild; handing the result back at bucket height; the pad-without-graphs switch.
• Replaced by main's mechanism:
  • its own prefill_graphs.py plus a forced CUDAGraphWrapper per piece → region_graph;
  • per-bucket pools → one shared pool;
  • its model allowlist → supports_token_graph;
  • ATOM_PREFILL_GRAPH_BUCKETS → ATOM_TOKEN_GRAPH_SIZES;
  • its per-kind output buffers → region_graph's copy into the recorded address;
  • fixed id / position / embedding buffers → _padded_inputs and TokenGraph's input buffers.
• Dropped:
  • the refusal list (DP, PP, PCP, TBO, EP, mrope): main settles one bucket across DP (fix: replay one prefill token-graph bucket across DP ranks; prune CI #2696, fix(dp): a decoding rank replays the group's prefill token-graph bucket #2697) and runs TBO prefill eagerly;
  • setting aiter's private ca_comm.enable_register_for_capturing: main records inside graph_capture();
  • ATOM_DUAL_STREAM_PIECEWISE;
  • the per-bucket step log.

Default

• On under the default TR_HYBRID mode, like DeepSeek-V4.1's, which matches the maintainers' preference for default-on.
• Off: ATOM_TOKEN_GRAPH_SIZES= (empty), an existing knob that records nothing; no new gate.
• Risk: padding drift. Kernel choices follow the padded height; for example, the all-reduce switches from bf16 to INT4 above about 585 tokens.
  • With the internal patch, accuracy held at concurrency 8.
  • With mono decode at concurrency 4, completions failed to agree. There, the padded eager run already disagreed with the plain run, so the padding, not the graphs, caused it. To be re-checked here with ATOM_TOKEN_GRAPH_EAGER_PAD (Test Plan).
  • If it fails again, we would ask whether Kimi-K3 should stay opt-in.

Not in this PR

• KV budget: counting the prefill token-graph pool in the KV budget. Main already reserves for decode token graphs, but the prefill buckets come out of the safety margin (docs/environment_variables.md). That is a framework change, so it is proposed separately.

Test Plan

• CPU unit tests (new): tests/test_kimi_k3_prefill_token_graphs.py, plus the related suites.

  python -m pytest tests/test_kimi_k3_prefill_token_graphs.py -q
  python -m pytest tests/test_region_graph.py tests/test_forward_mode.py tests/test_graph_holders.py tests/test_kda_layout_id.py tests/test_kda_checkpoint_slot_copy.py tests/test_state_checkpoint.py tests/test_gdn_midstep_checkpoints.py -q
  

• What the new tests cover:
  • cutting to the real rows and the group-major q_scale rebuild (ported from the internal patch's CPU test);
  • the hand-back at bucket height with zero pad rows, and its refusal of extra rows;
  • a region's shape end to end;
  • that only a marked step counts as padded;
  • the DP = 1 bucket choice and the off cases;
  • the eager-pad switch, both ways, and its default.
• Lint: black --check and ruff check on the touched files.
• GPU (MI350X and MI355X, TP8, Kimi-K3 recipe), to run:
  1. Capture: recording succeeds, including the capture builder's request-less prefill through Kimi-K3's build(), with the dual-stream MoE on. Record capture time and pool size.
  2. Agreement: eager against ATOM_TOKEN_GRAPH_EAGER_PAD=1 against replay, on the same prompts (completions and logprobs), at concurrency 8 and with mono decode at concurrency 4.
  3. Accuracy: GSM8K (.github/scripts/atom_test.sh accuracy), graphs on and off.
  4. Serving: a prefix-reuse benchmark, graphs on and off: TTFT p50 / p90, output tok/s and KV blocks. Also DSpark acceptance with graphs on.
  5. Sparse indexer: the MLA-indexer layers on a padded step.

Test Result

• New CPU tests: 19 passed.
• Related suites: 361 passed, 15 skipped. The skips are the CUDA-graph recording tests in test_region_graph.py and the two builder modules that import AITER.
• Lint: clean.
• Local run: Windows, torch 2.15 CPU, ATOM's declared dependencies, and a local stand-in for two POSIX-only names in atom.utils. ATOM's Linux CPU CI needs none of that.

Run	Result
GPU check	capture 10.5 s, +3.07 GB per GPU (KV blocks −5%); TTFT ×0.65 on prefills that fit a bucket; decode ×1.008
100-trajectory screen	output tok/s +5.7% over the window, +2.5% over the run; median TTFT −22%
Full run	output +5.0% / +4.7%; median TTFT 253 ms (−20%); accuracy 0.5904 against 0.5896

A prefill step at 512 / 1,024 / 1,536 / 2,048 tokens cost about 82 / 117 / 130 / 148 ms, against about 180 ms eager. At concurrency 128 the gain was close to neutral.

Submission Checklist

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

AI assistance: AI tools (Claude, in Cursor) helped write the change, its tests and this description. I reviewed every changed line and ran the tests above myself. The commit has a Co-authored-by: trailer.

[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 2960 --add-label <label>

[valar-cop]

some changes needed..

mono-decode agreement at concurrency 4, gsm8k on/off with the command, and a GPU run of the padded sparse-indexer step?

ci green.

[valar-cop]

groups_major hardcoded? fp4x2 producer returns token-major scales, this would permute them..

[valar-cop]

this sends every replay eager, not just padded steps? gate on the padded marker instead of graph.has()

[valar-cop]

mrope positions are (3, N), token axis last, so this returns them uncut while q is cut.. cut last axis or raise?

[valar-cop]

zeros bucket-high tensor per layer per step, ~2GB of fill on the TTFT path.. keep one zeroed buffer and overwrite real rows?

[valar-cop]

duplicates dummy_execution prefill.. share one constructor, and missing 0 < tokens bound?

[valar-cop]

this pins the over-broad switch, correct padded-only fix would fail it.. assert eager only for a marked step?