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 3, 22:45 UTC
[Bugfix] Fuse Qwen3.5 in_qkvz_proj forwarding with LoRA enabled (#37912)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Isotr0py committed
May 10, 2026 at 18:59 UTC
48698b1b9b30071952084ce12e8904b219193545
4 files changed
+84
-111
vllm/lora/layers/column_parallel_linear.py
+49
-11
@@ -195,9 +195,9 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
195
# There are two LoRA layers
196
# the output_sizes in MergedColumnParallelLinear is not sharded by tp
197
# we need to divide it by the tp_size to get correct slices size
198
- output_sizes = self.base_layer.output_sizes
198
+ self.output_sizes = self.base_layer.output_sizes
199
self.output_slices = tuple(
200
- divide(output_size, self.tp_size) for output_size in output_sizes
200
+ divide(output_size, self.tp_size) for output_size in self.output_sizes
201
)
202
self.n_slices = len(self.output_slices)
203
self.output_ids = (self.tp_rank,) * self.n_slices
@@ -261,6 +261,42 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
261
]
262
return sliced_lora_b
263
264
+ def expand_packed_lora(
265
+ self,
266
+ lora_a: list[torch.Tensor],
267
+ lora_b: list[torch.Tensor],
268
+ ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
269
+ """
270
+ Expand packed adapter groups when they don't match n_slices.
271
+ E.g. in_proj_qkv (covers Q+K+V) + in_proj_z
272
+ """
273
+ expanded_a: list[torch.Tensor] = []
274
+ expanded_b: list[torch.Tensor] = []
275
+ start_idx = 0
276
+ for a_i, b_i in zip(lora_a, lora_b):
277
+ # Determine which output slices this b_i covers.
278
+ b_rows, cu_rows, covered = b_i.shape[0], 0, 0
279
+ for i in range(start_idx, self.n_slices):
280
+ cu_rows += self.output_sizes[i]
281
+ if cu_rows == b_rows:
282
+ covered = i - start_idx + 1
283
+ break
284
+ else:
285
+ raise ValueError(
286
+ f"Cannot determine how to split lora_b with {b_rows} rows "
287
+ f"into {self.n_slices} slices with output sizes "
288
+ f"{self.output_sizes} starting from index {start_idx}."
289
+ )
290
+ # Split b_i into per-slice tensors and replicate a_i for each.
291
+ start = 0
292
+ for j in range(covered):
293
+ size = self.output_sizes[start_idx + j]
294
+ expanded_b.append(b_i[start : start + size, :])
295
+ expanded_a.append(a_i)
296
+ start += size
297
+ start_idx += covered
298
+ return expanded_a, expanded_b
299
+
300
def set_lora(
301
self,
302
index: int,
@@ -269,6 +305,12 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
305
):
306
self.reset_lora(index)
307
308
+ # Expand packed adapter groups when they don't match n_slices.
309
+ # E.g. in_proj_qkv (covers Q+K+V) + in_proj_z as 2 groups for a
310
+ # 4-slice layer: split b_qkv by output_sizes and replicate a_qkv.
311
+ if isinstance(lora_b, list) and len(lora_b) != self.n_slices:
312
+ lora_a, lora_b = self.expand_packed_lora(lora_a, lora_b)
313
+
314
if self.tp_size > 1:
315
lora_a = self.slice_lora_a(lora_a)
316
lora_b = self.slice_lora_b(lora_b)
@@ -497,18 +539,14 @@ class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLo
539
def slice_lora_a(
540
self, lora_a: list[torch.Tensor | None]
541
) -> list[torch.Tensor | None]:
500
- # NOTE: lora_a contains 2 subloras, and each sublora could be None.
542
output_shard_size = self.lora_a_stacked[0].shape[2]
543
output_start_idx = self.tp_rank * output_shard_size
503
- lora_a = [
504
- lora_a[0][output_start_idx : output_start_idx + output_shard_size, :]
505
- if lora_a[0] is not None
506
- else None,
507
- lora_a[1][output_start_idx : output_start_idx + output_shard_size, :]
508
- if lora_a[1] is not None
509
- else None,
544
+ return [
545
+ lora_a_i[output_start_idx : output_start_idx + output_shard_size, :]
546
+ if (lora_a_i := lora_a[i]) is not None
547
+ else None
548
+ for i in range(len(lora_a))
549
]
511
- return lora_a
550
551
def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
552
return _mcp_apply(x, bias, self)
vllm/lora/model_manager.py
+5
@@ -563,11 +563,16 @@ class LoRAModelManager:
563
else:
564
parts = module_name.split(".")
565
replacements = self.packed_modules_mapping[parts[-1]]
566
+ n_slices = getattr(module, "n_slices", len(replacements))
567
if module.__class__.__name__ == "FusedMoEWithLoRA":
568
replacements = replacements[
569
: len(module.lora_a_stacked) // self.lora_slots
570
]
571
subloras: list[LoRALayerWeights | None] = []
572
+ # HACK: overrides replacements for qkvz = qkv + z case.
573
+ # Any better methods to handle this case?
574
+ if n_slices != len(replacements):
575
+ replacements = [f"slice_{i}" for i in range(n_slices)]
576
for i, r in enumerate(replacements):
577
lora = LoRALayerWeights.create_dummy_lora_weights(
578
module_name + "." + r,
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+26
-57
@@ -263,7 +263,6 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
263
config: Qwen3NextConfig,
264
vllm_config: VllmConfig,
265
prefix: str = "",
266
- create_in_proj_qkvz: bool = True,
266
gqa_interleaved_layout=False,
267
) -> None:
268
super().__init__()
@@ -323,32 +322,14 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
322
# we need to create qkvz_proj adaptively here.
323
# When create_in_proj_qkvz is False (e.g. LoRA enabled in Qwen3.5),
324
# in_proj_qkv and in_proj_z are created separately instead.
326
- self.has_lora_projections = not create_in_proj_qkvz
327
- if create_in_proj_qkvz:
328
- self.in_proj_qkvz = self.create_qkvz_proj(
329
- hidden_size=self.hidden_size,
330
- key_dim=self.key_dim,
331
- value_dim=self.value_dim,
332
- quant_config=quant_config,
333
- prefix=f"{prefix}.in_proj_qkvz",
334
- )
335
- else:
336
- # LoRA case (Qwen3.5 only): keep q/k/v and z as separate modules
337
- # so that LoRA adapters can be applied independently.
338
- self.in_proj_qkv = MergedColumnParallelLinear(
339
- input_size=self.hidden_size,
340
- output_sizes=[self.key_dim, self.key_dim, self.value_dim],
341
- bias=False,
342
- quant_config=quant_config,
343
- prefix=f"{prefix}.in_proj_qkv",
344
- )
345
- self.in_proj_z = ColumnParallelLinear(
346
- input_size=self.hidden_size,
347
- output_size=self.value_dim,
348
- bias=False,
349
- quant_config=quant_config,
350
- prefix=f"{prefix}.in_proj_z",
351
- )
325
+ self.in_proj_qkvz = self.create_qkvz_proj(
326
+ hidden_size=self.hidden_size,
327
+ key_dim=self.key_dim,
328
+ value_dim=self.value_dim,
329
+ quant_config=quant_config,
330
+ prefix=f"{prefix}.in_proj_qkvz",
331
+ )
332
+
333
# ba_proj doesn't support blockwise fp8 quantization.
334
# Qwen3-Next and Qwen3.5 have different in_proj_ba checkpoint
335
# layouts, so we use a factory method to create the projection.
@@ -707,7 +688,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
688
):
689
"""ROCm forward using AITER Triton fused projection+attention when
690
available, otherwise falling back to the generic CUDA path."""
710
- if not self.has_lora_projections and GDN_AITER_TRITON_AVAILABLE:
691
+ if GDN_AITER_TRITON_AVAILABLE:
692
num_tokens = hidden_states.size(0)
693
projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
694
projected_states_ba, _ = self.in_proj_ba(hidden_states)
@@ -752,37 +733,27 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
733
# ============================================================
734
# Part 1: Input Projection
735
# ============================================================
755
- if self.has_lora_projections:
756
- # LoRA path (Qwen3.5 only): separate in_proj_qkv and in_proj_z
757
- mixed_qkv, _ = self.in_proj_qkv(hidden_states)
758
- ba, _ = self.in_proj_ba(hidden_states)
759
- z, _ = self.in_proj_z(hidden_states)
736
+ mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
737
+ ba, _ = self.in_proj_ba(hidden_states)
738
+
739
+ if self.gqa_interleaved_layout:
740
+ # Qwen3-Next: unpack the interleaved GQA layout
741
+ query, key, value, z, b, a = self.fix_query_key_value_ordering(
742
+ mixed_qkvz, ba
743
+ )
744
+ query, key, value = map(
745
+ lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value)
746
+ )
747
+ mixed_qkv = torch.cat((query, key, value), dim=-1)
748
+ else:
749
+ # Qwen3.5: weights are already in [q, k, v, z] and [b, a] order
750
+ qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
751
+ z_size = self.value_dim // self.tp_size
752
+ mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
753
z = z.reshape(z.size(0), -1, self.head_v_dim)
754
b, a = ba.chunk(2, dim=-1)
755
b = b.contiguous()
756
a = a.contiguous()
764
- else:
765
- mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
766
- ba, _ = self.in_proj_ba(hidden_states)
767
-
768
- if self.gqa_interleaved_layout:
769
- # Qwen3-Next: unpack the interleaved GQA layout
770
- query, key, value, z, b, a = self.fix_query_key_value_ordering(
771
- mixed_qkvz, ba
772
- )
773
- query, key, value = map(
774
- lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value)
775
- )
776
- mixed_qkv = torch.cat((query, key, value), dim=-1)
777
- else:
778
- # Qwen3.5: weights are already in [q, k, v, z] and [b, a] order
779
- qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
780
- z_size = self.value_dim // self.tp_size
781
- mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
782
- z = z.reshape(z.size(0), -1, self.head_v_dim)
783
- b, a = ba.chunk(2, dim=-1)
784
- b = b.contiguous()
785
- a = a.contiguous()
757
758
# ============================================================
759
# Part 2: Core Attention (Custom Op)
@@ -822,8 +793,6 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
793
"""
794
num_tokens = hidden_states.size(0)
795
825
- assert not self.has_lora_projections, "lora isn't supported on XPU."
826
-
796
# ============================================================
797
# Part 1: Input Projection
798
# ============================================================
vllm/model_executor/models/qwen3_5.py
+4
-43
@@ -138,7 +138,6 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
138
vllm_config=vllm_config,
139
prefix=f"{prefix}.linear_attn",
140
gqa_interleaved_layout=False,
141
- create_in_proj_qkvz=vllm_config.lora_config is None,
141
)
142
elif self.layer_type == "full_attention":
143
self.self_attn = Qwen3NextAttention(
@@ -217,7 +216,6 @@ class Qwen3_5Model(Qwen3NextModel):
216
self.num_redundant_experts = eplb_config.num_redundant_experts
217
218
self.config = config
220
- self.enable_lora = vllm_config.lora_config is not None
219
220
self.vocab_size = config.vocab_size
221
@@ -276,6 +274,9 @@ class Qwen3_5Model(Qwen3NextModel):
274
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
275
stacked_params_mapping = [
276
# (param_name, shard_name, shard_id)
277
+ # GDN
278
+ ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
279
+ ("in_proj_qkvz", "in_proj_z", 3),
280
# self attention
281
("qkv_proj", "q_proj", "q"),
282
("qkv_proj", "k_proj", "k"),
@@ -287,21 +288,6 @@ class Qwen3_5Model(Qwen3NextModel):
288
("in_proj_ba", "in_proj_a", 1),
289
]
290
290
- if self.enable_lora:
291
- stacked_params_mapping.extend(
292
- [
293
- ("in_proj_qkv", "in_proj_qkv", (0, 1, 2)),
294
- ("in_proj_z", "in_proj_z", 0),
295
- ]
296
- )
297
- else:
298
- stacked_params_mapping.extend(
299
- [
300
- ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
301
- ("in_proj_qkvz", "in_proj_z", 3),
302
- ]
303
- )
304
-
291
params_dict = dict(self.named_parameters())
292
loaded_params: set[str] = set()
293
expert_params_mapping = self.get_expert_mapping()
@@ -352,10 +338,7 @@ class Qwen3_5Model(Qwen3NextModel):
338
continue
339
param = params_dict[name]
340
weight_loader = param.weight_loader
355
- if param_name == "in_proj_z" and self.enable_lora:
356
- weight_loader(param, loaded_weight)
357
- else:
358
- weight_loader(param, loaded_weight, shard_id)
341
+ weight_loader(param, loaded_weight, shard_id)
342
break
343
else:
344
is_expert_weight = False
@@ -485,15 +468,6 @@ class Qwen3_5ForCausalLMBase(
468
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
469
)
470
488
- # When LoRA is enabled, GDN uses separate in_proj_qkv and in_proj_z
489
- # instead of merged in_proj_qkvz; pack mapping must match.
490
- if vllm_config.lora_config:
491
- base = getattr(Qwen3_5ForCausalLMBase, "packed_modules_mapping", {})
492
- self.packed_modules_mapping = {k: list(v) for k, v in base.items()}
493
- self.packed_modules_mapping.pop("in_proj_qkvz", None)
494
- self.packed_modules_mapping["in_proj_qkv"] = ["in_proj_qkv"]
495
- self.packed_modules_mapping["in_proj_z"] = ["in_proj_z"]
496
-
471
if get_pp_group().is_last_rank:
472
if config.tie_word_embeddings:
473
self.lm_head = self.model.embed_tokens
@@ -586,7 +560,6 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
560
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
561
# protocols have not __init__ method, so we need to use nn.Module.__init__
562
nn.Module.__init__(self)
589
- self.update_packed_mapping(enable_lora=vllm_config.lora_config is not None)
563
config: Qwen3_5Config = vllm_config.model_config.hf_config
564
quant_config = vllm_config.quant_config
565
multimodal_config = vllm_config.model_config.multimodal_config
@@ -614,17 +587,6 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
587
self.language_model.make_empty_intermediate_tensors
588
)
589
617
- def update_packed_mapping(self, enable_lora: bool):
618
- # When LoRA is enabled, GDN uses separate in_proj_qkv and in_proj_z
619
- if enable_lora:
620
- base = getattr(
621
- Qwen3_5ForConditionalGeneration, "packed_modules_mapping", {}
622
- )
623
- self.packed_modules_mapping = {k: list(v) for k, v in base.items()}
624
- self.packed_modules_mapping.pop("in_proj_qkvz", None)
625
- self.packed_modules_mapping["in_proj_qkv"] = ["in_proj_qkv"]
626
- self.packed_modules_mapping["in_proj_z"] = ["in_proj_z"]
627
-
590
def embed_input_ids(
591
self,
592
input_ids: torch.Tensor,
@@ -811,7 +773,6 @@ class Qwen3_5MoeForConditionalGeneration(
773
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
774
# protocols have not __init__ method, so we need to use nn.Module.__init__
775
nn.Module.__init__(self)
814
- self.update_packed_mapping(enable_lora=vllm_config.lora_config is not None)
776
config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config
777
quant_config = vllm_config.quant_config
778
multimodal_config = vllm_config.model_config.multimodal_config