[Test][Scheduler] Regression coverage for multi-module MTP prefill-lookahead chunking

[AnkitNakhawa]

What this is

Investigation into #53488 (prompt_logprobs corrupted for some requests under MTP speculative decoding). This is not a fix — it's a checkpoint: a regression test that closes a real coverage gap found while root-causing the issue, opened as draft while I continue reproducing on a GPU with a real multi-layer MTP checkpoint.

What I found

• Scheduler._reserve_prefill_lookahead only does real work when num_prefill_lookahead > 1, which only happens for multi-module MTP (num_speculative_tokens > 1 with a multi-layer draft config) — exactly the reporter's config (num_speculative_tokens=3).
• No existing test sets num_prefill_lookahead > 1; the two closest tests (test_mamba_align_eagle_schedules_encoder_at_boundary, one other) hardcode 1. So this lookahead-shortened-chunk regime is entirely untested.
• I traced the two most plausible root-cause candidates from the issue and ruled both out for the reporter's exact config:
  • draft_slots budget subtraction (max_num_new_slots_for_drafting) is provably 0 for plain method="mtp" — it can't be shrinking this request's budget.
  • The multi-module MTP speculator (vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py) uses its own private hidden-states buffer, not the base model's — no shared-buffer clobber.
• The new test here confirms the scheduler's own chunk-boundary bookkeeping (num_computed_tokens progression across a lookahead-shortened chunk) is also correct.

