[Bugfix] Fix race condition in MTP speculative decoding draft loop

[z1ying]

Summary

Fix #40756
A race condition in vllm/v1/spec_decode/llm_base_proposer.py that can trigger cudaErrorIllegalAddress under high concurrency when using MTP speculative decoding.

Root Cause

The draft token generation loop writes shared GPU buffers
(input_ids and hidden_states) and immediately launches the next
model forward pass using those buffers:

self.input_ids[:batch_size] = input_ids
self.hidden_states[:batch_size] = hidden_states

model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size]

Because CUDA operations are asynchronous, the buffer writes are not guaranteed to be visible to the next kernel before it begins execution. Under high concurrency this can cause FlashInfer attention kernels to observe stale or partially-written buffer contents, causing illegal memory access.

(EngineCore pid=196) ERROR ... SchedulerOutput(...),
total_num_scheduled_tokens=30,
scheduled_spec_decode_tokens={...: [-1, -1, -1, -1, -1], ...}

The issue disappears with CUDA_LAUNCH_BLOCKING=1 , suggesting a
missing synchronization between the buffer writes and the subsequent
draft model forward pass.

Fix

Add a stream-level synchronization after the buffer writes, before the model forward:

self.input_ids[:batch_size] = input_ids
self.hidden_states[:batch_size] = hidden_states
# Ensure buffer writes are committed before the draft model forward reads them
torch.accelerator.current_stream().synchronize()

Compared to torch.cuda.synchronize(), this approach is preferred because:

1. Stream-scoped synchronization
   Only synchronizes the current stream, avoiding full-device synchronization and reducing overhead.

2. Backend portability
   Uses the torch.accelerator abstraction, maintaining compatibility across CUDA, ROCm, and XPU.

Reproduction

Server startup command

vllm serve Qwen/Qwen3.5-9B \
  --dtype bfloat16 \
  --max-model-len 32768 \
  --max-num-seqs 64 \
  --gpu-memory-utilization 0.85 \
  --kv-cache-dtype bfloat16 \
  --attention-backend FLASHINFER \
  --enable-prefix-caching \
  --enable-chunked-prefill \
  --speculative-config '{"method": "mtp", "num_speculative_tokens": 4}' \
  --compilation-config '{"mode": 3}' \
  --language-model-only

Load test script (64 concurrent requests)

#!/usr/bin/env python3
"""vLLM MTP Crash Reproducer — 20 concurrent requests × 10 rounds"""
import asyncio, aiohttp, time, json, sys

URL = "http://localhost:8000/v1/chat/completions"
MODEL = "Qwen/Qwen3.5-9B"
CONCURRENT = 64
ROUNDS = 5

# ~4000 tokens 长 prompt，制造 KV cache 压力
PROMPT = ("Please provide an extremely detailed analysis covering "
          "every aspect including historical context, current state, "
          "and future implications of distributed computing. ") * 80

async def send(session, rid, rnd):
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": f"[{rnd}-{rid}] {PROMPT}"}],
        "max_tokens": 8192,
        "temperature": 0.9,
    }
    t0 = time.time()
    try:
        async with session.post(URL, json=payload,
                                timeout=aiohttp.ClientTimeout(total=600)) as r:
            body = await r.text()
            dt = time.time() - t0
            if r.status == 200:
                tok = json.loads(body)["usage"]["completion_tokens"]
                print(f"  [{rnd}-{rid:02d}] OK {dt:.1f}s {tok}tok", flush=True)
            else:
                print(f"  [{rnd}-{rid:02d}] ERR {r.status} {dt:.1f}s {body[:150]}", flush=True)
            return r.status
    except Exception as e:
        print(f"  [{rnd}-{rid:02d}] FAIL {time.time()-t0:.1f}s {e}", flush=True)
        return -1

async def run_round(rnd):
    print(f"\n{'='*50}\nRound {rnd}: {CONCURRENT} concurrent requests\n{'='*50}", flush=True)
    async with aiohttp.ClientSession(
        connector=aiohttp.TCPConnector(limit=CONCURRENT+5)
    ) as s:
        results = await asyncio.gather(
            *[send(s, i, rnd) for i in range(CONCURRENT)],
            return_exceptions=True
        )
    fails = sum(1 for r in results if r != 200)
    print(f"Round {rnd}: {CONCURRENT-fails} OK, {fails} failed", flush=True)
    return fails

async def main():
    print(f"Target: {URL}", flush=True)
    print(f"Model:  {MODEL}", flush=True)
    print(f"Concurrency: {CONCURRENT}", flush=True)
    print(f"Rounds: {ROUNDS}", flush=True)
    print(f"{'='*50}", flush=True)

    # warmup
    print("\n[Warmup] Single request...", flush=True)
    async with aiohttp.ClientSession() as s:
        r = await send(s, 0, 0)
        if r != 200:
            print("Warmup FAILED! Is vLLM running?", flush=True); sys.exit(1)
    print("[Warmup] OK\n", flush=True)
    await asyncio.sleep(2)

    for rnd in range(1, ROUNDS + 1):
        fails = await run_round(rnd)
        if fails > CONCURRENT // 2:
            print(f"\n!!! >50% requests failed in round {rnd}.", flush=True)
            print("!!! Server likely crashed. Check ~/vllm_server.log", flush=True)
            break
        await asyncio.sleep(3)

    print(f"\n{'='*50}", flush=True)
    print("Done. ", flush=True)
    

asyncio.run(main())

