[Mamba] Prefetch stochastic-rounding seeds on a side stream

tests/kernels/mamba/test_sr_seed_prefetch.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Identity tests for VLLM_MAMBA_SR_SEED_PREFETCH.
+
+The N per-layer stochastic-rounding seeds drawn up front on a side stream
+(StochasticRoundingSeedPrefetcher, event-joined per use) must equal the N
+per-layer `torch.randint(0, 2**32, (1,))` draws of stock, eagerly and across
+CUDA graph replays, with other work on the current stream between uses.
+"""
+
+import pytest
+import torch
+
+import vllm.model_executor.layers.mamba.ops.ssu_dispatch as ssu_dispatch
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    StochasticRoundingSeedPrefetcher,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+NUM_LAYERS = 40
+
+
+class _FakeForwardContext:
+    pass
+
+
+@pytest.fixture
+def forward_context(monkeypatch):
+    """A fresh fake ForwardContext per forward pass: call `.new()`."""
+
+    class _Holder:
+        current: _FakeForwardContext | None = None
+
+        def new(self) -> None:
+            self.current = _FakeForwardContext()
+
+    holder = _Holder()
+    monkeypatch.setattr(
+        ssu_dispatch, "is_forward_context_available", lambda: holder.current is not None
+    )
+    monkeypatch.setattr(ssu_dispatch, "get_forward_context", lambda: holder.current)
+    return holder
+
+
+def _stock(out, work, num_calls):
+    for k in range(num_calls):
+        seed = torch.randint(0, 2**32, (1,), device="cuda")
+        out[k].copy_(seed)
+        work.mul_(1.0)  # layer work between draws
+
+
+def _prefetched(prefetcher, out, work, num_calls):
+    # What FlashInferSSUBackend.__call__ does per SSU call.
+    for k in range(num_calls):
+        seed = prefetcher.next_seed(torch.device("cuda"))
+        if seed is None:
+            seed = torch.randint(0, 2**32, (1,), device="cuda")
+        out[k].copy_(seed)
+        work.mul_(1.0)
+
+
+def _run_eager(fn, seed, num_calls):
+    torch.manual_seed(seed)
+    out = torch.zeros(num_calls, 1, dtype=torch.int64, device="cuda")
+    fn(out)
+    torch.cuda.synchronize()
+    return out.clone()
+
+
+def _run_graph(fn, seed, num_calls, reps=3):
+    out = torch.zeros(num_calls, 1, dtype=torch.int64, device="cuda")
+    torch.manual_seed(seed)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            fn(out)
+    torch.cuda.current_stream().wait_stream(stream)
+    res = []
+    for _ in range(reps):
+        graph.replay()
+        torch.cuda.synchronize()
+        res.append(out.clone())
+    return res
+
+
+@pytest.mark.parametrize("extra_calls", [0, 2])
+def test_sr_seed_prefetch_eager_identical(forward_context, extra_calls):
+    work = torch.randn(4096, 4096, device="cuda")
+    prefetcher = StochasticRoundingSeedPrefetcher(NUM_LAYERS)
+    num_calls = NUM_LAYERS + extra_calls
+
+    def ours(out):
+        forward_context.new()
+        _prefetched(prefetcher, out, work, num_calls)
+
+    for seed in (123, 456):
+        ref = _run_eager(lambda out: _stock(out, work, num_calls), seed, num_calls)
+        got = _run_eager(ours, seed, num_calls)
+        assert torch.equal(ref, got)
+
+
+def test_sr_seed_prefetch_cuda_graph_identical(forward_context):
+    work = torch.randn(4096, 4096, device="cuda")
+    prefetcher = StochasticRoundingSeedPrefetcher(NUM_LAYERS)
+
+    def ours(out):
+        forward_context.new()
+        _prefetched(prefetcher, out, work, NUM_LAYERS)
+
+    # Create the side stream and events outside capture, as vLLM's eager
+    # warmup pass does before capturing.
+    _run_eager(ours, 0, NUM_LAYERS)
+    ref = _run_graph(lambda out: _stock(out, work, NUM_LAYERS), 7, NUM_LAYERS)
+    got = _run_graph(ours, 7, NUM_LAYERS)
+    for r, g in zip(ref, got):
+        assert torch.equal(r, g)
+    # Fresh seeds on every replay, as stock.
+    assert not torch.equal(got[0], got[-1])
+
+
+def test_sr_seed_prefetch_outside_forward_context_falls_back(forward_context):
+    prefetcher = StochasticRoundingSeedPrefetcher(NUM_LAYERS)
+    assert prefetcher.next_seed(torch.device("cuda")) is None

vllm/envs.py

@@ -296,6 +296,7 @@
     VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION: bool = False
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
     VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
+    VLLM_MAMBA_SR_SEED_PREFETCH: bool = False
     VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
     VLLM_USE_V2_MODEL_RUNNER: bool | None = None
     VLLM_LOG_MODEL_INSPECTION: bool = False
@@ -2024,6 +2025,15 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int(
         os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "1024")
     ),
+    # With the FlashInfer Mamba SSU backend and Mamba cache stochastic
+    # rounding enabled, draw the per-layer stochastic-rounding seeds of a
+    # forward pass up front on a side CUDA stream (same values, same order as
+    # the per-layer draws), instead of one small RNG kernel on the critical
+    # path before every SSU call. Each SSU call waits on a CUDA event for its
+    # seed.
+    "VLLM_MAMBA_SR_SEED_PREFETCH": lambda: bool(
+        int(os.getenv("VLLM_MAMBA_SR_SEED_PREFETCH", "0"))
+    ),
     # Format for saving torch.compile cache artifacts
     # - "binary": saves as binary file
     #     Safe for multiple vllm serve processes accessing the same torch compile cache.

vllm/model_executor/layers/mamba/ops/ssu_dispatch.py

@@ -9,15 +9,19 @@
 the backend defaults to 'cpu'.
 """
 
