@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 44af287ebe

[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>

Shanshan Shen committed Sep 27, 2026 at 23:48 UTC 44af287ebe38d6dc4e102948025f5e3e175aefd6

4 files changed +496 -20

tests/kernels/test_compressor_kv_cache.py

+266

@@ -2075,3 +2075,269 @@ def test_fused_kv_insert_split(num_tokens: int, kv_block_size: int):

2075 # RoPE (last 64): stored as bf16. The kernel recomputes the rotation, so it

2076 # is bf16-close to the reference rather than bit-exact (cf. test_cutedsl).

2077 torch.testing.assert_close(recovered[:, NOPE_DIM:], ref[:, NOPE_DIM:])

2078 +

2079 +

2080 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm stream coverage")

2081 +def test_v41_rocm_csa2_full_pipeline():

2082 + """_forward_csa2_full: one fork, then K write and q-side after the join.

2083 +

2084 + The default chain must run first on the current stream, the compressor

2085 + chain on aux 0 and the indexer weights projection on aux 1; the join must

2086 + publish both cache writes before _produce_k / forward_q run.

2087 + """

2088 + from vllm.forward_context import ForwardContext, override_forward_context

2089 + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention

2090 + from vllm.models.deepseek_v41.compressor import DeepseekCompressor

2091 +

2092 + torch.manual_seed(43)

2093 + num_tokens = 19

2094 + hidden_states = torch.randn(num_tokens, 1024, device="cuda")

2095 + positions = torch.arange(7, 26, device="cuda")

2096 + qr_kv = torch.zeros(num_tokens, 1536, device="cuda")

2097 + qr = torch.zeros(num_tokens, 4, device="cuda")

2098 + kv = torch.zeros(num_tokens, 512, device="cuda")

2099 + q = torch.zeros(num_tokens, 512, device="cuda")

2100 + latent = torch.ones(num_tokens, 512, dtype=torch.bfloat16, device="cuda")

2101 + weights = torch.randn(num_tokens, 32, device="cuda")

2102 + index_q = torch.full((num_tokens, 32, 128), 2.0, device="cuda")

2103 + index_q_scale = torch.full((num_tokens, 32), 3.0, device="cuda")

2104 + index_weights_out = torch.full((num_tokens, 32), 4.0, device="cuda")

2105 + main_cache = torch.full((1, 128, 584), 165, dtype=torch.uint8, device="cuda")

2106 + index_cache = torch.full((1, 128, 132), 165, dtype=torch.uint8, device="cuda")

2107 +

2108 + order: list[str] = []

2109 + streams: dict[str, torch.cuda.Stream] = {}

2110 +

2111 + compressor = DeepseekCompressor.__new__(DeepseekCompressor)

2112 + torch.nn.Module.__init__(compressor)

2113 + compressor.fused_wkv_wgate = SimpleNamespace(

2114 + weight=torch.randn(1024, 1024, device="cuda")

2115 + )

2116 +

2117 + def compressor_forward(kv_score, positions_):

2118 + order.append("compressor")

2119 + streams["compressor"] = torch.cuda.current_stream()

2120 + return latent

2121 +

2122 + def compressor_insert(latent_, positions_, rotary_emb):

2123 + order.append("insert_cache")

2124 + streams["insert_cache"] = torch.cuda.current_stream()

2125 + main_cache.fill_(5)

2126 +

2127 + compressor.forward = compressor_forward

2128 + compressor.insert_cache = compressor_insert

2129 +

2130 + def weights_proj(hidden_states_):

2131 + order.append("weights_proj")

2132 + streams["weights_proj"] = torch.cuda.current_stream()

2133 + return weights, None

2134 +

2135 + def produce_k(latent_, positions_, rotary_emb):

2136 + order.append("produce_k")

2137 + streams["produce_k"] = torch.cuda.current_stream()

2138 + index_cache.fill_(7)

2139 +

2140 + def forward_q(qr_, qr_scale, indexer_weights, positions_, rotary_emb):

2141 + order.append("forward_q")

2142 + streams["forward_q"] = torch.cuda.current_stream()

2143 + return index_q, index_q_scale, index_weights_out

2144 +

2145 + indexer = SimpleNamespace(

2146 + weights_proj=weights_proj,

2147 + _produce_k=produce_k,

2148 + forward_q=forward_q,

2149 + )

2150 + rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda"))

2151 + aux_streams = [torch.cuda.Stream(), torch.cuda.Stream()]

2152 +

2153 + def wq_b(qr_, qr_scale):

2154 + order.append("wq_b")

2155 + streams["wq_b"] = torch.cuda.current_stream()

2156 + return q

2157 +

2158 + def qnorm_insert(q_, kv_, positions_, attn_metadata):

2159 + order.append("qnorm_insert")

2160 + streams["qnorm_insert"] = torch.cuda.current_stream()

2161 + return q_

2162 +

2163 + attention = SimpleNamespace(

2164 + compressor=compressor,

2165 + indexer=indexer,

2166 + aux_stream_list=aux_streams,

2167 + ln_events=[torch.cuda.Event() for _ in range(4)],

2168 + rotary_emb=rotary,

2169 + indexer_rotary_emb=rotary,

2170 + n_local_heads=1,

2171 + head_dim=512,

2172 + _fused_wqa_wkv_gemm=lambda hidden_states_: qr_kv,

2173 + _split_qkv_and_norm=lambda qr_kv_: (qr, None, kv),

2174 + _wq_b_proj=wq_b,

2175 + _fused_qnorm_rope_kv_insert=qnorm_insert,

2176 + )

2177 +

2178 + context = ForwardContext({}, {}, {})

2179 + with override_forward_context(context):

2180 + result = DeepseekV41ROCMAiterMLAAttention._forward_csa2_full(

2181 + attention, hidden_states, positions

2182 + )

2183 +

2184 + q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result

2185 + assert torch.equal(q_out, q)

2186 + assert torch.equal(kv_out, kv)

2187 + assert torch.equal(index_q_out, index_q)

2188 + assert torch.equal(index_q_scale_out, index_q_scale)

2189 + assert torch.equal(index_weights_out_out, index_weights_out)

2190 + # Join publishes the aux-0 cache insert and orders the post-join K write.

2191 + assert torch.equal(main_cache, torch.full_like(main_cache, 5))

2192 + assert torch.equal(index_cache, torch.full_like(index_cache, 7))

2193 +

2194 + assert order == [

2195 + "wq_b",

2196 + "qnorm_insert",

2197 + "compressor",

2198 + "insert_cache",

2199 + "weights_proj",

2200 + "produce_k",

2201 + "forward_q",

2202 + ]

2203 + current = torch.cuda.current_stream()

2204 + for name in ("wq_b", "qnorm_insert", "produce_k", "forward_q"):

2205 + assert streams[name] == current, name

2206 + for name in ("compressor", "insert_cache"):

2207 + assert streams[name] == aux_streams[0], name

2208 + assert streams["weights_proj"] == aux_streams[1]

2209 +

2210 +

2211 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm stream coverage")

2212 +def test_v41_rocm_csa2_reindex_pipeline():

2213 + """_forward_csa2_reindex: serial projections, then SWA q path vs q-side."""

2214 + from vllm.forward_context import ForwardContext, override_forward_context

2215 + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention

2216 +

2217 + torch.manual_seed(45)

2218 + num_tokens = 19

2219 + hidden_states = torch.randn(num_tokens, 1024, device="cuda")

2220 + positions = torch.arange(7, 26, device="cuda")

2221 + qr_kv = torch.zeros(num_tokens, 1536, device="cuda")

2222 + qr = torch.zeros(num_tokens, 4, device="cuda")

2223 + kv = torch.zeros(num_tokens, 512, device="cuda")

2224 + q = torch.zeros(num_tokens, 512, device="cuda")

2225 + indexer_weights = torch.randn(num_tokens, 32, device="cuda")

2226 + index_q = torch.full((num_tokens, 32, 128), 2.0, device="cuda")

2227 + index_q_scale = torch.full((num_tokens, 32), 3.0, device="cuda")

2228 + index_weights_out = torch.full((num_tokens, 32), 4.0, device="cuda")

2229 +

2230 + order: list[str] = []

2231 + streams: dict[str, torch.cuda.Stream] = {}

2232 +

2233 + def forward_q(qr_, qr_scale, indexer_weights_, positions_, rotary_emb):

2234 + order.append("forward_q")

2235 + streams["forward_q"] = torch.cuda.current_stream()

2236 + return index_q, index_q_scale, index_weights_out

2237 +

2238 + indexer = SimpleNamespace(forward_q=forward_q)

2239 + rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda"))

2240 + aux0 = torch.cuda.Stream()

2241 +

2242 + def wq_b(qr_, qr_scale):

2243 + order.append("wq_b")

2244 + streams["wq_b"] = torch.cuda.current_stream()

2245 + return q

2246 +

2247 + def qnorm_insert(q_, kv_, positions_, attn_metadata):

2248 + order.append("qnorm_insert")

2249 + streams["qnorm_insert"] = torch.cuda.current_stream()

2250 + return q_

2251 +

2252 + attention = SimpleNamespace(

2253 + indexer=indexer,

2254 + aux_stream_list=[aux0],

2255 + ln_events=[torch.cuda.Event() for _ in range(2)],

2256 + indexer_rotary_emb=rotary,

2257 + n_local_heads=1,

2258 + head_dim=512,

2259 + _run_parallel_input_projections=lambda hidden_states_: (

2260 + qr_kv,

2261 + None,

2262 + indexer_weights,

2263 + ),

2264 + _split_qkv_and_norm=lambda qr_kv_: (qr, None, kv),

2265 + _wq_b_proj=wq_b,

2266 + _fused_qnorm_rope_kv_insert=qnorm_insert,

2267 + )

2268 +

2269 + context = ForwardContext({}, {}, {})

2270 + with override_forward_context(context):

2271 + result = DeepseekV41ROCMAiterMLAAttention._forward_csa2_reindex(

2272 + attention, hidden_states, positions

2273 + )

2274 +

2275 + q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result

2276 + assert torch.equal(q_out, q)

2277 + assert torch.equal(kv_out, kv)

2278 + assert torch.equal(index_q_out, index_q)

2279 + assert torch.equal(index_q_scale_out, index_q_scale)

2280 + assert torch.equal(index_weights_out_out, index_weights_out)

2281 +

2282 + assert order == ["wq_b", "qnorm_insert", "forward_q"]

2283 + current = torch.cuda.current_stream()

2284 + for name in ("wq_b", "qnorm_insert"):

2285 + assert streams[name] == current, name

2286 + assert streams["forward_q"] == aux0

2287 +

2288 +

2289 +@pytest.mark.parametrize("quant_tuple", [False, True])

2290 +def test_v41_indexer_forward_q(quant_tuple):

2291 + """forward_q runs wq_b then the fused RoPE/quant and unpacks both shapes."""

2292 + from unittest.mock import patch

2293 +

2294 + from vllm.models.deepseek_v41.attention import DeepseekV4Indexer

2295 +

2296 + num_tokens, n_head, head_dim = 5, 32, 128

2297 + qr = torch.randn(num_tokens, 1280)

2298 + q = torch.randn(num_tokens, n_head * head_dim)

2299 + positions = torch.arange(num_tokens)

2300 + indexer_weights = torch.randn(num_tokens, n_head)

2301 + rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64))

