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, 09:30 UTC
[Bugfix][LoRA] Fix Qwen35 LoRA (#36976)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Mar 20, 2026 at 11:09 UTC
8fbe3f303fbf995afe4c409309aa8e889235339e
5 files changed
+257
-46
.buildkite/test_areas/lora.yaml
+3
-2
@@ -8,7 +8,7 @@ steps:
8
- vllm/lora
9
- tests/lora
10
commands:
11
- - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py
11
+ - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemoel_lora.py
12
parallelism: 4
13
14
@@ -30,4 +30,5 @@ steps:
30
- pytest -v -s -x lora/test_llama_tp.py
31
- pytest -v -s -x lora/test_llm_with_multi_loras.py
32
- pytest -v -s -x lora/test_olmoe_tp.py
33
- - pytest -v -s -x lora/test_gptoss_tp.py
\ No newline at end of file
33
+ - pytest -v -s -x lora/test_gptoss_tp.py
34
+ - pytest -v -s -x lora/test_qwen35_densemoel_lora.py
\ No newline at end of file
tests/lora/conftest.py
+5
@@ -294,6 +294,11 @@ def whisper_lora_files():
294
return snapshot_download(repo_id="chengyili2005/whisper-small-mandarin-lora")
295
296
297
+@pytest.fixture(scope="session")
298
+def qwen35_dense_model_lora_files():
299
+ return snapshot_download(repo_id="jeeejeee/qwen35-4b-text-only-sql-lora")
300
+
301
+
302
@pytest.fixture
303
def reset_default_device():
304
"""
tests/lora/test_qwen35_densemoel_lora.py
new
+132
@@ -0,0 +1,132 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from transformers import AutoTokenizer
5
+
6
+import vllm
7
+import vllm.config
8
+from vllm.lora.request import LoRARequest
9
+
10
+from ..utils import create_new_process_for_each_test, multi_gpu_test
11
+
12
+MODEL_PATH = "Qwen/Qwen3.5-4B"
13
+
14
+PROMPT_TEMPLATE = """Write a SQL query for the given database.\nSchema:\nTables:\n - stadium(Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average)\n - singer(Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male)\n - concert(concert_ID, concert_Name, Theme, Stadium_ID, Year)\n - singer_in_concert(concert_ID, Singer_ID)\n\nQuestion:\n{query}""" # noqa: E501
15
+
16
+EXPECTED_LORA_OUTPUT = [
17
+ "SELECT count(*) FROM singer",
18
+ "SELECT avg(age) , min(age) , max(age) FROM singer WHERE country = 'France'",
19
+ "SELECT name FROM stadium WHERE stadium_id NOT IN (SELECT stadium_id FROM concert)",
20
+]
21
+
22
+
23
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
24
+
25
+
26
+def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
27
+ prompts = [
28
+ PROMPT_TEMPLATE.format(query="How many singers do we have?"),
29
+ PROMPT_TEMPLATE.format(
30
+ query=(
31
+ "What is the average, minimum, and maximum "
32
+ "age of all singers from France?"
33
+ )
34
+ ),
35
+ PROMPT_TEMPLATE.format(
36
+ query=("What are the names of the stadiums without any concerts?")
37
+ ),
38
+ ]
39
+ input_templates = []
40
+ for prmpt in prompts:
41
+ messages = [{"role": "user", "content": prmpt}]
42
+ prompt = tokenizer.apply_chat_template(
43
+ messages,
44
+ tokenize=False,
45
+ add_generation_prompt=True,
46
+ enable_thinking=False, # disable thinking
47
+ )
48
+ input_templates.append(prompt)
49
+ sampling_params = vllm.SamplingParams(temperature=0, max_tokens=512)
50
+ outputs = llm.generate(
51
+ input_templates,
52
+ sampling_params,
53
+ lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
54
+ )
55
+
56
+ generated_texts: list[str] = []
57
+ for output in outputs:
58
+ prompt = output.prompt
59
+ generated_text = output.outputs[0].text.strip()
60
+ generated_texts.append(generated_text)
61
+ print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
62
+ return generated_texts
63
+
64
+
65
+@create_new_process_for_each_test()
66
+def test_qwen35_dense_model_lora(qwen35_dense_model_lora_files):
67
+ llm = vllm.LLM(
68
+ MODEL_PATH,
69
+ max_model_len=512,
70
+ enable_lora=True,
71
+ max_loras=2,
72
+ max_num_seqs=16,
73
+ max_lora_rank=8,
74
+ trust_remote_code=True,
75
+ )
76
+
77
+ output1 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=1)
78
+ for i in range(len(EXPECTED_LORA_OUTPUT)):
79
+ assert output1[i] == EXPECTED_LORA_OUTPUT[i]
80
+ output2 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=2)
81
+ for i in range(len(EXPECTED_LORA_OUTPUT)):
82
+ assert output2[i] == EXPECTED_LORA_OUTPUT[i]
83
+
84
+
85
+@multi_gpu_test(num_gpus=4)
86
+def test_qwen35_dense_model_lora_tp4(qwen35_dense_model_lora_files):
87
+ llm = vllm.LLM(
88
+ MODEL_PATH,
89
+ max_model_len=1024,
90
+ enable_lora=True,
91
+ max_loras=2,
92
+ max_lora_rank=8,
93
+ max_num_seqs=16,
94
+ tensor_parallel_size=4,
95
+ trust_remote_code=True,
96
+ fully_sharded_loras=False,
97
+ compilation_config=vllm.config.CompilationConfig( # Avoid OOM
98
+ cudagraph_specialize_lora=False,
99
+ ),
100
+ )
101
+
102
+ output1 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=1)
103
+ print(output1)
104
+ for i in range(len(EXPECTED_LORA_OUTPUT)):
105
+ assert output1[i] == EXPECTED_LORA_OUTPUT[i]
106
+ output2 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=2)
107
+ for i in range(len(EXPECTED_LORA_OUTPUT)):
108
+ assert output2[i] == EXPECTED_LORA_OUTPUT[i]
109
+
110
+
111
+@multi_gpu_test(num_gpus=4)
112
+def test_qwen35_dense_model_lora_tp4_fully_sharded_loras(qwen35_dense_model_lora_files):
113
+ llm = vllm.LLM(
114
+ MODEL_PATH,
115
+ max_model_len=512,
116
+ enable_lora=True,
117
+ max_loras=2,
118
+ max_lora_rank=8,
119
+ tensor_parallel_size=4,
120
+ trust_remote_code=True,
121
+ fully_sharded_loras=True,
122
+ gpu_memory_utilization=0.8,
123
+ compilation_config=vllm.config.CompilationConfig( # Avoid OOM
124
+ cudagraph_specialize_lora=False,
125
+ ),
126
+ )
127
+ output1 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=1)
128
+ for i in range(len(EXPECTED_LORA_OUTPUT)):
129
+ assert output1[i] == EXPECTED_LORA_OUTPUT[i]
130
+ output2 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=2)
131
+ for i in range(len(EXPECTED_LORA_OUTPUT)):
132
+ assert output2[i] == EXPECTED_LORA_OUTPUT[i]
vllm/model_executor/models/qwen3_5.py
+100
-23
@@ -32,9 +32,7 @@ from einops import rearrange
32
from torch import nn
33
34
from vllm.compilation.decorators import support_torch_compile
35
-from vllm.config import (
36
- VllmConfig,
37
-)
35
+from vllm.config import VllmConfig
36
from vllm.distributed import (
37
get_pp_group,
38
)
@@ -42,7 +40,10 @@ from vllm.logger import init_logger
40
from vllm.model_executor.layers.layernorm import (
41
GemmaRMSNorm as Qwen3_5RMSNorm,
42
)
45
-from vllm.model_executor.layers.linear import MergedColumnParallelLinear
43
+from vllm.model_executor.layers.linear import (
44
+ ColumnParallelLinear,
45
+ MergedColumnParallelLinear,
46
+)
47
from vllm.model_executor.layers.logits_processor import LogitsProcessor
48
from vllm.model_executor.layers.mamba.mamba_utils import (
49
MambaStateCopyFunc,
@@ -130,6 +131,40 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
131
"Qwen3.5 Series dont need to fix query key value ordering"
132
)
133
134
+ def __init__(
135
+ self,
136
+ config: Qwen3_5Config,
137
+ vllm_config: VllmConfig,
138
+ prefix: str = "",
139
+ ) -> None:
140
+ create_in_proj_qkvz = vllm_config.lora_config is None
141
+ super().__init__(
142
+ config,
143
+ vllm_config=vllm_config,
144
+ prefix=prefix,
145
+ create_in_proj_qkvz=create_in_proj_qkvz,
146
+ )
147
+ if vllm_config.lora_config is not None:
148
+ # Separate in_proj_qkv (Q,K,V) and in_proj_z for LoRA compatibility.
149
+ # Use MergedColumnParallelLinear for in_proj_qkv because GDN can have
150
+ # linear_num_key_heads != linear_num_value_heads (e.g. 16 vs 32), so
151
+ # output sizes [key_dim, key_dim, value_dim] are not representable
152
+ # with a single QKVParallelLinear (which ties K and V head counts).
153
+ self.in_proj_qkv = MergedColumnParallelLinear(
154
+ input_size=self.hidden_size,
155
+ output_sizes=[self.key_dim, self.key_dim, self.value_dim],
156
+ bias=False,
157
+ quant_config=vllm_config.quant_config,
158
+ prefix=f"{prefix}.in_proj_qkv",
159
+ )
160
+ self.in_proj_z = ColumnParallelLinear(
161
+ input_size=self.hidden_size,
162
+ output_size=self.value_dim,
163
+ bias=False,
164
+ quant_config=vllm_config.quant_config,
165
+ prefix=f"{prefix}.in_proj_z",
166
+ )
167
+
168
def create_qkvz_proj(
169
self,
170
hidden_size: int,
@@ -180,15 +215,21 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
215
# ============================================================
216
# Part 1: Input Projection
217
# ============================================================
183
- mixed_qkvz, ba = torch.ops.vllm.gdn_in_proj(
184
- hidden_states,
185
- sum(self.in_proj_qkvz.output_sizes) // self.tp_size,
186
- sum(self.in_proj_ba.output_sizes) // self.tp_size,
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)
218
+ if hasattr(self, "in_proj_qkv"):
219
+ # LoRA path: separate in_proj_qkv and in_proj_z
220
+ mixed_qkv, _ = self.in_proj_qkv(hidden_states)
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
+ )
230
+ qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
231
+ z_size = self.value_dim // self.tp_size
232
+ mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
233
z = z.reshape(z.size(0), -1, self.head_v_dim)
234
b, a = ba.chunk(2, dim=-1)
235
@@ -240,18 +281,14 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
281
model_config = vllm_config.model_config
282
cache_config = vllm_config.cache_config
283
quant_config = vllm_config.quant_config
243
- speculative_config = vllm_config.speculative_config
284
285
self.layer_type = layer_type
286
self.layer_idx = extract_layer_index(prefix)
287
288
if self.layer_type == "linear_attention":
289
self.linear_attn = Qwen3_5GatedDeltaNet(
250
- config,
251
- model_config=model_config,
252
- cache_config=cache_config,
253
- quant_config=quant_config,
254
- speculative_config=speculative_config,
290
+ config=config,
291
+ vllm_config=vllm_config,
292
prefix=f"{prefix}.linear_attn",
293
)
294
elif self.layer_type == "full_attention":
@@ -331,6 +368,7 @@ class Qwen3_5Model(Qwen3NextModel):
368
self.num_redundant_experts = eplb_config.num_redundant_experts
369
370
self.config = config
371
+ self.enable_lora = vllm_config.lora_config is not None
372
373
self.vocab_size = config.vocab_size
374
@@ -396,13 +434,25 @@ class Qwen3_5Model(Qwen3NextModel):
434
# mlp
435
("gate_up_proj", "gate_proj", 0),
436
("gate_up_proj", "up_proj", 1),
399
- # GDN
400
- ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
401
- ("in_proj_qkvz", "in_proj_z", 3),
437
("in_proj_ba", "in_proj_b", 0),
438
("in_proj_ba", "in_proj_a", 1),
439
]
440
441
+ if self.enable_lora:
442
+ stacked_params_mapping.extend(
443
+ [
444
+ ("in_proj_qkv", "in_proj_qkv", (0, 1, 2)),
445
+ ("in_proj_z", "in_proj_z", 0),
446
+ ]
447
+ )
448
+ else:
449
+ stacked_params_mapping.extend(
450
+ [
451
+ ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
452
+ ("in_proj_qkvz", "in_proj_z", 3),
453
+ ]
454
+ )
455
+
456
params_dict = dict(self.named_parameters())
457
loaded_params: set[str] = set()
458
expert_params_mapping = self.get_expert_mapping()
@@ -450,7 +500,10 @@ class Qwen3_5Model(Qwen3NextModel):
500
continue
501
param = params_dict[name]
502
weight_loader = param.weight_loader
453
- weight_loader(param, loaded_weight, shard_id)
503
+ if param_name == "in_proj_z" and self.enable_lora:
504
+ weight_loader(param, loaded_weight)
505
+ else:
506
+ weight_loader(param, loaded_weight, shard_id)
507
break
508
else:
509
is_expert_weight = False
@@ -580,6 +633,15 @@ class Qwen3_5ForCausalLMBase(
633
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
634
)
635
636
+ # When LoRA is enabled, GDN uses separate in_proj_qkv and in_proj_z
637
+ # instead of merged in_proj_qkvz; pack mapping must match.
638
+ if vllm_config.lora_config:
639
+ base = getattr(Qwen3_5ForCausalLMBase, "packed_modules_mapping", {})
640
+ self.packed_modules_mapping = {k: list(v) for k, v in base.items()}
641
+ self.packed_modules_mapping.pop("in_proj_qkvz", None)
642
+ self.packed_modules_mapping["in_proj_qkv"] = ["in_proj_qkv"]
643
+ self.packed_modules_mapping["in_proj_z"] = ["in_proj_z"]
644
+
645
if get_pp_group().is_last_rank:
646
if config.tie_word_embeddings:
647
self.lm_head = self.model.embed_tokens
@@ -672,6 +734,7 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
734
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
735
# protocols have not __init__ method, so we need to use nn.Module.__init__
736
nn.Module.__init__(self)
737
+ self.update_packed_mapping(enable_lora=vllm_config.lora_config is not None)
738
config: Qwen3_5Config = vllm_config.model_config.hf_config
739
quant_config = vllm_config.quant_config
740
multimodal_config = vllm_config.model_config.multimodal_config
@@ -699,6 +762,16 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
762
self.language_model.make_empty_intermediate_tensors
763
)
764
765
+ def update_packed_mapping(self, enable_lora: bool):
766
+ # When LoRA is enabled, GDN uses separate in_proj_qkv and in_proj_z
767
+ if enable_lora:
768
+ base = getattr(
769
+ Qwen3_5ForConditionalGeneration, "packed_modules_mapping", {}
770
+ )
771
+ self.packed_modules_mapping = {k: list(v) for k, v in base.items()}
772
+ self.packed_modules_mapping.pop("in_proj_qkvz", None)
773
+ self.packed_modules_mapping["in_proj_qkv"] = ["in_proj_qkv"]
774
+
775
def embed_input_ids(
776
self,
777
input_ids: torch.Tensor,
@@ -879,9 +952,13 @@ class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):
952
class Qwen3_5MoeForConditionalGeneration(
953
Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts
954
):
955
+ # For MoE LoRA weights loading
956
+ is_3d_moe_weight: bool = True
957
+
958
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
959
# protocols have not __init__ method, so we need to use nn.Module.__init__
960
nn.Module.__init__(self)
961
+ self.update_packed_mapping(enable_lora=vllm_config.lora_config is not None)
962
config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config
963
quant_config = vllm_config.quant_config
964
multimodal_config = vllm_config.model_config.multimodal_config
vllm/model_executor/models/qwen3_next.py
+17
-21
@@ -15,7 +15,6 @@ from vllm.compilation.decorators import support_torch_compile
15
from vllm.config import (
16
CacheConfig,
17
ModelConfig,
18
- SpeculativeConfig,
18
VllmConfig,
19
get_current_vllm_config,
20
)
@@ -401,11 +400,9 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
400
def __init__(
401
self,
402
config: Qwen3NextConfig,
404
- model_config: ModelConfig | None = None,
405
- cache_config: CacheConfig | None = None,
406
- quant_config: QuantizationConfig | None = None,
407
- speculative_config: SpeculativeConfig | None = None,
403
+ vllm_config: VllmConfig,
404
prefix: str = "",
405
+ create_in_proj_qkvz: bool = True,
406
) -> None:
407
super().__init__()
408
self.tp_size = get_tensor_model_parallel_world_size()
@@ -432,10 +429,10 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
429
)
430
431
self.config = config
435
- self.model_config = model_config
436
- self.cache_config = cache_config
437
- self.quant_config = quant_config
438
- self.speculative_config = speculative_config
432
+ self.model_config = vllm_config.model_config
433
+ self.cache_config = vllm_config.cache_config
434
+ quant_config = vllm_config.quant_config
435
+ self.speculative_config = vllm_config.speculative_config
436
self.num_spec = (
437
self.speculative_config.num_speculative_tokens
438
if self.speculative_config
@@ -455,13 +452,16 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
452
# projection of the input hidden states
453
# Qwen3-Next and Qwen3.5 has a different qkv_proj layout,
454
# we need to create qkvz_proj adaptively here.
458
- self.in_proj_qkvz = self.create_qkvz_proj(
459
- hidden_size=self.hidden_size,
460
- key_dim=self.key_dim,
461
- value_dim=self.value_dim,
462
- quant_config=quant_config,
463
- prefix=f"{prefix}.in_proj_qkvz",
464
- )
455
+ # When create_in_proj_qkvz is False (e.g. LoRA enabled in Qwen3.5),
456
+ # the subclass creates in_proj_qkv and in_proj_z separately.
457
+ if create_in_proj_qkvz:
458
+ self.in_proj_qkvz = self.create_qkvz_proj(
459
+ hidden_size=self.hidden_size,
460
+ key_dim=self.key_dim,
461
+ value_dim=self.value_dim,
462
+ quant_config=quant_config,
463
+ prefix=f"{prefix}.in_proj_qkvz",
464
+ )
465
# ba_proj doesn't support blockwise fp8 quantization.
466
# Qwen3-Next and Qwen3.5 have different in_proj_ba checkpoint
467
# layouts, so we use a factory method to create the projection.
@@ -1207,7 +1207,6 @@ class Qwen3NextDecoderLayer(nn.Module):
1207
model_config = vllm_config.model_config
1208
cache_config = vllm_config.cache_config
1209
quant_config = vllm_config.quant_config
1210
- speculative_config = vllm_config.speculative_config
1210
1211
self.layer_type = layer_type
1212
self.layer_idx = extract_layer_index(prefix)
@@ -1215,10 +1214,7 @@ class Qwen3NextDecoderLayer(nn.Module):
1214
if self.layer_type == "linear_attention":
1215
self.linear_attn = Qwen3NextGatedDeltaNet(
1216
config,
1218
- model_config=model_config,
1219
- cache_config=cache_config,
1220
- quant_config=quant_config,
1221
- speculative_config=speculative_config,
1217
+ vllm_config=vllm_config,
1218
prefix=f"{prefix}.linear_attn",
1219
)
1220
elif self.layer_type == "full_attention":