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 5, 23:30 UTC
Disable dual stream execution of input projection for Qwen3 (#38152)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Mar 25, 2026 at 18:20 UTC
9704a5c3106a1d42c8f5b2dd51d8585ecc7db482
2 files changed
+5
-67
vllm/model_executor/models/qwen3_5.py
+2
-6
@@ -221,12 +221,8 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
221
ba, _ = self.in_proj_ba(hidden_states)
222
z, _ = self.in_proj_z(hidden_states)
223
else:
224
- mixed_qkvz, ba = torch.ops.vllm.gdn_in_proj(
225
- hidden_states,
226
- sum(self.in_proj_qkvz.output_sizes) // self.tp_size,
227
- sum(self.in_proj_ba.output_sizes) // self.tp_size,
228
- self.prefix,
229
- )
224
+ mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
225
+ ba, _ = self.in_proj_ba(hidden_states)
226
qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
227
z_size = self.value_dim // self.tp_size
228
mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
vllm/model_executor/models/qwen3_next.py
+3
-61
@@ -81,11 +81,7 @@ from vllm.platforms import current_platform
81
from vllm.sequence import IntermediateTensors
82
from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
83
from vllm.triton_utils import tl, triton
84
-from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
85
-from vllm.utils.torch_utils import (
86
- aux_stream,
87
- direct_register_custom_op,
88
-)
84
+from vllm.utils.torch_utils import direct_register_custom_op
85
from vllm.v1.attention.backend import AttentionMetadata
86
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
87
@@ -421,12 +417,6 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
417
self.act = ACT2FN[config.hidden_act]
418
self.layer_norm_epsilon = config.rms_norm_eps
419
self.prefix = prefix
424
- self.aux_stream = aux_stream()
425
- self.events = (
426
- [torch.cuda.Event(), torch.cuda.Event()]
427
- if current_platform.is_cuda_alike()
428
- else [None, None]
429
- )
420
421
self.config = config
422
self.model_config = vllm_config.model_config
@@ -659,12 +649,8 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
649
# ============================================================
650
# Part 1: Input Projection
651
# ============================================================
662
- projected_states_qkvz, projected_states_ba = torch.ops.vllm.gdn_in_proj(
663
- hidden_states,
664
- sum(self.in_proj_qkvz.output_sizes) // self.tp_size,
665
- sum(self.in_proj_ba.output_sizes) // self.tp_size,
666
- self.prefix,
667
- )
652
+ projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
653
+ projected_states_ba, _ = self.in_proj_ba(hidden_states)
654
query, key, value, z, b, a = self.fix_query_key_value_ordering(
655
projected_states_qkvz, projected_states_ba
656
)
@@ -804,18 +790,6 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
790
791
torch.accelerator.empty_cache()
792
807
- def _forward_in_proj(
808
- self, hidden_states: torch.Tensor
809
- ) -> tuple[torch.Tensor, torch.Tensor]:
810
- projected_states_qkvz, projected_states_ba = maybe_execute_in_parallel(
811
- lambda: self.in_proj_qkvz(hidden_states)[0],
812
- lambda: self.in_proj_ba(hidden_states)[0],
813
- self.events[0],
814
- self.events[1],
815
- self.aux_stream,
816
- )
817
- return projected_states_qkvz, projected_states_ba
818
-
793
def _forward_core(
794
self,
795
mixed_qkv: torch.Tensor,
@@ -1697,32 +1671,6 @@ class Qwen3NextForCausalLM(
1671
return self.model.get_expert_mapping()
1672
1673
1700
-def gdn_in_proj(
1701
- hidden_states: torch.Tensor,
1702
- qkvz_output_size: int,
1703
- ba_output_size: int,
1704
- layer_name: str,
1705
-) -> tuple[torch.Tensor, torch.Tensor]:
1706
- """
1707
- Custom op for the input projection.
1708
- """
1709
- forward_context: ForwardContext = get_forward_context()
1710
- self = forward_context.no_compile_layers[layer_name]
1711
- return self._forward_in_proj(hidden_states)
1712
-
1713
-
1714
-def gdn_in_proj_fake(
1715
- hidden_states: torch.Tensor,
1716
- qkvz_output_size: int,
1717
- ba_output_size: int,
1718
- layer_name: str,
1719
-) -> tuple[torch.Tensor, torch.Tensor]:
1720
- """Fake implementation for torch.compile."""
1721
- return hidden_states.new_empty(
1722
- hidden_states.shape[0], qkvz_output_size
1723
- ), hidden_states.new_empty(hidden_states.shape[0], ba_output_size)
1724
-
1725
-
1674
def gdn_attention_core(
1675
mixed_qkv: torch.Tensor,
1676
b: torch.Tensor,
@@ -1756,12 +1704,6 @@ def gdn_attention_core_fake(
1704
return
1705
1706
1759
-direct_register_custom_op(
1760
- op_name="gdn_in_proj",
1761
- op_func=gdn_in_proj,
1762
- fake_impl=gdn_in_proj_fake,
1763
-)
1764
-
1707
direct_register_custom_op(
1708
op_name="gdn_attention_core",
1709
op_func=gdn_attention_core,