2302 + fake = SimpleNamespace(

2303 + n_head=n_head,

2304 + head_dim=head_dim,

2305 + softmax_scale=head_dim**-0.5,

2306 + use_fp4_kv=False,

2307 + indexer_weights_dtype=torch.float32,

2308 + _wq_b_proj=lambda qr_, qr_scale: q,

2309 + )

2310 + q_quant = torch.randn(num_tokens, n_head, head_dim)

2311 + q_scale = torch.randn(num_tokens, n_head)

2312 + weights_out = torch.randn(num_tokens, n_head)

2313 + if quant_tuple:

2314 + fused_return = ((q_quant, q_scale), weights_out)

2315 + expected = (q_quant, q_scale, weights_out)

2316 + else:

2317 + fused_return = (q_quant, weights_out)

2318 + expected = (q_quant, None, weights_out)

2319 +

2320 + with patch(

2321 + "vllm.models.deepseek_v41.attention.fused_indexer_q_rope_quant",

2322 + return_value=fused_return,

2323 + ) as mock_fused:

2324 + result = DeepseekV4Indexer.forward_q(

2325 + fake, qr, None, indexer_weights, positions, rotary

2326 + )

2327 +

2328 + assert torch.equal(result[0], expected[0])

2329 + assert (result[1] is None) == (expected[1] is None)

