1958bda9b4 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/model_executor/models/arctic.py

+8 -1

@@ -75,7 +75,11 @@ class ArcticMLP(nn.Module):

75 )

76

77 self.w13 = MergedColumnParallelLinear(

78 - self.hidden_size, [self.ffn_dim] * 2, bias=False, quant_config=quant_config

78 + self.hidden_size,

79 + [self.ffn_dim] * 2,

80 + bias=False,

81 + quant_config=quant_config,

82 + prefix=f"{prefix}.w13",

83 )

84 self.w2 = RowParallelLinear(

85 self.ffn_dim,

@@ -83,6 +87,7 @@ class ArcticMLP(nn.Module):

87 bias=False,

88 reduce_results=reduce_results,

89 quant_config=quant_config,

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

91 )

92 if config.hidden_act != "silu":

93 raise ValueError(

@@ -297,6 +302,7 @@ class ArcticAttention(nn.Module):

302 self.total_num_kv_heads,

303 bias=False,

304 quant_config=quant_config,

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

306 )

307 self.o_proj = RowParallelLinear(

308 self.total_num_heads * self.head_dim,

@@ -304,6 +310,7 @@ class ArcticAttention(nn.Module):

310 bias=False,

311 reduce_results=True,

312 quant_config=quant_config,

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

314 )

315

316 self.rotary_emb = get_rope(

vllm/model_executor/models/baichuan.py

+14 -2

@@ -98,13 +98,22 @@ class BaiChuanMLP(nn.Module):

98 intermediate_size: int,

99 hidden_act: str,

100 quant_config: QuantizationConfig | None = None,

101 + prefix: str = "",

102 ):

103 super().__init__()

104 self.gate_up_proj = MergedColumnParallelLinear(

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

105 + hidden_size,

106 + [intermediate_size] * 2,

107 + bias=False,

108 + quant_config=quant_config,

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

110 )

111 self.down_proj = RowParallelLinear(

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

112 + intermediate_size,

113 + hidden_size,

114 + bias=False,

115 + quant_config=quant_config,

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

117 )

118 if hidden_act != "silu":

119 raise ValueError(

@@ -152,12 +161,14 @@ class BaiChuanAttention(nn.Module):

161 self.total_num_heads,

162 bias=False,

163 quant_config=quant_config,

164 + prefix=f"{prefix}.W_pack",

165 )

166 self.o_proj = RowParallelLinear(

167 self.total_num_heads * self.head_dim,

168 hidden_size,

169 bias=False,

170 quant_config=quant_config,

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

172 )

173 # Create the alibi slopes and slice them.

174 if self.position_embedding == "ALIBI":

@@ -235,6 +246,7 @@ class BaiChuanDecoderLayer(nn.Module):

246 intermediate_size=config.intermediate_size,

247 hidden_act=config.hidden_act,

248 quant_config=quant_config,

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

250 )

251 self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

252 self.post_attention_layernorm = RMSNorm(

vllm/model_executor/models/bamba.py

+11 -2

@@ -60,6 +60,7 @@ class BambaMLP(nn.Module):

60 config: BambaConfig,

61 quant_config: QuantizationConfig | None = None,

62 bias: bool = False,

63 + prefix: str = "",

64 ) -> None:

65 super().__init__()

66 self.gate_up_proj = MergedColumnParallelLinear(

@@ -67,12 +68,14 @@ class BambaMLP(nn.Module):

68 output_sizes=[config.intermediate_size] * 2,

69 bias=bias,

70 quant_config=quant_config,

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

72 )

73 self.down_proj = RowParallelLinear(

74 input_size=config.intermediate_size,

75 output_size=config.hidden_size,

76 bias=bias,

77 quant_config=quant_config,

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

79 )

80 if config.hidden_act != "silu":

81 raise ValueError(

@@ -118,7 +121,9 @@ class BambaMixerDecoderLayer(nn.Module):

121 prefix=f"{prefix}.mixer",

122 )

123

121 - self.feed_forward = BambaMLP(config, quant_config=quant_config)

124 + self.feed_forward = BambaMLP(

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

126 + )

127 self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

128 self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

129

@@ -202,12 +207,14 @@ class BambaAttentionDecoderLayer(nn.Module):

207 self.total_num_kv_heads,

208 bias=False,

209 quant_config=quant_config,

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

211 )

212 self.o_proj = RowParallelLinear(

213 self.total_num_heads * self.head_dim,

214 config.hidden_size,

215 bias=False,

216 quant_config=quant_config,

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

218 )

