@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5708297e4e

Fork 0

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 10, 17:45 UTC

[Misc][Model][Refactor] Pass the prefix into Linear layers (#31669)

Signed-off-by: Wang Kunpeng <1289706727@qq.com>

Wang Kunpeng committed Jan 6, 2026 at 04:03 UTC 5708297e4e7f27c4d06d6e4f62d75c3d2d1fc674

17 files changed +181 -40

vllm/model_executor/layers/mamba/mamba_mixer.py

+12 -2

@@ -82,6 +82,7 @@ class MambaMixer(MambaBase, CustomOp):

82 input_size=conv_kernel_size,

83 output_size=intermediate_size,

84 bias=use_conv_bias,

85 + prefix=f"{prefix}.conv1d",

86 )

87 # unsqueeze to fit conv1d weights shape into the linear weights shape.

88 # Can't do this in `weight_loader` since it already exists in

@@ -90,7 +91,10 @@ class MambaMixer(MambaBase, CustomOp):

91 self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)

93 self.in_proj = MergedColumnParallelLinear(

93 - hidden_size, [intermediate_size] * 2, bias=use_bias

94 + hidden_size,

95 + [intermediate_size] * 2,

96 + bias=use_bias,

97 + prefix=f"{prefix}.in_proj",

98 )

100 # selective projection used to make dt, B and C input dependent

@@ -98,12 +102,17 @@ class MambaMixer(MambaBase, CustomOp):

102 intermediate_size,

103 time_step_rank + ssm_state_size * 2,

104 bias=False,

105 + prefix=f"{prefix}.x_proj",

106 )

107 # time step projection (discretization) -

108 # In the forward we need to apply dt_proj without the bias,

109 # as the bias is added in the selective scan kernel.

110 self.dt_proj = ColumnParallelLinear(

106 - time_step_rank, intermediate_size, bias=True, skip_bias_add=True

111 + time_step_rank,

112 + intermediate_size,

113 + bias=True,

114 + skip_bias_add=True,

115 + prefix=f"{prefix}.dt_proj",

116 )

117

118 def weight_loader(param: Parameter, loaded_weight: torch.Tensor):

@@ -136,6 +145,7 @@ class MambaMixer(MambaBase, CustomOp):

145 hidden_size,

146 bias=use_bias,

147 input_is_parallel=True,

148 + prefix=f"{prefix}.out_proj",

149 )

150

151 self.dt_layernorm = (

vllm/model_executor/models/aria.py

+15 -5

@@ -127,11 +127,16 @@ class AriaProjectorMLP(nn.Module):

127 in_features: int,

128 hidden_features: int,

129 output_dim: int,

130 + prefix: str = "",

131 ) -> None:

132 super().__init__()

133

133 - self.linear_in = ColumnParallelLinear(in_features, hidden_features, bias=False)

134 - self.linear_out = RowParallelLinear(hidden_features, output_dim, bias=False)

134 + self.linear_in = ColumnParallelLinear(

135 + in_features, hidden_features, bias=False, prefix=f"{prefix}.linear_in"

136 + )

137 + self.linear_out = RowParallelLinear(

138 + hidden_features, output_dim, bias=False, prefix=f"{prefix}.linear_out"

139 + )

140 self.act = get_act_fn("gelu_new")

141

142 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

@@ -154,7 +159,7 @@ class AriaProjector(nn.Module):

159 A tensor with the shape of (batch_size, query_number, output_dim)