+import weakref
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from functools import cache
 
 import torch
 
+import vllm.envs as envs
 from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
+from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
 from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import current_stream
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
@@ -118,6 +122,83 @@ def reset_replayssm_ring_trackers(
     )
 
 
+class StochasticRoundingSeedPrefetcher:
+    """Draws a forward pass's stochastic-rounding seeds ahead, on a side stream.
+
+    Stock FlashInfer SSU draws one seed per Mamba layer per forward with
+    `torch.randint(0, 2**32, (1,))` right before the SSU kernel: a small,
+    non-overlapped kernel on the critical path of every Mamba layer.
+
+    Here the first SSU call of a forward pass records `_fork_event` on the
+    current stream, makes a dedicated side stream wait on it, and issues the
+    same `num_layers` randint calls (same shape, dtype, device and order) on
+    the side stream, recording `_events[k]` after the k-th. The k-th SSU call
+    of the forward makes the current stream wait on `_events[k]` and uses
+    seed k. The CUDA generator's (seed, offset) is advanced on the host at
+    each call, in call order, and no other user of the generator runs between
+    the Mamba layers of a forward, so seed k has exactly the value of the
+    stock k-th draw.
+
+    Ordering uses CUDA events only (graph edges under CUDA graph capture).
+    The seed tensors are allocated on the side stream; their blocks can only
+    be reused by a later side-stream allocation, i.e. by the next forward's
+    draws, which run after that forward's `_fork_event`, recorded on the
+    current stream after every consumer of the previous seeds was enqueued.
+
+    Forward passes are told apart by their ForwardContext. SSU calls beyond
+    `num_layers` in one forward, calls outside a forward context and calls
+    under DBO micro-batching return None, and the caller draws inline as
+    stock (the values stay identical). Every forward that calls the SSU must
+    call it at least `num_layers` times, otherwise the side stream is left
+    unjoined (an error under CUDA graph capture).
+    """
+
+    def __init__(self, num_layers: int):
+        assert num_layers > 0
+        self.num_layers = num_layers
+        self._stream: torch.cuda.Stream | None = None
+        self._fork_event: torch.cuda.Event | None = None
+        self._events: list[torch.cuda.Event] = []
+        self._seeds: list[torch.Tensor] = []
+        self._forward: weakref.ref | None = None
+        self._next = 0
+
+    def _draw(self, device: torch.device) -> None:
+        if self._stream is None:
+            self._stream = torch.cuda.Stream(device=device)
+            self._fork_event = torch.cuda.Event()
+            self._events = [torch.cuda.Event() for _ in range(self.num_layers)]
+        assert self._fork_event is not None
+        side = self._stream
+        self._fork_event.record(current_stream())
+        side.wait_event(self._fork_event)
+        seeds = []
+        with torch.cuda.stream(side):
+            for event in self._events:
+                seeds.append(torch.randint(0, 2**32, (1,), device=device))
+                event.record(side)
+        self._seeds = seeds
+
+    def next_seed(self, device: torch.device) -> torch.Tensor | None:
+        if not is_forward_context_available():
+            return None
+        from vllm.v1.worker.ubatching import dbo_enabled
+
+        if dbo_enabled():
+            return None
+        forward_context = get_forward_context()
+        if self._forward is None or self._forward() is not forward_context:
+            self._forward = weakref.ref(forward_context)
+            self._next = 0
+            self._draw(device)
+        k = self._next
+        if k >= self.num_layers:
+            return None
+        self._next = k + 1
+        current_stream().wait_event(self._events[k])
+        return self._seeds[k]
+
+
 class MambaSSUBackend(ABC):
     """Abstract base class for Mamba SSU backends."""
 
