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 8, 05:00 UTC
Revert "[Models] Fuse Qwen3.5 GDN's qkvz_proj and ba_proj" (#34683)
Jiangyun Zhu committed
Feb 17, 2026 at 17:29 UTC
1d65283e95f4d978c984df8585ca3f477166e651
3 files changed
+182
-87
vllm/model_executor/layers/linear.py
+6
-28
@@ -685,13 +685,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
685
self,
686
param: Parameter,
687
loaded_weight: torch.Tensor,
688
- loaded_shard_id: tuple[int, ...] | int | None = None,
688
+ loaded_shard_id: int | None = None,
689
):
690
- if isinstance(loaded_shard_id, tuple):
691
- raise NotImplementedError(
692
- "Shard id with multiple indices is not supported in weight_loader, "
693
- "please use weight_loader_v2 instead."
694
- )
690
# Special case for GGUF
691
# initialize GGUF param after we know the quantize type
692
is_gguf_weight = getattr(param, "is_gguf_weight", False)
@@ -830,10 +825,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
825
param_data.copy_(loaded_weight)
826
827
def _load_fused_module_from_checkpoint(
833
- self,
834
- param: BasevLLMParameter,
835
- loaded_weight: torch.Tensor,
836
- output_sizes: list[int] | None = None,
828
+ self, param: BasevLLMParameter, loaded_weight: torch.Tensor
829
):
830
"""
831
Handle special case for models where MLP layers are already
@@ -847,8 +839,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
839
840
current_shard_offset = 0
841
shard_offsets: list[tuple[int, int, int]] = []
850
- output_sizes = output_sizes or self.output_sizes
851
- for i, output_size in enumerate(output_sizes):
842
+ for i, output_size in enumerate(self.output_sizes):
843
shard_offsets.append((i, current_shard_offset, output_size))
844
current_shard_offset += output_size
845
@@ -873,30 +864,17 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
864
self,
865
param: BasevLLMParameter,
866
loaded_weight: torch.Tensor,
876
- loaded_shard_id: tuple[int, ...] | int | None = None,
867
+ loaded_shard_id: int | None = None,
868
):
878
- if loaded_shard_id is None or isinstance(loaded_shard_id, tuple):
869
+ if loaded_shard_id is None:
870
if isinstance(param, PerTensorScaleParameter):
871
param.load_merged_column_weight(loaded_weight=loaded_weight, shard_id=0)
872
return
873
elif type(param) in (RowvLLMParameter, BasevLLMParameter):
874
param.load_merged_column_weight(loaded_weight=loaded_weight)
875
return
885
- output_sizes = (
886
- [self.output_sizes[idx] for idx in loaded_shard_id]
887
- if loaded_shard_id
888
- else None
889
- )
890
- if isinstance(param, BlockQuantScaleParameter):
891
- weight_block_size = getattr(self, "weight_block_size", None)
892
- output_sizes = [
893
- adjust_block_scale_shard(weight_block_size, size, 0)[0]
894
- for size in (output_sizes or self.output_sizes)
895
- ]
876
# TODO: @dsikka - move to parameter.py
897
- self._load_fused_module_from_checkpoint(
898
- param, loaded_weight, output_sizes=output_sizes
899
- )
877
+ self._load_fused_module_from_checkpoint(param, loaded_weight)
878
return
879
880
assert loaded_shard_id < len(self.output_sizes)
vllm/model_executor/models/qwen3_5.py
+166
-32
@@ -30,20 +30,36 @@ from collections.abc import Callable, Iterable
30
import torch
31
from einops import rearrange
32
from torch import nn
33
+from transformers.activations import ACT2FN
34
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import (
37
+ CacheConfig,
38
+ ModelConfig,
39
+ SpeculativeConfig,
40
VllmConfig,
41
+ get_current_vllm_config,
42
)
43
from vllm.distributed import (
44
+ divide,
45
get_pp_group,
46
+ get_tensor_model_parallel_rank,
47
+ get_tensor_model_parallel_world_size,
48
)
49
from vllm.logger import init_logger
50
from vllm.model_executor.layers.layernorm import (
51
GemmaRMSNorm as Qwen3_5RMSNorm,
52
)
45
-from vllm.model_executor.layers.linear import MergedColumnParallelLinear
53
+from vllm.model_executor.layers.layernorm import RMSNormGated
54
+from vllm.model_executor.layers.linear import (
55
+ ColumnParallelLinear,
56
+ MergedColumnParallelLinear,
57
+ RowParallelLinear,
58
+)
59
from vllm.model_executor.layers.logits_processor import LogitsProcessor
60
+from vllm.model_executor.layers.mamba.mamba_mixer2 import (
61
+ mamba_v2_sharded_weight_loader,
62
+)
63
from vllm.model_executor.layers.mamba.mamba_utils import (
64
MambaStateCopyFunc,
65
MambaStateCopyFuncCalculator,
@@ -57,8 +73,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
73
)
74
from vllm.model_executor.model_loader.weight_utils import (
75
default_weight_loader,
76
+ sharded_weight_loader,
77
)
78
+from vllm.model_executor.utils import set_weight_attrs
79
from vllm.multimodal import MULTIMODAL_REGISTRY
80
+from vllm.platforms import current_platform
81
from vllm.sequence import IntermediateTensors
82
from vllm.transformers_utils.configs.qwen3_5 import (
83
Qwen3_5Config,
@@ -80,6 +99,7 @@ from .interfaces import (
99
)
100
from .qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
101
from .qwen3_next import (
102
+ ChunkGatedDeltaRule,
103
Qwen3NextAttention,
104
Qwen3NextDecoderLayer,
105
Qwen3NextGatedDeltaNet,
@@ -119,29 +139,152 @@ class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo):
139
140
141
class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
122
- def fix_query_key_value_ordering(
142
+ def __init__(
143
self,
124
- mixed_qkvz: torch.Tensor,
125
- mixed_ba: torch.Tensor,
126
- ):
127
- raise NotImplementedError(
128
- "Qwen3.5 Series dont need to fix query key value ordering"
144
+ config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig,
145
+ model_config: ModelConfig | None = None,
146
+ cache_config: CacheConfig | None = None,
147
+ quant_config: QuantizationConfig | None = None,
148
+ speculative_config: SpeculativeConfig | None = None,
149
+ prefix: str = "",
150
+ ) -> None:
151
+ super(Qwen3NextGatedDeltaNet, self).__init__()
152
+ self.tp_size = get_tensor_model_parallel_world_size()
153
+ self.tp_rank = get_tensor_model_parallel_rank()
154
+ self.hidden_size = config.hidden_size
155
+ self.num_v_heads = config.linear_num_value_heads
156
+ self.num_k_heads = config.linear_num_key_heads
157
+ self.head_k_dim = config.linear_key_head_dim
158
+ self.head_v_dim = config.linear_value_head_dim
159
+ self.key_dim = self.head_k_dim * self.num_k_heads
160
+ self.value_dim = self.head_v_dim * self.num_v_heads
161
+
162
+ self.conv_kernel_size = config.linear_conv_kernel_dim
163
+ self.layer_idx = extract_layer_index(prefix)
164
+ self.activation = config.hidden_act
165
+ self.act = ACT2FN[config.hidden_act]
166
+ self.layer_norm_epsilon = config.rms_norm_eps
167
+ self.prefix = prefix
168
+
169
+ self.config = config
170
+ self.model_config = model_config
171
+ self.cache_config = cache_config
172
+ self.quant_config = quant_config
173
+ self.speculative_config = speculative_config
174
+ self.num_spec = (
175
+ self.speculative_config.num_speculative_tokens
176
+ if self.speculative_config
177
+ else 0
178
)
179
131
- def create_qkvz_proj(
132
- self,
133
- hidden_size: int,
134
- key_dim: int,
135
- value_dim: int,
136
- quant_config: QuantizationConfig | None,
137
- prefix: str,
138
- ) -> MergedColumnParallelLinear:
139
- return MergedColumnParallelLinear(
140
- input_size=hidden_size,
141
- output_sizes=[key_dim, key_dim, value_dim, value_dim],
180
+ # QKV
181
+ self.conv_dim = self.key_dim * 2 + self.value_dim
182
+ self.conv1d = ColumnParallelLinear(
183
+ input_size=self.conv_kernel_size,
184
+ output_size=self.conv_dim,
185
+ bias=False,
186
+ prefix=f"{prefix}.conv1d",
187
+ )
188
+ self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
189
+
190
+ self.in_proj_qkv = MergedColumnParallelLinear(
191
+ input_size=self.hidden_size,
192
+ output_sizes=[self.key_dim, self.key_dim, self.value_dim],
193
+ bias=False,
194
+ quant_config=quant_config,
195
+ prefix=f"{prefix}.in_proj_qkv",
196
+ )
197
+ self.in_proj_z = ColumnParallelLinear(
198
+ input_size=self.hidden_size,
199
+ output_size=self.value_dim,
200
+ bias=False,
201
+ quant_config=quant_config,
202
+ prefix=f"{prefix}.in_proj_z",
203
+ )
204
+ self.in_proj_b = ColumnParallelLinear(
205
+ input_size=self.hidden_size,
206
+ output_size=self.num_v_heads,
207
+ bias=False,
208
+ quant_config=quant_config,
209
+ prefix=f"{prefix}.in_proj_b",
210
+ )
211
+ self.in_proj_a = ColumnParallelLinear(
212
+ input_size=self.hidden_size,
213
+ output_size=self.num_v_heads,
214
bias=False,
215
quant_config=quant_config,
144
- prefix=prefix,
216
+ prefix=f"{prefix}.in_proj_a",
217
+ )
218
+
219
+ query_key_settings = (self.key_dim, 0, False)
220
+ value_settings = (self.value_dim, 0, False)
221
+
222
+ delattr(self.conv1d.weight, "weight_loader")
223
+ set_weight_attrs(
224
+ self.conv1d.weight,
225
+ {
226
+ "weight_loader": mamba_v2_sharded_weight_loader(
227
+ [
228
+ query_key_settings,
229
+ query_key_settings,
230
+ value_settings,
231
+ ],
232
+ self.tp_size,
233
+ self.tp_rank,
234
+ )
235
+ },
236
+ )
237
+
238
+ # selective projection used to make dt, B and C input dependant
239
+
240
+ # time step projection (discretization)
241
+ # instantiate once and copy inv_dt in init_weights of PretrainedModel
242
+ self.dt_bias = nn.Parameter(
243
+ torch.ones(self.num_v_heads // self.tp_size),
244
+ )
245
+ self.A_log = nn.Parameter(
246
+ torch.empty(
247
+ divide(self.num_v_heads, self.tp_size),
248
+ )
249
+ )
250
+
251
+ set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
252
+ set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
253
+
254
+ self.norm = RMSNormGated(
255
+ self.head_v_dim,
256
+ eps=self.layer_norm_epsilon,
257
+ group_size=None,
258
+ norm_before_gate=True,
259
+ device=current_platform.current_device(),
260
+ dtype=config.dtype,
261
+ )
262
+
263
+ self.out_proj = RowParallelLinear(
264
+ self.value_dim,
265
+ self.hidden_size,
266
+ bias=False,
267
+ input_is_parallel=True,
268
+ quant_config=quant_config,
269
+ prefix=f"{prefix}.out_proj",
270
+ )
271
+
272
+ self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
273
+
274
+ compilation_config = get_current_vllm_config().compilation_config
275
+ if prefix in compilation_config.static_forward_context:
276
+ raise ValueError(f"Duplicate layer name: {prefix}")
277
+ compilation_config.static_forward_context[prefix] = self
278
+
279
+ def fix_query_key_value_ordering(
280
+ self,
281
+ mixed_qkv,
282
+ z,
283
+ b,
284
+ a,
285
+ ):
286
+ raise NotImplementedError(
287
+ "Qwen3.5 Series dont need to fix query key value ordering"
288
)
289
290
def forward(
@@ -160,13 +303,11 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
303
# ============================================================
304
# Part 1: Input Projection
305
# ============================================================
163
- mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
164
- qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
165
- z_size = self.value_dim // self.tp_size
166
- mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
306
+ mixed_qkv, _ = self.in_proj_qkv(hidden_states)
307
+ z, _ = self.in_proj_z(hidden_states)
308
z = z.reshape(z.size(0), -1, self.head_v_dim)
168
- ba, _ = self.in_proj_ba(hidden_states)
169
- b, a = ba.chunk(2, dim=-1)
309
+ b, _ = self.in_proj_b(hidden_states)
310
+ a, _ = self.in_proj_a(hidden_states)
311
312
b = b.contiguous()
313
a = a.contiguous()
@@ -365,18 +506,11 @@ class Qwen3_5Model(Qwen3NextModel):
506
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
507
stacked_params_mapping = [
508
# (param_name, shard_name, shard_id)
368
- # self attention
509
("qkv_proj", "q_proj", "q"),
510
("qkv_proj", "k_proj", "k"),
511
("qkv_proj", "v_proj", "v"),
372
- # mlp
512
("gate_up_proj", "gate_proj", 0),
513
("gate_up_proj", "up_proj", 1),
375
- # GDN
376
- ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
377
- ("in_proj_qkvz", "in_proj_z", 3),
378
- ("in_proj_ba", "in_proj_b", 0),
379
- ("in_proj_ba", "in_proj_a", 1),
514
]
515
516
params_dict = dict(self.named_parameters())
vllm/model_executor/models/qwen3_next.py
+10
-27
@@ -44,7 +44,6 @@ from vllm.model_executor.layers.layernorm import (
44
from vllm.model_executor.layers.layernorm import RMSNormGated
45
from vllm.model_executor.layers.linear import (
46
ColumnParallelLinear,
47
- MergedColumnParallelLinear,
47
QKVParallelLinear,
48
ReplicatedLinear,
49
RowParallelLinear,
@@ -407,19 +406,19 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
406
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
407
408
# projection of the input hidden states
410
- # Qwen3-Next and Qwen3.5 has a different qkv_proj layout,
411
- # we need to create qkvz_proj adaptively here.
412
- self.in_proj_qkvz = self.create_qkvz_proj(
413
- hidden_size=self.hidden_size,
414
- key_dim=self.key_dim,
415
- value_dim=self.value_dim,
409
+ self.projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2
410
+ self.projection_size_ba = self.num_v_heads * 2
411
+ self.in_proj_qkvz = ColumnParallelLinear(
412
+ input_size=self.hidden_size,
413
+ output_size=self.projection_size_qkvz,
414
+ bias=False,
415
quant_config=quant_config,
416
prefix=f"{prefix}.in_proj_qkvz",
417
)
418
# ba_proj doesn't support blockwise fp8 quantization.
420
- self.in_proj_ba = MergedColumnParallelLinear(
419
+ self.in_proj_ba = ColumnParallelLinear(
420
input_size=self.hidden_size,
422
- output_sizes=[self.num_v_heads] * 2,
421
+ output_size=self.projection_size_ba,
422
bias=False,
423
quant_config=quant_config,
424
prefix=f"{prefix}.in_proj_ba",
@@ -485,26 +484,10 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
484
raise ValueError(f"Duplicate layer name: {prefix}")
485
compilation_config.static_forward_context[prefix] = self
486
488
- def create_qkvz_proj(
489
- self,
490
- hidden_size: int,
491
- key_dim: int,
492
- value_dim: int,
493
- quant_config: QuantizationConfig | None,
494
- prefix: str,
495
- ) -> MergedColumnParallelLinear:
496
- return MergedColumnParallelLinear(
497
- input_size=hidden_size,
498
- output_sizes=[sum((key_dim, key_dim, value_dim)), value_dim],
499
- bias=False,
500
- quant_config=quant_config,
501
- prefix=f"{prefix}.in_proj_qkvz",
502
- )
503
-
487
def fix_query_key_value_ordering(
488
self,
506
- mixed_qkvz: torch.Tensor,
507
- mixed_ba: torch.Tensor,
489
+ mixed_qkvz,
490
+ mixed_ba,
491
):
492
"""
493
Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.