So the corruption most likely lives in how _get_prompt_logprobs_dict (or the V2 runner's compute_prompt_logprobs_with_chunking) reads hidden states for a chunk shaped by this lookahead reservation — but confirming the exact line needs a real GPU run against a multi-layer MTP checkpoint, which I don't have in this environment.

Test plan

.venv/bin/python -m pytest tests/v1/core/test_scheduler.py -k test_multi_module_mtp_prefill_lookahead_chunk_accounting -v
.venv/bin/python -m pytest tests/v1/core/test_scheduler.py -q   # full file, 151 passed
pre-commit run --files tests/v1/core/test_scheduler.py          # clean

Duplicate check

gh issue view 53488 --repo vllm-project/vllm --comments and gh pr list --repo vllm-project/vllm --state open --search "53488 in:body" — no open PR addresses this issue yet.

AI assistance

Investigation and this test were written with AI assistance (Claude). I've reviewed the diff and reasoning above myself; I'm continuing the GPU-side root-cause work before this is ready for review.

[nyhhome]

Reporter of #53488 here. Your coverage finding stands on its own, but I think one of your ruled-out candidates was scoped to the wrong buffer, and I have a measurement that discriminates between your remaining hypothesis and the fix that has since been proposed.

On timing, not awareness: your dup check was correct when you ran it. You opened this at 2026-08-24T03:26Z; @hoobnn cross-referenced #53520 onto the issue at 04:31Z, an hour later. So #53520 ("Fix prompt_logprobs corruption from CUDA graph output buffer reuse by padded drafter") may simply postdate your search.

The buffer you ruled out is not the buffer #53520 fixes. You wrote:

The multi-module MTP speculator (.../multi_module_mtp/speculator.py) uses its own private hidden-states buffer, not the base model's — no shared-buffer clobber.

That is correct about the speculator's private hidden-states buffer. #53520 is about a different object: the CUDA graph output buffer that the padded drafter writes over. A negative result on the first does not exclude the second — and I think that is why your remaining hypothesis points one layer too high (at _get_prompt_logprobs_dict / chunking) rather than at graph replay.

A discriminator, from the runs I did while validating #53520 on Jetson AGX Thor (sm110a, aarch64):

Every corrupted row I have runs with CUDA graphs on. Two earlier speculative K=3 rows that were run with --enforce-eager scored clean on the same corpus and the same checkpoint.

I want to be precise about how much that carries: those eager rows are not strict single-variable controls — they also differ in kernel path (triton vs the fused CUDA path) and were built on a different day, so I record them as corroboration, not proof. What does not depend on that caveat is the scheduler-side reasoning: --enforce-eager does not touch _reserve_prefill_lookahead, and it does not change the reduced per-step budget (max_num_scheduled_tokens = 2048/(K+1) = 512). A lookahead-shaped-chunk explanation therefore predicts corruption in the eager rows too. There is none.

Two more data points pointing the same way:

• The boundary is at 512 polluted / 513 clean — the same adjacent pair I had bisected on a GB10, reproduced exactly on the Thor. On the speculative rows the engine reports max_cudagraph_capture_size: 512; on the non-speculative row it is 256. A 512 boundary is capture-size shaped.
• After applying [Bugfix][Spec Decode] Fix prompt_logprobs corruption from CUDA graph output buffer reuse by padded drafter #53520, an 11-point length sweep is 11/11 clean, and the 512↔513 discontinuity goes from 8x (15.878 → 1.939 mean NLL) to 0.008% (1.94532 → 1.94516) — the discontinuity is flattened, not moved.

One caveat I want to state rather than let you infer: I did not run a knob-isolating experiment to establish which setting fixes the number at 512. In this configuration 2048/(K+1), max_num_seqs 128x(K+1) and max_cudagraph_capture_size all land on 512, so the fact that the boundary sits there is consistent with a capture-size origin but does not establish one. The CUDA-graph dependence itself rests on the eager rows above, with the caveat I gave for them.

Your test is not made redundant by #53520. The gap you found is real and orthogonal: num_prefill_lookahead > 1 has no coverage, the two nearest tests hardcode 1, and #53520 does not add any. The scheduler bookkeeping you verified is correct — that is a useful negative result to have in the tree, not a discarded one.

Offer: you mention you lack a GPU with a real multi-layer MTP checkpoint. I have that — Jetson AGX Thor (sm110a) and a GB10, with three Qwen3.8-27B checkpoints (NVFP4 ModelOpt, NVFP4 compressed-tensors, W4A16) and the scoring harness from the issue. Root-causing may be moot now, but if it is useful I am happy to run your test on top of #53520's head and report back, or to re-run any specific configuration you want to see.

Disclosure: this analysis was done with AI assistance; the runs, the hardware and the numbers above are real and reproducible, and the account owner reviewed the content before posting. Raw artifacts available on request.

[nyhhome]

Clarification on my earlier comment, and one new measurement that bears on the root cause.

Clarification. I wrote "…num_prefill_lookahead > 1 has no coverage, the two nearest tests hardcode 1, and #53520 does not add any." The "any" referred to coverage for num_prefill_lookahead > 1, not to tests in general — #53520 has carried its own unit tests (tests/v1/worker/test_gpu_model_runner.py, +54) since its first push on 2026-08-24, five days before I wrote that. Restated precisely: #53520 adds no coverage for num_prefill_lookahead > 1; its two tests exercise the model-runner bookkeeping/clone path via GPUModelRunner.__new__ with a hand-built SimpleNamespace, while this PR's test exercises scheduler chunk accounting in tests/v1/core/test_scheduler.py. The files do not overlap — test_gpu_model_runner.py at head contains no occurrence of "lookahead" at all, and num_prefill_lookahead has exactly one test site in the repo, the one this PR adds. So the substance stands: your test is not made redundant by #53520. Sorry for the ambiguous phrasing.

New measurement. I have since run the single-variable experiment I owed on my earlier claim, and it lands on the CUDA-graph side rather than the scheduler side. Same box, same image, same unpatched gpu_model_runner.py, --max-num-seqs 128, num_speculative_tokens=3, CUDA graphs on; the only knob changed is an explicit max_cudagraph_capture_size = 256 instead of the default 512:

prompt tokens	mean NLL	verdict
200	13.544	polluted
256	13.577	polluted
257	1.990	clean
300 / 480 / 512 / 513	2.089 / 1.985 / 1.945 / 1.936	clean

The boundary moved from 512│513 to 256│257, tracking the capture limit exactly, while max_num_seqs × (K+1) stayed at 128 × 4 = 512 and 512 tokens became clean. Details in my comment on #53520.

That is consistent with the graph-replay mechanism and, I think, with why the lookahead-chunking hypothesis did not pan out: the scheduler-side bookkeeping you verified is correct, and the corruption is downstream of it. Your coverage gap is still real and still unfilled by #53520 — it just isn't where this particular bug lives.

Offer stands: if it is useful I can run your test on top of #53520's current head (78e87980be8d) on the same hardware and report back.

Disclosure: this analysis and the runs behind it were done with AI assistance; the hardware, the runs and the numbers are real and reproducible, and the account owner reviewed the content before posting.