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

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 14, 23:45 UTC

[Model]: Fused MoE for nomic-embed-text-v2-moe (#18321)

Signed-off-by: isotr0py <2037008807@qq.com> Signed-off-by: Isotr0py <2037008807@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Isotr0py committed Jul 29, 2025 at 18:13 UTC a4528f0cac5d2857ccc56d2a2e1a1c43142643ce

2 files changed +142 -113

vllm/model_executor/layers/fused_moe/fused_moe.py

+36 -11

@@ -7,6 +7,7 @@ import os

7 from typing import Any, Callable, Optional

9 import torch

10 +import torch.nn.functional as F

12 import vllm.envs as envs

13 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

@@ -1001,6 +1002,7 @@ def inplace_fused_experts(hidden_states: torch.Tensor,

1002 topk_weights: torch.Tensor,

1003 topk_ids: torch.Tensor,

1004 activation: str = "silu",

1005 + is_act_and_mul: bool = True,

1006 apply_router_weight_on_input: bool = False,

1007 use_fp8_w8a8: bool = False,

1008 use_int8_w8a8: bool = False,

@@ -1018,7 +1020,8 @@ def inplace_fused_experts(hidden_states: torch.Tensor,

1020 a2_scale: Optional[torch.Tensor] = None,

1021 block_shape: Optional[list[int]] = None) -> None:

1022 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,

1021 - activation, apply_router_weight_on_input, use_fp8_w8a8,

1023 + activation, is_act_and_mul,

1024 + apply_router_weight_on_input, use_fp8_w8a8,

1025 use_int8_w8a8, use_int8_w8a16, use_int4_w4a16,

1026 use_mxfp4_w4a4, per_channel_quant, global_num_experts,

1027 expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,

@@ -1032,6 +1035,7 @@ def inplace_fused_experts_fake(

1035 topk_weights: torch.Tensor,

1036 topk_ids: torch.Tensor,

1037 activation: str = "silu",

1038 + is_act_and_mul: bool = True,

1039 apply_router_weight_on_input: bool = False,

1040 use_fp8_w8a8: bool = False,

1041 use_int8_w8a8: bool = False,

@@ -1167,6 +1171,7 @@ def outplace_fused_experts(

1171 topk_weights: torch.Tensor,

1172 topk_ids: torch.Tensor,

1173 activation: str = "silu",

1174 + is_act_and_mul: bool = True,

1175 apply_router_weight_on_input: bool = False,

1176 use_fp8_w8a8: bool = False,

1177 use_int8_w8a8: bool = False,

@@ -1183,13 +1188,12 @@ def outplace_fused_experts(

1188 a1_scale: Optional[torch.Tensor] = None,

1189 a2_scale: Optional[torch.Tensor] = None,

1190 block_shape: Optional[list[int]] = None) -> torch.Tensor:

1186 - return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,

1187 - False, activation, apply_router_weight_on_input,

1188 - use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,

1189 - use_int4_w4a16, use_mxfp4_w4a4,

1190 - per_channel_quant, global_num_experts,

1191 - expert_map, w1_scale, w2_scale, w1_zp, w2_zp,

1192 - a1_scale, a2_scale, block_shape)

1191 + return fused_experts_impl(

1192 + hidden_states, w1, w2, topk_weights, topk_ids, False, activation,

1193 + is_act_and_mul, apply_router_weight_on_input, use_fp8_w8a8,

1194 + use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, use_mxfp4_w4a4,

1195 + per_channel_quant, global_num_experts, expert_map, w1_scale, w2_scale,

1196 + w1_zp, w2_zp, a1_scale, a2_scale, block_shape)

1197

1198

1199 def outplace_fused_experts_fake(

@@ -1199,6 +1203,7 @@ def outplace_fused_experts_fake(

1203 topk_weights: torch.Tensor,

1204 topk_ids: torch.Tensor,

1205 activation: str = "silu",

1206 + is_act_and_mul: bool = True,

1207 use_fp8_w8a8: bool = False,

1208 use_int8_w8a8: bool = False,

1209 use_int8_w8a16: bool = False,

@@ -1253,6 +1258,7 @@ def fused_experts(

1258 topk_ids: torch.Tensor,

1259 inplace: bool = False,

1260 activation: str = "silu",

1261 + is_act_and_mul: bool = True,

1262 apply_router_weight_on_input: bool = False,

1263 use_fp8_w8a8: bool = False,

1264 use_int8_w8a8: bool = False,

@@ -1283,6 +1289,8 @@ def fused_experts(

1289 or is_blackwell_deep_gemm_used())

1290 if (allow_deep_gemm and use_fp8_w8a8 and should_use_deep_gemm):

1291 assert apply_router_weight_on_input is False

1292 + assert is_act_and_mul, (

1293 + "DeepGemm only supports is_act_and_mul=True for now.")

1294 return deep_gemm_moe_fp8(

1295 hidden_states=hidden_states,

1296 w1=w1,

@@ -1319,6 +1327,7 @@ def fused_experts(

1327 topk_weights=topk_weights,

1328 topk_ids=topk_ids,

1329 activation=activation,

1330 + is_act_and_mul=is_act_and_mul,

1331 apply_router_weight_on_input=apply_router_weight_on_input,

1332 use_fp8_w8a8=use_fp8_w8a8,

1333 use_int8_w8a8=use_int8_w8a8,

@@ -1345,6 +1354,7 @@ def fused_experts_impl(

1354 topk_ids: torch.Tensor,

1355 inplace: bool = False,

1356 activation: str = "silu",

1357 + is_act_and_mul: bool = True,

1358 apply_router_weight_on_input: bool = False,

1359 use_fp8_w8a8: bool = False,

1360 use_int8_w8a8: bool = False,

@@ -1503,14 +1513,21 @@ def fused_experts_impl(

1513 per_channel_quant=per_channel_quant,

1514 block_shape=block_shape)

1515

1506 - if activation == "silu":

1516 + # Activation function with multiplication

1517 + if activation == "silu" and is_act_and_mul:

1518 torch.ops._C.silu_and_mul(intermediate_cache2,

1519 intermediate_cache1.view(-1, N))

1509 - elif activation == "gelu":

1520 + elif activation == "gelu" and is_act_and_mul:

1521 torch.ops._C.gelu_and_mul(intermediate_cache2,

1522 intermediate_cache1.view(-1, N))

1523 + # Activation function without multiplication

1524 + elif activation == "silu":

1525 + intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N))

1526 + elif activation == "gelu":

1527 + intermediate_cache2 = F.gelu(intermediate_cache1.view(-1, N))

1528 else:

1513 - raise ValueError(f"Unsupported FusedMoe activation: {activation}")

1529 + raise ValueError(f"Unsupported FusedMoe activation: {activation}, "

1530 + f"with is_act_and_mul={is_act_and_mul}.")

1531

1532 qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(

1533 A=intermediate_cache2,

@@ -1555,6 +1572,7 @@ def fused_moe(

1572 renormalize: bool,

1573 inplace: bool = False,

1574 activation: str = "silu",

1575 + is_act_and_mul: bool = True,

1576 use_grouped_topk: bool = False,

1577 num_expert_group: Optional[int] = None,

1578 topk_group: Optional[int] = None,

@@ -1591,6 +1609,9 @@ def fused_moe(

1609 Defaults to False.

1610 - activation (str): The activation function to apply after the first

1611 MoE layer.

1612 + - is_act_and_mul (bool): If True, use activation-and-mul function for

1613 + activation (self-gated activation), otherwise use activation function

1614 + for activation (ungated activation).

1615 - num_expert_group: Optional[int]: additional parameter for grouped_topk

1616 - topk_group: Optional[int]: additional parameter for grouped_topk

1617 - use_grouped_topk: If True, use grouped_topk instead of fused_topk

@@ -1627,6 +1648,9 @@ def fused_moe(

1648 Returns:

1649 - torch.Tensor: The output tensor after applying the MoE layer.

1650 """

1651 + if not is_act_and_mul:

1652 + assert inplace is False, (

1653 + "is_act_and_mul=False is not supported with inplace=True")

1654

1655 if use_grouped_topk:

1656 assert num_expert_group is not None and topk_group is not None

@@ -1647,6 +1671,7 @@ def fused_moe(

1671 topk_ids,

1672 inplace=inplace,

1673 activation=activation,

1674 + is_act_and_mul=is_act_and_mul,

1675 use_fp8_w8a8=use_fp8_w8a8,

1676 use_int8_w8a8=use_int8_w8a8,

1677 use_int8_w8a16=use_int8_w8a16,

vllm/model_executor/models/bert_with_rope.py

+106 -102

@@ -10,9 +10,12 @@ from transformers import PretrainedConfig

10 from vllm.attention import Attention, AttentionType

11 from vllm.compilation.decorators import support_torch_compile

12 from vllm.config import CacheConfig, VllmConfig

13 -from vllm.distributed import get_tensor_model_parallel_world_size

13 +from vllm.distributed import (divide, get_tensor_model_parallel_rank,

14 + get_tensor_model_parallel_world_size,

15 + tensor_model_parallel_all_reduce)

16 from vllm.model_executor.layers.activation import (get_act_and_mul_fn,

17 get_act_fn)

18 +from vllm.model_executor.layers.fused_moe import fused_moe

19 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

20 MergedColumnParallelLinear,

21 QKVParallelLinear,

@@ -26,6 +29,8 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader

29 from vllm.model_executor.models import SupportsV0Only

30 from vllm.model_executor.models.interfaces import SupportsQuant

31 from vllm.model_executor.models.utils import WeightsMapper

32 +from vllm.model_executor.utils import set_weight_attrs

33 +from vllm.platforms import current_platform

34 from vllm.sequence import IntermediateTensors

@@ -201,114 +206,101 @@ class BertWithRopeMLP(nn.Module):

206 return hidden_states

207

208

204 -class NomicRouter(nn.Module):

209 +class NomicMoE(nn.Module):

210

206 - def __init__(self, hidden_size: int, moe_num_experts: int, moe_top_k: int):

211 + def __init__(

212 + self,

213 + num_experts: int,

214 + top_k: int,

215 + hidden_size: int,

216 + intermediate_size: int,

217 + hidden_act: str,

218 + params_dtype: Optional[torch.dtype] = None,

219 + tp_size: Optional[int] = None,

220 + ):

221 super().__init__()

208 - self.moe_top_k = moe_top_k

209 - self.layer = ReplicatedLinear(hidden_size, moe_num_experts, bias=False)

222

211 - def forward(

212 - self, x: torch.Tensor

213 - ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:

214 - weights = self.layer(x.view(-1, x.shape[-1]))[0].softmax(

215 - dim=-1, dtype=torch.float32)

216 - top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)

217 - weights = weights.to(x.dtype)

218 - top_weights = top_weights.to(x.dtype)

219 - return weights, top_weights, top_experts # type: ignore

220 -

221 -

222 -class NomicExpertMLP(nn.Module):

223 -

224 - def __init__(self, hidden_size: int, ffn_hidden_size: int,

225 - moe_num_experts: int, ffn_act_fn: str):

226 - super().__init__()

223 + self.tp_size = tp_size or get_tensor_model_parallel_world_size()

224 + self.num_total_experts = num_experts

225 + self.top_k = top_k

226 self.hidden_size = hidden_size

228 - self.ffn_hidden_size = ffn_hidden_size

229 - self.moe_num_experts = moe_num_experts

227 + self.total_intermediate_size = intermediate_size

228 + self.intermediate_size = divide(intermediate_size, self.tp_size)

229 + self.hidden_act = hidden_act

230

231 + if params_dtype is None:

232 + params_dtype = torch.get_default_dtype()

233 + self.params_dtype = params_dtype

234 +

235 + self.router = ReplicatedLinear(self.hidden_size,

236 + self.num_total_experts,

237 + bias=False)

238 self.w1 = nn.Parameter(

232 - torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))

239 + torch.empty(self.num_total_experts,

240 + self.intermediate_size,

241 + self.hidden_size,

242 + device=current_platform.device_type,

243 + dtype=self.params_dtype))

244 self.w2 = nn.Parameter(

234 - torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))

235 - self.activation_fn = get_act_fn(ffn_act_fn)

236 -

237 - def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor:

238 - expert_w1 = self.w1.view(self.moe_num_experts, self.ffn_hidden_size,

239 - self.hidden_size)[expert_idx]

240 - expert_w2 = self.w2.view(self.moe_num_experts, self.ffn_hidden_size,

241 - self.hidden_size)[expert_idx]

242 -

243 - x1 = x.matmul(expert_w1.t())

244 - act_out = self.activation_fn(x1)

245 - x2 = act_out.matmul(expert_w2)

246 - return x2

247 -

248 -

249 -class NomicExperts(nn.Module):

250 -

251 - def __init__(self, config, hidden_size: int, ffn_hidden_size: int,

252 - moe_num_experts: int):

253 - super().__init__()

254 - self.moe_num_experts = moe_num_experts

255 -

256 - self.mlp = NomicExpertMLP(hidden_size=config.n_embd,

257 - ffn_hidden_size=config.n_inner,

258 - moe_num_experts=moe_num_experts,

259 - ffn_act_fn=config.hidden_act)

260 - self.bias = nn.Parameter(torch.zeros(config.n_embd))

261 -

262 - def forward(self, x: torch.Tensor, weights: torch.Tensor,

263 - top_weights: torch.Tensor,

264 - top_experts: torch.LongTensor) -> torch.Tensor:

265 - q_len, hidden_size = x.shape

266 - x = x.view(-1, hidden_size)

267 - out = torch.zeros_like(x)

268 -

269 - expert_mask = nn.functional.one_hot(

270 - top_experts, num_classes=self.moe_num_experts).permute(2, 1, 0)

271 - for expert_idx in range(0, self.moe_num_experts):

272 - topk_idx, token_idx = torch.where(expert_mask[expert_idx])

273 - if token_idx.shape[0] == 0:

274 - continue

275 -

276 - token_list = token_idx.tolist()

277 - topk_list = topk_idx.tolist()

278 -

279 - expert_tokens = x[None, token_list].reshape(-1, hidden_size)

280 - expert_out = self.mlp(

281 - expert_tokens, expert_idx) * top_weights[token_list, topk_list,

282 - None]

283 -

284 - out.index_add_(0, token_idx, expert_out)

285 -

286 - out = out.reshape(q_len, hidden_size)

287 - return out + self.bias

288 -

289 -

290 -class NomicMoELayer(nn.Module):

291 -

292 - def __init__(self, config: PretrainedConfig):

293 - super().__init__()

294 -

295 - self.router = NomicRouter(

296 - config.n_embd,

297 - moe_num_experts=config.num_experts,

298 - moe_top_k=config.moe_top_k,

299 - )

245 + torch.empty(self.num_total_experts,

246 + self.hidden_size,

247 + self.intermediate_size,

248 + device=current_platform.device_type,

249 + dtype=self.params_dtype))

250 + self.bias = nn.Parameter(torch.zeros(self.hidden_size))

251 + set_weight_attrs(self.w1, {

252 + "weight_loader": self.weight_loader,

253 + })

254 + set_weight_attrs(self.w2, {

255 + "weight_loader": self.weight_loader,

256 + })

257

301 - self.experts = NomicExperts(

302 - config,

303 - hidden_size=config.n_embd,

304 - ffn_hidden_size=config.n_inner,

305 - moe_num_experts=config.num_experts,

306 - )

258 + def weight_loader(

259 + self,

260 + param: nn.Parameter,

261 + loaded_weight: torch.Tensor,

262 + weight_name: str,

263 + ):

264 + # NOTE: Nomic-MoE has fused experts weights with shape

265 + # (num_experts * intermediate_size, hidden_size)

266 + tp_rank = get_tensor_model_parallel_rank()

267 + param_data = param.data

268 + shard_size = self.intermediate_size

269 + shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)

270 + if weight_name.endswith("w1"):

271 + loaded_weight = loaded_weight.reshape(

272 + self.num_total_experts,

273 + self.total_intermediate_size,

274 + self.hidden_size,

275 + )[:, shard]

276 + if weight_name.endswith("w2"):

277 + loaded_weight = loaded_weight.reshape(

278 + self.num_total_experts,

279 + self.total_intermediate_size,

280 + self.hidden_size,

281 + )[:, shard].transpose(1, 2)

282 + param_data.copy_(loaded_weight)

283

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

309 - weights, top_weights, top_experts = self.router(x)

310 - out = self.experts(x, weights, top_weights, top_experts)

311 - return out

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

285 + num_tokens, hidden_size = hidden_states.shape

286 + hidden_states = hidden_states.view(-1, self.hidden_size)

287 + # router_logits: (num_tokens, n_experts)

288 + router_logits, _ = self.router(hidden_states)

289 + final_hidden_states = fused_moe(hidden_states,

290 + self.w1,

291 + self.w2,

292 + router_logits,

293 + self.top_k,

294 + renormalize=False,

295 + inplace=False,

296 + activation=self.hidden_act,

297 + is_act_and_mul=False)

298 +

299 + if self.tp_size > 1:

300 + final_hidden_states = tensor_model_parallel_all_reduce(

301 + final_hidden_states)

302 +

303 + return final_hidden_states.view(num_tokens, hidden_size) + self.bias

304

305

306 class BertWithRopeBlock(nn.Module):

@@ -332,7 +324,11 @@ class BertWithRopeBlock(nn.Module):

324 prefix=f"{prefix}.attention")

325

326 if moe:

335 - self.mlp = NomicMoELayer(config=config, )

327 + self.mlp = NomicMoE(num_experts=config.num_experts,

328 + top_k=config.moe_top_k,

329 + hidden_size=config.hidden_size,

330 + intermediate_size=config.intermediate_size,

331 + hidden_act=config.hidden_act)

332 else:

333 if config.hidden_act in ["silu", "geglu"]:

334 self.mlp = BertWithRopeGatedMLP(

@@ -463,7 +459,11 @@ class BertWithRope(nn.Module, SupportsV0Only, SupportsQuant):

459 param = params_dict[name]

460 weight_loader = getattr(param, "weight_loader",

461 default_weight_loader)

466 - weight_loader(param, loaded_weight)

462 + if name.endswith((".w1", ".w2")):

463 + # Nomic-MoE has fused experts weights

464 + weight_loader(param, loaded_weight, name)

465 + else:

466 + weight_loader(param, loaded_weight)

467 loaded_params.add(name)

468 return loaded_params

469

@@ -481,6 +481,10 @@ class NomicBertModel(BertWithRope):

481 "mlp.fc12": "mlp.gate_proj",

482 "mlp.fc2": "mlp.down_proj",

483 "norm2": "mlp_ln",

484 + # MoE mapping

485 + "experts.mlp.": "",

486 + "experts.": "",

487 + "router.layer": "router",

488 })

489

490