[MTP] Fix negative pad-row KV length in draft decode metadata

[whx-sjtu]

Problem

_enter_decode_metadata rebases kv_indptr for the scheduled rows only:

kv_indptr[1 : scheduled_bs + 1] -= torch.cumsum(num_reject_tokens, dim=0)

A padded row therefore keeps the range whatever batch last occupied that row left behind. The comment above this line assumed those stale entries are harmless because they are "one of its own valid ranges, so their reads stay in bounds" — they are not. Once the real rows total more tokens than the stale value, the tail of the indptr runs backwards, so the pad row's length comes out negative. mtp_prepare_decode_mla_kernel turns that into a negative sparse count, and sparse decode derives kv_start/kv_end from the negative extent and faults on an address unrelated to any live allocation:

Memory access fault by GPU node-6 ... on address 0x783dda000000. Reason: Unknown.

Per-row KV lengths from the batch that faulted (bs=5, padded to running_bs=8):

kvcnt=[3477, 3493, 3596, 3725, 3201, -11147, 980, 1058]
spcnt=[2048, 2048, 2048, 2048, 2048, -11147, 980, 1058]

The five real rows are correct (spcnt = min(ctx, index_topk)); the first pad row is negative.

Reaching it needs all three of:

1. a padded draft batch — so cudagraph mode, since positions is only staged up to running_bs when self.step is not None;
2. num_speculative_tokens >= 2 — a draft step past 0 must exist to rebuild the decode metadata;
3. a long request landing on a size bucket a shorter batch left behind — otherwise the stale value still exceeds the real total and the tail stays monotonic.

That combination is why it shows up as an intermittent crash under mixed long/short traffic rather than in any single-length run.

Fix

Close the tail in the backend's own prepare_mtp_decode, which is where the base contract already puts this:

Backends that distinguish the two mark the padded tail so downstream kernels skip it; those that do not simply never read bs.

It receives bs as the scheduled count and reads running_bs off the row buffer, so both are already in hand — this is the same place and the same if running_bs > bs: shape DeepSeek-V4 uses to mark its own padded tail.

The value is the last real end repeated, matching what a draft PREFILL's widened rows already get from _pad_prefill_mla_draft_tail. It has to precede the fused kernel, which reads the diffs; the += cu_seqlens_q that kernel applies right after then leaves each pad row holding the same single token it leaves every real one, so no kernel sees an empty range. The row's output stays discarded exactly as before, and its writes remain neutralized by the existing -1 slot path.

TritonMLAMetadataBuilder inherits this through its super().prepare_mtp_decode(...) call. MHA and GDN never take the rebase that strands the tail, and DeepSeek-V4 rebuilds its indptr over the full running_bs.

Validation

DeepSeek-V3.2 + MTP-3, TP8 on MI355X, cudagraph_mode=FULL (draft graph enabled), full GSM8K at 64 concurrency. 5-shot and 20-shot were driven against one server simultaneously, which is the mixed long/short scenario that reproduces the fault.

	before	after
memory access fault	8/8 ranks, < 1 min in	none
GSM8K 5-shot (n=1319)	crash	0.9538 ± 0.0058
GSM8K 20-shot (n=1319)	crash	0.9522 ± 0.0059
MTP acceptance	—	68.86%, 3.07 toks/fwd
per-position acceptance	—	21.6% / 31.9% / 40.5%

Without the fix the same workload faulted on every rank within a minute; --num-speculative-tokens 1 or an eager draft avoided it, matching the conditions above.

Localized by bisecting with per-kernel torch.cuda.synchronize() markers, which pinned the fault inside the draft decode replay, then by dumping per-row KV and sparse counts at the faulting batch.

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

[valarLip]

should be kv_indptr[scheduled_bs + 1 : running_bs + 1]=kv_indptr[scheduled_bs + 1] ?

[whx-sjtu]

Adopted, with the anchor one slot earlier — and moved the whole thing out of here into the backend's prepare_mtp_decode, since marking the padded tail is what the base contract already asks a backend to do:

Backends that distinguish the two mark the padded tail so downstream kernels skip it; those that do not simply never read bs.

It gets bs as the scheduled count and reads running_bs off the row buffer, so this now reads the same way DeepSeek-V4 marks its own tail a few files over:

if running_bs > bs:
    kv_indptr[bs + 1 : running_bs + 1] = kv_indptr[bs]

On the index itself: it has to be kv_indptr[bs], not kv_indptr[bs + 1]. bs + 1 is the first stale entry, so assigning from it starts with a self-assign and never touches the pair that goes negative. Replaying kv_indptr from the batch that faulted — 5 real rows padded to the captured bucket of 8, entries 6..8 left by the previous occupant:

no fix                 [0, 3476, 6968, 10563, 14287, 17487,  6339,  7318,  8375]
                       -> lengths [3476, 3492, 3595, 3724, 3200, -11148, 979, 1057]

= kv_indptr[bs + 1]    [0, 3476, 6968, 10563, 14287, 17487,  6339,  6339,  6339]
                       -> lengths [3476, 3492, 3595, 3724, 3200, -11148,   0,    0]   <- still negative

= kv_indptr[bs]        [0, 3476, 6968, 10563, 14287, 17487, 17487, 17487, 17487]
                       -> lengths [3476, 3492, 3595, 3724, 3200,      0,   0,    0]

The negative is the 17487 -> 6339 drop at the boundary; only anchoring on kv_indptr[bs] closes it.

(Those are pre-kernel values. The probe that captured them runs after mtp_prepare_decode_mla_kernel has applied kv_indptr += cu_seqlens_q — an arange, +i on entry i — so the array above is the printed one minus that arange, i.e. what the assignment sees. The leading 0 is pinned by the measurement rather than assumed: the printed per-row lengths sum to 8383, which is also the printed kv_indptr[8], and that sum is by definition kv_indptr[8] - kv_indptr[0].)

That same arange is why the zero lengths above are not what any kernel sees: the assignment lands before the fused kernel, which then gives every row one more token, pad rows included. So each pad row reaches the sparse decode holding exactly one token, on a slot kv_indices_generate_triton fills from block_tables — no empty range to handle.

Re-ran the full validation on this version: GSM8K 5-shot and 20-shot driven at one server together, 64 concurrency, cudagraph FULL — 0.9538 / 0.9522 exact_match, no fault, MTP acceptance 68.86% (3.07 toks/fwd). Description updated.