@@ -224,6 +305,9 @@ def __init__(self, mamba_config: MambaConfig):
             ) from e
         logger.info_once("Using FlashInfer Mamba SSU algorithm: %s", self._algorithm)
         self._kernel = _fi_ssu
+        # Set by initialize_mamba_ssu_backend when VLLM_MAMBA_SR_SEED_PREFETCH
+        # is enabled.
+        self.sr_seed_prefetcher: StochasticRoundingSeedPrefetcher | None = None
 
     @property
     def _algorithm(self) -> MambaSSUAlgorithm:
@@ -253,11 +337,12 @@ def __call__(
         cu_seqlens: torch.Tensor | None = None,
         is_blackwell: bool = False,
     ) -> None:
-        rand_seed = (
-            torch.randint(0, 2**32, (1,), device=state.device)
-            if self._mamba_config.enable_stochastic_rounding
-            else None
-        )
+        rand_seed = None
+        if self._mamba_config.enable_stochastic_rounding:
+            if self.sr_seed_prefetcher is not None:
+                rand_seed = self.sr_seed_prefetcher.next_seed(state.device)
+            if rand_seed is None:
+                rand_seed = torch.randint(0, 2**32, (1,), device=state.device)
         self._kernel(
             state,
             x,
@@ -460,6 +545,54 @@ def selective_state_update_replayssm_flashinfer(
     return result
 
 
+def _maybe_make_sr_seed_prefetcher(
+    mamba_config: MambaConfig,
+    kv_cache_config: KVCacheConfig,
+    use_replayssm: bool,
+) -> StochasticRoundingSeedPrefetcher | None:
+    """Build the seed prefetcher for VLLM_MAMBA_SR_SEED_PREFETCH, if usable.
+
+    The number of seeds per forward is the number of Mamba-2 layers in the KV
+    cache config: each MambaMixer2 decode makes one SSU call per forward.
+    """
+    if not envs.VLLM_MAMBA_SR_SEED_PREFETCH:
+        return None
+    from vllm.platforms import current_platform
+
+    if (
+        not mamba_config.enable_stochastic_rounding
+        or use_replayssm
+        or not current_platform.is_cuda()
+    ):
+        logger.info_once(
+            "VLLM_MAMBA_SR_SEED_PREFETCH has no effect: it needs CUDA and "
+            "Mamba cache stochastic rounding, without ReplaySSM."
+        )
+        return None
+    num_layers = 0
+    for group in kv_cache_config.kv_cache_groups:
+        spec = group.kv_cache_spec
+        if not isinstance(spec, MambaSpec):
+            continue
+        if spec.mamba_type == MambaAttentionBackendEnum.MAMBA1:
+            # Mamba-1 layers also call the SSU backend; not counted here.
+            logger.info_once(
+                "VLLM_MAMBA_SR_SEED_PREFETCH is disabled for models with "
+                "Mamba-1 layers."
+            )
+            return None
+        if spec.mamba_type == MambaAttentionBackendEnum.MAMBA2:
+            num_layers += len(group.layer_names)
+    if num_layers == 0:
+        return None
+    logger.info_once(
+        "Prefetching Mamba stochastic-rounding seeds for %d layers per "
+        "forward on a side stream.",
+        num_layers,
+    )
+    return StochasticRoundingSeedPrefetcher(num_layers)
+
+
 def initialize_mamba_ssu_backend(
     mamba_config: MambaConfig,
     kv_cache_config: KVCacheConfig,
@@ -504,6 +637,11 @@ def initialize_mamba_ssu_backend(
         _mamba_ssu_backend = backend_cls(mamba_config)
         logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
 
+    if isinstance(_mamba_ssu_backend, FlashInferSSUBackend):
+        _mamba_ssu_backend.sr_seed_prefetcher = _maybe_make_sr_seed_prefetcher(
+            mamba_config, kv_cache_config, use_replayssm=use_replayssm
+        )
+
     _flashinfer_replayssm_kernel = None
     if use_replayssm and backend == MambaBackendEnum.FLASHINFER:
         try: