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

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, 05:30 UTC

Add routed_scaling_factor to MoE grouped topk (#23123)

Signed-off-by: Xin Yang <xyangx@amazon.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Xin Yang committed Aug 29, 2025 at 21:36 UTC 8fb85b7bb67408c725474c31fe2a8f980c250277

19 files changed +77 -4

vllm/model_executor/layers/fused_moe/cpu_fused_moe.py

+12

@@ -21,6 +21,7 @@ def grouped_topk(

21 num_expert_group: int = 0,

22 topk_group: int = 0,

23 scoring_func: str = "softmax",

24 + routed_scaling_factor: float = 1.0,

25 e_score_correction_bias: Optional[torch.Tensor] = None

26 ) -> tuple[torch.Tensor, torch.Tensor]:

27 assert hidden_states.shape[0] == gating_output.shape[0], (

@@ -65,6 +66,8 @@ def grouped_topk(

66 if renormalize:

67 topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)

69 + if routed_scaling_factor != 1.0:

70 + topk_weights = topk_weights * routed_scaling_factor

71 return topk_weights, topk_ids.to(torch.int32)

@@ -78,6 +81,7 @@ def select_experts(

81 num_expert_group: Optional[int] = None,

82 custom_routing_function: Optional[Callable] = None,

83 scoring_func: str = "softmax",

84 + routed_scaling_factor: float = 1.0,

85 e_score_correction_bias: Optional[torch.Tensor] = None,

86 ) -> tuple[torch.Tensor, torch.Tensor]:

87 if use_grouped_topk:

@@ -90,6 +94,7 @@ def select_experts(

94 num_expert_group=num_expert_group,

95 topk_group=topk_group,

96 scoring_func=scoring_func,

97 + routed_scaling_factor=routed_scaling_factor,

98 e_score_correction_bias=e_score_correction_bias)

99 elif custom_routing_function is None:

100 assert scoring_func == "softmax"

@@ -131,12 +136,15 @@ class IPEXFusedMOE:

136 expert_map: Optional[torch.Tensor] = None,

137 custom_routing_function: Optional[Callable] = None,

138 scoring_func: str = "softmax",

139 + routed_scaling_factor: float = 1.0,

140 e_score_correction_bias: Optional[torch.Tensor] = None,

141 apply_router_weight_on_input: bool = False,

142 activation: str = "silu",

143 ) -> torch.Tensor:

144 assert activation == "silu", f"{activation} is not supported."

145 assert not apply_router_weight_on_input

146 + assert routed_scaling_factor == 1.0, \

147 + f"routed_scaling_factor {routed_scaling_factor} is not supported."

148 return layer.ipex_fusion(

149 x,

150 use_grouped_topk,

@@ -170,6 +178,7 @@ class SGLFusedMOE:

178 expert_map: Optional[torch.Tensor] = None,

179 custom_routing_function: Optional[Callable] = None,

180 scoring_func: str = "softmax",

181 + routed_scaling_factor: float = 1.0,

182 e_score_correction_bias: Optional[torch.Tensor] = None,

183 apply_router_weight_on_input: bool = False,

184 activation: str = "silu",

@@ -186,6 +195,7 @@ class SGLFusedMOE:

195 num_expert_group=num_expert_group,

196 custom_routing_function=custom_routing_function,

197 scoring_func=scoring_func,

198 + routed_scaling_factor=routed_scaling_factor,

199 e_score_correction_bias=e_score_correction_bias,

200 )

201

@@ -227,6 +237,7 @@ class CPUFusedMOE:

237 expert_map: Optional[torch.Tensor] = None,

238 custom_routing_function: Optional[Callable] = None,

239 scoring_func: str = "softmax",

240 + routed_scaling_factor: float = 1.0,

241 e_score_correction_bias: Optional[torch.Tensor] = None,

242 apply_router_weight_on_input: bool = False,

243 activation: str = "silu",

@@ -243,6 +254,7 @@ class CPUFusedMOE:

254 num_expert_group=num_expert_group,

255 custom_routing_function=custom_routing_function,

256 scoring_func=scoring_func,

257 + routed_scaling_factor=routed_scaling_factor,

258 e_score_correction_bias=e_score_correction_bias,

259 )

