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 14, 05:30 UTC
[Feature] models: pass layer prefix to replace_linear_class for per-layer quantization routing. Addresses #23239 (#23556)
Signed-off-by: Shrey Gupta <shreyg1303@gmail.com>
Shrey Gupta committed
Aug 28, 2025 at 08:42 UTC
1b7b161a09289214eea41e17895a68a7ccd4b1dc
2 files changed
+18
-8
vllm/model_executor/models/deepseek_vl2.py
+8
-4
@@ -408,13 +408,17 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
408
if isinstance(module, nn.Linear):
409
parent, attr_name = self._get_parent_and_attr(vit, name)
410
if isinstance(parent, timm.layers.Mlp) and attr_name == "fc1":
411
- new_linear = replace_linear_class(module, "colwise",
412
- quant_config)
411
+ new_linear = replace_linear_class(module,
412
+ "colwise",
413
+ quant_config,
414
+ prefix=name)
415
setattr(parent, attr_name, new_linear)
416
elif isinstance(parent,
417
timm.layers.Mlp) and attr_name == "fc2":
416
- new_linear = replace_linear_class(module, "rowwise",
417
- quant_config)
418
+ new_linear = replace_linear_class(module,
419
+ "rowwise",
420
+ quant_config,
421
+ prefix=name)
422
setattr(parent, attr_name, new_linear)
423
424
return vit
vllm/model_executor/models/transformers.py
+10
-4
@@ -106,8 +106,11 @@ def can_enable_torch_compile(vllm_config: VllmConfig) -> bool:
106
107
108
def replace_linear_class(
109
- linear: nn.Linear, style: Literal["colwise", "rowwise"],
110
- quant_config: QuantizationConfig
109
+ linear: nn.Linear,
110
+ style: Literal["colwise", "rowwise"],
111
+ quant_config: QuantizationConfig,
112
+ *,
113
+ prefix: str = "",
114
) -> Union[ColumnParallelLinear, RowParallelLinear, ReplicatedLinear]:
115
"""
116
Replace nn.Linear with one of vLLM's tensor parallel linear classes.
@@ -141,6 +144,7 @@ def replace_linear_class(
144
output_size=linear.out_features,
145
bias=linear.bias is not None,
146
quant_config=quant_config,
147
+ prefix=prefix,
148
return_bias=False,
149
**vllm_linear_kwargs,
150
)
@@ -557,8 +561,10 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
561
generator = (p for p in tp_plan if re.match(p, qual_name))
562
pattern = next(generator, None)
563
style = tp_plan.get(pattern, "replicate")
560
- new_module = replace_linear_class(child_module, style,
561
- self.quant_config)
564
+ new_module = replace_linear_class(child_module,
565
+ style,
566
+ self.quant_config,
567
+ prefix=qual_name)
568
setattr(module, child_name, new_module)
569
log_replacement(qual_name, child_module, new_module)
570
else: