A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 13:15 UTC
[Perf] Enable dual stream execution of input projection for Qwen3 (#36795)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Mar 17, 2026 at 20:13 UTC
f1740006e47d580656668ba5a9253a4e4340e198
3 files changed
+115
-5
vllm/model_executor/models/qwen3_5.py
+6
-2
@@ -180,12 +180,16 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
180
# ============================================================
181
# Part 1: Input Projection
182
# ============================================================
183
- mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
183
+ mixed_qkvz, ba = torch.ops.vllm.gdn_in_proj(
184
+ hidden_states,
185
+ self.in_proj_qkvz.weight.shape[0],
186
+ self.in_proj_ba.weight.shape[0],
187
+ self.prefix,
188
+ )
189
qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
190
z_size = self.value_dim // self.tp_size
191
mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
192
z = z.reshape(z.size(0), -1, self.head_v_dim)
188
- ba, _ = self.in_proj_ba(hidden_states)
193
b, a = ba.chunk(2, dim=-1)
194
195
b = b.contiguous()
vllm/model_executor/models/qwen3_next.py
+61
-3
@@ -82,7 +82,11 @@ from vllm.platforms import current_platform
82
from vllm.sequence import IntermediateTensors
83
from vllm.transformers_utils.configs import Qwen3NextConfig
84
from vllm.triton_utils import tl, triton
85
-from vllm.utils.torch_utils import direct_register_custom_op
85
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
86
+from vllm.utils.torch_utils import (
87
+ aux_stream,
88
+ direct_register_custom_op,
89
+)
90
from vllm.v1.attention.backend import AttentionMetadata
91
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
92
@@ -419,6 +423,12 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
423
self.act = ACT2FN[config.hidden_act]
424
self.layer_norm_epsilon = config.rms_norm_eps
425
self.prefix = prefix
426
+ self.aux_stream = aux_stream()
427
+ self.events = (
428
+ [torch.cuda.Event(), torch.cuda.Event()]
429
+ if current_platform.is_cuda()
430
+ else [None, None]
431
+ )
432
433
self.config = config
434
self.model_config = model_config
@@ -647,8 +657,12 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
657
# ============================================================
658
# Part 1: Input Projection
659
# ============================================================
650
- projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
651
- projected_states_ba, _ = self.in_proj_ba(hidden_states)
660
+ projected_states_qkvz, projected_states_ba = torch.ops.vllm.gdn_in_proj(
661
+ hidden_states,
662
+ self.in_proj_qkvz.weight.shape[0],
663
+ self.in_proj_ba.weight.shape[0],
664
+ self.prefix,
665
+ )
666
query, key, value, z, b, a = self.fix_query_key_value_ordering(
667
projected_states_qkvz, projected_states_ba
668
)
@@ -783,6 +797,18 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
797
798
torch.accelerator.empty_cache()
799
800
+ def _forward_in_proj(
801
+ self, hidden_states: torch.Tensor
802
+ ) -> tuple[torch.Tensor, torch.Tensor]:
803
+ projected_states_qkvz, projected_states_ba = maybe_execute_in_parallel(
804
+ lambda: self.in_proj_qkvz(hidden_states)[0],
805
+ lambda: self.in_proj_ba(hidden_states)[0],
806
+ self.events[0],
807
+ self.events[1],
808
+ self.aux_stream,
809
+ )
810
+ return projected_states_qkvz, projected_states_ba
811
+
812
def _forward_core(
813
self,
814
mixed_qkv: torch.Tensor,
@@ -1670,6 +1696,32 @@ class Qwen3NextForCausalLM(
1696
return self.model.get_expert_mapping()
1697
1698
1699
+def gdn_in_proj(
1700
+ hidden_states: torch.Tensor,
1701
+ qkvz_output_size: int,
1702
+ ba_output_size: int,
1703
+ layer_name: str,
1704
+) -> tuple[torch.Tensor, torch.Tensor]:
1705
+ """
1706
+ Custom op for the input projection.
1707
+ """
1708
+ forward_context: ForwardContext = get_forward_context()
1709
+ self = forward_context.no_compile_layers[layer_name]
1710
+ return self._forward_in_proj(hidden_states)
1711
+
1712
+
1713
+def gdn_in_proj_fake(
1714
+ hidden_states: torch.Tensor,
1715
+ qkvz_output_size: int,
1716
+ ba_output_size: int,
1717
+ layer_name: str,
1718
+) -> tuple[torch.Tensor, torch.Tensor]:
1719
+ """Fake implementation for torch.compile."""
1720
+ return hidden_states.new_empty(
1721
+ hidden_states.shape[0], qkvz_output_size
1722
+ ), hidden_states.new_empty(hidden_states.shape[0], ba_output_size)
1723
+
1724
+
1725
def gdn_attention_core(
1726
mixed_qkv: torch.Tensor,
1727
b: torch.Tensor,
@@ -1703,6 +1755,12 @@ def gdn_attention_core_fake(
1755
return
1756
1757
1758
+direct_register_custom_op(
1759
+ op_name="gdn_in_proj",
1760
+ op_func=gdn_in_proj,
1761
+ fake_impl=gdn_in_proj_fake,
1762
+)
1763
+
1764
direct_register_custom_op(
1765
op_name="gdn_attention_core",
1766
op_func=gdn_attention_core,
vllm/utils/multi_stream_utils.py
new
+48
@@ -0,0 +1,48 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Callable
5
+from typing import Any
6
+
7
+import torch
8
+
9
+
10
+def maybe_execute_in_parallel(
11
+ fn0: Callable[[], Any],
12
+ fn1: Callable[[], Any],
13
+ event0: torch.cuda.Event,
14
+ event1: torch.cuda.Event,
15
+ aux_stream: torch.cuda.Stream | None = None,
16
+) -> tuple[Any, Any]:
17
+ """Run two functions potentially in parallel on separate CUDA streams.
18
+
19
+ When aux_stream is provided, fn0 runs on the current (default) stream and
20
+ fn1 runs on aux_stream, synchronized via CUDA events. When aux_stream is
21
+ None, both functions execute sequentially on the current stream.
22
+
23
+ This design follows TensorRT-LLM's maybe_execute_in_parallel pattern
24
+ (tensorrt_llm/_torch/modules/multi_stream_utils.py).
25
+
26
+ Args:
27
+ fn0: Callable for the default stream.
28
+ fn1: Callable for the auxiliary stream.
29
+ event0: CUDA event recorded before fn0 so aux_stream can wait.
30
+ event1: CUDA event recorded after fn1 so default stream can wait.
31
+ aux_stream: The second CUDA stream for fn1.
32
+ Multi-stream is disabled when aux_stream is None.
33
+
34
+ Returns:
35
+ Tuple of (fn0_result, fn1_result).
36
+ """
37
+ if aux_stream is not None:
38
+ event0.record()
39
+ result0 = fn0()
40
+ with torch.cuda.stream(aux_stream):
41
+ event0.wait()
42
+ result1 = fn1()
43
+ event1.record()
44
+ event1.wait()
45
+ else:
46
+ result0 = fn0()
47
+ result1 = fn1()
48
+ return (result0, result1)