260

vllm/model_executor/layers/fused_moe/fused_moe.py

+4 -3

@@ -1011,7 +1011,8 @@ def grouped_topk(

1011 if renormalize:

1012 topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)

1013

1014 - topk_weights = topk_weights * routed_scaling_factor

1014 + if routed_scaling_factor != 1.0:

1015 + topk_weights = topk_weights * routed_scaling_factor

1016 return topk_weights.to(torch.float32), topk_ids.to(torch.int32)

1017

1018

@@ -1790,8 +1791,8 @@ def fused_moe(

1791 Defaults to False.

1792 - global_num_experts (int): The total number of experts in the global

1793 expert space.

1793 - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices

1794 - from the global expert space to the local expert space of the expert

1794 + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices

1795 + from the global expert space to the local expert space of the expert

1796 parallel shard.

1797 - w1_scale (Optional[torch.Tensor]): Optional scale to be used for

1798 w1.

vllm/model_executor/layers/fused_moe/layer.py

+18

@@ -244,6 +244,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):

244 expert_map: Optional[torch.Tensor] = None,

245 custom_routing_function: Optional[Callable] = None,

246 scoring_func: str = "softmax",

247 + routed_scaling_factor: float = 1.0,

248 e_score_correction_bias: Optional[torch.Tensor] = None,

249 apply_router_weight_on_input: bool = False,

250 activation: str = "silu",

@@ -400,6 +401,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

401 expert_map: Optional[torch.Tensor] = None,

402 custom_routing_function: Optional[Callable] = None,

403 scoring_func: str = "softmax",

404 + routed_scaling_factor: float = 1.0,

405 e_score_correction_bias: Optional[torch.Tensor] = None,

406 apply_router_weight_on_input: bool = False,

407 activation: str = "silu",

@@ -427,6 +429,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

429 expert_map=expert_map,

430 custom_routing_function=custom_routing_function,

431 scoring_func=scoring_func,

432 + routed_scaling_factor=routed_scaling_factor,

433 e_score_correction_bias=e_score_correction_bias,

434 activation=activation,

435 apply_router_weight_on_input=apply_router_weight_on_input,

@@ -450,6 +453,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

453 expert_map: Optional[torch.Tensor] = None,

454 custom_routing_function: Optional[Callable] = None,

455 scoring_func: str = "softmax",

456 + routed_scaling_factor: float = 1.0,

457 e_score_correction_bias: Optional[torch.Tensor] = None,

458 apply_router_weight_on_input: bool = False,

459 activation: str = "silu",

@@ -469,6 +473,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

473 num_expert_group=num_expert_group,

474 custom_routing_function=custom_routing_function,

475 scoring_func=scoring_func,

476 + routed_scaling_factor=routed_scaling_factor,

477 e_score_correction_bias=e_score_correction_bias,

478 indices_type=self.topk_indices_dtype,

479 enable_eplb=enable_eplb,

@@ -534,6 +539,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

539 expert_map: Optional[torch.Tensor] = None,

540 custom_routing_function: Optional[Callable] = None,

541 scoring_func: str = "softmax",

542 + routed_scaling_factor: float = 1.0,

543 e_score_correction_bias: Optional[torch.Tensor] = None,

544 apply_router_weight_on_input: bool = False,

545 activation: str = "silu",

@@ -560,6 +566,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

566 expert_map,

567 custom_routing_function,

568 scoring_func,

569 + routed_scaling_factor,

570 e_score_correction_bias,

571 apply_router_weight_on_input,

572 activation,

@@ -579,6 +586,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

586 expert_map: Optional[torch.Tensor] = None,

587 custom_routing_function: Optional[Callable] = None,

588 scoring_func: str = "softmax",

589 + routed_scaling_factor: float = 1.0,

590 e_score_correction_bias: Optional[torch.Tensor] = None,

591 apply_router_weight_on_input: bool = False,

592 activation: str = "silu",