219

220 self.attn = Attention(

@@ -219,7 +226,9 @@ class BambaAttentionDecoderLayer(nn.Module):

226 prefix=f"{prefix}.attn",

227 )

228

222 - self.feed_forward = BambaMLP(config, quant_config=quant_config)

229 + self.feed_forward = BambaMLP(

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

231 + )

232 self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

233 self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

234

vllm/model_executor/models/bloom.py

+6 -1

@@ -108,12 +108,14 @@ class BloomAttention(nn.Module):

108 self.total_num_heads,

109 bias=True,

110 quant_config=quant_config,

111 + prefix=f"{prefix}.query_key_value",

112 )

113 self.dense = RowParallelLinear(

114 self.hidden_size,

115 self.hidden_size,

116 bias=True,

117 quant_config=quant_config,

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

119 )

120

121 # Create the alibi slopes and slice them.

@@ -152,6 +154,7 @@ class BloomMLP(nn.Module):

154 self,

155 config: BloomConfig,

156 quant_config: QuantizationConfig | None = None,

157 + prefix: str = "",

158 ):

159 super().__init__()

160 hidden_size = config.hidden_size

@@ -159,12 +162,14 @@ class BloomMLP(nn.Module):

162 hidden_size,

163 4 * hidden_size,

164 quant_config=quant_config,

165 + prefix=f"{prefix}.dense_h_to_4h",

166 )

167 self.gelu_impl = get_act_fn("gelu")

168 self.dense_4h_to_h = RowParallelLinear(

169 4 * hidden_size,

170 hidden_size,

171 quant_config=quant_config,

172 + prefix=f"{prefix}.dense_4h_to_h",

173 )

174

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

@@ -192,7 +197,7 @@ class BloomBlock(nn.Module):

197 self.post_attention_layernorm = nn.LayerNorm(

198 hidden_size, eps=config.layer_norm_epsilon

199 )

195 - self.mlp = BloomMLP(config, quant_config)

200 + self.mlp = BloomMLP(config, quant_config, prefix=f"{prefix}.mlp")

201 self.apply_residual_connection_post_layernorm = (

202 config.apply_residual_connection_post_layernorm

203 )

vllm/model_executor/models/chameleon.py

+7

@@ -227,6 +227,7 @@ class ChameleonMLP(nn.Module):

227 hidden_act: str,

228 quant_config: QuantizationConfig | None = None,

229 bias: bool = False,

230 + prefix: str = "",

231 ) -> None:

232 super().__init__()

233 self.gate_up_proj = MergedColumnParallelLinear(

@@ -234,12 +235,14 @@ class ChameleonMLP(nn.Module):

235 output_sizes=[intermediate_size] * 2,

236 bias=bias,

237 quant_config=quant_config,

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

239 )

240 self.down_proj = RowParallelLinear(

241 input_size=intermediate_size,

242 output_size=hidden_size,

243 bias=bias,

244 quant_config=quant_config,

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

246 )

247 if hidden_act != "silu":

248 raise ValueError(

@@ -299,12 +302,14 @@ class ChameleonAttention(nn.Module):

302 total_num_kv_heads=self.total_num_kv_heads,

303 bias=bias,

304 quant_config=quant_config,

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

306 )

307 self.o_proj = RowParallelLinear(

308 input_size=self.total_num_heads * self.head_dim,

309 output_size=hidden_size,

310 bias=bias,

311 quant_config=quant_config,

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

313 )

314 self.q_norm = ChameleonLayerNorm((self.num_heads, self.head_dim))

315 self.k_norm = ChameleonLayerNorm((self.num_kv_heads, self.head_dim))

@@ -393,6 +398,7 @@ class ChameleonDecoderLayer(nn.Module):

398 hidden_act=config.hidden_act,

399 quant_config=quant_config,

400 bias=getattr(config, "mlp_bias", False),

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

402 )

403 self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

404 self.post_attention_layernorm = RMSNorm(

@@ -462,6 +468,7 @@ class ChameleonSwinDecoderLayer(nn.Module):

468 hidden_act=config.hidden_act,

469 quant_config=quant_config,

470 bias=getattr(config, "mlp_bias", False),

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

472 )

473 self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

474 self.post_attention_layernorm = RMSNorm(

vllm/model_executor/models/dbrx.py

+2

@@ -209,12 +209,14 @@ class DbrxAttention(nn.Module):

209 self.total_num_kv_heads,

210 bias=False,

211 quant_config=quant_config,

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

213 )

