[Bugfix] Record aux-stream MoE allocations on the consuming stream

[rchalamala]

Purpose

Aux-stream MoE allocations that escape to the main stream are never handed to Tensor.record_stream(), so the caching allocator can recycle a block while the main stream is still reading it.

The sites all look correctly synchronised, and for execution order they are — events and wait_stream do make the consumer run after the producer. But ordering is not lifetime. A tensor allocated inside with torch.cuda.stream(aux) is owned by the allocator's aux pool; when its Python reference dies, the block is returned to that pool with only aux's completion tracked. The next aux allocation can take the block while the main-stream consumer's kernel is still reading it. record_stream is the mechanism that defers reuse until the consuming stream is done, and this repo already documents exactly that in vllm/v1/worker/gpu/spec_decode/utils.py:40-42:

without record_stream, the caching allocator may reuse its memory

Four escaping allocations were missing it:

# vllm/utils/multi_stream_utils.py -- maybe_execute_in_parallel
with torch.cuda.stream(aux_stream):
    event0.wait()
    result1 = fn1()          # allocated on aux_stream
    event1.record()
event1.wait()                # orders the consumer -- does not pin the block
+record_stream_if_safe(result1, torch.cuda.current_stream())

# vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py
record_stream_if_safe(shared_output, shared_expert_stream)   # tags the INPUT
with torch.cuda.stream(shared_expert_stream):
    shared_output = tensor_model_parallel_all_reduce(shared_output)  # REBOUND: new aux tensor
main.wait_stream(shared_expert_stream)
+record_stream_if_safe(shared_output, main)   # the rebound tensor was never tagged

SharedExperts.forward (shared_experts.py) and execute_in_parallel have the same shape. shared_experts.py also carried a comment asserting the tag was unnecessary, which conflates the two guarantees:

-# NOTE: We don't need shared_output.record_stream(current_stream())
-# because we synch the streams before using shared_output.

Symptom class: intermittent, scheduling-dependent corruption or unmapped reads (XID 31) in a bystander kernel downstream of the MoE — the faulting kernel is whichever one happens to consume the recycled block, not the buggy one. Every site is normally rescued by the next iteration's wait_stream/event.record, which is why it is rare rather than deterministic.

Scope, and what this deliberately does not do. The helper no-ops during CUDA graph capture (graph-private pool allocations are static, not recycled) so this is an eager-path fix. An earlier draft also suppressed multi-stream overlap during capture on the theory that captured side-stream work gets no overlap at replay. That theory is wrong and was dropped — measured on B300 (2×4096² bf16 GEMMs, 16 per graph, 20 replays):

captured graph	replay
single-stream	5.700 ms
multi-stream	4.856 ms (1.174x)

A graph captured from two streams replays its DAG concurrently, so suppressing overlap under capture would cost real decode latency on small-batch decode. Behaviour under capture is therefore unchanged.

Affects every model using shared experts / the multi-stream helpers (Kimi K3 latent-MoE, DeepSeek shared experts, the DeepSeek-V4 attention path via execute_in_parallel). Found while investigating a K3 DFlash warmup crash on 8×B300 — see the honesty note under Test Result.

Test Plan

# lint (repo pre-commit, all hooks)
uv venv --python 3.12 && uv pip install -r requirements/lint.txt && pre-commit install
pre-commit run

# new unit test, on 8x B300 (sm_103), torch 2.11
.venv/bin/python -m pytest tests/utils_/test_multi_stream_utils.py -v

# standalone harness against real capture + replay, exercising the patched
# module directly (the two measurements quoted above and below)
python3 test_aux_stream_fix.py
python3 test_capture_overlap.py

Test Result

pre-commit run: all hooks pass (ruff, mypy-3.10, SPDX, check-torch-cuda-call, lazy-import, docstring/config validators).

Standalone harness on B300, exercising the patched multi_stream_utils.py against real capture and replay — 8/8:

== 1. record_stream_if_safe semantics ==
  [PASS] outside capture: tolerates tensor / tuple / None / non-tensor
  [PASS] inside capture: detected as capturing
  [PASS] inside capture: no-op, does not raise
== 2. maybe_execute_in_parallel: eager multi-stream ==
  [PASS] numerics match sequential
== 3. multi-stream still ACTIVE under capture (deliberately) ==
  [PASS] captured graph replays correctly
== 4. breakable-capture fallback preserved (unchanged behaviour) ==
  [PASS] breakable capture -> sequential, correct
== 5. execute_in_parallel (N-ary, deepseek_v4 attention path) ==
  [PASS] N-ary correct, None slots skipped, no raise on tagging
  [PASS] N-ary sequential fallback correct

Also verified directly that an unguarded record_stream() inside capture does not raise on torch 2.11 — so the guard is for correctness-of-meaning, not to avoid an exception.

Pending: pytest tests/utils_/test_multi_stream_utils.py has not run yet — it needs a compiled tree (vllm._C_stable_libtorch) and the source build of this branch is still in progress. The logic under test is what the standalone harness above already covers on the same GPUs. I will post the pytest output when the build lands.

Honesty note on the motivating crash. This came out of a source review of a K3 DFlash k=16 warmup crash (XID 31 MMU fault → CUBLAS_STATUS_EXECUTION_FAILED in the drafter's fc GEMM). I have not reproduced that crash against this patch, so I am not claiming this fixes it. The reported producer/consumer pair was not confirmed either: vllm/v1/worker/gpu/spec_decode/dflash/ and qwen3_dflash.py contain no stream or event operations at all, so the DFlash combine_hidden_states path is single-stream. What is defensible from source is the lifetime hazard fixed here, and the fact that warmup is the phase most exposed to it — all four multi-stream gates are upper-bound token thresholds (256/256/512), so multi-stream is ON for warmup's 1-2 request batches, and warmup.py:329-341 deliberately enumerates mixed spec/non-spec batches whose shapes miss the captured uniform-decode graphs and run eager.

Not a duplicate. Checked open PRs (gh pr list --repo vllm-project/vllm --state open --search "record_stream" / "shared experts stream" / "multi_stream_utils"). Nearest neighbours are vllm-project#33225 (replaces with torch.cuda.stream() by the torch.Stream context-manager form — pure syntax, no lifetime change, but it edits adjacent lines in the older fused_moe/layer.py shared-experts block, so expect a trivial conflict if both land) and vllm-project#45452 (opt-in K2.6-NVFP4 two-stream perf flags, default-off, adds overlap rather than fixing tagging). Neither adds record_stream to an escaping aux allocation.

AI assistance was used for this change: the source review, patch, tests and measurements were produced by Devin. A human must review every line and defend it before this goes upstream.

Upstreaming. This branch targets k3-pr50000. Two of the three files — vllm/utils/multi_stream_utils.py and .../runner/shared_experts.py — are byte-identical on vllm-project/vllm main (checked at 15d65f8) and carry the bug verbatim, so those two hunks are directly upstreamable as a standalone PR. latent_moe_runner.py does not exist on main and arrives with the K3 drop, so that hunk belongs on PR vllm-project#50000.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Link to Devin session: https://modal.devinenterprise.com/sessions/857595e915f9407db96cd4e8140caa24
Requested by: @rchalamala

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR that start with 'DevinAI' or '@devin'.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.