@@ -617,6 +625,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

625 expert_map: Optional[torch.Tensor] = None,

626 custom_routing_function: Optional[Callable] = None,

627 scoring_func: str = "softmax",

628 + routed_scaling_factor: float = 1.0,

629 e_score_correction_bias: Optional[torch.Tensor] = None,

630 apply_router_weight_on_input: bool = False,

631 activation: str = "silu",

@@ -637,6 +646,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

646 raise NotImplementedError(

647 "Expert score correction bias is not supported for TPU.")

648 assert activation == "silu", f"{activation} is not supported for TPU."

649 + assert routed_scaling_factor == 1.0, \

650 + f"routed_scaling_factor {routed_scaling_factor} is not supported " \

651 + f"for TPU."

652 if enable_eplb is not False or expert_load_view is not None or \

653 logical_to_physical_map is not None or \

654 logical_replica_count is not None:

@@ -766,6 +778,7 @@ class FusedMoE(CustomOp):

778 prefix: str = "",

779 custom_routing_function: Optional[Callable] = None,

780 scoring_func: str = "softmax",

781 + routed_scaling_factor: float = 1.0,

782 e_score_correction_bias: Optional[torch.Tensor] = None,

783 apply_router_weight_on_input: bool = False,

784 activation: str = "silu",

@@ -848,6 +861,7 @@ class FusedMoE(CustomOp):

861 self.topk_group = topk_group

862 self.custom_routing_function = custom_routing_function

863 self.scoring_func = scoring_func

864 + self.routed_scaling_factor = routed_scaling_factor

865 self.e_score_correction_bias = e_score_correction_bias

866 self.apply_router_weight_on_input = apply_router_weight_on_input

867 self.activation = activation

@@ -1416,6 +1430,7 @@ class FusedMoE(CustomOp):

1430 num_expert_group: Optional[int] = None,

1431 custom_routing_function: Optional[Callable] = None,

1432 scoring_func: str = "softmax",

1433 + routed_scaling_factor: float = 1.0,

1434 e_score_correction_bias: Optional[torch.Tensor] = None,

1435 indices_type: Optional[torch.dtype] = None,

1436 enable_eplb: bool = False,

@@ -1460,6 +1475,7 @@ class FusedMoE(CustomOp):

1475 num_expert_group=num_expert_group,

1476 topk_group=topk_group,

1477 scoring_func=scoring_func,

1478 + routed_scaling_factor=routed_scaling_factor,

1479 e_score_correction_bias=e_score_correction_bias)

1480 if indices_type is not None:

1481 topk_ids = topk_ids.to(dtype=indices_type)

@@ -1627,6 +1643,7 @@ class FusedMoE(CustomOp):

1643 num_expert_group=self.num_expert_group,

1644 custom_routing_function=self.custom_routing_function,

1645 scoring_func=self.scoring_func,

1646 + routed_scaling_factor=self.routed_scaling_factor,

1647 e_score_correction_bias=self.e_score_correction_bias,

1648 activation=self.activation,

1649 enable_eplb=self.enable_eplb,

@@ -1695,6 +1712,7 @@ class FusedMoE(CustomOp):

1712 num_expert_group=self.num_expert_group,

1713 custom_routing_function=self.custom_routing_function,

1714 scoring_func=self.scoring_func,

1715 + routed_scaling_factor=self.routed_scaling_factor,

1716 e_score_correction_bias=self.e_score_correction_bias,

1717 activation=self.activation,

1718 apply_router_weight_on_input=self.apply_router_weight_on_input,

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

@@ -267,6 +267,7 @@ def rocm_aiter_grouped_topk(

267 num_expert_group: int = 0,

268 topk_group: int = 0,

269 scoring_func: str = "softmax",

270 + routed_scaling_factor: float = 1.0,

271 e_score_correction_bias: Optional[torch.Tensor] = None

272 ) -> tuple[torch.Tensor, torch.Tensor]:

273 token = hidden_states.shape[0]

@@ -298,6 +299,8 @@ def rocm_aiter_grouped_topk(

299 scoring_func,

300 )

301