214 self.out_proj = RowParallelLinear(

215 self.d_model,

216 self.d_model,

217 bias=False,

218 quant_config=quant_config,

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

220 )

221 self.rotary_emb = get_rope(

222 self.head_dim,

vllm/model_executor/models/deepseek.py

+8 -1

@@ -82,7 +82,11 @@ class DeepseekMLP(nn.Module):

82 ) -> None:

83 super().__init__()

84 self.gate_up_proj = MergedColumnParallelLinear(

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

85 + hidden_size,

86 + [intermediate_size] * 2,

87 + bias=False,

88 + quant_config=quant_config,

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

90 )

91 self.down_proj = RowParallelLinear(

92 intermediate_size,

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

94 bias=False,

95 quant_config=quant_config,

96 reduce_results=reduce_results,

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

98 )

99 if hidden_act != "silu":

100 raise ValueError(

@@ -239,6 +244,7 @@ class DeepseekAttention(nn.Module):

244 self.total_num_kv_heads,

245 bias=False,

246 quant_config=quant_config,

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

248 )

249

250 self.o_proj = RowParallelLinear(

@@ -246,6 +252,7 @@ class DeepseekAttention(nn.Module):

252 hidden_size,

253 bias=False,

254 quant_config=quant_config,

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

256 )

257

258 self.rotary_emb = get_rope(

vllm/model_executor/models/dots1.py

+2

@@ -240,6 +240,7 @@ class Dots1Attention(nn.Module):

240 self.total_num_kv_heads,

241 bias=attention_bias,

242 quant_config=quant_config,

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

244 )

245

246 self.o_proj = RowParallelLinear(

@@ -247,6 +248,7 @@ class Dots1Attention(nn.Module):

248 hidden_size,

249 bias=False,

250 quant_config=quant_config,

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

252 )

253

254 self.rotary_emb = get_rope(

vllm/model_executor/models/falcon.py

+6 -1

@@ -137,6 +137,7 @@ class FalconAttention(nn.Module):

137 bias=config.bias,

138 skip_bias_add=True,

139 quant_config=quant_config,

140 + prefix=f"{prefix}.query_key_value",

141 )

142 self.q_size = self.num_heads * self.head_dim

143 self.kv_size = self.num_kv_heads * self.head_dim

@@ -153,6 +154,7 @@ class FalconAttention(nn.Module):

154 skip_bias_add=True,

155 quant_config=quant_config,

156 reduce_results=self.reduce_row_parallel_results,

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

158 )

159

160 self.use_rotary = config.rotary

@@ -227,6 +229,7 @@ class FalconMLP(nn.Module):

229 self,

230 config: FalconConfig,

231 quant_config: QuantizationConfig | None = None,

232 + prefix: str = "",

233 ):

234 super().__init__()

235 hidden_size = config.hidden_size

@@ -237,6 +240,7 @@ class FalconMLP(nn.Module):

240 bias=config.bias,

241 skip_bias_add=True,

242 quant_config=quant_config,

243 + prefix=f"{prefix}.dense_h_to_4h",

244 )

245 self.act = get_act_fn("gelu")

246 self.reduce_row_parallel_results = not (

@@ -249,6 +253,7 @@ class FalconMLP(nn.Module):

253 skip_bias_add=True,

254 reduce_results=self.reduce_row_parallel_results,

255 quant_config=quant_config,

256 + prefix=f"{prefix}.dense_4h_to_h",

257 )

258

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

@@ -275,7 +280,7 @@ class FalconDecoderLayer(nn.Module):

280 self.self_attention = FalconAttention(

281 config, cache_config, quant_config, prefix=f"{prefix}.self_attention"

282 )

278 - self.mlp = FalconMLP(config, quant_config)

283 + self.mlp = FalconMLP(config, quant_config, prefix=f"{prefix}.mlp")

284 self.config = config

285

286 if not hasattr(config, "num_ln_in_parallel_attn"):

vllm/model_executor/models/falcon_h1.py

+4 -1

@@ -59,6 +59,7 @@ class FalconH1MLP(nn.Module):

59 config: FalconH1Config,

60 quant_config: QuantizationConfig | None = None,

61 bias: bool = False,

62 + prefix: str = "",

63 ) -> None:

64 super().__init__()

65 self.gate_up_proj = MergedColumnParallelLinear(

@@ -66,12 +67,14 @@ class FalconH1MLP(nn.Module):

67 output_sizes=[config.intermediate_size] * 2,

68 bias=bias,

69 quant_config=quant_config,

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

71 )

72 self.down_proj = RowParallelLinear(

73 input_size=config.intermediate_size,

74 output_size=config.hidden_size,

75 bias=bias,

76 quant_config=quant_config,

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

78 )

79 self.tp_size = get_tensor_model_parallel_world_size()

80 self.intermediate_size = config.intermediate_size

@@ -365,7 +368,7 @@ class FalconH1ParallelHybrid(nn.Module):

368 self.attention_in_multiplier = config.attention_in_multiplier

369 self.attn_out_multiplier = config.attention_out_multiplier

370

368 - self.feed_forward = FalconH1MLP(config)

371 + self.feed_forward = FalconH1MLP(config, prefix=f"{prefix}.feed_forward")

372

373 self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

374 self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

vllm/model_executor/models/gemma2.py

+14 -2

@@ -66,13 +66,22 @@ class Gemma2MLP(nn.Module):

66 hidden_act: str,

67 hidden_activation: str,

68 quant_config: QuantizationConfig | None = None,

69 + prefix: str = "",

70 ) -> None:

71 super().__init__()

72 self.gate_up_proj = MergedColumnParallelLinear(

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

73 + hidden_size,

74 + [intermediate_size] * 2,

75 + bias=False,

76 + quant_config=quant_config,

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

78 )

79 self.down_proj = RowParallelLinear(

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

80 + intermediate_size,

81 + hidden_size,

82 + bias=False,

83 + quant_config=quant_config,

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

85 )

86 if not (hidden_act == hidden_activation == "gelu_pytorch_tanh"):

87 raise ValueError(

@@ -134,12 +143,14 @@ class Gemma2Attention(nn.Module):

143 self.total_num_kv_heads,

144 bias=config.attention_bias,

145 quant_config=quant_config,

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

147 )

148 self.o_proj = RowParallelLinear(

149 self.total_num_heads * self.head_dim,

150 hidden_size,

151 bias=config.attention_bias,

152 quant_config=quant_config,

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

154 )

155 self.rotary_emb = get_rope(

156 self.head_dim,

@@ -208,6 +219,7 @@ class Gemma2DecoderLayer(nn.Module):

219 hidden_act=config.hidden_act,

220 hidden_activation=config.hidden_activation,

221 quant_config=quant_config,

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

223 )

224 self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

225 self.post_attention_layernorm = GemmaRMSNorm(

vllm/model_executor/models/gpt_j.py

+6 -1

@@ -78,12 +78,14 @@ class GPTJAttention(nn.Module):

78 self.total_num_heads,

79 bias=False,

80 quant_config=quant_config,

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

82 )

83 self.out_proj = RowParallelLinear(

84 config.hidden_size,

85 config.hidden_size,

86 bias=False,

87 quant_config=quant_config,

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

89 )

90

91 tp_world_size = get_tensor_model_parallel_world_size()

@@ -130,6 +132,7 @@ class GPTJMLP(nn.Module):

132 intermediate_size: int,

133 config: GPTJConfig,

134 quant_config: QuantizationConfig | None = None,

135 + prefix: str = "",

136 ):

137 super().__init__()

138 hidden_size = config.n_embd

@@ -137,11 +140,13 @@ class GPTJMLP(nn.Module):

140 hidden_size,

141 intermediate_size,

142 quant_config=quant_config,

143 + prefix=f"{prefix}.fc_in",

144 )

145 self.fc_out = RowParallelLinear(

146 intermediate_size,

147 hidden_size,

148 quant_config=quant_config,

149 + prefix=f"{prefix}.fc_out",

150 )

151 self.act = get_act_fn(config.activation_function)

152

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

171 self.attn = GPTJAttention(

172 config, cache_config, quant_config, prefix=f"{prefix}.attn"

173 )

169 - self.mlp = GPTJMLP(inner_dim, config, quant_config)

174 + self.mlp = GPTJMLP(inner_dim, config, quant_config, prefix=f"{prefix}.mlp")

175

176 def forward(

177 self,

vllm/model_executor/models/gpt_neox.py

+5

@@ -80,12 +80,14 @@ class GPTNeoXAttention(nn.Module):

80 self.total_num_heads,

81 bias=self.bias,

82 quant_config=quant_config,

83 + prefix=f"{prefix}.query_key_value",

84 )

85 self.dense = RowParallelLinear(

86 config.hidden_size,

87 config.hidden_size,

88 bias=self.bias,

89 quant_config=quant_config,

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

91 )

92 scaling = self.head_size**-0.5

93 rotary_dim = int(self.head_size * config.rotary_pct)

@@ -125,17 +127,20 @@ class GPTNeoXMLP(nn.Module):

127 self,

128 config: GPTNeoXConfig,

129 quant_config: QuantizationConfig | None = None,

130 + prefix: str = "",

131 ):

132 super().__init__()

133 self.dense_h_to_4h = ColumnParallelLinear(

134 config.hidden_size,

135 config.intermediate_size,

136 quant_config=quant_config,

137 + prefix=f"{prefix}.dense_h_to_4h",

138 )

139 self.dense_4h_to_h = RowParallelLinear(

140 config.intermediate_size,

141 config.hidden_size,

142 quant_config=quant_config,

143 + prefix=f"{prefix}.dense_4h_to_h",

144 )

145 self.act = get_act_fn(config.hidden_act)

146

vllm/model_executor/models/jais.py

+7 -1

@@ -107,12 +107,14 @@ class JAISAttention(nn.Module):

107 total_num_heads,

108 bias=True,

109 quant_config=quant_config,

110 + prefix=f"{prefix}.c_attn",

111 )

112 self.c_proj = RowParallelLinear(

113 self.hidden_size,

114 self.hidden_size,

115 bias=True,

116 quant_config=quant_config,

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

118 )

119

120 tp_rank = get_tensor_model_parallel_rank()

@@ -147,6 +149,7 @@ class JAISMLP(nn.Module):

149 intermediate_size: int,

150 config: JAISConfig,

151 quant_config: QuantizationConfig | None = None,

152 + prefix: str = "",

153 ):

154 super().__init__()

155 hidden_size = config.hidden_size

@@ -156,6 +159,7 @@ class JAISMLP(nn.Module):

159 intermediate_size,

160 bias=True,

161 quant_config=quant_config,

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

163 )

164 self.c_fc2 = (

165 ColumnParallelLinear(

@@ -163,6 +167,7 @@ class JAISMLP(nn.Module):

167 intermediate_size,

168 bias=True,

169 quant_config=quant_config,

170 + prefix=f"{prefix}.c_fc2",

171 )

172 if self.swiglu

173 else None

@@ -172,6 +177,7 @@ class JAISMLP(nn.Module):

177 hidden_size,

178 bias=True,

179 quant_config=quant_config,

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

181 )

182

183 self.act = SwiGLUActivation()

@@ -206,7 +212,7 @@ class JAISBlock(nn.Module):

212 config, cache_config, quant_config, prefix=f"{prefix}.attn"

213 )

214 self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)

209 - self.mlp = JAISMLP(inner_dim, config, quant_config)

215 + self.mlp = JAISMLP(inner_dim, config, quant_config, prefix=f"{prefix}.mlp")

216

217 def forward(

218 self,

vllm/model_executor/models/jamba.py

+2

@@ -220,12 +220,14 @@ class JambaAttentionDecoderLayer(nn.Module):

220 self.total_num_kv_heads,

221 bias=False,

222 quant_config=quant_config,

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

224 )

225 self.o_proj = RowParallelLinear(

226 self.total_num_heads * self.head_dim,

227 config.hidden_size,

228 bias=False,

229 quant_config=quant_config,

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

231 )

232

233 self.attn = Attention(

vllm/model_executor/models/minicpm.py

+13 -2

@@ -191,13 +191,22 @@ class MiniCPMMLP(nn.Module):

191 hidden_act: str,

192 hidden_act_param: float,

193 quant_config: QuantizationConfig | None = None,

194 + prefix: str = "",

195 ) -> None:

196 super().__init__()

197 self.gate_up_proj = MergedColumnParallelLinear(

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

198 + hidden_size,

199 + [intermediate_size] * 2,

200 + bias=False,

201 + quant_config=quant_config,

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

203 )

204 self.down_proj = RowParallelLinear(

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

205 + intermediate_size,

206 + hidden_size,

207 + bias=False,

208 + quant_config=quant_config,

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

210 )

211 if hidden_act == "silu":

212 self.act_fn = SiluAndMul()

@@ -259,12 +268,14 @@ class MiniCPMAttention(nn.Module):

268 self.total_num_kv_heads,

269 bias=False,

270 quant_config=quant_config,

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

272 )

273 self.o_proj = RowParallelLinear(

274 self.total_num_heads * self.head_dim,

275 hidden_size,

276 bias=False,

277 quant_config=quant_config,

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

279 )

280

281 self.rotary_emb = get_rope(

vllm/model_executor/models/minicpm3.py

+4

@@ -96,6 +96,7 @@ class MiniCPM3Attention(nn.Module):

96 self.num_heads * self.qk_head_dim,

97 bias=False,

98 quant_config=quant_config,

99 + prefix=f"{prefix}.q_b_proj",

100 )

101

102 self.kv_a_proj_with_mqa = ReplicatedLinear(

@@ -103,6 +104,7 @@ class MiniCPM3Attention(nn.Module):

104 self.kv_lora_rank + self.qk_rope_head_dim,

105 bias=False,

106 quant_config=quant_config,

107 + prefix=f"{prefix}.kv_a_proj_with_mqa",

108 )

109 self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)

110 self.kv_b_proj = ColumnParallelLinear(

@@ -110,6 +112,7 @@ class MiniCPM3Attention(nn.Module):

112 self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),

113 bias=False,

114 quant_config=quant_config,

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

116 )

117 # O projection.

118 self.o_proj = RowParallelLinear(

@@ -117,6 +120,7 @@ class MiniCPM3Attention(nn.Module):

120 self.hidden_size,

121 bias=False,

122 quant_config=quant_config,

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

124 )

125

126 self.rotary_emb = get_rope(

vllm/model_executor/models/mpt.py

+6 -1

@@ -83,6 +83,7 @@ class MPTAttention(nn.Module):

83 self.total_num_kv_heads,

84 bias=not config.no_bias,

85 quant_config=quant_config,

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

87 )

88 if self.qk_ln:

89 self.q_ln = nn.LayerNorm(self.d_model)

@@ -92,6 +93,7 @@ class MPTAttention(nn.Module):

93 self.d_model,

94 bias=not config.no_bias,

95 quant_config=quant_config,

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

97 )

98

99 tp_world_size = get_tensor_model_parallel_world_size()

@@ -152,6 +154,7 @@ class MPTMLP(nn.Module):

154 self,

155 config: MptConfig,

156 quant_config: QuantizationConfig | None = None,

157 + prefix: str = "",

158 ):

159 super().__init__()

160 hidden_size = config.d_model

@@ -162,6 +165,7 @@ class MPTMLP(nn.Module):

165 intermediate_size,

166 bias=not config.no_bias,

167 quant_config=quant_config,

168 + prefix=f"{prefix}.up_proj",

169 )

170 self.act = get_act_fn("gelu")

171 self.down_proj = RowParallelLinear(

@@ -169,6 +173,7 @@ class MPTMLP(nn.Module):

173 hidden_size,

174 bias=not config.no_bias,

175 quant_config=quant_config,

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

177 )

178

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

@@ -193,7 +198,7 @@ class MPTBlock(nn.Module):

198 config, cache_config, quant_config, prefix=f"{prefix}.attn"

199 )

200 self.norm_2 = nn.LayerNorm(hidden_size)

196 - self.ffn = MPTMLP(config, quant_config)

201 + self.ffn = MPTMLP(config, quant_config, prefix=f"{prefix}.ffn")

202

203 def forward(

204 self,

vllm/model_executor/models/olmoe.py

+2

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

158 self.total_num_kv_heads,

159 bias=False,

160 quant_config=quant_config,

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

162 )

163 self.tp_size = tp_size

164 self.tp_rank = get_tensor_model_parallel_rank()

@@ -168,6 +169,7 @@ class OlmoeAttention(nn.Module):

169 self.hidden_size,

170 bias=False,

171 quant_config=quant_config,

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

173 )

174

175 self.rotary_emb = get_rope(

vllm/model_executor/models/orion.py

+14 -2

@@ -52,13 +52,22 @@ class OrionMLP(nn.Module):

52 intermediate_size: int,

53 hidden_act: str,

54 quant_config: QuantizationConfig | None = None,

55 + prefix: str = "",

56 ) -> None:

57 super().__init__()

58 self.gate_up_proj = MergedColumnParallelLinear(

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

59 + hidden_size,

60 + [intermediate_size] * 2,

61 + bias=False,

62 + quant_config=quant_config,

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

64 )

65 self.down_proj = RowParallelLinear(

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

66 + intermediate_size,

67 + hidden_size,

68 + bias=False,

69 + quant_config=quant_config,

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

71 )

72 if hidden_act != "silu":

