[AMD] Fix DSv4 draft extend taking the target compression path during prefill

[Duyi-Wang]

Motivation

On MI355X (gfx950), DeepSeek-V4-Pro with EAGLE/MTP at TP 4 with DP attention aborts within
90 s of the server becoming ready — on the first batch of real prefills, before a single decode
step completes:

HSA_STATUS_ERROR_EXCEPTION: An HSAIL operation resulted in a hardware exception. code: 0x1016
torch.AcceleratorError: HIP error: unspecified launch failure  (hipErrorLaunchFailure)

  eagle_worker_v2.py         _draft_extend_for_prefill
    deepseek_v4_backend_hip_radix.py  init_forward_metadata_prefill
      _attach_unified_kv_prefill_meta -> torch.repeat_interleave

Deterministic, reproduced on clean nodes across several configurations. The reporting rank
varies (DP1 and DP2 observed), so it is not rank-specific.

Root cause

Two call sites reach the prefill branch of init_forward_metadata for a draft model, and only
one rewrites forward_mode:

path	call site	forward_mode
decode	eagle_worker_common.prepare_for_draft_extend()	sets ForwardMode.DRAFT_EXTEND_V2
prefill	eagle_worker_v2._draft_extend_for_prefill()	hands over the target's batch untouched → stays ForwardMode.EXTEND

The backend decides compression from forward_mode alone (need_compress = not is_draft), so on
the prefill path a draft model is sent down the compression path. It owns no compression pool to
begin with — the two pools initialise as

draft:  c4_size=0        c128_size=0
target: c4_size=1832064  c128_size=57252

so this is not a draft/target size mismatch but compression metadata computed against a pool that
does not exist. The resulting out-of-bounds access is reported asynchronously and surfaces at the
next synchronisation point — the repeat_interleave in _attach_unified_kv_prefill_meta, which
omits output_size and therefore implies a sum() D2H copy. That line is where the fault is
observed, not where it occurs.

Instrumenting the branch confirms it at runtime: on the failing ranks the draft worker is handed
ForwardMode.EXTEND, and the write offsets it computes (c4_out_loc up to 64) point into a pool
of size 0.

Modifications

The backend already reads model_runner.is_draft_worker in __init__ (for the dspark check).
Keep it as a named attribute and use it as a second, mode-independent source of truth:

is_draft = (
    forward_batch.forward_mode.is_draft_extend_v2() or self.is_draft_worker
)

A draft worker now never takes the compression path regardless of the forward_mode it is
handed. Only draft-worker backend instances are affected; the target is unchanged.

Fixing it here rather than by rewriting forward_mode in _draft_extend_for_prefill keeps the
blast radius to this backend — forward_mode is read by the logits processor, the CUDA-graph
runners and several other backends.

Tests

Full agentic replay (long-context, tool-heavy, multi-turn) on a single MI355X node, TP 4 + DP
attention, EAGLE MTP (3 steps, topk 1, 4 draft tokens), hierarchical cache, concurrency 32.

	before	after
hipErrorLaunchFailure	crash 57 s after ready	0
Prefill batch	stops at the crash	915
Decode batch	0	reached
client warmup	returned=1/354, then stalled	returned=330/354, errors=0

Same image, same node class, same workload; the only difference is this patch.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #33732720236
Latest PR Test (Extra): ❌ Run #33735198934
Latest PR Test (AMD ROCm 7.2): ❌ Run #33732720244

[kkHuang-amd]

LGTM.

One question, this issue should be not related to TP size, right?