302 + if routed_scaling_factor != 1.0:

303 + topk_weights = topk_weights * routed_scaling_factor

304 return topk_weights, topk_ids

305

306

vllm/model_executor/layers/quantization/awq_marlin.py

@@ -497,6 +497,7 @@ class AWQMoEMethod(FusedMoEMethodBase):

497 expert_map: Optional[torch.Tensor] = None,

498 custom_routing_function: Optional[Callable] = None,

499 scoring_func: str = "softmax",

500 + routed_scaling_factor: float = 1.0,

501 e_score_correction_bias: Optional[torch.Tensor] = None,

502 apply_router_weight_on_input: bool = False,

503 activation: str = "silu",

@@ -523,6 +524,7 @@ class AWQMoEMethod(FusedMoEMethodBase):

524 num_expert_group=num_expert_group,

525 custom_routing_function=custom_routing_function,

526 scoring_func=scoring_func,

527 + routed_scaling_factor=routed_scaling_factor,

528 e_score_correction_bias=e_score_correction_bias,

529 indices_type=self.topk_indices_dtype)

530

vllm/model_executor/layers/quantization/bitsandbytes.py

@@ -466,6 +466,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

466 expert_map: Optional[torch.Tensor] = None,

467 custom_routing_function: Optional[Callable] = None,

468 scoring_func: str = "softmax",

469 + routed_scaling_factor: float = 1.0,

470 e_score_correction_bias: Optional[torch.Tensor] = None,

471 apply_router_weight_on_input: bool = False,

472 activation: str = "silu",

@@ -490,6 +491,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

491 num_expert_group=num_expert_group,

492 custom_routing_function=custom_routing_function,

493 scoring_func=scoring_func,

494 + routed_scaling_factor=routed_scaling_factor,

495 e_score_correction_bias=e_score_correction_bias,

496 indices_type=self.topk_indices_dtype)

497 if self.quant_config.load_in_8bit:

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+10

@@ -350,6 +350,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

350 expert_map: Optional[torch.Tensor] = None,

351 custom_routing_function: Optional[Callable] = None,

352 scoring_func: str = "softmax",

353 + routed_scaling_factor: float = 1.0,

354 e_score_correction_bias: Optional[torch.Tensor] = None,

355 apply_router_weight_on_input: bool = False,

356 activation: str = "silu",

@@ -375,6 +376,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

376 num_expert_group=num_expert_group,

377 custom_routing_function=custom_routing_function,

378 scoring_func=scoring_func,

379 + routed_scaling_factor=routed_scaling_factor,

380 e_score_correction_bias=e_score_correction_bias,

381 indices_type=self.topk_indices_dtype,

382 )

@@ -809,6 +811,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

811 expert_map: Optional[torch.Tensor] = None,

812 custom_routing_function: Optional[Callable] = None,

813 scoring_func: str = "softmax",

814 + routed_scaling_factor: float = 1.0,

815 e_score_correction_bias: Optional[torch.Tensor] = None,

816 apply_router_weight_on_input: bool = False,

817 activation: str = "silu",

@@ -832,6 +835,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

835 num_expert_group=num_expert_group,

836 custom_routing_function=custom_routing_function,

837 scoring_func=scoring_func,

838 + routed_scaling_factor=routed_scaling_factor,

839 e_score_correction_bias=e_score_correction_bias,

840 indices_type=self.topk_indices_dtype,

841 )

@@ -1057,6 +1061,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1061 expert_map: Optional[torch.Tensor] = None,

1062 custom_routing_function: Optional[Callable] = None,

1063 scoring_func: str = "softmax",

1064 + routed_scaling_factor: float = 1.0,

1065 e_score_correction_bias: Optional[torch.Tensor] = None,

1066 apply_router_weight_on_input: bool = False,

1067 activation: str = "silu",

@@ -1084,6 +1089,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1089 num_expert_group=num_expert_group,

1090 custom_routing_function=custom_routing_function,

1091 scoring_func=scoring_func,

1092 + routed_scaling_factor=routed_scaling_factor,

1093 e_score_correction_bias=e_score_correction_bias,

