@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fdcd42350b

[ROCm][Perf] Enable HCA dual-stream overlap for DeepSeek-V4 (#56853)

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

Shanshan Shen committed Sep 17, 2026 at 14:46 UTC fdcd42350bd2475abffc01bb73eec5d7c7a1e549

1 file changed +35 -9

vllm/models/deepseek_v4/amd/rocm.py

+35 -9

@@ -558,14 +558,17 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

558 self._fused_compressor_split_sizes: tuple[int, int] | None = None

559

560 if self.indexer is None:

561 - # Only enable multi-stream overlap for CSA layer now.

562 - self.aux_stream_list = None

561 + # Dense layers have no compressor work to overlap; HCA layers

562 + # (compressor, no indexer) keep the streams for the dual-stream

563 + # fork below.

564 + if self.compressor is None:

565 + self.aux_stream_list = None

566 else:

567 # Disable indexer inner overlap.

568 self.indexer.aux_stream = None

569

567 - def _enable_csa_multi_stream(self) -> bool:

568 - """All CSA multi-stream gates: env var, streams, and capture region.

570 + def _enable_multi_stream_overlap(self) -> bool:

571 + """ROCm multi-stream gates: streams and capture region.

572

573 Dict metadata marks piecewise cudagraph, whose eager breaks rebuild

574 the attention inputs on the owning stream. Forking side streams

@@ -573,7 +576,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

576 this overlap on ROCm (event waits can hang), so multi-stream only

577 runs where the fork/join becomes static graph edges: inside capture,

578 or with non-dict metadata (full cudagraph or the profile run), which

576 - has no eager breaks.

579 + has no eager breaks. Covers both the HCA and CSA forks.

580 """

581 attn_metadata = get_forward_context().attn_metadata

582 return self.aux_stream_list is not None and (

@@ -624,7 +627,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

627 device=hidden_states.device,

628 )

629

627 - if self._enable_csa_multi_stream():

630 + if self._enable_multi_stream_overlap():

631 # The ROCm override consumes these sentinels inside the capture

632 # boundary, moving the stream fan-out ahead of the projections.

633 self._prepare_and_attn_fn(

@@ -658,7 +661,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

661 positions: torch.Tensor,

662 o_padded: torch.Tensor,

663 ) -> None:

661 - """Run the ROCm CSA fork/join inside the graph capture boundary."""

664 + """Run the ROCm fork/join (HCA or CSA) inside the capture boundary."""

665 aux_streams = self.aux_stream_list

666 # The sequential pipeline disables aux_stream_list before calling

667 # back with real projection inputs; aux_streams is None ends that

@@ -684,13 +687,13 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

687

688 # Re-check: forward's gate ran inside a captured segment that

689 # _prepare_and_attn_eager (MRV1) then broke, making this region eager.

687 - if not self._enable_csa_multi_stream():

690 + if not self._enable_multi_stream_overlap():

691 self._run_sequential_pipeline(hidden_states, positions, o_padded)

692 return

693

694 indexer = self.indexer

695 compressor = self.compressor

693 - assert indexer is not None and compressor is not None

696 + assert compressor is not None

697

698 def default_chain():

699 qr_kv = self._fused_wqa_wkv_gemm(hidden_states)

@@ -710,6 +713,24 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

713 )

714 compressor(score, positions, self.rotary_emb)

715

716 + if indexer is None:

717 + # HCA dual-stream: the main compressor runs on aux stream 0 while

718 + # the default stream produces q and inserts KV into the SWA cache.

719 + # Both branches only read hidden_states, so the join merely has to

720 + # precede the sparse attention that consumes the compressed KV.

721 + (q, _qr_out, _qr_scale_out, kv_out), _ = execute_in_parallel(

722 + default_chain,

723 + [main_compressor_chain],

724 + self.ln_events[0],

725 + [self.ln_events[1]],

726 + aux_streams[:1],

727 + enable=True,

728 + )

729 + self._sparse_indexer_and_attn(

730 + hidden_states, None, None, None, q, kv_out, positions, o_padded

731 + )

732 + return

733 +

734 def indexer_compressor_chain() -> None:

735 score = torch.mm(

736 hidden_states,

@@ -718,6 +739,11 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

739 )

740 indexer.compressor(score, positions, self.indexer_rotary_emb)

741

742 + # CSA three-stream: the main and indexer compressors run on aux

743 + # streams 0 and 1 while the default stream produces q and inserts KV

744 + # into the SWA cache. Every branch only reads hidden_states plus its

745 + # own state, so the join merely has to precede the indexer op and the

746 + # sparse attention, which consume the compressed KV caches.

747 (q, qr_out, qr_scale_out, kv_out), _ = execute_in_parallel(

748 default_chain,

749 [main_compressor_chain, indexer_compressor_chain],