[Bugfix] MTP: emit one spec-decode row per decode seq to fix IndexError

[yhl-amd]

Motivation

The scheduled DeepSeek-V4-Pro DPA MTP3 8k1k benchmark crashes during decode:

File "atom/model_engine/model_runner.py", line 524, in prepare_input_ids
    draft_tokens = batch.scheduled_spec_decode_tokens[new_curr_indices]
IndexError: index 127 is out of bounds for axis 0 with size 127

(all 8 DP ranks die → benchmark step fails; e.g. run 28461535343).

Technical Details

Scheduler only records a scheduled_spec_decode_tokens entry for decode seqs
whose spec_token_ids is non-empty:

if seq.spec_token_ids.size > 0:
    scheduled_spec_decode_tokens[seq.id] = seq.spec_token_ids

ScheduledBatch.__init__ then turns that dict into a positional array via
np.asarray(list(values())), so its length is number of drafting seqs, not
number of decode seqs. But TokenIDProcessor.prepare_input_ids indexes it by
full-batch sequence positions (new_curr_indices) in the [deferred | new]
layout. A decode seq without drafts — e.g. one that just transitioned from
prefill and sits at the batch tail — makes the array shorter than the batch, and
the positional index runs off the end.

Failure flow

                  decode batch: N=128 seqs (in batch order)
        +----+----+----+- -+----+----+----+
        |seq0|seq1|seq2|...| 125| 126| 127|   seq127 just finished prefill:
        +----+----+----+- -+----+----+----+   first decode, no draft proposed yet
                                        |
                                        v
                          spec_token_ids.size == 0 ?
                                        |
                    no (seq0..126)      |      yes (seq127)
                        |               |          |
                        v               |          v
        =============================   |   +---------------------+
        SCHEDULER (scheduler.py:980)    |   | skipped, no dict row |
        dict[id] = spec_token_ids       |   |  X seq127 dropped     |
        =============================   |   +---------------------+
                        |
                        v
        np.asarray(list(dict.values()))          # scheduler.py:315
                        |
                        v
        +----+----+----+- -+----+----+
   arr: |row0|row1|row2|...| 125| 126|      shape = (127, 3)   # K=127 < N=128
        +----+----+----+- -+----+----+      valid rows 0..126
                        |
                        |   (decode-only batch, prev step was also decode)
                        v
        ==========================================
        MODEL_RUNNER.prepare_input_ids
        get_token_locations() ->
          new_curr_indices = [127]              # seq127's batch position
        ==========================================
                        |
                        v
              takes the [deferred | new] branch
                        |
                        v
        draft = arr[new_curr_indices] = arr[127]   # model_runner.py:524
                        |
              arr only has rows 0..126
                        v
        +=======================================================+
        | IndexError: index 127 is out of bounds                |
        |            for axis 0 with size 127                   |
        |   -> all 8 DP ranks crash, benchmark step exit code 1 |
        +=======================================================+

Root mismatch: the scheduler numbers the array by drafting-seq order
(compacted, 0..126), while the model runner indexes it by batch position
(127). The dropped seq sits at the tail, so the lookup overruns.

Regression history (bisected)

This is a latent defect from #219 exposed by a later scheduling change,
not a same-PR break. The DeepSeek-V4-Pro DPA MTP3 8k1k / c=1024 job passed on
2026-06-26/27/28 and started failing on 2026-06-30:

 06-26 pass   06-27 pass   06-28 pass  |  #1318 (2026-06-29)  |  06-30 fail   07-01 fail

• mtp refine #219 (mtp refine) planted the fragile contract: it converted
  scheduled_spec_decode_tokens from a dict (consumed missing-safe via
  .get(req_id, [])) into a positional numpy array while keeping the size > 0
  filter that drops rows. A draft-less decode seq is harmless unless it lands
  at the batch tail.
• [Feature] OFFLOAD: standalone LMCache CPU/NVMe KV-offload connector #1318 (OFFLOAD connector) is the only commit touching scheduler.py /
  model_runner.py in the passing→failing window. Its scheduler refactor
  changed prefill/decode admission ordering (new skipped_partial_prefills +
  running.extendleft handling), which — with 8k chunked prefill — routes a
  freshly-prefilled, draft-less seq to the [deferred | new] tail position and
  into the positional-index path, triggering the latent bug.
• (Pinpointing the exact ordering line inside [Feature] OFFLOAD: standalone LMCache CPU/NVMe KV-offload connector #1318 would need an e2e serve
  bisect across the two SHAs; the fix here is at the contract layer so it holds
  regardless of scheduling order.)

The --enable-dp-attention --method mtp + long-prefill (8k) + high-concurrency
combination is what makes the tail layout common (all three are necessary; the
same job with 1k input, or without MTP, passes).

Prior-art note: the pre-#219 code may have been latently broken too

While bisecting we also looked at the original (pre-#219) implementation, and it
looks like the draft-less decode seq was never handled fully correctly there
either -- it just wasn't reachable:

• The scheduler has always scheduled a fixed mtp_k + 1 tokens for every
  decode seq, regardless of whether it has drafts (num_new_tokens = self.mtp_k + 1
  is unconditional, outside the if seq.spec_token_ids gate). So the compute
  width per decode seq is identical across all versions; this fix adds no extra
  compute for the placeholder seq -- those token slots were already scheduled.
• Pre-mtp refine #219, scheduled_spec_decode_tokens was a dict consumed via
  .get(req_id, []), and the [deferred | new] branch did
  new_token_ids.extend([tokens[-1]] + draft_tokens). For a draft-less seq that
  yields 1 token, while the scheduler reserved mtp_k + 1 slots for it -- a
  latent length mismatch. It appears this never fired because the
  [deferred | new] tail layout with a draft-less seq did not occur until
  [Feature] OFFLOAD: standalone LMCache CPU/NVMe KV-offload connector #1318's scheduling change (consistent with the bisect: the job was green
  through 06-28).

In other words, "a draft-less decode seq flowing through the [deferred | new]
path" was never actually exercised before; #219 turned the dict into a positional
array (making the missing row fatal instead of silently .get-defaulted), and
#1318 first made that path reachable. This fix handles the case explicitly at the
contract layer, so it is correct regardless of which path (front/tail) or
scheduling order produces a draft-less decode seq.

Fix

Give every decode seq a row. A seq with no drafts does not accept draft
tokens this step (num_new_token + num_rejected == 1), so the draft-slot values
are don't-cares; reuse the seq's own scheduled decode tokens (guaranteed valid
token ids). Rows for seqs that already had drafts are unchanged, so the common
path is byte-for-byte identical and the array is now dense
(num_decode_seqs, mtp_k).

Test Plan / Result

• Reproduced the exact IndexError: index 127 ... size 127 against the real
  ScheduledBatch on CPU (compacted (127, 3) array indexed at position 127).
• Added tests/test_scheduler_spec_decode_dense.py driving the real
  Scheduler.schedule() decode path:
  • before fix: array is (1, 3) → AssertionError (the compaction bug)
  • after fix: array is (2, 3), positional indexing of every batch
    position succeeds, drafting seqs' rows unchanged.
• pytest tests/test_scheduler.py tests/test_arg_utils_spec.py tests/test_sequence.py
  → 66 passed (no regressions).
• Note: full 8-GPU end-to-end serve of DSV4-Pro DPA MTP3 not yet run (cluster
  GPUs were fully occupied); the crashing operation is pure CPU numpy and is
  covered by the repro + regression test above. Happy to attach an e2e serve log
  once GPUs free up.