diff --git a/tests/compile/test_side_stream.py b/tests/compile/test_side_stream.py new file mode 100644 index 000000000000..83a8d6386fec --- /dev/null +++ b/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() diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index dc12acbaf4a8..2d267ae08802 100644 --- a/vllm/compilation/backends.py +++ b/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,6 +651,7 @@ 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 @@ -650,6 +659,7 @@ def wrap_with_cudagraph_if_needed( 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,22 +877,29 @@ 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 @@ -889,6 +907,7 @@ def collect_standalone_compile_artifacts( 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? diff --git a/vllm/compilation/caching.py b/vllm/compilation/caching.py index 62da2d9de35b..2d3a04bcf3a3 100644 --- a/vllm/compilation/caching.py +++ b/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 diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 742ec55e6efa..2290bd847685 100644 --- a/vllm/compilation/compiler_interface.py +++ b/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() diff --git a/vllm/compilation/side_stream.py b/vllm/compilation/side_stream.py new file mode 100644 index 000000000000..72f12b1a9dc3 --- /dev/null +++ b/vllm/compilation/side_stream.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared high-priority CUDA side streams.""" + +import torch + +from vllm.platforms import current_platform +from vllm.utils.torch_utils import current_stream + +_streams: dict[int, torch.cuda.Stream] = {} + + +def get_side_stream() -> torch.cuda.Stream | None: + """Return the current device's shared high-priority CUDA stream.""" + if not current_platform.is_cuda(): + return None + + main_stream = current_stream() + device_index = main_stream.device.index + assert device_index is not None + stream = _streams.get(device_index) + if stream is None: + _, high_priority = torch.cuda.Stream.priority_range() + stream = torch.cuda.Stream(device=device_index, priority=high_priority) + _streams[device_index] = stream + return stream + + +def wait_side_stream() -> None: + """Make the current stream wait for work issued on the side stream. + + No-op on platforms without a CUDA side stream. + """ + main_stream = current_stream() + stream = get_side_stream() + if stream is not None: + main_stream.wait_stream(stream) + + +def register_side_stream() -> None: + """Register streams used by PyTorch's compiled stream operators.""" + from torch._dynamo.graph_bytecode_inputs import ( + CURRENT_STREAM_INDEX, + index_to_external_object_weakref, + set_external_object_by_index, + ) + + main_stream = current_stream() + side_stream = get_side_stream() + if side_stream is None: + return + + for index, stream in ( + (CURRENT_STREAM_INDEX, main_stream), + (CURRENT_STREAM_INDEX + 1, side_stream), + ): + stream_ref = index_to_external_object_weakref.get(index) + if stream_ref is None or stream_ref() is not stream: + set_external_object_by_index(index, stream) + + +def _patch_external_object_getter() -> None: + import torch._dynamo.graph_bytecode_inputs as graph_inputs + + getter = graph_inputs.get_external_object_by_index + if getattr(getter, "_vllm_side_stream_patched", False): + return + + def get_external_object_by_index(index: int): + if index not in index_to_external_object_weakref and index in (0, 1): # type: ignore[name-defined] # noqa: F821 + _vllm_register_side_stream() # type: ignore[name-defined] # noqa: F821 + if index not in index_to_external_object_weakref: # type: ignore[name-defined] # noqa: F821 + raise AssertionError("Index not registered in index_to_user_object_weakref") + obj = index_to_external_object_weakref[index]() # type: ignore[name-defined] # noqa: F821 + if obj is None: + raise AssertionError("User object is no longer alive") + return obj + + graph_inputs._vllm_register_side_stream = register_side_stream + getter.__code__ = get_external_object_by_index.__code__ + getter._vllm_side_stream_patched = True + + +_patch_external_object_getter() diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index b8f197b31614..4621fdb46b53 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -201,6 +201,7 @@ from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.compilation.side_stream import wait_side_stream from vllm.config import ( CacheConfig, ModelConfig, @@ -511,7 +512,6 @@ def __init__( and _vllm_config.parallel_config.decode_context_parallel_size > 1 and _vllm_config.parallel_config.dcp_comm_backend == "a2a" ) - # Initialize q/k/v range constants. self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32) self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32) @@ -556,6 +556,7 @@ def forward( kv_c_normed: torch.Tensor, k_pe: torch.Tensor, output_shape: torch.Size | None = None, + indexer_dummy_dep: torch.Tensor | None = None, ) -> torch.Tensor: if self.calculate_kv_scales: torch.ops.vllm.maybe_calc_kv_scales( @@ -599,6 +600,7 @@ def forward( self_kv_cache, attn_metadata, output=output, + indexer_dummy_dep=indexer_dummy_dep, ) return output else: @@ -618,6 +620,7 @@ def forward( output, encoded, kv_cache_dummy_dep=kv_cache_dummy_dep, + indexer_dummy_dep=indexer_dummy_dep, ) return output @@ -635,6 +638,7 @@ def forward_impl( quant_scale_ue8m0: bool | None = None, quant_col_major: bool | None = None, quant_tma_aligned: bool | None = None, + indexer_dummy_dep: torch.Tensor | None = None, ) -> torch.Tensor: assert output is not None, "Output tensor must be provided." @@ -809,7 +813,14 @@ def forward_impl( # call decode attn if not is_sparse_impl: assert attn_metadata.decode is not None - attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined] + if indexer_dummy_dep is not None: + wait_side_stream() + attn_out, lse = self.impl.forward_mqa( # type: ignore[attr-defined] + mqa_q, + kv_cache, + attn_metadata, + self, + ) # correct dcp attn_out with lse. if self.impl.dcp_world_size > 1: @@ -1092,6 +1103,7 @@ def unified_mla_attention_with_output( quant_scale_ue8m0: bool | None = None, quant_col_major: bool | None = None, quant_tma_aligned: bool | None = None, + indexer_dummy_dep: torch.Tensor | None = None, ) -> None: # kv_cache_dummy_dep is not used but accepting it creates a data dependency # that ensures torch.compile preserves ordering between KV cache update and @@ -1112,6 +1124,7 @@ def unified_mla_attention_with_output( quant_scale_ue8m0=quant_scale_ue8m0, quant_col_major=quant_col_major, quant_tma_aligned=quant_tma_aligned, + indexer_dummy_dep=indexer_dummy_dep, ) @@ -1128,7 +1141,9 @@ def unified_mla_attention_with_output_fake( quant_scale_ue8m0: bool | None = None, quant_col_major: bool | None = None, quant_tma_aligned: bool | None = None, + indexer_dummy_dep: torch.Tensor | None = None, ) -> None: + del indexer_dummy_dep return diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 66a95b43c71f..0910e422f2f3 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -85,6 +85,11 @@ def __init__( self.rotary_emb = mla_modules.rotary_emb self.o_proj = mla_modules.o_proj self.indexer = mla_modules.indexer + self.indexer_side_stream = ( + getattr(self.indexer, "indexer_side_stream", None) + if self.indexer is not None + else None + ) self.indexer_rope_emb = mla_modules.indexer_rotary_emb self.is_sparse = mla_modules.is_sparse @@ -114,7 +119,6 @@ def __init__( indexer=self.indexer, topk_indices_buffer=mla_modules.topk_indices_buffer, ) - self.prefix = prefix def forward( @@ -125,6 +129,8 @@ def forward( ) -> torch.Tensor: q_c = None kv_lora = None + run_indexer = self.indexer and self.is_sparse and not self.skip_topk + indexer_dummy_dep = None if self.q_lora_rank is not None: assert self.fused_qkv_a_proj is not None, ( @@ -143,6 +149,11 @@ def forward( dim=-1, ) q_c = self.q_a_layernorm(q_c) + if run_indexer: + assert self.indexer is not None + indexer_dummy_dep = self.indexer( + hidden_states, q_c, positions, self.indexer_rope_emb + ) q = self.q_b_proj(q_c)[0] else: assert self.kv_a_proj_with_mqa is not None, ( @@ -151,6 +162,11 @@ def forward( assert self.q_proj is not None, ( "q_proj is required when q_lora_rank is None" ) + if run_indexer: + assert self.indexer is not None + indexer_dummy_dep = self.indexer( + hidden_states, q_c, positions, self.indexer_rope_emb + ) kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0] q = self.q_proj(hidden_states)[0] @@ -166,17 +182,17 @@ def forward( positions, q[..., self.qk_nope_head_dim :], k_pe ) - if self.indexer and self.is_sparse and not self.skip_topk: - self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) - if llama_4_scaling is not None: q *= llama_4_scaling + if indexer_dummy_dep is not None and self.indexer_side_stream is not None: + torch.accelerator.current_stream().wait_stream(self.indexer_side_stream) attn_out = self.mla_attn( q, kv_c_normed, k_pe, output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim), + indexer_dummy_dep=indexer_dummy_dep, ) return self.o_proj(attn_out)[0] diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index ceb52e5d329a..8a9fcee28281 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -134,14 +134,9 @@ def _fused_indexer_q_rope_quant_kernel( q_fp8, q_fp8_s0, q_fp8_s1, - weights, - weights_s0, - weights_s1, - weights_out, - weights_out_s0, - weights_out_s1, - softmax_scale, - head_scale, + q_scale_out, + q_scale_out_s0, + q_scale_out_s1, fp8_min: tl.constexpr, fp8_max: tl.constexpr, is_neox: tl.constexpr, @@ -199,30 +194,25 @@ def _fused_indexer_q_rope_quant_kernel( tl.clamp(q_nope / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), ) - weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32) - tl.store( - weights_out + token * weights_out_s0 + head * weights_out_s1, - weight * q_scale * softmax_scale * head_scale, - ) + tl.store(q_scale_out + token * q_scale_out_s0 + head * q_scale_out_s1, q_scale) def fused_indexer_q_rope_quant( positions: torch.Tensor, q: torch.Tensor, cos_sin_cache: torch.Tensor, - weights: torch.Tensor, - softmax_scale: float, - head_scale: float, is_neox: bool, ) -> tuple[torch.Tensor, torch.Tensor]: + """Fused rope + per-head fp8 quant for the indexer q; returns the raw + per-head q scale so callers fold it into the indexer weights, letting + this kernel issue before the wk/weights GEMM.""" assert current_platform.is_cuda() assert q.dtype == torch.bfloat16 assert q.shape[-1] == 128 assert cos_sin_cache.shape[-1] == 64 - assert weights.shape == q.shape[:2] q_fp8 = torch.empty_like(q, dtype=current_platform.fp8_dtype()) - weights_out = torch.empty_like(weights, dtype=torch.float32) + q_scale = torch.empty(q.shape[:2], device=q.device, dtype=torch.float32) fp8_min, fp8_max = get_fp8_min_max() _fused_indexer_q_rope_quant_kernel[(q.shape[0], q.shape[1])]( positions, @@ -234,20 +224,15 @@ def fused_indexer_q_rope_quant( q_fp8, q_fp8.stride(0), q_fp8.stride(1), - weights, - weights.stride(0), - weights.stride(1), - weights_out, - weights_out.stride(0), - weights_out.stride(1), - softmax_scale, - head_scale, + q_scale, + q_scale.stride(0), + q_scale.stride(1), fp8_min=fp8_min, fp8_max=fp8_max, is_neox=is_neox, num_warps=1, ) - return q_fp8, weights_out + return q_fp8, q_scale def _gather_workspace_shapes( @@ -364,7 +349,7 @@ def sparse_attn_indexer( num_decode_tokens = attn_metadata_narrowed.num_decode_tokens # q_scale is required iff the FP4 cache path is enabled; the FP8 path - # folds the Q scale into `weights` inside fused_indexer_q_rope_quant. + # folds the Q scale into `weights` before calling this op. if use_fp4_cache: assert q_scale is not None, "use_fp4_cache=True requires q_scale" else: diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index ac187151a0a6..ad0421adaf58 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -35,6 +35,7 @@ import vllm._custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile +from vllm.compilation.side_stream import get_side_stream from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( get_ep_group, @@ -713,6 +714,11 @@ def __init__( ) self.is_inplace_rope = is_inplace_rope + # Run the indexer on a high-priority side stream so it overlaps the + # main stream's path to attention. Applies to all CUDA sparse-MLA + # decode, not just DCP. + self.use_indexer_side_stream = current_platform.is_cuda() + self.indexer_side_stream = get_side_stream() self.n_head_scale = self.n_head**-0.5 self.use_fused_indexer_q = ( current_platform.is_cuda() @@ -722,9 +728,13 @@ def __init__( and self.scale_fmt is not None ) - def forward( - self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb - ) -> torch.Tensor: + def _prepare_indexer_inputs( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + positions, + rotary_emb, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: q, _ = self.wq_b(qr) q = q.view(-1, self.n_head, self.head_dim) @@ -746,6 +756,15 @@ def forward( positions, q[..., : self.rope_dim], k[..., : self.rope_dim].unsqueeze(1) ) elif self.use_fused_indexer_q and q.dtype == torch.bfloat16: + # q rope+quant first so it is not serialized behind the wk GEMM; + # the q scale is folded into the weights afterwards + q_fp8, q_scale = fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + rotary_emb.is_neox_style, + ) + # fused wk + weights_proj: one GEMM, then split kw, _ = self.wk_weights_proj(hidden_states) k = kw[:, : self.head_dim] @@ -756,16 +775,6 @@ def forward( k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - q_fp8, weights = fused_indexer_q_rope_quant( - positions, - q, - rotary_emb.cos_sin_cache, - weights, - self.softmax_scale, - self.n_head_scale, - rotary_emb.is_neox_style, - ) - # rotate only the MQA K k_pe = k_pe.unsqueeze(1) q_dummy = torch.empty_like(k_pe) @@ -773,7 +782,9 @@ def forward( k_pe = k_pe.reshape(-1, self.rope_dim) k = torch.cat([k_pe, k_nope], dim=-1) - return self.indexer_op(hidden_states, q_fp8, k, weights) + weights = weights * q_scale * self.softmax_scale * self.n_head_scale + + return q_fp8, k, weights else: q_pe, q_nope = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 @@ -813,6 +824,23 @@ def forward( weights = weights * q_scale * self.softmax_scale * self.n_head_scale + return q_fp8, k, weights + + def forward( + self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb + ) -> torch.Tensor: + if self.use_indexer_side_stream: + assert self.indexer_side_stream is not None + self.indexer_side_stream.wait_stream(torch.accelerator.current_stream()) + with self.indexer_side_stream: + q_fp8, k, weights = self._prepare_indexer_inputs( + hidden_states, qr, positions, rotary_emb + ) + topk_indices = self.indexer_op(hidden_states, q_fp8, k, weights) + return topk_indices + q_fp8, k, weights = self._prepare_indexer_inputs( + hidden_states, qr, positions, rotary_emb + ) return self.indexer_op(hidden_states, q_fp8, k, weights) diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py index dcf955ad59ba..8adbebd44dd8 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/vllm/models/deepseek_v32/nvidia/attention.py @@ -400,6 +400,7 @@ def _fused_attention( if self.indexer is not None: has_indexer = True + assert self.indexer_rope_emb is not None indexer_k_norm_w = self.indexer.k_norm.weight indexer_k_norm_bias = self.indexer.k_norm.bias indexer_k_norm_eps = self.indexer.k_norm.eps @@ -467,7 +468,7 @@ def _fused_attention( q_pe, self.rotary_emb.cos_sin_cache, index_q, - self.indexer_rope_emb.cos_sin_cache if has_indexer else None, + indexer_k_rope_cos_sin_cache if has_indexer else None, ql_nope, self._q_scale, index_weights, diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index bef811f6e61d..a49faeb47a0f 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -445,7 +445,7 @@ def compile_flash_attn_varlen_func_from_specs( raise NotImplementedError("FA4 compile-only wrapper does not support dropout") del deterministic - from vllm.vllm_flash_attn.cute.interface import ( + from vllm.vllm_flash_attn.cute.interface import ( # type: ignore[attr-defined] compile_flash_attn_varlen_func_from_specs as _fa4_compile_flash_attn_varlen_func_from_specs, )