[Attention] Overlap sparse MLA indexer with native streams

tests/compile/test_side_stream.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import tempfile
+from contextlib import contextmanager
+
+import pytest
+import torch
+from torch._dynamo.testing import EagerAndRecordGraphs
+
+import vllm.compilation.side_stream as side_stream
+from vllm.compilation.backends import graph_uses_stream_ops
+from vllm.compilation.decorators import support_torch_compile
+from vllm.compilation.side_stream import get_side_stream
+from vllm.config import (
+    CompilationConfig,
+    CompilationMode,
+    VllmConfig,
+    set_current_vllm_config,
+)
+from vllm.envs import disable_envs_cache
+from vllm.forward_context import set_forward_context
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import is_torch_equal_or_newer
+
+
+@contextmanager
+def use_vllm_config(vllm_config: VllmConfig):
+    with set_forward_context({}, vllm_config), set_current_vllm_config(vllm_config):
+        yield
+
+
+@support_torch_compile
+class NativeSideStreamModule(torch.nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.side_stream = get_side_stream()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert self.side_stream is not None
+        self.side_stream.wait_stream(torch.accelerator.current_stream())
+        with self.side_stream:
+            side = x + 1
+        main = x * 2
+        torch.accelerator.current_stream().wait_stream(self.side_stream)
+        return main + side
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only test")
+def test_side_stream_uses_native_compile_context() -> None:
+    stream = get_side_stream()
+    assert stream is not None
+    backend = EagerAndRecordGraphs()
+
+    def run(x: torch.Tensor) -> torch.Tensor:
+        stream.wait_stream(torch.accelerator.current_stream())
+        with stream:
+            side = x + 1
+        torch.accelerator.current_stream().wait_stream(stream)
+        return side
+
+    x = torch.zeros(4, device="cuda")
+    actual = torch.compile(run, backend=backend, fullgraph=True)(x)
+    assert torch.equal(actual, x + 1)
+    assert len(backend.graphs) == 1
+    assert graph_uses_stream_ops(backend.graphs[0])
+
+    annotated_nodes = [
+        node
+        for node in backend.graphs[0].graph.nodes
+        if node.meta.get("custom", {}).get("stream") not in (None, 0)
+    ]
+    assert annotated_nodes
+    assert all(
+        "vllm.side_stream" not in str(node.target)
+        for node in backend.graphs[0].graph.nodes
+    )
+    wait_stream_nodes = [
+        node
+        for node in backend.graphs[0].graph.nodes
+        if "streams.wait_stream" in str(node.target)
+    ]
+    assert len(wait_stream_nodes) == 2
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only test")
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.13.0"), reason="requires torch 2.13")
+def test_side_stream_aot_cache_round_trip(monkeypatch: pytest.MonkeyPatch) -> None:
+    with tempfile.TemporaryDirectory() as cache_dir, monkeypatch.context() as m:
+        m.setenv("VLLM_CACHE_ROOT", cache_dir)
+        m.setenv("VLLM_USE_AOT_COMPILE", "1")
+        m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
+        m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
+        disable_envs_cache()
+
+        vllm_config = VllmConfig(
+            compilation_config=CompilationConfig(
+                mode=CompilationMode.VLLM_COMPILE,
+                backend="inductor",
+            )
+        )
+        x = torch.randn(16, device="cuda")
+        expected = 3 * x + 1
+        with use_vllm_config(vllm_config):
+            compiled_module = NativeSideStreamModule(vllm_config=vllm_config)
+            torch.testing.assert_close(compiled_module(x), expected)
+
+        disable_envs_cache()
+        m.setenv("VLLM_FORCE_AOT_LOAD", "1")
+        vllm_config = VllmConfig(
+            compilation_config=CompilationConfig(
+                mode=CompilationMode.VLLM_COMPILE,
+                backend="inductor",
+            )
+        )
+        with use_vllm_config(vllm_config):
+            cached_module = NativeSideStreamModule(vllm_config=vllm_config)
+            from torch._dynamo.graph_bytecode_inputs import reset_user_object_tracking
+
+            reset_user_object_tracking()
+            side_stream._streams.clear()
+            torch.testing.assert_close(cached_module(x), expected)
+
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                captured = cached_module(x)
+            graph.replay()
+            torch.testing.assert_close(captured, expected)
+
+        assert cached_module.was_aot_compile_fn_loaded_from_disk
+        disable_envs_cache()

vllm/compilation/backends.py

@@ -56,6 +56,13 @@
 logger = init_logger(__name__)
 
 
+def graph_uses_stream_ops(graph: torch.fx.GraphModule) -> bool:
+    return any(
+        node.op == "call_function" and str(node.target).startswith("streams.")
+        for node in graph.graph.nodes
+    )
+
+
 def make_copy_and_call(
     sym_tensor_indices: list[int],
     input_buffers: list[torch.Tensor | None],
@@ -631,6 +638,7 @@ def wrap_with_cudagraph_if_needed(
     compilation_config: CompilationConfig,
     is_first_graph: bool,
     is_last_graph: bool,
+    uses_stream_ops: bool = False,
 ) -> Any:
     """
     Wrap a piecewise backend with CUDA graph wrapper if needed.
@@ -643,13 +651,15 @@ def wrap_with_cudagraph_if_needed(
         compilation_config: The compilation configuration
         is_first_graph: Whether this is the first graph in the sequence
         is_last_graph: Whether this is the last graph in the sequence
+        uses_stream_ops: Whether the graph switches between CUDA streams
 
     Returns:
         The wrapped backend if CUDA graphs are enabled, otherwise the original backend
     """
     if (
         not compilation_config.cudagraph_mode.has_piecewise_cudagraphs()
         or compilation_config.use_inductor_graph_partition
+        or uses_stream_ops
     ):
         return piecewise_backend
 
@@ -764,6 +774,7 @@ def call_module(
                 self.compilation_config,
                 piecewise_backend.is_first_graph,
                 piecewise_backend.is_last_graph,
+                graph_uses_stream_ops(submod),
             )
 
             compilation_counter.num_piecewise_capturable_graphs_seen += 1
@@ -866,29 +877,37 @@ def __init__(
 
     def collect_standalone_compile_artifacts(
         self,
-    ) -> tuple[Any, dict[str, list[int]] | None, dict[str, bool] | None]:
+    ) -> tuple[
+        Any,
+        dict[str, list[int]] | None,
+        dict[str, bool] | None,
+        dict[str, bool] | None,
+    ]:
         """Collect inductor cache artifacts from all piecewise backends.
 
         Returns:
             tuple: (standalone_compile_artifacts, sym_shape_indices_map,
-                    returns_tuple_map)
+                    returns_tuple_map, uses_stream_ops_map)
                 - standalone_compile_artifacts: StandaloneCompiledArtifacts
                   with compiled artifacts
                 - sym_shape_indices_map: dict mapping submod_name to
                   sym_shape_indices
                 - returns_tuple_map: dict mapping submod_name to
                   returns_tuple
+                - uses_stream_ops_map: dict mapping submod_name to whether
+                  the graph contains native stream operations
         """
 
         if not envs.VLLM_USE_MEGA_AOT_ARTIFACT:
-            return None, None, None
+            return None, None, None, None
 
         from .caching import StandaloneCompiledArtifacts
         from .piecewise_backend import PiecewiseBackend
 
         standalone_compile_artifacts = StandaloneCompiledArtifacts()
         sym_shape_indices_map = {}
         returns_tuple_map = {}
+        uses_stream_ops_map = {}
 
         for name, _ in self.split_gm.named_children():
             # get the actual attribute (shadowed by PiecewiseBackend in __dict__)
@@ -902,6 +921,8 @@ def collect_standalone_compile_artifacts(
             submod_name = name
             sym_shape_indices_map[submod_name] = piecewise_backend.sym_shape_indices
             returns_tuple_map[submod_name] = piecewise_backend.returns_tuple
+            original_submod = self.split_gm._modules[name]
+            uses_stream_ops_map[submod_name] = graph_uses_stream_ops(original_submod)
 
             for shape_str, bytes_data in piecewise_backend.to_bytes().items():
                 standalone_compile_artifacts.insert(submod_name, shape_str, bytes_data)
@@ -924,7 +945,12 @@ def collect_standalone_compile_artifacts(
             list(standalone_compile_artifacts.submodule_bytes.keys()),
         )
 
-        return standalone_compile_artifacts, sym_shape_indices_map, returns_tuple_map
+        return (
+            standalone_compile_artifacts,
+            sym_shape_indices_map,
+            returns_tuple_map,
+            uses_stream_ops_map,
+        )
 
     def configure_post_pass(self) -> None:
         # TODO proper PassManager?

vllm/compilation/caching.py

@@ -214,6 +214,9 @@ def __init__(
             self.shape_env = sym_input.node.shape_env
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        from vllm.compilation.side_stream import register_side_stream
+
+        register_side_stream()
         return self.optimized_call(*args, **kwargs)
 
     @classmethod
@@ -260,11 +263,25 @@ def serialize_compile_artifacts(
         for node in state["graph_module"].graph.nodes:
             node.meta.pop("source_fn_stack", None)
             node.meta.pop("nn_module_stack", None)
+            if (
+                getattr(node.target, "__module__", None)
+                == "torch._dynamo.graph_bytecode_inputs"
+                and getattr(node.target, "__name__", None)
+                == "get_external_object_by_index"
+            ):
+                node.meta.pop("example_value", None)
         for name, submod in state["graph_module"].named_children():
             if hasattr(submod, "graph"):
                 for node in submod.graph.nodes:
                     node.meta.pop("source_fn_stack", None)
                     node.meta.pop("nn_module_stack", None)
+                    if (
+                        getattr(node.target, "__module__", None)
+                        == "torch._dynamo.graph_bytecode_inputs"
+                        and getattr(node.target, "__name__", None)
+                        == "get_external_object_by_index"
+                    ):
+                        node.meta.pop("example_value", None)
 
         if state.get("sym_tensor_indices"):
             # put tensor inputs on meta device since their data
@@ -290,10 +307,12 @@ def serialize_compile_artifacts(
                 standalone_compile_artifacts,
                 sym_shape_indices_map,
                 returns_tuple_map,
+                uses_stream_ops_map,
             ) = compiled_fn.vllm_backend.collect_standalone_compile_artifacts()
             state["standalone_compile_artifacts"] = standalone_compile_artifacts
             state["sym_shape_indices_map"] = sym_shape_indices_map
             state["returns_tuple_map"] = returns_tuple_map
+            state["uses_stream_ops_map"] = uses_stream_ops_map
         return pickle.dumps(state)
 
     @classmethod
@@ -309,6 +328,7 @@ def deserialize_compile_artifacts(cls, data: bytes) -> "VllmSerializableFunction
         standalone_compile_artifacts = state.pop("standalone_compile_artifacts", None)
         sym_shape_indices_map = state.pop("sym_shape_indices_map", {})
         returns_tuple_map = state.pop("returns_tuple_map", {})
+        uses_stream_ops_map = state.pop("uses_stream_ops_map", {})
 
         saved_aot_autograd_config = state["aot_autograd_config"]
         if saved_aot_autograd_config is not None:
@@ -329,6 +349,7 @@ def deserialize_compile_artifacts(cls, data: bytes) -> "VllmSerializableFunction
                     vllm_config=get_current_vllm_config(),
                     sym_shape_indices_map=sym_shape_indices_map,
                     returns_tuple_map=returns_tuple_map,
+                    uses_stream_ops_map=uses_stream_ops_map,
                     fake_mode=fake_mode,
                 )
 
@@ -414,6 +435,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
     vllm_config: VllmConfig,
     sym_shape_indices_map: dict[str, list[int]],
     returns_tuple_map: dict[str, bool],
+    uses_stream_ops_map: dict[str, bool],
     fake_mode: FakeTensorMode,
 ) -> "VllmSerializableFunction":
     """Construct a VllmSerializableFunction from cached inductor artifacts.
@@ -444,6 +466,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
         vllm_config: The vLLM configuration.
         sym_shape_indices_map: Mapping from submod_name to sym_shape_indices.
         returns_tuple_map: Mapping from submod_name to returns_tuple.
+        uses_stream_ops_map: Mapping from submod_name to whether it uses stream ops.
 
     Returns:
         A VllmSerializableFunction that can be called directly.
@@ -516,6 +539,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
             compilation_config,
             is_first,
             is_last,
+            uses_stream_ops_map.get(submod_name, False),
         )
 
         submod_callables[submod_name] = wrapped_backend

vllm/compilation/compiler_interface.py

@@ -24,6 +24,13 @@
 logger = init_logger(__name__)
 
 
+def _uses_non_default_stream(graph: fx.GraphModule) -> bool:
+    return any(
+        node.meta.get("custom", {}).get("stream") not in (None, 0)
+        for node in graph.graph.nodes
+    )
+
+
 class CompilerInterface:
     """
     The interface for a compiler that can be used by vLLM.
@@ -290,6 +297,10 @@ def compile(
         current_config = {}
         if compiler_config is not None:
             current_config.update(compiler_config)
+        if _uses_non_default_stream(graph):
+            # PyTorch 2.13's compile-time autotune wrapper does not initialize
+            # the raw handle for kernels assigned to a non-default stream.
+            current_config["triton.autotune_at_compile_time"] = False
         set_inductor_config(current_config, compile_range)
         set_functorch_config()