2330 + if expected[1] is not None:

2331 + assert torch.equal(result[1], expected[1])

2332 + assert torch.equal(result[2], expected[2])

2333 + args = mock_fused.call_args.args

2334 + assert args[0] is positions

2335 + torch.testing.assert_close(args[1], q.view(-1, n_head, head_dim))

2336 + assert args[2] is rotary.cos_sin_cache

2337 + assert args[3] is indexer_weights

2338 + assert args[4] == head_dim**-0.5

2339 + assert args[5] == n_head**-0.5

2340 + assert mock_fused.call_args.kwargs == {

2341 + "use_fp4": False,

2342 + "weights_out_dtype": torch.float32,

2343 + }

vllm/models/deepseek_v41/amd/rocm.py

+189

@@ -29,6 +29,7 @@ from vllm.models.deepseek_v41.sparse_mla import (

29 from vllm.platforms import current_platform

30 from vllm.platforms.rocm import _ON_GFX950

31 from vllm.triton_utils import tl, triton

32 +from vllm.utils.multi_stream_utils import execute_in_parallel

33 from vllm.v1.attention.backend import (

34 CommonAttentionMetadata,

35 )

@@ -529,6 +530,9 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):

530 "this rank; PP splits inside a v4.1 index-sharing group "

531 "are not supported."