73 raise ValueError(

@@ -116,12 +125,14 @@ class OrionAttention(nn.Module):

125 self.total_num_kv_heads,

126 bias=False,

127 quant_config=quant_config,

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

129 )

130 self.o_proj = RowParallelLinear(

131 self.total_num_heads * self.head_dim,

132 hidden_size,

133 bias=False,

134 quant_config=quant_config,

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

136 )

137

138 self.rotary_emb = get_rope(

@@ -183,6 +194,7 @@ class OrionDecoderLayer(nn.Module):

194 intermediate_size=config.intermediate_size,

195 hidden_act=config.hidden_act,

196 quant_config=quant_config,

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

198 )

199

200 self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)

vllm/model_executor/models/persimmon.py

+19 -4

@@ -62,14 +62,23 @@ from .utils import (

62

63 class PersimmonMLP(nn.Module):

64 def __init__(

65 - self, config: PersimmonConfig, quant_config: QuantizationConfig | None = None

65 + self,

66 + config: PersimmonConfig,

67 + quant_config: QuantizationConfig | None = None,

68 + prefix: str = "",

69 ):

70 super().__init__()

71 self.dense_h_to_4h = ColumnParallelLinear(

69 - config.hidden_size, config.intermediate_size, quant_config=quant_config

72 + config.hidden_size,

73 + config.intermediate_size,

74 + quant_config=quant_config,

75 + prefix=f"{prefix}.dense_h_to_4h",

76 )

77 self.dense_4h_to_h = RowParallelLinear(

72 - config.intermediate_size, config.hidden_size, quant_config=quant_config

78 + config.intermediate_size,

79 + config.hidden_size,

80 + quant_config=quant_config,

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

82 )

83 self.act = get_act_fn(config.hidden_act)

84

@@ -110,12 +119,14 @@ class PersimmonAttention(nn.Module):

119 self.total_num_heads,

120 bias=True,

121 quant_config=quant_config,

122 + prefix=f"{prefix}.query_key_value",

123 )

124 self.dense = RowParallelLinear(

125 self.total_num_heads * self.head_dim,

126 self.hidden_size,

127 bias=True,

128 quant_config=quant_config,

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

130 )

131 self.is_qk_layernorm = config.qk_layernorm

132

@@ -192,7 +203,11 @@ class PersimmonDecoderLayer(nn.Module):

203 quant_config=quant_config,

204 prefix=f"{prefix}.self_attn",

205 )

195 - self.mlp = PersimmonMLP(config, quant_config=quant_config)

206 + self.mlp = PersimmonMLP(

207 + config,

208 + quant_config=quant_config,

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

210 + )

211 self.input_layernorm = nn.LayerNorm(

212 config.hidden_size, eps=config.layer_norm_eps

213 )

vllm/model_executor/models/phi.py

+9 -2

@@ -99,11 +99,13 @@ class PhiAttention(nn.Module):

99 self.total_num_heads,

100 bias=True,

101 quant_config=quant_config,

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

103 )

104 self.dense = RowParallelLinear(

105 self.hidden_size,

106 self.hidden_size,

107 quant_config=quant_config,

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

109 )

110

111 scaling = self.head_size**-0.5

@@ -148,7 +150,10 @@ class PhiAttention(nn.Module):

150

151 class PhiMLP(nn.Module):

152 def __init__(

151 - self, config: PhiConfig, quant_config: QuantizationConfig | None = None

153 + self,

154 + config: PhiConfig,

155 + quant_config: QuantizationConfig | None = None,

156 + prefix: str = "",

157 ):

158 super().__init__()

159

@@ -159,11 +164,13 @@ class PhiMLP(nn.Module):

164 config.hidden_size,

165 n_inner,

166 quant_config=quant_config,

167 + prefix=f"{prefix}.fc1",

168 )

169 self.fc2 = RowParallelLinear(

170 n_inner,

171 config.hidden_size,

172 quant_config=quant_config,

173 + prefix=f"{prefix}.fc2",

174 )

175 self.act = get_act_fn(config.hidden_act)

176

@@ -189,7 +196,7 @@ class PhiLayer(nn.Module):

196 self.self_attn = PhiAttention(

197 config, cache_config, quant_config, prefix=f"{prefix}.self_attn"

198 )

192 - self.mlp = PhiMLP(config, quant_config)

199 + self.mlp = PhiMLP(config, quant_config, prefix=f"{prefix}.mlp")

200

