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

tests/utils_/test_multi_stream_utils.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+
+from vllm.utils.multi_stream_utils import (
+    execute_in_parallel,
+    maybe_execute_in_parallel,
+    record_stream_if_safe,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+
+@pytest.fixture
+def inputs():
+    x = torch.randn(64, 64, device="cuda")
+    w = torch.randn(64, 64, device="cuda")
+    return x, w
+
+
+def test_record_stream_if_safe_ignores_non_tensors():
+    stream = torch.cuda.Stream()
+    tensor = torch.randn(8, device="cuda")
+
+    record_stream_if_safe(tensor, stream)
+    record_stream_if_safe((tensor, None, "not a tensor"), stream)
+    record_stream_if_safe([tensor], stream)
+    record_stream_if_safe(None, stream)
+
+
+def test_record_stream_if_safe_is_noop_during_capture():
+    """record_stream is meaningless for graph-pool blocks, so skip it."""
+    stream = torch.cuda.Stream()
+    tensor = torch.randn(8, device="cuda")
+
+    # Warm up on a side stream, as CUDA graph capture requires.
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        tensor.mul(2.0)
+    torch.cuda.current_stream().wait_stream(side)
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = tensor * 2.0
+        assert torch.cuda.is_current_stream_capturing()
+        record_stream_if_safe(captured, stream)
+
+    graph.replay()
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(captured, tensor * 2.0)
+
+
+def test_maybe_execute_in_parallel_matches_sequential(inputs):
+    x, w = inputs
+    aux = torch.cuda.Stream()
+    event0, event1 = torch.cuda.Event(), torch.cuda.Event()
+
+    parallel = maybe_execute_in_parallel(
+        lambda: x @ w, lambda: x * 2.0, event0, event1, aux
+    )
+    sequential = maybe_execute_in_parallel(
+        lambda: x @ w, lambda: x * 2.0, event0, event1, None
+    )
+    torch.accelerator.synchronize()
+
+    for got, want in zip(parallel, sequential):
+        torch.testing.assert_close(got, want)
+
+
+def test_execute_in_parallel_matches_sequential(inputs):
+    x, w = inputs
+    start = torch.cuda.Event()
+    done = [torch.cuda.Event(), torch.cuda.Event()]
+    streams = [torch.cuda.Stream(), torch.cuda.Stream()]
+    fns = [lambda: x * 2.0, None]
+
+    default, aux = execute_in_parallel(
+        lambda: x @ w, fns, start, done, streams, enable=True
+    )
+    torch.accelerator.synchronize()
+
+    torch.testing.assert_close(default, x @ w)
+    torch.testing.assert_close(aux[0], x * 2.0)
+    assert aux[1] is None

vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py

@@ -12,6 +12,7 @@
     flashinfer_trtllm_fused_allreduce_norm,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.utils.multi_stream_utils import record_stream_if_safe
 from vllm.utils.torch_utils import aux_stream, current_stream
 
 from .moe_runner import MoERunner, _unpack
@@ -200,12 +201,15 @@ def _fused_forward(
         if shared_expert_stream is not None:
             # overlap shared expert allreduce with latent up_proj
             main = current_stream()
-            shared_output.record_stream(shared_expert_stream)
+            record_stream_if_safe(shared_output, shared_expert_stream)
             shared_expert_stream.wait_stream(main)
             with torch.cuda.stream(shared_expert_stream):
                 shared_output = tensor_model_parallel_all_reduce(shared_output)
             result = torch.mm(fused_latent, transform.up_proj.weight.t())
             main.wait_stream(shared_expert_stream)
+            # The all-reduce output is a new aux-stream allocation consumed
+            # below on the main stream.
+            record_stream_if_safe(shared_output, main)
         else:
             shared_output = tensor_model_parallel_all_reduce(shared_output)
             result = torch.mm(fused_latent, transform.up_proj.weight.t())

vllm/model_executor/layers/fused_moe/runner/shared_experts.py

@@ -11,6 +11,7 @@
     FusedMoEConfig,
 )
 from vllm.platforms import current_platform
+from vllm.utils.multi_stream_utils import record_stream_if_safe
 from vllm.utils.torch_utils import (
     aux_stream,
     current_stream,
@@ -120,9 +121,7 @@ def maybe_sync_shared_experts_stream(
             # Record that the clone will be used by shared_experts_stream
             # to avoid gc issue from deallocation of hidden_states_clone
             # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
-            # NOTE: We don't need shared_output.record_stream(current_stream())
-            # because we synch the streams before using shared_output.
-            shared_experts_input.record_stream(self._stream)
+            record_stream_if_safe(shared_experts_input, self._stream)
 
             # Mark sync start point for the aux stream since we will
             # run in parallel with router/gate.
@@ -138,6 +137,10 @@ def _run_in_aux_stream(
         with torch.cuda.stream(self._stream):
             output = self._layer(shared_experts_input)
         current_stream().wait_stream(self._stream)
+        # output was allocated on the aux stream; stream ordering alone does
+        # not stop the allocator from recycling the block underneath the
+        # main-stream consumer.
+        record_stream_if_safe(output, current_stream())
 
         return output
 

vllm/utils/multi_stream_utils.py

@@ -8,6 +8,31 @@
 import torch
 
 
+def record_stream_if_safe(value: Any, stream: torch.cuda.Stream) -> None:
+    """Mark aux-stream allocations as also used on ``stream``.
+
+    A tensor produced inside ``with torch.cuda.stream(aux)`` is owned by the
+    caching allocator's ``aux`` pool. Stream ordering (events / ``wait_stream``)
+    makes the consumer run after the producer, but it does not stop the
+    allocator from handing the block to a later ``aux`` allocation while the
+    consumer is still reading it. ``record_stream`` is what defers that reuse.
+
+    Skipped during CUDA graph capture, where allocations come from the
+    graph-private pool and are not recycled.
+
+    Args:
+        value: A tensor, or a tuple/list that may contain tensors. Non-tensor
+            entries are ignored.
+        stream: The stream the tensors are consumed on.
+    """
+    if not torch.cuda.is_available() or torch.cuda.is_current_stream_capturing():
+        return
+    values = value if isinstance(value, (tuple, list)) else (value,)
+    for v in values:
+        if isinstance(v, torch.Tensor):
+            v.record_stream(stream)
+
+
 class AuxStreamType(Enum):
     Attention = 1
 
@@ -60,6 +85,7 @@ def maybe_execute_in_parallel(
             result1 = fn1()
             event1.record()
         event1.wait()
+        record_stream_if_safe(result1, torch.cuda.current_stream())
     else:
         result0 = fn0()
         result1 = fn1()
@@ -133,4 +159,8 @@ def execute_in_parallel(
     for ev in pending:
         ev.wait()
 
+    current = torch.cuda.current_stream()
+    for res in aux_results:
+        record_stream_if_safe(res, current)
+
     return default_result, aux_results