From 44af287ebe38d6dc4e102948025f5e3e175aefd6 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Sun, 27 Sep 2026 23:48:01 +0800 Subject: [PATCH] [ROCm][Perf] Enable layer-aware CSA2 multi-stream overlap for DeepSeek-V4.1-Flash (#57407) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- tests/kernels/test_compressor_kv_cache.py | 266 ++++++++++++++++++++++ vllm/models/deepseek_v41/amd/rocm.py | 189 +++++++++++++++ vllm/models/deepseek_v41/attention.py | 50 ++-- vllm/utils/multi_stream_utils.py | 11 +- 4 files changed, 496 insertions(+), 20 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 06897f04c0..278c0ef5d0 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -2075,3 +2075,269 @@ def test_fused_kv_insert_split(num_tokens: int, kv_block_size: int): # RoPE (last 64): stored as bf16. The kernel recomputes the rotation, so it # is bf16-close to the reference rather than bit-exact (cf. test_cutedsl). torch.testing.assert_close(recovered[:, NOPE_DIM:], ref[:, NOPE_DIM:]) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm stream coverage") +def test_v41_rocm_csa2_full_pipeline(): + """_forward_csa2_full: one fork, then K write and q-side after the join. + + The default chain must run first on the current stream, the compressor + chain on aux 0 and the indexer weights projection on aux 1; the join must + publish both cache writes before _produce_k / forward_q run. + """ + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention + from vllm.models.deepseek_v41.compressor import DeepseekCompressor + + torch.manual_seed(43) + num_tokens = 19 + hidden_states = torch.randn(num_tokens, 1024, device="cuda") + positions = torch.arange(7, 26, device="cuda") + qr_kv = torch.zeros(num_tokens, 1536, device="cuda") + qr = torch.zeros(num_tokens, 4, device="cuda") + kv = torch.zeros(num_tokens, 512, device="cuda") + q = torch.zeros(num_tokens, 512, device="cuda") + latent = torch.ones(num_tokens, 512, dtype=torch.bfloat16, device="cuda") + weights = torch.randn(num_tokens, 32, device="cuda") + index_q = torch.full((num_tokens, 32, 128), 2.0, device="cuda") + index_q_scale = torch.full((num_tokens, 32), 3.0, device="cuda") + index_weights_out = torch.full((num_tokens, 32), 4.0, device="cuda") + main_cache = torch.full((1, 128, 584), 165, dtype=torch.uint8, device="cuda") + index_cache = torch.full((1, 128, 132), 165, dtype=torch.uint8, device="cuda") + + order: list[str] = [] + streams: dict[str, torch.cuda.Stream] = {} + + compressor = DeepseekCompressor.__new__(DeepseekCompressor) + torch.nn.Module.__init__(compressor) + compressor.fused_wkv_wgate = SimpleNamespace( + weight=torch.randn(1024, 1024, device="cuda") + ) + + def compressor_forward(kv_score, positions_): + order.append("compressor") + streams["compressor"] = torch.cuda.current_stream() + return latent + + def compressor_insert(latent_, positions_, rotary_emb): + order.append("insert_cache") + streams["insert_cache"] = torch.cuda.current_stream() + main_cache.fill_(5) + + compressor.forward = compressor_forward + compressor.insert_cache = compressor_insert + + def weights_proj(hidden_states_): + order.append("weights_proj") + streams["weights_proj"] = torch.cuda.current_stream() + return weights, None + + def produce_k(latent_, positions_, rotary_emb): + order.append("produce_k") + streams["produce_k"] = torch.cuda.current_stream() + index_cache.fill_(7) + + def forward_q(qr_, qr_scale, indexer_weights, positions_, rotary_emb): + order.append("forward_q") + streams["forward_q"] = torch.cuda.current_stream() + return index_q, index_q_scale, index_weights_out + + indexer = SimpleNamespace( + weights_proj=weights_proj, + _produce_k=produce_k, + forward_q=forward_q, + ) + rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda")) + aux_streams = [torch.cuda.Stream(), torch.cuda.Stream()] + + def wq_b(qr_, qr_scale): + order.append("wq_b") + streams["wq_b"] = torch.cuda.current_stream() + return q + + def qnorm_insert(q_, kv_, positions_, attn_metadata): + order.append("qnorm_insert") + streams["qnorm_insert"] = torch.cuda.current_stream() + return q_ + + attention = SimpleNamespace( + compressor=compressor, + indexer=indexer, + aux_stream_list=aux_streams, + ln_events=[torch.cuda.Event() for _ in range(4)], + rotary_emb=rotary, + indexer_rotary_emb=rotary, + n_local_heads=1, + head_dim=512, + _fused_wqa_wkv_gemm=lambda hidden_states_: qr_kv, + _split_qkv_and_norm=lambda qr_kv_: (qr, None, kv), + _wq_b_proj=wq_b, + _fused_qnorm_rope_kv_insert=qnorm_insert, + ) + + context = ForwardContext({}, {}, {}) + with override_forward_context(context): + result = DeepseekV41ROCMAiterMLAAttention._forward_csa2_full( + attention, hidden_states, positions + ) + + q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result + assert torch.equal(q_out, q) + assert torch.equal(kv_out, kv) + assert torch.equal(index_q_out, index_q) + assert torch.equal(index_q_scale_out, index_q_scale) + assert torch.equal(index_weights_out_out, index_weights_out) + # Join publishes the aux-0 cache insert and orders the post-join K write. + assert torch.equal(main_cache, torch.full_like(main_cache, 5)) + assert torch.equal(index_cache, torch.full_like(index_cache, 7)) + + assert order == [ + "wq_b", + "qnorm_insert", + "compressor", + "insert_cache", + "weights_proj", + "produce_k", + "forward_q", + ] + current = torch.cuda.current_stream() + for name in ("wq_b", "qnorm_insert", "produce_k", "forward_q"): + assert streams[name] == current, name + for name in ("compressor", "insert_cache"): + assert streams[name] == aux_streams[0], name + assert streams["weights_proj"] == aux_streams[1] + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm stream coverage") +def test_v41_rocm_csa2_reindex_pipeline(): + """_forward_csa2_reindex: serial projections, then SWA q path vs q-side.""" + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention + + torch.manual_seed(45) + num_tokens = 19 + hidden_states = torch.randn(num_tokens, 1024, device="cuda") + positions = torch.arange(7, 26, device="cuda") + qr_kv = torch.zeros(num_tokens, 1536, device="cuda") + qr = torch.zeros(num_tokens, 4, device="cuda") + kv = torch.zeros(num_tokens, 512, device="cuda") + q = torch.zeros(num_tokens, 512, device="cuda") + indexer_weights = torch.randn(num_tokens, 32, device="cuda") + index_q = torch.full((num_tokens, 32, 128), 2.0, device="cuda") + index_q_scale = torch.full((num_tokens, 32), 3.0, device="cuda") + index_weights_out = torch.full((num_tokens, 32), 4.0, device="cuda") + + order: list[str] = [] + streams: dict[str, torch.cuda.Stream] = {} + + def forward_q(qr_, qr_scale, indexer_weights_, positions_, rotary_emb): + order.append("forward_q") + streams["forward_q"] = torch.cuda.current_stream() + return index_q, index_q_scale, index_weights_out + + indexer = SimpleNamespace(forward_q=forward_q) + rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda")) + aux0 = torch.cuda.Stream() + + def wq_b(qr_, qr_scale): + order.append("wq_b") + streams["wq_b"] = torch.cuda.current_stream() + return q + + def qnorm_insert(q_, kv_, positions_, attn_metadata): + order.append("qnorm_insert") + streams["qnorm_insert"] = torch.cuda.current_stream() + return q_ + + attention = SimpleNamespace( + indexer=indexer, + aux_stream_list=[aux0], + ln_events=[torch.cuda.Event() for _ in range(2)], + indexer_rotary_emb=rotary, + n_local_heads=1, + head_dim=512, + _run_parallel_input_projections=lambda hidden_states_: ( + qr_kv, + None, + indexer_weights, + ), + _split_qkv_and_norm=lambda qr_kv_: (qr, None, kv), + _wq_b_proj=wq_b, + _fused_qnorm_rope_kv_insert=qnorm_insert, + ) + + context = ForwardContext({}, {}, {}) + with override_forward_context(context): + result = DeepseekV41ROCMAiterMLAAttention._forward_csa2_reindex( + attention, hidden_states, positions + ) + + q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result + assert torch.equal(q_out, q) + assert torch.equal(kv_out, kv) + assert torch.equal(index_q_out, index_q) + assert torch.equal(index_q_scale_out, index_q_scale) + assert torch.equal(index_weights_out_out, index_weights_out) + + assert order == ["wq_b", "qnorm_insert", "forward_q"] + current = torch.cuda.current_stream() + for name in ("wq_b", "qnorm_insert"): + assert streams[name] == current, name + assert streams["forward_q"] == aux0 + + +@pytest.mark.parametrize("quant_tuple", [False, True]) +def test_v41_indexer_forward_q(quant_tuple): + """forward_q runs wq_b then the fused RoPE/quant and unpacks both shapes.""" + from unittest.mock import patch + + from vllm.models.deepseek_v41.attention import DeepseekV4Indexer + + num_tokens, n_head, head_dim = 5, 32, 128 + qr = torch.randn(num_tokens, 1280) + q = torch.randn(num_tokens, n_head * head_dim) + positions = torch.arange(num_tokens) + indexer_weights = torch.randn(num_tokens, n_head) + rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64)) + fake = SimpleNamespace( + n_head=n_head, + head_dim=head_dim, + softmax_scale=head_dim**-0.5, + use_fp4_kv=False, + indexer_weights_dtype=torch.float32, + _wq_b_proj=lambda qr_, qr_scale: q, + ) + q_quant = torch.randn(num_tokens, n_head, head_dim) + q_scale = torch.randn(num_tokens, n_head) + weights_out = torch.randn(num_tokens, n_head) + if quant_tuple: + fused_return = ((q_quant, q_scale), weights_out) + expected = (q_quant, q_scale, weights_out) + else: + fused_return = (q_quant, weights_out) + expected = (q_quant, None, weights_out) + + with patch( + "vllm.models.deepseek_v41.attention.fused_indexer_q_rope_quant", + return_value=fused_return, + ) as mock_fused: + result = DeepseekV4Indexer.forward_q( + fake, qr, None, indexer_weights, positions, rotary + ) + + assert torch.equal(result[0], expected[0]) + assert (result[1] is None) == (expected[1] is None) + if expected[1] is not None: + assert torch.equal(result[1], expected[1]) + assert torch.equal(result[2], expected[2]) + args = mock_fused.call_args.args + assert args[0] is positions + torch.testing.assert_close(args[1], q.view(-1, n_head, head_dim)) + assert args[2] is rotary.cos_sin_cache + assert args[3] is indexer_weights + assert args[4] == head_dim**-0.5 + assert args[5] == n_head**-0.5 + assert mock_fused.call_args.kwargs == { + "use_fp4": False, + "weights_out_dtype": torch.float32, + } diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index e1588b55e7..9de09bad8c 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -29,6 +29,7 @@ from vllm.models.deepseek_v41.sparse_mla import ( from vllm.platforms import current_platform from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton +from vllm.utils.multi_stream_utils import execute_in_parallel from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) @@ -529,6 +530,9 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): "this rank; PP splits inside a v4.1 index-sharing group " "are not supported." ) + if self.compressor is None and self.indexer is None: + # Dense layers have nothing to overlap; keep the base serial path. + self.aux_stream_list = None @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: @@ -600,8 +604,193 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): def _run_parallel_input_projections( self, hidden_states: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + if self.indexer is not None and self.compressor is None: + # Serial: qr must sit on the default stream before the q-side + # fork below can consume it (ROCm). + aux_streams = self.aux_stream_list + self.aux_stream_list = None + try: + return super()._run_parallel_input_projections(hidden_states) + finally: + self.aux_stream_list = aux_streams return super()._run_parallel_input_projections(hidden_states) + def _enable_multi_stream_overlap(self) -> bool: + """ROCm multi-stream gates: streams and capture region. + + Dict metadata marks piecewise cudagraph, whose eager breaks rebuild + the attention inputs on the owning stream. Forking side streams + there would rely on runtime HIP event sync, which is unreliable in + this overlap on ROCm (event waits can hang), so multi-stream only + runs where the fork/join becomes static graph edges: inside capture, + or with non-dict metadata (full cudagraph or the profile run), which + has no eager breaks. + """ + attn_metadata = get_forward_context().attn_metadata + return self.aux_stream_list is not None and ( + torch.cuda.is_current_stream_capturing() + or not isinstance(attn_metadata, dict) + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + if ( + self.indexer is None + or self._prepare_and_attn_fn == self._prepare_and_attn_eager + or not self._enable_multi_stream_overlap() + ): + # Sequential fallback: no forks outside the capture region, where + # HIP event sync is unreliable; MRV1 keeps the input prep in its + # wide eager region (#51430). + aux_streams = self.aux_stream_list + self.aux_stream_list = None + try: + return super().forward(positions, hidden_states, llama_4_scaling) + finally: + self.aux_stream_list = aux_streams + + attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states) + if self.compressor is not None: + q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_full( + hidden_states, positions + ) + else: + q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_reindex( + hidden_states, positions + ) + self._sparse_indexer_and_attn( + hidden_states, + index_q, + index_q_scale, + index_weights, + q, + kv, + positions, + attn_out, + ) + return self._o_proj(attn_out, positions) + + def _forward_csa2_full( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + """CSA layers: one fork replaces the base pipeline's three. + + Each fork/join is a HIP event pair per side stream on ROCm, so the + base three-fork pipeline is expensive there; the compressor chain and + the indexer weights projection run fully in parallel with the default + chain. The indexer K write and q-side stay serial after the join: + ``wk`` consumes the aux-0 latent and the q-side's wq_b consumes the + default-stream qr. + """ + attn_metadata = get_forward_context().attn_metadata + compressor = self.compressor + indexer = self.indexer + assert compressor is not None and indexer is not None + aux_streams = self.aux_stream_list + assert aux_streams is not None and len(aux_streams) >= 2 + + def default_chain(): + qr_kv = self._fused_wqa_wkv_gemm(hidden_states) + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) + q = self._wq_b_proj(qr, qr_scale).view( + -1, self.n_local_heads, self.head_dim + ) + return ( + self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata), + qr, + qr_scale, + kv, + ) + + def compressor_chain(): + kv_score = torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + latent = compressor(kv_score, positions) + compressor.insert_cache(latent, positions, self.rotary_emb) + return latent + + def indexer_weights_chain(): + # ReplicatedLinear returns (output, bias); bias is None. + weights, _ = indexer.weights_proj(hidden_states) + return weights + + (q, qr, qr_scale, kv), (latent, indexer_weights) = execute_in_parallel( + default_chain, + [compressor_chain, indexer_weights_chain], + self.ln_events[0], + self.ln_events[1:3], + aux_streams[:2], + enable=True, + default_first=True, + ) + + indexer._produce_k(latent, positions, self.indexer_rotary_emb) + index_q, index_q_scale, index_weights_out = indexer.forward_q( + qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb + ) + return q, kv, index_q, index_q_scale, index_weights_out + + def _forward_csa2_reindex( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + """Indexer-only layers: the q-side overlaps the SWA q path. + + The input projections run serial first: the q-side consumes qr, so + it can only overlap the SWA q projection and KV insert. + """ + indexer = self.indexer + assert indexer is not None + attn_metadata = get_forward_context().attn_metadata + # Serial on these layers via the _run_parallel_input_projections + # override: qr must sit on the default stream before the fork. + qr_kv, _, indexer_weights = self._run_parallel_input_projections(hidden_states) + assert indexer_weights is not None + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) + + def swa_q_chain(): + q = self._wq_b_proj(qr, qr_scale).view( + -1, self.n_local_heads, self.head_dim + ) + return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + def indexer_q_chain(): + return indexer.forward_q( + qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb + ) + + aux_streams = self.aux_stream_list + assert aux_streams is not None + q, aux_results = execute_in_parallel( + swa_q_chain, + [indexer_q_chain], + self.ln_events[0], + self.ln_events[1:2], + aux_streams[:1], + enable=True, + default_first=True, + ) + index_q, index_q_scale, index_weights_out = aux_results[0] + return q, kv, index_q, index_q_scale, index_weights_out + @functools.cached_property def _wq_b_uses_aiter_block_scaled(self) -> bool: """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel. diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 6d961f0055..c27a8d9baf 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -478,7 +478,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # graph and MRV1 produces garbage (#51430). self._prepare_and_attn_fn = self._prepare_and_attn_eager - # Will be None on ROCm for now. self.aux_stream_list = aux_stream_list # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins @@ -1381,6 +1380,37 @@ class DeepseekV4Indexer(nn.Module): self.use_fp4_kv, ) + def forward_q( + self, + qr: torch.Tensor | QuantizedActivation, + qr_scale: torch.Tensor | None, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + rotary_emb: nn.Module, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Build the indexer queries: wq_b over qr plus fused RoPE/quant. + + Split out so the ROCm layer can schedule it apart from + ``_produce_k`` (e.g. on an aux stream once ``qr`` is ready). + """ + q = self._wq_b_proj(qr, qr_scale) + q = q.view(-1, self.n_head, self.head_dim) + q_quant, weights = fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + use_fp4=self.use_fp4_kv, + weights_out_dtype=self.indexer_weights_dtype, + ) + if isinstance(q_quant, tuple): + q, q_scale = q_quant + else: + q, q_scale = q_quant, None + return q, q_scale, weights + def forward( self, qr: torch.Tensor | QuantizedActivation, @@ -1422,23 +1452,7 @@ class DeepseekV4Indexer(nn.Module): # (skip_k_cache_insert=True). self._produce_k(latent, positions, rotary_emb) - q = self._wq_b_proj(qr, qr_scale) - q = q.view(-1, self.n_head, self.head_dim) - q_quant, weights = fused_indexer_q_rope_quant( - positions, - q, - rotary_emb.cos_sin_cache, - indexer_weights, - self.softmax_scale, - self.n_head**-0.5, - use_fp4=self.use_fp4_kv, - weights_out_dtype=self.indexer_weights_dtype, - ) - if isinstance(q_quant, tuple): - q, q_scale = q_quant - else: - q, q_scale = q_quant, None - return q, q_scale, weights + return self.forward_q(qr, qr_scale, indexer_weights, positions, rotary_emb) def _wq_b_proj( self, diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py index 6190b76991..5cb7f00f76 100644 --- a/vllm/utils/multi_stream_utils.py +++ b/vllm/utils/multi_stream_utils.py @@ -70,6 +70,7 @@ def execute_in_parallel( done_events: list[torch.cuda.Event], aux_streams: list[torch.cuda.Stream] | None = None, enable: bool = False, + default_first: bool = False, ) -> tuple[Any, list[Any]]: """Run default_fn on the current stream and aux_fns concurrently on aux_streams. @@ -97,6 +98,10 @@ def execute_in_parallel( so callers that pass aux_streams must also pass enable=True (typically gated by an env var) to actually overlap. When False, execution falls back to sequential on the current stream. + default_first: When True, enqueue the default chain before the aux + chains. The default chain is typically the longest, so it starts + first; the aux streams still only wait on start_event, so they + overlap the default chain. Defaults to False. Returns: Tuple of (default_result, aux_results) where aux_results[i] is the @@ -117,6 +122,8 @@ def execute_in_parallel( pending: list[torch.cuda.Event] = [] start_event.record() + if default_first: + default_result = default_fn() for i, fn in enumerate(aux_fns): if fn is None: continue @@ -125,8 +132,8 @@ def execute_in_parallel( aux_results[i] = fn() done_events[i].record() pending.append(done_events[i]) - - default_result = default_fn() + if not default_first: + default_result = default_fn() for ev in pending: ev.wait()