[yifehuan]

Fixes #3177.

Problem

On gfx1250, DeepSeek-V4.1-Flash crashed on the first decode step that the single-launch kernel served:

File "atom/model_engine/model_runner.py", line 3791, in postprocess
    1 : hidden_states.shape[0] + 1
AttributeError: 'NoneType' object has no attribute 'shape'

Cause

The bug came from two PRs meeting on main, #3109 and #3136. #3091 did not cause it.

• [DSV4.1][gfx1250] Single-launch decode step (FlyDSL) #3109 (03b584b8, the gfx1250 single-launch decode step) runs the whole decode step in one FlyDSL launch per rank: target, sampling, verification and the next drafts. It plugs into ModelRunner through StepOverrideModel / _step_override. When it serves a step, run_model returns (None, None). The drafter's propose is wrapped, so on a served step it returns the drafts the kernel already drew and never reads the hidden states.
• refactor(tbo): stage, publish and finish attention metadata once per micro-batch #3136 (bf657561, TBO metadata refactor) landed on main before [DSV4.1][gfx1250] Single-launch decode step (FlyDSL) #3109 merged. It changed the drafter's input ids in postprocess (and in the align-only draft call) from input_ids.gpu[1 : batch.total_tokens_num + 1] to input_ids.gpu[1 : hidden_states.shape[0] + 1].

#3109 was developed and validated on a base without #3136, so 03b584b8 as merged already crashes on every served step. The 5b4fbfde in the report is a later main.

#3091 (gfx950 high-interactivity monokernels) is not involved. Its served rounds (served_by_round, compute_head → round/post_verify) keep the target forward returning hidden states, and its drafts run through the normal propose_draft_token_ids. Its compute_head exists only on the hi model, and consumes_device_verdict is False without a hi decode round, so neither is active on gfx1250.

Fix

ModelRunner._draft_input_ids(batch, hidden_states) returns the ids the drafter reads, one per row the forward ran:

• With hidden states: it uses their row count, so every existing path keeps its behavior. This includes a PCP rank's own shard on a pure middle chunk, where the rows are not running_tokens.
• Without hidden states (a step the model served whole): it uses the step's running_tokens, from the forward context's ForwardMode.

All three call sites (deferred, undeferred and align-only drafts) use it now.

Note: two mechanisms for a model taking over part of a decode step

ATOM now has two parallel ways for a model to take over part of a decode step, and each hooks into run_model / postprocess differently:

	Who decides	What the runner gets
#3091	The attention builder (served_by_round)	Hidden states, plus the head's ids and verdict (compute_head)
#3109	The model wrapper (StepOverrideModel)	(None, None); outputs read from _served

This crash happened because postprocess assumed only the first. Unifying the two is a larger change and should be discussed with both PRs' authors. It is not done here.

Test

• black and ruff pass on the changed file.
• GPU, on gfx1250 (MI455 ×4, TP4):
  • Setup: DSpark 5, BF16 KV, FP8 index cache, c1, an agentic client.
  • The run on main reproduces the crash at the first served step.
  • The run with this branch serves without error.

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