1094 indices_type=self.topk_indices_dtype)

1095

@@ -1361,6 +1367,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1367 expert_map: Optional[torch.Tensor] = None,

1368 custom_routing_function: Optional[Callable] = None,

1369 scoring_func: str = "softmax",

1370 + routed_scaling_factor: float = 1.0,

1371 e_score_correction_bias: Optional[torch.Tensor] = None,

1372 apply_router_weight_on_input: bool = False,

1373 activation: str = "silu",

@@ -1389,6 +1396,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1396 num_expert_group=num_expert_group,

1397 custom_routing_function=custom_routing_function,

1398 scoring_func=scoring_func,

1399 + routed_scaling_factor=routed_scaling_factor,

1400 e_score_correction_bias=e_score_correction_bias,

1401 indices_type=self.topk_indices_dtype)

1402

@@ -1592,6 +1600,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1600 expert_map: Optional[torch.Tensor] = None,

1601 custom_routing_function: Optional[Callable] = None,

1602 scoring_func: str = "softmax",

1603 + routed_scaling_factor: float = 1.0,

1604 e_score_correction_bias: Optional[torch.Tensor] = None,

1605 apply_router_weight_on_input: bool = False,

1606 activation: str = "silu",

@@ -1618,6 +1627,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1627 num_expert_group=num_expert_group,

1628 custom_routing_function=custom_routing_function,

1629 scoring_func=scoring_func,

1630 + routed_scaling_factor=routed_scaling_factor,

1631 e_score_correction_bias=e_score_correction_bias,

1632 indices_type=self.topk_indices_dtype)

1633

vllm/model_executor/layers/quantization/experts_int8.py

@@ -120,6 +120,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

120 expert_map: Optional[torch.Tensor] = None,

121 custom_routing_function: Optional[Callable] = None,

122 scoring_func: str = "softmax",

123 + routed_scaling_factor: float = 1.0,

124 e_score_correction_bias: Optional[torch.Tensor] = None,

125 apply_router_weight_on_input: bool = False,

126 activation: str = "silu",

@@ -146,6 +147,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

147 num_expert_group=num_expert_group,

148 custom_routing_function=custom_routing_function,

149 scoring_func=scoring_func,

150 + routed_scaling_factor=routed_scaling_factor,

151 e_score_correction_bias=e_score_correction_bias,

152 indices_type=self.topk_indices_dtype)

153

vllm/model_executor/layers/quantization/fp8.py

+3 -1

@@ -955,6 +955,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

955 expert_map: Optional[torch.Tensor] = None,

956 custom_routing_function: Optional[Callable] = None,

957 scoring_func: str = "softmax",

958 + routed_scaling_factor: float = 1.0,

959 e_score_correction_bias: Optional[torch.Tensor] = None,

960 apply_router_weight_on_input: bool = False,

961 activation: str = "silu",

@@ -994,7 +995,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

995 expert_offset=layer.ep_rank * layer.local_num_experts,

996 local_num_experts=layer.local_num_experts,

997 block_shape=self.quant_config.weight_block_size,

997 - routed_scaling=1.0,

998 + routed_scaling=routed_scaling_factor,

999 )

1000 else:

1001 assert (not renormalize

@@ -1020,6 +1021,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1021 num_expert_group=num_expert_group,

1022 custom_routing_function=custom_routing_function,

1023 scoring_func=scoring_func,

1024 + routed_scaling_factor=routed_scaling_factor,

1025 e_score_correction_bias=e_score_correction_bias,

1026 indices_type=self.topk_indices_dtype,

1027 enable_eplb=enable_eplb,

vllm/model_executor/layers/quantization/gguf.py

@@ -532,6 +532,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):

532 expert_map: Optional[torch.Tensor] = None,

533 custom_routing_function: Optional[Callable] = None,

534 scoring_func: str = "softmax",

535 + routed_scaling_factor: float = 1.0,

536 e_score_correction_bias: Optional[torch.Tensor] = None,

537 apply_router_weight_on_input: bool = False,

538 activation: str = "silu",