201 def forward(

202 self,

vllm/model_executor/models/phimoe.py

+2

@@ -343,12 +343,14 @@ class PhiMoEAttention(nn.Module):

343 self.total_num_kv_heads,

344 bias=True,

345 quant_config=quant_config,

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

347 )

348 self.o_proj = RowParallelLinear(

349 self.total_num_heads * self.head_dim,

350 hidden_size,

351 bias=True,

352 quant_config=quant_config,

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

354 )

355 self.rotary_emb = get_rope(

356 self.head_dim,

vllm/model_executor/models/plamo2.py

+2

@@ -567,12 +567,14 @@ class Plamo2AttentionMixer(nn.Module):

567 self.total_num_kv_heads,

568 bias=False,

569 quant_config=quant_config,

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

571 )

572 self.o_proj = RowParallelLinear(

573 self.total_num_heads * self.head_dim,

574 config.hidden_size,

575 bias=False,

576 quant_config=quant_config,

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

578 )

579

580 self.rope_theta = config.rope_theta if hasattr(config, "rope_theta") else 10000

vllm/model_executor/models/qwen.py

+2

@@ -102,12 +102,14 @@ class QWenAttention(nn.Module):

102 self.total_num_heads,

103 bias=True,

104 quant_config=quant_config,

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

106 )

107 self.c_proj = RowParallelLinear(

108 self.total_num_heads * self.head_dim,

109 hidden_size,

110 bias=False,

111 quant_config=quant_config,

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

113 )

114 self.scaling = self.head_dim**-0.5

115

vllm/model_executor/models/zamba2.py

+15 -1

@@ -75,7 +75,12 @@ class Zamba2LoRA(nn.Module):

75 super().__init__()

76

77 self.A = ColumnParallelLinear(

78 - input_dim, rank, bias=False, quant_config=quant_config, gather_output=True

78 + input_dim,

79 + rank,

80 + bias=False,

81 + quant_config=quant_config,

82 + gather_output=True,

83 + prefix=f"{prefix}.A",

84 )

85

86 if isinstance(output_dim, list):

@@ -150,12 +155,14 @@ class Zamba2Attention(nn.Module):

155 self.total_num_attention_heads,

156 bias=False,

157 quant_config=quant_config,

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

159 )

160 self.o_proj = RowParallelLinear(

161 self.attention_hidden_size,

162 config.hidden_size,

163 bias=False,

164 quant_config=quant_config,

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

166 )

167

168 # Even though in Zamba2 weights are shared between attention layers, KV

@@ -197,18 +204,21 @@ class Zamba2Attention(nn.Module):

204 config.adapter_rank,

205 self.attention_hidden_size,

206 quant_config=quant_config,

207 + prefix=f"{prefix}.linear_q_adapter",

208 )

209 linear_k_adapter = Zamba2LoRA(

210 self.attention_hidden_size,

211 config.adapter_rank,

212 self.attention_hidden_size,

213 quant_config=quant_config,

214 + prefix=f"{prefix}.linear_k_adapter",

215 )

216 linear_v_adapter = Zamba2LoRA(

217 self.attention_hidden_size,

218 config.adapter_rank,

219 self.attention_hidden_size,

220 quant_config=quant_config,

221 + prefix=f"{prefix}.linear_v_adapter",

222 )

223 else:

224 linear_q_adapter = nn.Identity()

@@ -312,6 +322,7 @@ class Zamba2MLP(nn.Module):

322 2 * [self.intermediate_size], # 2x for gate and input projections

323 bias=self.config.add_bias_linear,

324 quant_config=quant_config,

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

326 )

327

328 self.down_proj = RowParallelLinear(

@@ -319,6 +330,7 @@ class Zamba2MLP(nn.Module):

330 self.hidden_size,

331 bias=self.config.add_bias_linear,

332 quant_config=quant_config,

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

334 )

335

336 # Only allow GELU activations

@@ -418,6 +430,7 @@ class Zamba2AttentionDecoderLayer(nn.Module):

430 bare_block_idx=bare_block_idx,

431 num_hybrid_layers=num_hybrid_layers,

432 quant_config=quant_config,

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

434 )

435

436 # Initialize layer normalizations

@@ -599,6 +612,7 @@ class Zamba2HybridLayer(nn.Module):

612 config.hidden_size,

613 bias=False,

614 quant_config=quant_config,

615 + prefix=f"{prefix}.linear",

616 )

617 self.mamba_decoder = Zamba2MambaDecoderLayer(

618 config,

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