160 """

161

157 - def __init__(self, config: AriaConfig) -> None:

162 + def __init__(self, config: AriaConfig, prefix: str = "") -> None:

163 super().__init__()

164

165 self.patch_to_query_dict = config.projector_patch_to_query_dict

@@ -174,7 +179,10 @@ class AriaProjector(nn.Module):

179

180 self.layer_norm = nn.LayerNorm(self.in_features)

181 self.feed_forward = AriaProjectorMLP(

177 - self.in_features, self.hidden_features, self.output_dim

182 + self.in_features,

183 + self.hidden_features,

184 + self.output_dim,

185 + prefix=f"{prefix}.feed_forward",

186 )

187

188 def forward(

@@ -536,7 +544,9 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

544 quant_config=quant_config,

545 prefix=f"{prefix}.vision_tower",

546 )

539 - self.multi_modal_projector = AriaProjector(config)

547 + self.multi_modal_projector = AriaProjector(

548 + config, prefix=maybe_prefix(prefix, "multi_modal_projector")

549 + )

550 self.vocab_size = config.text_config.vocab_size

551 self.language_model = AriaTextModel(

552 vllm_config=vllm_config.with_hf_config(config.text_config),

vllm/model_executor/models/gpt_neox.py

+1 -1

@@ -166,7 +166,7 @@ class GPTNeoXLayer(nn.Module):

166 self.attention = GPTNeoXAttention(

167 config, cache_config, quant_config, prefix=f"{prefix}.attention"

168 )

169 - self.mlp = GPTNeoXMLP(config, quant_config)

169 + self.mlp = GPTNeoXMLP(config, quant_config, prefix=f"{prefix}.mlp")

170

171 def forward(

172 self,

vllm/model_executor/models/hunyuan_v1.py

@@ -427,6 +427,7 @@ class HunYuanSparseMoeBlock(nn.Module):

427 hidden_act=config.hidden_act,

428 quant_config=quant_config,

429 reduce_results=False,

430 + prefix=f"{prefix}.shared_mlp",

431 )

432 else:

433 self.shared_mlp = None

vllm/model_executor/models/jamba.py

@@ -78,6 +78,7 @@ class JambaMoE(nn.Module):

78 bias=False,

79 quant_config=None,

80 params_dtype=params_dtype,

81 + prefix=f"{prefix}.router",

82 )

84 self.experts = FusedMoE(

vllm/model_executor/models/jina_vl.py

+14 -4

@@ -27,15 +27,23 @@ logger = init_logger(__name__)

29 class JinaVLScorer(nn.Module):

30 - def __init__(self, model_config: "ModelConfig"):

30 + def __init__(self, model_config: "ModelConfig", prefix: str = ""):

31 super().__init__()

32 config = model_config.hf_config.get_text_config()

33 head_dtype = model_config.head_dtype

34 self.dense = ColumnParallelLinear(

35 - config.hidden_size, config.hidden_size, params_dtype=head_dtype, bias=True

35 + config.hidden_size,

36 + config.hidden_size,

37 + params_dtype=head_dtype,

38 + bias=True,

39 + prefix=f"{prefix}.dense",

40 )

41 self.out_proj = RowParallelLinear(

38 - config.hidden_size, config.num_labels, params_dtype=head_dtype, bias=True

42 + config.hidden_size,

43 + config.num_labels,

44 + params_dtype=head_dtype,

45 + bias=True,

46 + prefix=f"{prefix}.out_proj",

47 )

49 def forward(self, x, **kwargs):

@@ -94,7 +102,9 @@ class JinaVLForSequenceClassification(

102 pooler_config = vllm_config.model_config.pooler_config

103 assert pooler_config is not None

104

97 - self.score = JinaVLScorer(vllm_config.model_config)

105 + self.score = JinaVLScorer(

106 + vllm_config.model_config, prefix=maybe_prefix(prefix, "score")

107 + )

108 self.pooler = DispatchPooler(

109 {

110 "token_classify": Pooler.for_token_classify(

vllm/model_executor/models/minicpm.py

@@ -90,6 +90,7 @@ class MiniCPMMoE(nn.Module):

90 intermediate_size: int,

91 params_dtype: torch.dtype | None = None,

92 tp_size: int | None = None,

93 + prefix: str = "",

94 ):

95 super().__init__()

96 self.tp_size = tp_size or get_tensor_model_parallel_world_size()

@@ -108,6 +109,7 @@ class MiniCPMMoE(nn.Module):

109 bias=False,

110 params_dtype=self.params_dtype,

111 quant_config=None,

112 + prefix=f"{prefix}.gate",

113 )

114

115 self.ws = nn.Parameter(

@@ -352,6 +354,7 @@ class MiniCPMDecoderLayer(nn.Module):

354 hidden_act=self.config.hidden_act,

355 hidden_act_param=getattr(self.config, "hidden_act_param", 0.0),

356 quant_config=self.quant_config,

357 + prefix=f"{self.prefix}.mlp",

358 )

359 else:

360 self.mlp = MiniCPMMoE(

@@ -359,6 +362,7 @@ class MiniCPMDecoderLayer(nn.Module):

362 top_k=self.config.num_experts_per_tok,

363 hidden_size=self.config.hidden_size,

364 intermediate_size=self.config.intermediate_size,

365 + prefix=f"{self.prefix}.mlp",

366 )

367

368 def forward(

vllm/model_executor/models/minicpm_eagle.py

@@ -108,6 +108,7 @@ class EagleMiniCPMDecoderLayer(nn.Module):

108 top_k=self.config.num_experts_per_tok,

109 hidden_size=self.config.hidden_size,

110 intermediate_size=self.config.intermediate_size,

111 + prefix=f"{self.prefix}.mlp",

112 )

113

114 def forward(

vllm/model_executor/models/mistral_large_3_eagle.py

@@ -67,6 +67,7 @@ class EagleMistralLarge3Model(DeepseekV2Model):

67 input_is_parallel=False,

68 quant_config=quant_config,

69 return_bias=False,

70 + prefix=maybe_prefix(prefix, "fc"),

71 )

72 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

73 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

vllm/model_executor/models/modernbert.py

+25 -8

@@ -63,7 +63,9 @@ class ModernBertEmbeddings(nn.Module):

65 class ModernBertAttention(nn.Module):

66 - def __init__(self, config: ModernBertConfig, layer_id: int | None = None):

66 + def __init__(

67 + self, config: ModernBertConfig, layer_id: int | None = None, prefix: str = ""

68 + ):

69 super().__init__()

70 self.config = config

71 self.hidden_size = config.hidden_size

@@ -80,6 +82,7 @@ class ModernBertAttention(nn.Module):

82 self.head_dim,

83 self.num_heads,

84 bias=config.attention_bias,

85 + prefix=f"{prefix}.Wqkv",

86 )

88 if layer_types := getattr(config, "layer_types", None):

@@ -117,7 +120,10 @@ class ModernBertAttention(nn.Module):

120 per_layer_sliding_window=sliding_window,

121 )

122 self.Wo = RowParallelLinear(

120 - config.hidden_size, config.hidden_size, bias=config.attention_bias

123 + config.hidden_size,

124 + config.hidden_size,

125 + bias=config.attention_bias,

126 + prefix=f"{prefix}.Wo",

127 )

128

129 def forward(

@@ -135,7 +141,7 @@ class ModernBertAttention(nn.Module):

141

142

143 class ModernBertMLP(nn.Module):

138 - def __init__(self, config: ModernBertConfig):

144 + def __init__(self, config: ModernBertConfig, prefix: str = ""):

145 super().__init__()

146 self.config = config

147 self.Wi = nn.Linear(

@@ -143,7 +149,10 @@ class ModernBertMLP(nn.Module):

149 )

150 self.act = nn.GELU()

151 self.Wo = RowParallelLinear(

146 - config.intermediate_size, config.hidden_size, bias=config.mlp_bias

152 + config.intermediate_size,

153 + config.hidden_size,

154 + bias=config.mlp_bias,

155 + prefix=f"{prefix}.Wo",

156 )

157

158 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

@@ -163,11 +172,13 @@ class ModernBertLayer(nn.Module):

172 self.attn_norm = nn.LayerNorm(

173 config.hidden_size, eps=config.norm_eps, bias=config.norm_bias

174 )

166 - self.attn = ModernBertAttention(config=config, layer_id=layer_id)

175 + self.attn = ModernBertAttention(

176 + config=config, layer_id=layer_id, prefix=f"{prefix}.attn"

177 + )

178 self.mlp_norm = nn.LayerNorm(

179 config.hidden_size, eps=config.norm_eps, bias=config.norm_bias

180 )

170 - self.mlp = ModernBertMLP(config)

181 + self.mlp = ModernBertMLP(config, prefix=f"{prefix}.mlp")

182

183 def forward(

184 self,

@@ -189,7 +200,11 @@ class ModernBertEncoderLayer(nn.Module):

200 config = vllm_config.model_config.hf_config

201 self.layers = nn.ModuleList(

202 [

192 - ModernBertLayer(config=config, layer_id=layer_id)

203 + ModernBertLayer(

204 + config=config,

205 + layer_id=layer_id,

206 + prefix=f"{prefix}.layers.{layer_id}",

207 + )

208 for layer_id in range(config.num_hidden_layers)

209 ]

210 )

@@ -220,7 +235,9 @@ class ModernBertModel(nn.Module):

235 config = vllm_config.model_config.hf_config

236 self.config = config

237 self.embeddings = ModernBertEmbeddings(config)

223 - self.encoder_layer = ModernBertEncoderLayer(vllm_config)

238 + self.encoder_layer = ModernBertEncoderLayer(

239 + vllm_config, prefix=f"{prefix}.encoder_layer"

240 + )

241 self.final_norm = nn.LayerNorm(

242 config.hidden_size, eps=config.norm_eps, bias=config.norm_bias

243 )

vllm/model_executor/models/molmo.py

+50 -9

@@ -142,6 +142,7 @@ class ViTMLP(nn.Module):

142 self,

143 config: VisionBackboneConfig,

144 quant_config: QuantizationConfig | None = None,

145 + prefix: str = "",

146 ):

147 super().__init__()

148 self.w1 = ColumnParallelLinear(

@@ -149,6 +150,7 @@ class ViTMLP(nn.Module):

150 config.image_mlp_dim,

151 bias=True,

152 quant_config=quant_config,

153 + prefix=f"{prefix}.w1",

154 )

155 # Activation function.

156 assert config.image_mlp_activations == "quick_gelu"

@@ -158,6 +160,7 @@ class ViTMLP(nn.Module):

160 config.image_emb_dim,

161 bias=True,

162 quant_config=quant_config,

163 + prefix=f"{prefix}.w2",

164 )

165

166 def forward(self, x: torch.Tensor) -> torch.Tensor:

@@ -176,6 +179,7 @@ class MultiHeadDotProductAttention(nn.Module):

179 use_bias: bool = True,

180 nlayers: int = 1,

181 quant_config: QuantizationConfig | None = None,

182 + prefix: str = "",

183 ):

184 super().__init__()

185

@@ -202,24 +206,28 @@ class MultiHeadDotProductAttention(nn.Module):

206 self.total_num_heads * self.head_dim,

207 bias=use_bias,

208 quant_config=quant_config,

209 + prefix=f"{prefix}.wq",

210 )

211 self.wk = ColumnParallelLinear(

212 nlayers * self.hidden_size,

213 self.total_num_kv_heads * self.head_dim,

214 bias=use_bias,

215 quant_config=quant_config,

216 + prefix=f"{prefix}.wk",

217 )

218 self.wv = ColumnParallelLinear(

219 nlayers * self.hidden_size,

220 self.total_num_kv_heads * self.head_dim,

221 bias=use_bias,

222 quant_config=quant_config,

223 + prefix=f"{prefix}.wv",

224 )

225 self.wo = RowParallelLinear(

226 self.total_num_heads * self.head_dim,

227 self.hidden_size,

228 bias=use_bias,

229 quant_config=quant_config,

230 + prefix=f"{prefix}.wo",

231 )

232

233 self.scale = self.head_dim**-0.5

@@ -254,10 +262,15 @@ class ResidualAttentionBlock(nn.Module):

262 self,

263 config: VisionBackboneConfig,

264 quant_config: QuantizationConfig | None = None,

265 + prefix: str = "",

266 ):

267 super().__init__()

259 - self.attention = MultiHeadDotProductAttention(config, quant_config=quant_config)

260 - self.feed_forward = ViTMLP(config, quant_config)

268 + self.attention = MultiHeadDotProductAttention(

269 + config, quant_config=quant_config, prefix=f"{prefix}.attention"

270 + )

271 + self.feed_forward = ViTMLP(

272 + config, quant_config, prefix=f"{prefix}.feed_forward"

273 + )

274 self.attention_norm = nn.LayerNorm(

275 config.image_emb_dim,

276 eps=config.image_norm_eps,

@@ -280,12 +293,15 @@ class BlockCollection(nn.Module):

293 self,

294 config: VisionBackboneConfig,

295 quant_config: QuantizationConfig | None = None,

296 + prefix: str = "",

297 ):

298 super().__init__()

299 self.resblocks = nn.ModuleList(

300 [

287 - ResidualAttentionBlock(config, quant_config)

288 - for _ in range(config.image_num_layers)

301 + ResidualAttentionBlock(

302 + config, quant_config, prefix=f"{prefix}.resblocks.{i}"

303 + )

304 + for i in range(config.image_num_layers)

305 ]

306 )

307

@@ -308,6 +324,7 @@ class VisionTransformer(nn.Module):

324 self,

325 config: VisionBackboneConfig,

326 quant_config: QuantizationConfig | None = None,

327 + prefix: str = "",

328 ):

329 super().__init__()

330 scale = config.image_emb_dim**-0.5

@@ -324,7 +341,9 @@ class VisionTransformer(nn.Module):

341 bias=False,

342 )

343 self.pre_ln = nn.LayerNorm(config.image_emb_dim, eps=config.image_norm_eps)

327 - self.transformer = BlockCollection(config, quant_config)

344 + self.transformer = BlockCollection(

345 + config, quant_config, prefix=f"{prefix}.transformer"

346 + )

347

348 def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor:

349 cls_emb = self.positional_embedding[0:1]

@@ -419,6 +438,7 @@ class MolmoAttention(nn.Module):

438 self.total_num_kv_heads,

439 bias=config.qkv_bias,

440 quant_config=quant_config,

441 + prefix=f"{prefix}.qkv_proj",

442 )

443

444 self.tp_rank: int | None = None

@@ -454,6 +474,7 @@ class MolmoAttention(nn.Module):

474 self.hidden_size,

475 bias=False,

476 quant_config=quant_config,

477 + prefix=f"{prefix}.o_proj",

478 )

479

480 def _apply_qk_norm(

@@ -493,6 +514,7 @@ class LanguageModelMLP(nn.Module):

514 config: PretrainedConfig,

515 input_dim: int | None = None,

516 quant_config: QuantizationConfig | None = None,

517 + prefix: str = "",

518 ) -> None:

519 super().__init__()

520 self.hidden_size = config.hidden_size

@@ -503,6 +525,7 @@ class LanguageModelMLP(nn.Module):

525 [self.intermediate_size] * 2,

526 bias=False,

527 quant_config=quant_config,

528 + prefix=f"{prefix}.gate_up_proj",

529 )

530 # Activation function.

531 self.act_fn = MulAndSilu()

@@ -512,6 +535,7 @@ class LanguageModelMLP(nn.Module):

535 self.hidden_size,

536 bias=False,

537 quant_config=quant_config,

538 + prefix=f"{prefix}.down_proj",

539 )

540

541 def forward(

@@ -532,6 +556,7 @@ class ImageProjectorMLP(nn.Module):

556 config: PretrainedConfig,

557 input_dim: int | None = None,

558 quant_config: QuantizationConfig | None = None,

559 + prefix: str = "",

560 ) -> None:

561 super().__init__()

562 self.hidden_size = config.hidden_size

@@ -542,6 +567,7 @@ class ImageProjectorMLP(nn.Module):

567 [self.intermediate_size] * 2,

568 bias=False,

569 quant_config=quant_config,

570 + prefix=f"{prefix}.merged_linear",

571 )

572 # Activation function.

573 self.act_fn = SiluAndMul()

@@ -552,6 +578,7 @@ class ImageProjectorMLP(nn.Module):

578 self.hidden_size,

579 bias=False,

580 quant_config=quant_config,

581 + prefix=f"{prefix}.down_proj",

582 )

583

584 def forward(

@@ -579,7 +606,9 @@ class MolmoDecoderLayer(nn.Module):

606 )

607

608 # MLP block.

582 - self.mlp = LanguageModelMLP(config, quant_config=quant_config)

609 + self.mlp = LanguageModelMLP(

610 + config, quant_config=quant_config, prefix=f"{prefix}.mlp"

611 + )

612

613 # LayerNorm

614 assert config.layer_norm_type == "rms"

@@ -643,6 +672,7 @@ class MolmoVisionBackbone(nn.Module, SupportsQuant):

672 config: PretrainedConfig,

673 vision_config: VisionBackboneConfig,

674 quant_config: QuantizationConfig | None = None,

675 + prefix: str = "",

676 ) -> None:

677 super().__init__()

678 self.vit_layers = VIT_LAYERS

@@ -651,18 +681,24 @@ class MolmoVisionBackbone(nn.Module, SupportsQuant):

681 (self.image_num_patch[0] + 1) // POOLING_SIZE,

682 (self.image_num_patch[1] + 1) // POOLING_SIZE,

683 )

654 - self.image_vit = VisionTransformer(vision_config, quant_config=quant_config)

684 + self.image_vit = VisionTransformer(

685 + vision_config, quant_config=quant_config, prefix=f"{prefix}.image_vit"

686 + )

687 self.num_prefix_tokens = self.image_vit.num_prefix_tokens

688 assert self.num_prefix_tokens in {0, 1}, (

689 "Only 0 or 1 prefix tokens are supported"

690 )

691 self.image_pooling_2d = MultiHeadDotProductAttention(

660 - vision_config, nlayers=len(self.vit_layers), quant_config=quant_config

692 + vision_config,

693 + nlayers=len(self.vit_layers),

694 + quant_config=quant_config,

695 + prefix=f"{prefix}.image_pooling_2d",

696 )

697 self.image_projector = ImageProjectorMLP(

698 config,

699 input_dim=vision_config.image_emb_dim,

700 quant_config=quant_config,

701 + prefix=f"{prefix}.image_projector",

702 )

703

704 image_dim = vision_config.image_emb_dim * len(self.vit_layers)

@@ -1405,7 +1441,12 @@ class MolmoForCausalLM(

1441 self.multimodal_config = multimodal_config

1442

1443 vision_config = VisionBackboneConfig()

1408 - self.vision_backbone = MolmoVisionBackbone(config, vision_config, quant_config)

1444 + self.vision_backbone = MolmoVisionBackbone(

1445 + config,

1446 + vision_config,

1447 + quant_config,

1448 + prefix=maybe_prefix(prefix, "vision_backbone"),

1449 + )

1450 self.model = MolmoModel(

1451 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

1452 )

vllm/model_executor/models/olmoe.py

+5 -1

@@ -86,7 +86,11 @@ class OlmoeMoE(nn.Module):

87 # Gate always runs at half / full precision for now.

88 self.gate = ReplicatedLinear(

89 - hidden_size, num_experts, bias=False, quant_config=None

89 + hidden_size,

90 + num_experts,

91 + bias=False,

92 + quant_config=None,

93 + prefix=f"{prefix}.gate",

94 )

96 self.experts = FusedMoE(

vllm/model_executor/models/phimoe.py

@@ -272,6 +272,7 @@ class PhiMoE(nn.Module):

272 bias=False,

273 params_dtype=params_dtype,

274 quant_config=None,

275 + prefix=f"{prefix}.gate",

276 )

277

278 self.experts = FusedMoE(

vllm/model_executor/models/qwen.py

+15 -3

@@ -56,13 +56,22 @@ class QWenMLP(nn.Module):

56 intermediate_size: int,

57 hidden_act: str = "silu",

58 quant_config: QuantizationConfig | None = None,

59 + prefix: str = "",

60 ):

61 super().__init__()

62 self.gate_up_proj = MergedColumnParallelLinear(

62 - hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config

63 + hidden_size,

64 + [intermediate_size] * 2,

65 + bias=False,

66 + quant_config=quant_config,

67 + prefix=f"{prefix}.gate_up_proj",

68 )

69 self.c_proj = RowParallelLinear(

65 - intermediate_size, hidden_size, bias=False, quant_config=quant_config

70 + intermediate_size,

71 + hidden_size,

72 + bias=False,

73 + quant_config=quant_config,

74 + prefix=f"{prefix}.c_proj",

75 )

76 if hidden_act != "silu":

77 raise ValueError(

@@ -163,7 +172,10 @@ class QWenBlock(nn.Module):

172 self.ln_2 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

173

174 self.mlp = QWenMLP(

166 - config.hidden_size, config.intermediate_size // 2, quant_config=quant_config

175 + config.hidden_size,

176 + config.intermediate_size // 2,

177 + quant_config=quant_config,

178 + prefix=f"{prefix}.mlp",

179 )

180

181 def forward(

vllm/model_executor/models/qwen3_next.py

@@ -864,6 +864,7 @@ class Qwen3NextDecoderLayer(nn.Module):

864 intermediate_size=config.intermediate_size,

865 hidden_act=config.hidden_act,

866 quant_config=quant_config,

867 + prefix=f"{prefix}.mlp",

868 )

869

870 self.input_layernorm = Qwen3NextRMSNorm(

vllm/model_executor/models/qwen_vl.py

+26 -6

@@ -109,6 +109,7 @@ class VisualAttention(nn.Module):

109 bias: bool = True,

110 kdim: int | None = None,

111 vdim: int | None = None,

112 + prefix: str = "",

113 ):

114 super().__init__()

115 self.embed_dim = embed_dim

@@ -128,8 +129,12 @@ class VisualAttention(nn.Module):

129 assert self._qkv_same_embed_dim, (

130 "Visual Attention implementation only supports self-attention"

131 )

131 - self.in_proj = ReplicatedLinear(embed_dim, 3 * embed_dim)

132 - self.out_proj = ReplicatedLinear(embed_dim, embed_dim)

132 + self.in_proj = ReplicatedLinear(

133 + embed_dim, 3 * embed_dim, prefix=f"{prefix}.in_proj"

134 + )

135 + self.out_proj = ReplicatedLinear(

136 + embed_dim, embed_dim, prefix=f"{prefix}.out_proj"

137 + )

138 self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)

139

140 def forward(

@@ -214,10 +219,15 @@ class QwenVLMLP(nn.Module):

219 hidden_size: int,

220 intermediate_size: int,

221 quant_config: QuantizationConfig | None = None,

222 + prefix: str = "",

223 ):

224 super().__init__()

225 self.c_fc = ColumnParallelLinear(

220 - hidden_size, intermediate_size, bias=True, quant_config=quant_config

226 + hidden_size,

227 + intermediate_size,

228 + bias=True,

229 + quant_config=quant_config,

230 + prefix=f"{prefix}.c_fc",

231 )

232 self.act_fn = get_act_fn("gelu")

233 self.c_proj = RowParallelLinear(

@@ -225,6 +235,7 @@ class QwenVLMLP(nn.Module):

235 hidden_size,

236 bias=True,

237 quant_config=quant_config,

238 + prefix=f"{prefix}.c_proj",

239 )

240

241 def forward(self, x):

@@ -242,17 +253,19 @@ class VisualAttentionBlock(nn.Module):

253 mlp_ratio: float = 4.0,

254 norm_layer: Callable[[int], nn.Module] = nn.LayerNorm,

255 quant_config: QuantizationConfig | None = None,

256 + prefix: str = "",

257 ):

258 super().__init__()

259

260 self.ln_1 = norm_layer(d_model)

261 self.ln_2 = norm_layer(d_model)

262 mlp_width = int(d_model * mlp_ratio)

251 - self.attn = VisualAttention(d_model, n_head)

263 + self.attn = VisualAttention(d_model, n_head, prefix=f"{prefix}.attn")

264 self.mlp = QwenVLMLP(

265 hidden_size=d_model,

266 intermediate_size=mlp_width,

267 quant_config=quant_config,

268 + prefix=f"{prefix}.mlp",

269 )

270

271 def attention(

@@ -282,6 +295,7 @@ class TransformerBlock(nn.Module):

295 mlp_ratio: float = 4.0,

296 norm_layer: Callable[[int], nn.Module] = nn.LayerNorm,

297 quant_config: QuantizationConfig | None = None,

298 + prefix: str = "",

299 ):

300 super().__init__()

301 self.width = width

@@ -295,8 +309,9 @@ class TransformerBlock(nn.Module):

309 mlp_ratio,

310 norm_layer=norm_layer,

311 quant_config=quant_config,

312 + prefix=f"{prefix}.resblocks.{i}",

313 )

299 - for _ in range(layers)

314 + for i in range(layers)

315 ]

316 )

317

@@ -327,6 +342,7 @@ class VisionTransformer(nn.Module):

342 output_dim: int = 512,

343 image_start_id: int = 151857,

344 quant_config: QuantizationConfig | None = None,

345 + prefix: str = "",

346 **kwargs,

347 ):

348 super().__init__()

@@ -356,6 +372,7 @@ class VisionTransformer(nn.Module):

372 mlp_ratio,

373 norm_layer=norm_layer,

374 quant_config=quant_config,

375 + prefix=f"{prefix}.transformer",

376 )

377

378 self.attn_pool = Resampler2(

@@ -366,6 +383,7 @@ class VisionTransformer(nn.Module):

383 norm_layer=norm_layer,

384 adaptive=False,

385 do_post_projection=False,

386 + prefix=f"{prefix}.attn_pool",

387 ).to(

388 device=self.positional_embedding.device,

389 dtype=self.positional_embedding.dtype,

@@ -413,7 +431,9 @@ class QwenVLModel(QWenModel):

431 config = vllm_config.model_config.hf_config

432 quant_config = vllm_config.quant_config

433

416 - self.visual = VisionTransformer(**config.visual, quant_config=quant_config)

434 + self.visual = VisionTransformer(

435 + **config.visual, quant_config=quant_config, prefix=f"{prefix}.visual"

436 + )

437

438

439 @lru_cache(maxsize=1)

vllm/model_executor/models/zamba2.py

+8 -1

@@ -86,7 +86,13 @@ class Zamba2LoRA(nn.Module):

86 B_class = MergedColumnParallelLinear

87 else:

88 B_class = ColumnParallelLinear

89 - self.B = B_class(rank, output_dim, bias=False, quant_config=quant_config)

89 + self.B = B_class(

90 + rank,

91 + output_dim,

92 + bias=False,

93 + quant_config=quant_config,

94 + prefix=f"{prefix}.B",

95 + )

97 def forward(

98 self,

@@ -346,6 +352,7 @@ class Zamba2MLP(nn.Module):

352 config.adapter_rank,

353 2 * [self.intermediate_size],

354 quant_config,

355 + prefix=f"{prefix}.gate_up_proj_adapter_list.{block_idx}",

356 )

357 else:

358 gate_up_proj_adapter = nn.Identity()