532 )

533 + if self.compressor is None and self.indexer is None:

534 + # Dense layers have nothing to overlap; keep the base serial path.

535 + self.aux_stream_list = None

536

537 @classmethod

538 def get_padded_num_q_heads(cls, num_heads: int) -> int:

@@ -600,8 +604,193 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):

604 def _run_parallel_input_projections(

605 self, hidden_states: torch.Tensor

606 ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:

607 + if self.indexer is not None and self.compressor is None:

608 + # Serial: qr must sit on the default stream before the q-side

609 + # fork below can consume it (ROCm).

610 + aux_streams = self.aux_stream_list

611 + self.aux_stream_list = None

612 + try:

613 + return super()._run_parallel_input_projections(hidden_states)

614 + finally:

615 + self.aux_stream_list = aux_streams

616 return super()._run_parallel_input_projections(hidden_states)

617

618 + def _enable_multi_stream_overlap(self) -> bool:

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

620 +

621 + Dict metadata marks piecewise cudagraph, whose eager breaks rebuild

622 + the attention inputs on the owning stream. Forking side streams

623 + there would rely on runtime HIP event sync, which is unreliable in

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

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

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

627 + has no eager breaks.

628 + """

629 + attn_metadata = get_forward_context().attn_metadata

630 + return self.aux_stream_list is not None and (

631 + torch.cuda.is_current_stream_capturing()

632 + or not isinstance(attn_metadata, dict)

633 + )

634 +

635 + def forward(

636 + self,

637 + positions: torch.Tensor,

638 + hidden_states: torch.Tensor,

639 + llama_4_scaling: torch.Tensor | None = None,

640 + ) -> torch.Tensor:

641 + if (

642 + self.indexer is None

643 + or self._prepare_and_attn_fn == self._prepare_and_attn_eager

644 + or not self._enable_multi_stream_overlap()

645 + ):

646 + # Sequential fallback: no forks outside the capture region, where

647 + # HIP event sync is unreliable; MRV1 keeps the input prep in its

648 + # wide eager region (#51430).

649 + aux_streams = self.aux_stream_list

650 + self.aux_stream_list = None

651 + try:

652 + return super().forward(positions, hidden_states, llama_4_scaling)

653 + finally:

654 + self.aux_stream_list = aux_streams

655 +

656 + attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states)

657 + if self.compressor is not None:

658 + q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_full(

659 + hidden_states, positions

660 + )

661 + else:

662 + q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_reindex(

663 + hidden_states, positions

664 + )

665 + self._sparse_indexer_and_attn(

666 + hidden_states,

667 + index_q,

668 + index_q_scale,

669 + index_weights,

670 + q,

671 + kv,

672 + positions,

673 + attn_out,

674 + )

675 + return self._o_proj(attn_out, positions)

676 +

677 + def _forward_csa2_full(

678 + self, hidden_states: torch.Tensor, positions: torch.Tensor

679 + ) -> tuple[

680 + torch.Tensor,

681 + torch.Tensor,

682 + torch.Tensor | None,

683 + torch.Tensor | None,

684 + torch.Tensor | None,

685 + ]:

686 + """CSA layers: one fork replaces the base pipeline's three.

687 +

688 + Each fork/join is a HIP event pair per side stream on ROCm, so the

689 + base three-fork pipeline is expensive there; the compressor chain and

690 + the indexer weights projection run fully in parallel with the default

691 + chain. The indexer K write and q-side stay serial after the join:

692 + ``wk`` consumes the aux-0 latent and the q-side's wq_b consumes the

693 + default-stream qr.

694 + """

695 + attn_metadata = get_forward_context().attn_metadata

696 + compressor = self.compressor

697 + indexer = self.indexer

698 + assert compressor is not None and indexer is not None

699 + aux_streams = self.aux_stream_list

700 + assert aux_streams is not None and len(aux_streams) >= 2

701 +

702 + def default_chain():

703 + qr_kv = self._fused_wqa_wkv_gemm(hidden_states)

704 + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)

705 + q = self._wq_b_proj(qr, qr_scale).view(

706 + -1, self.n_local_heads, self.head_dim

707 + )

708 + return (

709 + self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata),

710 + qr,

711 + qr_scale,

712 + kv,

713 + )

714 +

715 + def compressor_chain():

716 + kv_score = torch.mm(

717 + hidden_states,

718 + compressor.fused_wkv_wgate.weight.T,

719 + out_dtype=torch.float32,

720 + )

721 + latent = compressor(kv_score, positions)

722 + compressor.insert_cache(latent, positions, self.rotary_emb)

723 + return latent

724 +

725 + def indexer_weights_chain():

726 + # ReplicatedLinear returns (output, bias); bias is None.

727 + weights, _ = indexer.weights_proj(hidden_states)

728 + return weights

729 +

730 + (q, qr, qr_scale, kv), (latent, indexer_weights) = execute_in_parallel(

731 + default_chain,

732 + [compressor_chain, indexer_weights_chain],

733 + self.ln_events[0],

734 + self.ln_events[1:3],

735 + aux_streams[:2],

736 + enable=True,

737 + default_first=True,

738 + )

739 +

740 + indexer._produce_k(latent, positions, self.indexer_rotary_emb)

741 + index_q, index_q_scale, index_weights_out = indexer.forward_q(

742 + qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb

743 + )

744 + return q, kv, index_q, index_q_scale, index_weights_out

745 +

746 + def _forward_csa2_reindex(

747 + self, hidden_states: torch.Tensor, positions: torch.Tensor

748 + ) -> tuple[

749 + torch.Tensor,

750 + torch.Tensor,

751 + torch.Tensor | None,

752 + torch.Tensor | None,

753 + torch.Tensor | None,

754 + ]:

755 + """Indexer-only layers: the q-side overlaps the SWA q path.

756 +

757 + The input projections run serial first: the q-side consumes qr, so

758 + it can only overlap the SWA q projection and KV insert.

759 + """

760 + indexer = self.indexer

761 + assert indexer is not None

762 + attn_metadata = get_forward_context().attn_metadata

763 + # Serial on these layers via the _run_parallel_input_projections

764 + # override: qr must sit on the default stream before the fork.

765 + qr_kv, _, indexer_weights = self._run_parallel_input_projections(hidden_states)

766 + assert indexer_weights is not None

767 + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)

768 +

769 + def swa_q_chain():

770 + q = self._wq_b_proj(qr, qr_scale).view(

771 + -1, self.n_local_heads, self.head_dim

772 + )

773 + return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)

774 +

775 + def indexer_q_chain():

776 + return indexer.forward_q(

777 + qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb

778 + )

779 +

780 + aux_streams = self.aux_stream_list

781 + assert aux_streams is not None

782 + q, aux_results = execute_in_parallel(

783 + swa_q_chain,

784 + [indexer_q_chain],

785 + self.ln_events[0],

786 + self.ln_events[1:2],

787 + aux_streams[:1],

788 + enable=True,

789 + default_first=True,

790 + )

791 + index_q, index_q_scale, index_weights_out = aux_results[0]

792 + return q, kv, index_q, index_q_scale, index_weights_out

793 +

794 @functools.cached_property

795 def _wq_b_uses_aiter_block_scaled(self) -> bool:

796 """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.

vllm/models/deepseek_v41/attention.py

+32 -18

@@ -478,7 +478,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

478 # graph and MRV1 produces garbage (#51430).

479 self._prepare_and_attn_fn = self._prepare_and_attn_eager

480

481 - # Will be None on ROCm for now.

481 self.aux_stream_list = aux_stream_list

482 # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;

483 # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins

@@ -1381,6 +1380,37 @@ class DeepseekV4Indexer(nn.Module):

1380 self.use_fp4_kv,

1381 )