@@ -562,6 +563,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):

563 num_expert_group=num_expert_group,

564 custom_routing_function=custom_routing_function,

565 scoring_func=scoring_func,

566 + routed_scaling_factor=routed_scaling_factor,

567 e_score_correction_bias=e_score_correction_bias,

568 indices_type=self.topk_indices_dtype)

569 return fused_moe_gguf(x, layer.w13_qweight, layer.w2_qweight,

vllm/model_executor/layers/quantization/gptq_marlin.py

@@ -643,6 +643,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

643 expert_map: Optional[torch.Tensor] = None,

644 custom_routing_function: Optional[Callable] = None,

645 scoring_func: str = "softmax",

646 + routed_scaling_factor: float = 1.0,

647 e_score_correction_bias: Optional[torch.Tensor] = None,

648 apply_router_weight_on_input: bool = False,

649 activation: str = "silu",

@@ -669,6 +670,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

670 num_expert_group=num_expert_group,

671 custom_routing_function=custom_routing_function,

672 scoring_func=scoring_func,

673 + routed_scaling_factor=routed_scaling_factor,

674 e_score_correction_bias=e_score_correction_bias,

675 indices_type=self.topk_indices_dtype)

676

vllm/model_executor/layers/quantization/modelopt.py

@@ -483,6 +483,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

483 expert_map: Optional[torch.Tensor] = None,

484 custom_routing_function: Optional[Callable] = None,

485 scoring_func: str = "softmax",

486 + routed_scaling_factor: float = 1.0,

487 e_score_correction_bias: Optional[torch.Tensor] = None,

488 apply_router_weight_on_input: bool = False,

489 activation: str = "silu",

@@ -521,6 +522,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

522 num_expert_group=num_expert_group,

523 custom_routing_function=custom_routing_function,

524 scoring_func=scoring_func,

525 + routed_scaling_factor=routed_scaling_factor,

526 e_score_correction_bias=e_score_correction_bias,

527 indices_type=self.topk_indices_dtype,

528 )

@@ -1356,6 +1358,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1358 expert_map: Optional[torch.Tensor] = None,

1359 custom_routing_function: Optional[Callable] = None,

1360 scoring_func: str = "softmax",

1361 + routed_scaling_factor: float = 1.0,

1362 e_score_correction_bias: Optional[torch.Tensor] = None,

1363 apply_router_weight_on_input: bool = False,

1364 activation: str = "silu",

@@ -1434,6 +1437,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1437 num_expert_group=num_expert_group,

1438 custom_routing_function=custom_routing_function,

1439 scoring_func=scoring_func,

1440 + routed_scaling_factor=routed_scaling_factor,

1441 e_score_correction_bias=e_score_correction_bias,

1442 indices_type=self.topk_indices_dtype)

1443

vllm/model_executor/layers/quantization/moe_wna16.py

@@ -297,6 +297,7 @@ class MoeWNA16Method(FusedMoEMethodBase):

297 expert_map: Optional[torch.Tensor] = None,

298 custom_routing_function: Optional[Callable] = None,

299 scoring_func: str = "softmax",

300 + routed_scaling_factor: float = 1.0,

301 e_score_correction_bias: Optional[torch.Tensor] = None,

302 apply_router_weight_on_input: bool = False,

303 activation: str = "silu",

@@ -322,6 +323,7 @@ class MoeWNA16Method(FusedMoEMethodBase):

323 num_expert_group=num_expert_group,

324 custom_routing_function=custom_routing_function,

325 scoring_func=scoring_func,

326 + routed_scaling_factor=routed_scaling_factor,

327 e_score_correction_bias=e_score_correction_bias,

328 indices_type=self.topk_indices_dtype)

329

vllm/model_executor/layers/quantization/mxfp4.py

@@ -546,6 +546,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

546 expert_map: Optional[torch.Tensor] = None,

547 custom_routing_function: Optional[Callable] = None,

548 scoring_func: str = "softmax",

549 + routed_scaling_factor: float = 1.0,

550 e_score_correction_bias: Optional[torch.Tensor] = None,