Crashes within minutes without the fix. Add CUDA_LAUNCH_BLOCKING=1 to confirm race condition — crash disappears.

Evidence

Before patch — crash log

After patch — load test passes

Community Validation

Independently verified by two users across different hardware and vLLM versions:

Validator 1 — RTX PRO 6000 96GB (Blackwell, sm_120), vLLM 0.20.1, Qwen3.6-27B-FP8, PyTorch 2.11.0+cu130, CUDA 13.2, 40 concurrent agents × 2048 max_tokens:

MTP	Throughput	Avg Latency	Crashes
2	1429.6 tok/s	55.77s	0 ✅
3	1163.2 tok/s	67.98s	0 ✅
4	1167.8 tok/s	67.96s	0 ✅
5	690.7 tok/s	113.41s	0 ✅

After the patch: All configurations (MTP=2/3/4/5) remained completely stable under 40-concurrent stress testing with zero crashes. The race condition has been effectively resolved by the explicit synchronize() call.

Additional finding: GPU utilization averaged 91% during testing, VRAM increase was only +202MB (0.2%) under 40 concurrent requests — no measurable overhead introduced by the synchronize() call.

We are now running MTP=2 in production with this patch.

Validator 2 — 4× RTX 3090 (TP=4), vLLM 0.20.2, Qwen3.6-27B-FP16:

MTP	Prefill (tok/s)	Decode (tok/s)	Crashes
2	1383.4	25.9	0 ✅
3	1438.7	27.9	0 ✅
4	1354.6	27.9	0 ✅
5	1363.5	27.3	0 ✅

Bug confirmed in v0.20.1 and v0.20.2. Fix confirmed across Blackwell and Ampere architectures, single-GPU and TP=4, FP8 and FP16.

Test Results

Regression Tests (48 passed, 7 skipped, 0 failed)

Tests run against all SpecDecodeBaseProposer subclasses that inherit the modified propose() method:

File	Passed	Skipped	Failed
test_mtp.py	2 ✅	0	0
test_eagle.py	45 ✅	6 (TRITON_ATTN unsupported on platform)	0
test_speculators_dflash.py	1 ✅	1 (requires 40GB GPU memory)	0
Total	48	7	0

python -m pytest \
  tests/v1/spec_decode/test_mtp.py \
  tests/v1/spec_decode/test_eagle.py \
  tests/v1/spec_decode/test_speculators_dflash.py \
  -v

All tests passed

[gemini-code-assist]

Code Review

This pull request introduces a stream synchronization call in llm_base_proposer.py to ensure that input data writes to CUDA graph buffers are completed before downstream kernels are launched. This change aims to maintain data consistency with minimal CPU overhead compared to global synchronization. There were no review comments provided, so I have no additional feedback to offer.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[MatthewBonanni]

These operations are all on the same stream, no? So they should operate sequentially, i.e. no race condition. If there were actually a race condition somewhere, I suspect that this change just artificially slows down the execution enough for the copy to finish.

You mention #41162 is a related fix and this is a workaround "while the upstream fix propagates". 41162 has landed -- does it fix this issue? Are you still seeing this failure on the latest ToT? Were you using model runner V2? 41162 only touches model runner V2.

Finally, in the future, please disclose when you are using AI for PR contributions.

[z1ying]

Hi @MatthewBonanni,

Thanks for reviewing! You are right—there’s another factor I need to consider. I will dive deeper to gather more persuasive evidence and make the PR clearer. Just sharing some findings I discovered earlier:

1. During debugging, I found that adding torch.cuda.synchronize() at this point eliminates the crash.
   The issue is suspected to be a race condition, based on attempts to reproduce it across multiple GPUs (3080Ti, 3090, 5090). It is difficult to trigger, and more aggressive configurations lead to OOM.

2. On a 3090, under the same conditions, MTP=1/2/3 do not crash, while MTP=4 does.

Additional observations:

• Single-request warm-up does not trigger the crash.
• CUDA_LAUNCH_BLOCKING=1 does not trigger the crash.
• --no-async-scheduling still crashes.

Added probes in gpu_model_runner: input_fits_in_drafter is always true, and _draft_token_ids never contains negative values. Later, in llm_base_proposer.py, adding torch.cuda.synchronize() at the cudagraph input buffer assignment eliminates the crash.

self.input_ids[:batch_size] = input_ids
self.hidden_states[:batch_size] = hidden_states

3. Enabling FLASHINFER may cause operations not to execute sequentially on the same stream.
   This looks like similar to the issue mentioned in [Model Runner V2] Rebuild attn metadata between draft decode steps #41162 — the attention metadata builder depends on positional state, which must be updated synchronously with each position advance.

I’m still investigating the FLASHINFER source to confirm its actual execution order; this will take some time.
At least, this patch is effective and has only minor performance impact. It has resolved the crashes reported by two users in the #42603 comments, and I included their post-patch data in the PR. I will work on finding the root cause with evidence to support it and update the commit message with help from Claude Sonnet 4.6 later.

[benchislett]

While I don't doubt that there may be a bug, forcing a synchronization is an unacceptable fix. Please migrate to a github issue, creating one if needed, until a root-cause is identified.

[schopra8]

@benchislett Do you have a suggestion for how to tackle this issue? Seems to be similar to another speculative decoding issue I'm seeing with Gemma MTP #42572

[schopra8]

@benchislett Do you have a suggestion for how to tackle this issue? Seems to be similar to another speculative decoding issue I'm seeing with Gemma MTP #42572

Just to clarify, I'm referring to an issue in vLLM v1