1382

1383 + def forward_q(

1384 + self,

1385 + qr: torch.Tensor | QuantizedActivation,

1386 + qr_scale: torch.Tensor | None,

1387 + indexer_weights: torch.Tensor,

1388 + positions: torch.Tensor,

1389 + rotary_emb: nn.Module,

1390 + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:

1391 + """Build the indexer queries: wq_b over qr plus fused RoPE/quant.

1392 +

1393 + Split out so the ROCm layer can schedule it apart from

1394 + ``_produce_k`` (e.g. on an aux stream once ``qr`` is ready).

1395 + """

1396 + q = self._wq_b_proj(qr, qr_scale)

1397 + q = q.view(-1, self.n_head, self.head_dim)

1398 + q_quant, weights = fused_indexer_q_rope_quant(

1399 + positions,

1400 + q,

1401 + rotary_emb.cos_sin_cache,

1402 + indexer_weights,

1403 + self.softmax_scale,

1404 + self.n_head**-0.5,

1405 + use_fp4=self.use_fp4_kv,

1406 + weights_out_dtype=self.indexer_weights_dtype,

1407 + )

1408 + if isinstance(q_quant, tuple):

1409 + q, q_scale = q_quant

1410 + else:

1411 + q, q_scale = q_quant, None

1412 + return q, q_scale, weights

1413 +

1414 def forward(

1415 self,

1416 qr: torch.Tensor | QuantizedActivation,

@@ -1422,23 +1452,7 @@ class DeepseekV4Indexer(nn.Module):

1452 # (skip_k_cache_insert=True).

1453 self._produce_k(latent, positions, rotary_emb)

1454

1425 - q = self._wq_b_proj(qr, qr_scale)

1426 - q = q.view(-1, self.n_head, self.head_dim)

1427 - q_quant, weights = fused_indexer_q_rope_quant(

1428 - positions,

1429 - q,

1430 - rotary_emb.cos_sin_cache,

1431 - indexer_weights,

1432 - self.softmax_scale,

1433 - self.n_head**-0.5,

1434 - use_fp4=self.use_fp4_kv,

1435 - weights_out_dtype=self.indexer_weights_dtype,

1436 - )

1437 - if isinstance(q_quant, tuple):

1438 - q, q_scale = q_quant

1439 - else:

1440 - q, q_scale = q_quant, None

1441 - return q, q_scale, weights

1455 + return self.forward_q(qr, qr_scale, indexer_weights, positions, rotary_emb)

1456

1457 def _wq_b_proj(

1458 self,

vllm/utils/multi_stream_utils.py

+9 -2

@@ -70,6 +70,7 @@ def execute_in_parallel(

70 done_events: list[torch.cuda.Event],

71 aux_streams: list[torch.cuda.Stream] | None = None,

72 enable: bool = False,

73 + default_first: bool = False,

74 ) -> tuple[Any, list[Any]]:

75 """Run default_fn on the current stream and aux_fns concurrently on

76 aux_streams.

@@ -97,6 +98,10 @@ def execute_in_parallel(

98 so callers that pass aux_streams must also pass enable=True

99 (typically gated by an env var) to actually overlap. When False,

100 execution falls back to sequential on the current stream.

101 + default_first: When True, enqueue the default chain before the aux

102 + chains. The default chain is typically the longest, so it starts

103 + first; the aux streams still only wait on start_event, so they

104 + overlap the default chain. Defaults to False.

105

106 Returns:

107 Tuple of (default_result, aux_results) where aux_results[i] is the

@@ -117,6 +122,8 @@ def execute_in_parallel(

122 pending: list[torch.cuda.Event] = []

123

124 start_event.record()

125 + if default_first:

126 + default_result = default_fn()

127 for i, fn in enumerate(aux_fns):

128 if fn is None:

129 continue

@@ -125,8 +132,8 @@ def execute_in_parallel(

132 aux_results[i] = fn()

133 done_events[i].record()

134 pending.append(done_events[i])

128 -

129 - default_result = default_fn()

135 + if not default_first:

136 + default_result = default_fn()

137

138 for ev in pending:

139 ev.wait()