551 apply_router_weight_on_input: bool = False,

552 activation: str = "silu",

@@ -569,6 +570,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

570 num_expert_group=num_expert_group,

571 custom_routing_function=custom_routing_function,

572 scoring_func=scoring_func,

573 + routed_scaling_factor=routed_scaling_factor,

574 e_score_correction_bias=e_score_correction_bias)

575

576 return torch.ops.vllm.fused_marlin_moe(

vllm/model_executor/layers/quantization/quark/quark_moe.py

@@ -218,6 +218,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

218 expert_map: Optional[torch.Tensor] = None,

219 custom_routing_function: Optional[Callable] = None,

220 scoring_func: str = "softmax",

221 + routed_scaling_factor: float = 1.0,

222 e_score_correction_bias: Optional[torch.Tensor] = None,

223 apply_router_weight_on_input: bool = False,

224 activation: str = "silu",

@@ -244,6 +245,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

245 num_expert_group=num_expert_group,

246 custom_routing_function=custom_routing_function,

247 scoring_func=scoring_func,

248 + routed_scaling_factor=routed_scaling_factor,

249 e_score_correction_bias=e_score_correction_bias,

250 indices_type=self.topk_indices_dtype)

251

@@ -380,6 +382,7 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):

382 expert_map: Optional[torch.Tensor] = None,

383 custom_routing_function: Optional[Callable] = None,

384 scoring_func: str = "softmax",

385 + routed_scaling_factor: float = 1.0,

386 e_score_correction_bias: Optional[torch.Tensor] = None,

387 apply_router_weight_on_input: bool = False,

388 activation: str = "silu",

@@ -406,6 +409,7 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):

409 num_expert_group=num_expert_group,

410 custom_routing_function=custom_routing_function,

411 scoring_func=scoring_func,

412 + routed_scaling_factor=routed_scaling_factor,

413 e_score_correction_bias=e_score_correction_bias,

414 indices_type=self.topk_indices_dtype)

415

vllm/model_executor/layers/quantization/rtn.py

@@ -283,6 +283,7 @@ class RTNMoEMethod(FusedMoEMethodBase):

283 expert_map: Optional[torch.Tensor] = None,

284 custom_routing_function: Optional[Callable] = None,

285 scoring_func: str = "softmax",

286 + routed_scaling_factor: float = 1.0,

287 e_score_correction_bias: Optional[torch.Tensor] = None,

288 apply_router_weight_on_input: bool = False,

289 activation: str = "silu",

@@ -309,6 +310,7 @@ class RTNMoEMethod(FusedMoEMethodBase):

310 num_expert_group=num_expert_group,

311 custom_routing_function=custom_routing_function,

312 scoring_func=scoring_func,

313 + routed_scaling_factor=routed_scaling_factor,

314 e_score_correction_bias=e_score_correction_bias,

315 indices_type=self.topk_indices_dtype)

316

vllm/model_executor/models/deepseek_v2.py

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

160 topk_group=config.topk_group,

161 prefix=f"{prefix}.experts",

162 scoring_func=config.scoring_func,

163 + routed_scaling_factor=self.routed_scaling_factor,

164 e_score_correction_bias=self.gate.e_score_correction_bias,

165 enable_eplb=self.enable_eplb,

166 num_redundant_experts=self.n_redundant_experts)

vllm/model_executor/models/dots1.py

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

137 topk_group=config.topk_group,

138 prefix=f"{prefix}.experts",

139 scoring_func=config.scoring_func,

140 + routed_scaling_factor=self.routed_scaling_factor,

141 e_score_correction_bias=self.gate.e_score_correction_bias)

142

143 if config.n_shared_experts is not None:

vllm/model_executor/models/glm4_moe.py

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

159 topk_group=config.topk_group,

160 prefix=f"{prefix}.experts",

161 scoring_func="sigmoid",

162 + routed_scaling_factor=self.routed_scaling_factor,

163 e_score_correction_bias=self.gate.e_score_correction_bias,

164 enable_eplb=self.enable_eplb,

165 num_redundant_experts=self.n_redundant_experts)