de6674b002 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/model_executor/layers/fused_moe/router/gate_linear.py

+53 -13

@@ -5,7 +5,11 @@ from torch.nn.parameter import Parameter

5

6 import vllm._custom_ops as ops

7 from vllm.model_executor.custom_op import PluggableLayer

8 -from vllm.model_executor.layers.linear import ReplicatedLinear

8 +from vllm.model_executor.layers.linear import (

9 + ReplicatedLinear,

10 + UnquantizedLinearMethod,

11 +)

12 +from vllm.model_executor.layers.quantization import QuantizationConfig

13 from vllm.platforms import current_platform

14 from vllm.utils.torch_utils import direct_register_custom_op

15

@@ -25,12 +29,18 @@ class GateLinear(ReplicatedLinear):

29 The ``out_dtype`` attribute is mutable and can be set after init

30 (e.g. when the required dtype depends on the expert quantization

31 method which is only known later).

32 +

33 + A ``quant_config`` that actually quantizes the gate disables every

34 + specialized tier, leaving plain ``ReplicatedLinear`` behavior.

35 """

36

37 # (hidden_size, num_experts) pairs with an instantiated fp32 kernel:

38 # (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3

39 FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)}

40 FP32_MAX_TOKENS = 32

41 + # Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses

42 + # whose router shape stays ahead further up the range may raise it.

43 + LL_BF16_MAX_TOKENS = 16

44

45 def __init__(

46 self,

@@ -40,7 +50,10 @@ class GateLinear(ReplicatedLinear):

50 out_dtype: torch.dtype | None = None,

51 params_dtype: torch.dtype | None = None,

52 force_fp32_compute: bool = False,

53 + skip_bias_add: bool = False,

54 + quant_config: QuantizationConfig | None = None,

55 prefix: str = "",

56 + return_bias: bool = True,

57 ):

58 is_hopper = current_platform.is_device_capability((9, 0))

59 is_blackwell = current_platform.is_device_capability_family(100)

@@ -72,17 +85,25 @@ class GateLinear(ReplicatedLinear):

85 input_size,

86 output_size,

87 bias=bias,

88 + skip_bias_add=skip_bias_add,

89 params_dtype=params_dtype,

76 - quant_config=None,

90 + quant_config=quant_config,

91 prefix=prefix,

92 + return_bias=return_bias,

93 )

94 self.out_dtype = out_dtype

95

96 + # A quantized gate exposes no plain ``weight``, so every specialized

97 + # tier below is disabled and forward falls back to ReplicatedLinear.

98 + self.is_unquantized = isinstance(self.quant_method, UnquantizedLinearMethod)

99 + can_use_specialized_kernels &= self.is_unquantized

100 +

101 self.allow_specialized_router_gemm = can_use_specialized_kernels

102

103 # fp32 specialized kernel eligibility (exact dims, fp32 weight)

104 self.allow_fp32_router_gemm = (

85 - not bias

105 + self.is_unquantized

106 + and not bias

107 and self.weight.dtype == torch.float32

108 and (

109 (

@@ -94,7 +115,8 @@ class GateLinear(ReplicatedLinear):

115 )

116 )

117 self.allow_bf16x3_router_gemm = (

97 - not bias

118 + self.is_unquantized

119 + and not bias

120 and self.weight.dtype == torch.float32

121 and current_platform.is_cuda()

122 and is_blackwell

@@ -108,7 +130,7 @@ class GateLinear(ReplicatedLinear):

130 # applies on any CUDA-alike device (no bias, since torch.mm has no bias

131 # term). The specialized-kernel gate above excludes family-120 Blackwell

132 # (GB10 / DGX Spark), which this tier still covers. See #49921.

111 - self._router_gemm_no_bias = not bias

133 + self._router_gemm_no_bias = self.is_unquantized and not bias

134 self._router_gemm_cublas_capable = (

135 current_platform.is_cuda() or current_platform.is_rocm()

136 ) and self._router_gemm_no_bias

@@ -162,17 +184,26 @@ class GateLinear(ReplicatedLinear):

184 and is_available()

185 )

186

187 + def _return(

188 + self, output: torch.Tensor

189 + ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:

190 + return output if not self.return_bias else (output, None)

191 +

192 def forward(

193 self, x: torch.Tensor

194 ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:

195 # Tier 1: cuteDSL ll_bf16_gemm (SM90+, any dims)

169 - if self.allow_ll_bf16_gemm and x.shape[0] <= 16 and x.dtype == torch.bfloat16:

196 + if (

197 + self.allow_ll_bf16_gemm

198 + and x.shape[0] <= self.LL_BF16_MAX_TOKENS

199 + and x.dtype == torch.bfloat16

200 + ):

201 from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (

202 ll_bf16_gemm,

203 )

204

205 output = ll_bf16_gemm(x, self.weight)

175 - return output, None

206 + return self._return(output)

207

208 # Tier 2: fp32 specialized kernel (model-specific shapes, M<=32)

209 # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph

@@ -184,7 +215,7 @@ class GateLinear(ReplicatedLinear):

215 output = torch.ops.vllm.fp32_router_gemm_dispatch(

216 x, self.weight, self.allow_bf16x3_router_gemm

217 )

187 - return output, None

218 + return self._return(output)

219

220 # Tier 3: bf16x3 CuteDSL kernel for fp32 router weights

221 if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:

@@ -193,20 +224,29 @@ class GateLinear(ReplicatedLinear):

224 )

225

226 output = bf16x3_router_gemm(x, self.weight)

196 - return output, None

227 + return self._return(output)

228

229 # Tier 4: cuBLAS bf16→fp32

230 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:

231 output = torch.mm(x, self.weight.T, out_dtype=torch.float32)

201 - return output, None

232 + return self._return(output)

233

234 # Tier 5: F.linear (ReplicatedLinear)

204 - if self.out_dtype is not None and x.dtype != self.weight.dtype:

235 + if (

236 + self.out_dtype is not None

237 + and self.is_unquantized

238 + and x.dtype != self.weight.dtype

239 + ):

240 x = x.to(self.weight.dtype)

206 - output, output_bias = super().forward(x)

241 + base_output = super().forward(x)

242 + output_bias: Parameter | None = None

243 + if isinstance(base_output, tuple):

244 + output, output_bias = base_output

245 + else:

246 + output = base_output

247 if self.out_dtype is not None and output.dtype != self.out_dtype:

248 output = output.to(self.out_dtype)

209 - return output, output_bias

249 + return output if not self.return_bias else (output, output_bias)

250

251

252 _FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS

vllm/model_executor/models/AXK1.py

+2 -3

@@ -44,6 +44,7 @@ from vllm.logger import init_logger

44 from vllm.model_executor.layers.attention import Attention

45 from vllm.model_executor.layers.fused_moe import (

46 FusedMoEFactory,

47 + GateLinear,

48 fused_moe_make_expert_params_mapping,

49 )

50 from vllm.model_executor.layers.fused_moe.utils import (

@@ -128,11 +129,9 @@ class AXK1MoE(nn.Module):

129 "Only silu is supported for now."

130 )

131

131 - self.gate = ReplicatedLinear(

132 + self.gate = GateLinear(

133 config.hidden_size,

134 config.n_routed_experts,

134 - bias=False,

135 - quant_config=None,

135 prefix=f"{prefix}.gate",

136 )

137 if config.topk_method == "noaux_tc":

vllm/model_executor/models/afmoe.py

+6 -4

@@ -19,6 +19,7 @@ from vllm.logger import init_logger

19 from vllm.model_executor.layers.attention import Attention

20 from vllm.model_executor.layers.fused_moe import (

21 FusedMoEFactory,

22 + GateLinear,

23 MoERunner,

24 )

25 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -82,11 +83,12 @@ class AfmoeMoE(nn.Module):

83 )

84

85 # Router gate

85 - self.gate = nn.Linear(

86 + self.gate = GateLinear(

87 config.hidden_size,

88 config.num_experts,

88 - bias=False,

89 - dtype=torch.float32,

89 + out_dtype=torch.float32,

90 + params_dtype=torch.float32,

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

92 )

93 self.expert_bias = nn.Parameter(

94 torch.empty(config.num_experts, dtype=torch.float32)

@@ -140,7 +142,7 @@ class AfmoeMoE(nn.Module):

142 num_tokens, hidden_dim = hidden_states.shape

143 hidden_states = hidden_states.view(-1, hidden_dim)

144

143 - router_logits = self.gate(hidden_states.to(dtype=torch.float32))

145 + router_logits, _ = self.gate(hidden_states)

146

147 final_hidden_states = self.experts(

148 hidden_states=hidden_states, router_logits=router_logits

vllm/model_executor/models/bailing_moe.py

+9 -5

@@ -41,7 +41,10 @@ from vllm.distributed import (

41 )

42 from vllm.model_executor.layers.activation import SiluAndMul

43 from vllm.model_executor.layers.attention import Attention

44 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

44 +from vllm.model_executor.layers.fused_moe import (

45 + FusedMoEFactory,

46 + GateLinear,

47 +)

48 from vllm.model_executor.layers.layernorm import RMSNorm

49 from vllm.model_executor.layers.linear import (

50 MergedColumnParallelLinear,

@@ -243,11 +246,12 @@ class BailingMoE(nn.Module):

246 else:

247 self.router_dtype = torch.bfloat16

248

246 - self.gate = nn.Linear(

249 + self.gate = GateLinear(

250 self.hidden_size,

251 self.num_experts,

249 - bias=False,

250 - dtype=self.router_dtype,

252 + out_dtype=self.router_dtype,

253 + params_dtype=self.router_dtype,

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

255 )

256

257 if getattr(config, "moe_router_enable_expert_bias", False):

@@ -312,7 +316,7 @@ class BailingMoE(nn.Module):

316 hidden_states = hidden_states.view(-1, hidden_size)

317

318 # router_logits: (num_tokens, n_experts)

315 - router_logits = self.gate(hidden_states.to(self.router_dtype))

319 + router_logits, _ = self.gate(hidden_states)

320 router_logits = router_logits.to(hidden_states.dtype)

321

322 final_hidden_states = self.experts(

vllm/model_executor/models/bailing_moe_linear.py

+12 -35

@@ -5,7 +5,6 @@ from collections.abc import Callable, Iterable

5

6 import torch

7 import torch.nn as nn

8 -import torch.nn.functional as F

8 from transformers.configuration_utils import PretrainedConfig

9

10 from vllm.compilation.decorators import support_torch_compile

@@ -19,6 +18,7 @@ from vllm.forward_context import get_forward_context

18 from vllm.logger import init_logger

19 from vllm.model_executor.layers.fused_moe import (

20 FusedMoEFactory,

21 + GateLinear,

22 fused_moe_make_expert_params_mapping,

23 )

24 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -229,37 +229,6 @@ class BailingMoeV25MLAAttention(nn.Module):

229 return self.mla_attn(positions, hidden_states)

230

231

232 -class BailingMoEGate(nn.Module):

233 - def __init__(

234 - self,

235 - config: PretrainedConfig,

236 - params_dtype: torch.dtype | None = None,

237 - prefix: str = "",

238 - ):

239 - super().__init__()

240 - if params_dtype is None:

241 - params_dtype = torch.get_default_dtype()

242 - self.params_dtype = params_dtype

243 - self.weight = nn.Parameter(

244 - torch.empty(

245 - (config.num_experts, config.hidden_size),

246 - dtype=self.params_dtype,

247 - ),

248 - )

249 - if getattr(config, "moe_router_enable_expert_bias", False):

250 - self.expert_bias = nn.Parameter(

251 - torch.empty((config.num_experts,), dtype=torch.float32),

252 - )

253 - else:

254 - self.expert_bias = None

255 -

256 - def forward(self, hidden_states):

257 - logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None).to(

258 - hidden_states.dtype

259 - )

260 - return logits

261 -

262 -

232 class BailingMoeV25(nn.Module):

233 """Bailing MoE v2.5 - standalone implementation for linear attention model."""

234

@@ -296,11 +265,19 @@ class BailingMoeV25(nn.Module):

265 self.router_dtype = torch.bfloat16

266

267 # Gate for routing

299 - self.gate = BailingMoEGate(

300 - config=config,

268 + self.gate = GateLinear(

269 + self.hidden_size,

270 + self.num_experts,

271 + out_dtype=self.router_dtype,

272 params_dtype=self.router_dtype,

273 prefix=f"{prefix}.gate",

274 )

275 + if getattr(config, "moe_router_enable_expert_bias", False):

276 + self.gate.expert_bias = nn.Parameter(

277 + torch.empty((self.num_experts,), dtype=torch.float32),

278 + )

279 + else:

280 + self.gate.expert_bias = None

281 correction_bias = (

282 self.gate.expert_bias if self.gate.expert_bias is not None else None

283 )

@@ -357,7 +334,7 @@ class BailingMoeV25(nn.Module):

334 hidden_states = hidden_states.contiguous().view(-1, hidden_size)

335

336 # router_logits: (num_tokens, n_experts)

360 - router_logits = self.gate(hidden_states.to(self.router_dtype))

337 + router_logits, _ = self.gate(hidden_states)

338 router_logits = router_logits.to(hidden_states.dtype)

339

340 final_hidden_states = self.experts(

vllm/model_executor/models/bailing_moe_v3.py

+12 -25

@@ -16,7 +16,6 @@ from typing import TypeGuard

16 import regex as re

17 import torch

18 import torch.nn as nn

19 -import torch.nn.functional as F

19 from einops import rearrange

20 from transformers.configuration_utils import PretrainedConfig

21

@@ -32,6 +31,7 @@ from vllm.model_executor.custom_op import PluggableLayer

31 from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul

32 from vllm.model_executor.layers.fused_moe import (

33 FusedMoEFactory,

34 + GateLinear,

35 fused_moe_make_expert_params_mapping,

36 )

37 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -1118,28 +1118,6 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):

1118 core_attn_out[0, :num_actual_tokens] = out[0, :num_actual_tokens]

1119

1120

1121 -class BailingMoeV3Gate(nn.Module):

1122 - def __init__(

1123 - self,

1124 - config: PretrainedConfig,

1125 - params_dtype: torch.dtype | None = None,

1126 - ) -> None:

1127 - super().__init__()

1128 - if params_dtype is None:

1129 - params_dtype = torch.float32

1130 - self.weight = nn.Parameter(

1131 - torch.empty((config.num_experts, config.hidden_size), dtype=params_dtype)

1132 - )

1133 - self.expert_bias = nn.Parameter(

1134 - torch.empty((config.num_experts,), dtype=torch.float32)

1135 - )

1136 -

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

1138 - return F.linear(hidden_states.to(self.weight.dtype), self.weight).to(

1139 - hidden_states.dtype

1140 - )

1141 -

1142 -

1121 class BailingMoeV3MoE(nn.Module):

1122 def __init__(

1123 self,

@@ -1153,7 +1131,16 @@ class BailingMoeV3MoE(nn.Module):

1131 self.top_k = config.num_experts_per_tok

1132 self.hidden_size = config.hidden_size

1133 self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)

1156 - self.gate = BailingMoeV3Gate(config)

1134 + self.gate = GateLinear(

1135 + config.hidden_size,

1136 + config.num_experts,

1137 + out_dtype=torch.float32,

1138 + params_dtype=torch.float32,

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

1140 + )

1141 + self.gate.expert_bias = nn.Parameter(

1142 + torch.empty((config.num_experts,), dtype=torch.float32)

1143 + )

1144 self.expert_swiglu_limit = _get_layer_swiglu_limit(

1145 getattr(config, "expert_swiglu_limit_list", None), layer_id

1146 )

@@ -1196,7 +1183,7 @@ class BailingMoeV3MoE(nn.Module):

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

1184 num_tokens, hidden_size = hidden_states.shape

1185 hidden_states = hidden_states.contiguous().view(-1, hidden_size)

1199 - router_logits = self.gate(hidden_states.to(torch.float32))

1186 + router_logits, _ = self.gate(hidden_states)

1187 hidden_states = self.experts(

1188 hidden_states=hidden_states, router_logits=router_logits

1189 )

vllm/model_executor/models/bert_with_rope.py

+6 -4

@@ -18,7 +18,11 @@ from vllm.model_executor.layers.activation import get_act_and_mul_fn, get_act_fn

18 from vllm.model_executor.layers.attention import (

19 EncoderOnlyAttention,

20 )

21 -from vllm.model_executor.layers.fused_moe import activation_without_mul, fused_topk

21 +from vllm.model_executor.layers.fused_moe import (

22 + GateLinear,

23 + activation_without_mul,

24 + fused_topk,

25 +)

26 from vllm.model_executor.layers.linear import (

27 ColumnParallelLinear,

28 MergedColumnParallelLinear,

@@ -254,9 +258,7 @@ class NomicMoE(nn.Module):

258 params_dtype = torch.get_default_dtype()

259 self.params_dtype = params_dtype

260

257 - self.router = ReplicatedLinear(

258 - self.hidden_size, self.num_total_experts, bias=False

259 - )

261 + self.router = GateLinear(self.hidden_size, self.num_total_experts)

262 self.w1 = nn.Parameter(

263 torch.empty(

264 self.num_total_experts,

vllm/model_executor/models/cohere2_moe.py

+5 -5

@@ -16,11 +16,13 @@ from vllm.distributed import (

16 )

17 from vllm.model_executor.layers.activation import SiluAndMul

18 from vllm.model_executor.layers.attention import Attention

19 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

19 +from vllm.model_executor.layers.fused_moe import (

20 + FusedMoEFactory,

21 + GateLinear,

22 +)

23 from vllm.model_executor.layers.linear import (

24 MergedColumnParallelLinear,

25 QKVParallelLinear,

23 - ReplicatedLinear,

26 RowParallelLinear,

27 )

28 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -275,12 +277,10 @@ class Cohere2Moe(nn.Module):

277 else:

278 self.custom_routing_function = None

279

278 - self.gate = ReplicatedLinear(

280 + self.gate = GateLinear(

281 config.hidden_size,

282 config.num_experts,

281 - bias=False,

283 params_dtype=params_dtype,

283 - quant_config=None,

284 prefix=f"{prefix}.gate",

285 )

286

vllm/model_executor/models/dbrx.py

+2 -4

@@ -18,11 +18,11 @@ from vllm.distributed import (

18 from vllm.model_executor.layers.attention import Attention

19 from vllm.model_executor.layers.fused_moe import (

20 FusedMoEFactory,

21 + GateLinear,

22 RoutedExperts,

23 )

24 from vllm.model_executor.layers.linear import (

25 QKVParallelLinear,

25 - ReplicatedLinear,

26 RowParallelLinear,

27 )

28 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -62,12 +62,10 @@ class DbrxRouter(nn.Module):

62 self.tp_size = get_tensor_model_parallel_world_size()

63 self.num_total_experts = config.ffn_config.moe_num_experts

64 self.d_model = config.d_model

65 - self.layer = ReplicatedLinear(

65 + self.layer = GateLinear(

66 self.d_model,

67 self.num_total_experts,

68 - bias=False,

68 params_dtype=params_dtype,

70 - quant_config=None,

69 )

70

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

vllm/model_executor/models/ernie45_moe.py

+8 -6

@@ -41,12 +41,15 @@ from vllm.distributed import (

41 from vllm.logger import init_logger

42 from vllm.model_executor.layers.activation import SiluAndMul

43 from vllm.model_executor.layers.attention import Attention

44 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner

44 +from vllm.model_executor.layers.fused_moe import (

45 + FusedMoEFactory,

46 + GateLinear,

47 + MoERunner,

48 +)

49 from vllm.model_executor.layers.layernorm import RMSNorm

50 from vllm.model_executor.layers.linear import (

51 MergedColumnParallelLinear,

52 QKVParallelLinear,

49 - ReplicatedLinear,

53 RowParallelLinear,

54 )

55 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -150,12 +153,11 @@ class Ernie4_5_MoeMoE(nn.Module):

153 f"the number of experts {config.moe_num_experts}."

154 )

155

153 - self.gate = ReplicatedLinear(

156 + self.gate = GateLinear(

157 config.hidden_size,

158 config.moe_num_experts,

156 - bias=False,

159 + out_dtype=torch.float32,

160 params_dtype=torch.float32,

158 - quant_config=None,

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

162 )

163

@@ -199,7 +201,7 @@ class Ernie4_5_MoeMoE(nn.Module):

201 hidden_dim = hidden_states.shape[-1]

202 hidden_states = hidden_states.view(-1, hidden_dim)

203

202 - router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))

204 + router_logits, _ = self.gate(hidden_states)

205

206 final_hidden_states = self.experts(

207 hidden_states=hidden_states, router_logits=router_logits

vllm/model_executor/models/exaone_moe.py

+5 -5

@@ -30,11 +30,13 @@ from vllm.distributed import (

30 get_tensor_model_parallel_world_size,

31 )

32 from vllm.model_executor.layers.attention import Attention

33 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

33 +from vllm.model_executor.layers.fused_moe import (

34 + FusedMoEFactory,

35 + GateLinear,

36 +)

37 from vllm.model_executor.layers.layernorm import RMSNorm

38 from vllm.model_executor.layers.linear import (

39 QKVParallelLinear,

37 - ReplicatedLinear,

40 RowParallelLinear,

41 )

42 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -84,11 +86,9 @@ class ExaoneMoe(nn.Module):

86 f"the number of experts {config.num_experts}."

87 )

88

87 - self.gate = ReplicatedLinear(

89 + self.gate = GateLinear(

90 config.hidden_size,

91 config.num_experts,

90 - bias=False,

91 - quant_config=None,

92 prefix=f"{prefix}.gate",

93 )

94

vllm/model_executor/models/glm4_moe.py

+6 -7

@@ -43,6 +43,7 @@ from vllm.model_executor.layers.activation import SiluAndMul

43 from vllm.model_executor.layers.attention import Attention

44 from vllm.model_executor.layers.fused_moe import (

45 FusedMoEFactory,

46 + GateLinear,

47 )

48 from vllm.model_executor.layers.fused_moe.utils import (

49 is_model_fused_shared_expert_compatible,

@@ -137,14 +138,12 @@ class Glm4MoE(nn.Module):

138 f"Unsupported activation: {config.hidden_act}. "

139 "Only silu is supported for now."

140 )

140 - # NOTE In the transformers implementation, the gate isn't an nn.Linear,

141 - # so we cannot use ReplicatedLinear here.

142 - # See: https://github.com/huggingface/transformers/blob/v4.55.1/src/transformers/models/glm4_moe/modeling_glm4_moe.py#L260

143 - self.gate = nn.Linear(

141 + self.gate = GateLinear(

142 config.hidden_size,

143 config.n_routed_experts,

146 - bias=False,

147 - dtype=torch.float32,

144 + out_dtype=torch.float32,

145 + params_dtype=torch.float32,

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

147 )

148 self.gate.e_score_correction_bias = nn.Parameter(

149 torch.empty(config.n_routed_experts, dtype=torch.float32)

@@ -212,7 +211,7 @@ class Glm4MoE(nn.Module):

211 hidden_states = hidden_states.view(-1, hidden_dim)

212

213 # router_logits: (num_tokens, n_experts)

215 - router_logits = self.gate(hidden_states.to(dtype=torch.float32))

214 + router_logits, _ = self.gate(hidden_states)

215

216 final_hidden_states = self.experts(

217 hidden_states=hidden_states, router_logits=router_logits

vllm/model_executor/models/gpt_oss.py

+2 -3

@@ -22,6 +22,7 @@ from vllm.distributed import (

22 from vllm.model_executor.layers.attention import Attention

23 from vllm.model_executor.layers.fused_moe import (

24 FusedMoEFactory,

25 + GateLinear,

26 fused_moe_make_expert_params_mapping,

27 )

28 from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig

@@ -29,7 +30,6 @@ from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts

30 from vllm.model_executor.layers.layernorm import RMSNorm

31 from vllm.model_executor.layers.linear import (

32 QKVParallelLinear,

32 - ReplicatedLinear,

33 RowParallelLinear,

34 )

35 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -377,11 +377,10 @@ class MLPBlock(torch.nn.Module):

377 self.hidden_size = config.hidden_size

378 self.experts_per_token = config.num_experts_per_tok

379 self.world_size = dist.get_world_size() if dist.is_initialized() else 1

380 - self.router = ReplicatedLinear(

380 + self.router = GateLinear(

381 config.hidden_size,

382 config.num_local_experts,

383 bias=True,

384 - quant_config=None,

384 prefix=f"{prefix}.router",

385 return_bias=False,

386 )

vllm/model_executor/models/granitemoe.py

+5 -5

@@ -39,11 +39,13 @@ from vllm.distributed import (

39 tensor_model_parallel_all_gather,

40 )

41 from vllm.model_executor.layers.attention import Attention

42 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

42 +from vllm.model_executor.layers.fused_moe import (

43 + FusedMoEFactory,

44 + GateLinear,

45 +)

46 from vllm.model_executor.layers.layernorm import RMSNorm

47 from vllm.model_executor.layers.linear import (

48 QKVParallelLinear,

46 - ReplicatedLinear,

49 RowParallelLinear,

50 )

51 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -94,12 +96,10 @@ class GraniteMoeMoE(nn.Module):

96 self.is_sequence_parallel = is_sequence_parallel

97

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

97 - self.gate = ReplicatedLinear(

99 + self.gate = GateLinear(

100 hidden_size,

101 num_experts,

100 - bias=False,

102 params_dtype=params_dtype,

102 - quant_config=None,

103 prefix=f"{prefix}.gate",

104 )

105

vllm/model_executor/models/interns1_pro.py

+2 -3

@@ -43,12 +43,12 @@ from vllm.model_executor.layers.activation import SiluAndMul

43 from vllm.model_executor.layers.attention import Attention

44 from vllm.model_executor.layers.fused_moe import (

45 FusedMoEFactory,

46 + GateLinear,

47 )

48 from vllm.model_executor.layers.layernorm import RMSNorm

49 from vllm.model_executor.layers.linear import (

50 MergedColumnParallelLinear,

51 QKVParallelLinear,

51 - ReplicatedLinear,

52 RowParallelLinear,

53 )

54 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -181,10 +181,9 @@ class InternS1ProMoeSparseMoeBlock(nn.Module):

181 custom_routing_function=self._custom_routing_function,

182 )

183

184 - self.gate = ReplicatedLinear(

184 + self.gate = GateLinear(

185 config.hidden_size,

186 config.num_experts,

187 - bias=False,

187 prefix=f"{prefix}.gate",

188 )

189

vllm/model_executor/models/interns2_mobius.py

+5 -4

@@ -17,7 +17,10 @@ from vllm.distributed import (

17 tensor_model_parallel_all_gather,

18 tensor_model_parallel_reduce_scatter,

19 )

20 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

20 +from vllm.model_executor.layers.fused_moe import (

21 + FusedMoEFactory,

22 + GateLinear,

23 +)

24 from vllm.model_executor.layers.linear import ReplicatedLinear

25 from vllm.model_executor.layers.logits_processor import LogitsProcessor

26 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (

@@ -93,11 +96,9 @@ class InternS2MobiusMetaMoeBlock(nn.Module):

96 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

97 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

98

96 - self.gate = ReplicatedLinear(

99 + self.gate = GateLinear(

100 config.hidden_size,

101 config.num_experts,

99 - bias=False,

100 - quant_config=None,

102 prefix=f"{prefix}.gate",

103 )

104 self.experts = FusedMoEFactory(

vllm/model_executor/models/jamba.py

+2 -4

@@ -16,11 +16,11 @@ from vllm.distributed.parallel_state import get_pp_group

16 from vllm.model_executor.layers.attention import Attention

17 from vllm.model_executor.layers.fused_moe import (

18 FusedMoEFactory,

19 + GateLinear,

20 )

21 from vllm.model_executor.layers.layernorm import RMSNorm

22 from vllm.model_executor.layers.linear import (

23 QKVParallelLinear,

23 - ReplicatedLinear,

24 RowParallelLinear,

25 )

26 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -74,11 +74,9 @@ class JambaMoE(nn.Module):

74 self.intermediate_size = config.intermediate_size

75

76 if self.num_total_experts > 1:

77 - self.router = ReplicatedLinear(

77 + self.router = GateLinear(

78 self.hidden_size,

79 self.num_total_experts,

80 - bias=False,

81 - quant_config=None,

80 params_dtype=params_dtype,

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

82 )

vllm/model_executor/models/k2_horizon.py

+2 -1

@@ -49,6 +49,7 @@ from vllm.model_executor.layers.activation import SiluAndMul

49 from vllm.model_executor.layers.attention import Attention

50 from vllm.model_executor.layers.fused_moe import (

51 FusedMoEFactory,

52 + GateLinear,

53 fused_moe_make_expert_params_mapping,

54 )

55 from vllm.model_executor.layers.fused_moe.config import (

@@ -398,7 +399,7 @@ class K2HorizonSparseMoeBlock(nn.Module):

399 self.physical_expert_start + self.n_local_physical_experts

400 )

401

401 - self.gate = ReplicatedLinear(

402 + self.gate = GateLinear(

403 config.hidden_size,

404 config.num_experts,

405 bias=config.moe_gate_bias,

vllm/model_executor/models/laguna.py

+5 -5

@@ -19,12 +19,14 @@ from vllm.distributed import (

19 )

20 from vllm.logger import init_logger

21 from vllm.model_executor.layers.attention import Attention

22 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

22 +from vllm.model_executor.layers.fused_moe import (

23 + FusedMoEFactory,

24 + GateLinear,

25 +)

26 from vllm.model_executor.layers.layernorm import RMSNorm

27 from vllm.model_executor.layers.linear import (

28 ColumnParallelLinear,

29 QKVParallelLinear,

27 - ReplicatedLinear,

30 RowParallelLinear,

31 )

32 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -171,11 +173,9 @@ class LagunaMoE(nn.Module):

173 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

174 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

175 # Router gate

174 - self.gate = ReplicatedLinear(

176 + self.gate = GateLinear(

177 config.hidden_size,

178 config.num_experts,

177 - bias=False,

178 - quant_config=None,

179 prefix=f"{prefix}.gate",

180 )

181

vllm/model_executor/models/lfm2_moe.py

+2 -3

@@ -17,12 +17,12 @@ from vllm.model_executor.layers.activation import SiluAndMul

17 from vllm.model_executor.layers.attention import Attention

18 from vllm.model_executor.layers.fused_moe import (

19 FusedMoEFactory,

20 + GateLinear,

21 )

22 from vllm.model_executor.layers.layernorm import RMSNorm

23 from vllm.model_executor.layers.linear import (

24 MergedColumnParallelLinear,

25 QKVParallelLinear,

25 - ReplicatedLinear,

26 RowParallelLinear,

27 )

28 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -125,10 +125,9 @@ class Lfm2MoeSparseMoeBlock(nn.Module):

125 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

126 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

127

128 - self.gate = ReplicatedLinear(

128 + self.gate = GateLinear(

129 config.hidden_size,

130 config.num_experts,

131 - bias=False,

131 quant_config=quant_config,

132 prefix=f"{prefix}.gate",

133 )

vllm/model_executor/models/llama4.py

+2 -4

@@ -38,12 +38,12 @@ from vllm.model_executor.layers.attention import (

38 )

39 from vllm.model_executor.layers.fused_moe import (

40 FusedMoEFactory,

41 + GateLinear,

42 fused_moe_make_expert_params_mapping,

43 )

44 from vllm.model_executor.layers.layernorm import RMSNorm

45 from vllm.model_executor.layers.linear import (

46 QKVParallelLinear,

46 - ReplicatedLinear,

47 RowParallelLinear,

48 )

49 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -98,11 +98,9 @@ class Llama4MoE(nn.Module):

98 self.ep_size = self.ep_group.size()

99

100 intermediate_size_moe = config.intermediate_size

101 - self.router = ReplicatedLinear(

101 + self.router = GateLinear(

102 config.hidden_size,

103 config.num_local_experts,

104 - bias=False,

105 - quant_config=None,

104 prefix=f"{prefix}.router",

105 )

106

vllm/model_executor/models/longcat_flash.py

+4 -6

@@ -48,12 +48,12 @@ from vllm.logger import init_logger

48 from vllm.model_executor.layers.activation import SiluAndMul

49 from vllm.model_executor.layers.fused_moe import (

50 FusedMoEFactory,

51 + GateLinear,

52 fused_moe_make_expert_params_mapping,

53 )

54 from vllm.model_executor.layers.layernorm import RMSNorm

55 from vllm.model_executor.layers.linear import (

56 MergedColumnParallelLinear,

56 - ReplicatedLinear,

57 RowParallelLinear,

58 )

59 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -254,12 +254,12 @@ class LongcatRouter(nn.Module):

254 else config.num_experts[0]

255 )

256 self.n_routed_experts = self.n_routed_experts + zero_expert_num

257 - self.classifier = ReplicatedLinear(

257 + self.classifier = GateLinear(

258 config.hidden_size,

259 self.n_routed_experts,

260 bias=config.router_bias,

261 + out_dtype=router_params_dtype,

262 params_dtype=router_params_dtype,

262 - quant_config=None,

263 prefix=f"{prefix}.classifier",

264 )

265 self.e_score_correction_bias = nn.Parameter(

@@ -331,9 +331,7 @@ class LongcatMoe(nn.Module):

331 else:

332 hidden_states_padded = hidden_states

333

334 - router_logits_full = self.router(

335 - hidden_states_padded.to(self.router_params_dtype)

336 - )

334 + router_logits_full = self.router(hidden_states_padded)

335

336 # MoERunner handles routing memoization and zero expert computation

337 # internally. Pass full router_logits (including zero experts) so that

vllm/model_executor/models/minicpm.py

+6 -5

@@ -43,12 +43,15 @@ from vllm.distributed import (

43 )

44 from vllm.model_executor.layers.activation import FatreluAndMul, SiluAndMul

45 from vllm.model_executor.layers.attention import Attention

46 -from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk

46 +from vllm.model_executor.layers.fused_moe import (

47 + GateLinear,

48 + fused_experts,

49 + fused_topk,

50 +)

51 from vllm.model_executor.layers.layernorm import RMSNorm

52 from vllm.model_executor.layers.linear import (

53 MergedColumnParallelLinear,

54 QKVParallelLinear,

51 - ReplicatedLinear,

55 RowParallelLinear,

56 )

57 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -109,12 +112,10 @@ class MiniCPMMoE(nn.Module):

112 params_dtype = torch.get_default_dtype()

113 self.params_dtype = params_dtype

114

112 - self.gate = ReplicatedLinear(

115 + self.gate = GateLinear(

116 self.hidden_size,

117 self.num_total_experts,

115 - bias=False,

118 params_dtype=self.params_dtype,

117 - quant_config=None,

119 prefix=f"{prefix}.gate",

120 )

121

vllm/model_executor/models/mixtral.py

+5 -5

@@ -39,11 +39,13 @@ from vllm.distributed import (

39 get_tensor_model_parallel_world_size,

40 )

41 from vllm.model_executor.layers.attention import Attention

42 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

42 +from vllm.model_executor.layers.fused_moe import (

43 + FusedMoEFactory,

44 + GateLinear,

45 +)

46 from vllm.model_executor.layers.layernorm import RMSNorm

47 from vllm.model_executor.layers.linear import (

48 QKVParallelLinear,

46 - ReplicatedLinear,

49 RowParallelLinear,

50 )

51 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -106,12 +108,10 @@ class MixtralMoE(nn.Module):

108 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

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

110

109 - self.gate = ReplicatedLinear(

111 + self.gate = GateLinear(

112 hidden_size,

113 num_experts,

112 - bias=False,

114 params_dtype=params_dtype,

114 - quant_config=None,

115 prefix=f"{prefix}.gate",

116 )

117

vllm/model_executor/models/moondream3.py

+6 -4

@@ -28,12 +28,15 @@ from vllm.model_executor.layers.attention import Attention

28 from vllm.model_executor.layers.attention.mm_encoder_attention import (

29 MMEncoderAttention,

30 )

31 -from vllm.model_executor.layers.fused_moe import MoEActivation, fused_experts

31 +from vllm.model_executor.layers.fused_moe import (

32 + GateLinear,

33 + MoEActivation,

34 + fused_experts,

35 +)

36 from vllm.model_executor.layers.fused_moe.config import biased_moe_quant_config

37 from vllm.model_executor.layers.linear import (

38 ColumnParallelLinear,

39 QKVParallelLinear,

36 - ReplicatedLinear,

40 RowParallelLinear,

41 )

42 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -453,11 +456,10 @@ class Moondream3TextMoE(nn.Module):

456 self.num_local_experts = self.experts_per_rank

457

458 # Router (gate) - use ReplicatedLinear for compatibility

456 - self.gate = ReplicatedLinear(

459 + self.gate = GateLinear(

460 hidden_size,

461 num_experts,

462 bias=True,

460 - quant_config=None,

463 prefix=f"{prefix}.gate",

464 )

465

vllm/model_executor/models/olmoe.py

+5 -5

@@ -32,11 +32,13 @@ from vllm.distributed import (

32 from vllm.distributed.utils import split_tensor_along_last_dim

33 from vllm.logger import init_logger

34 from vllm.model_executor.layers.attention import Attention

35 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

35 +from vllm.model_executor.layers.fused_moe import (

36 + FusedMoEFactory,

37 + GateLinear,

38 +)

39 from vllm.model_executor.layers.layernorm import RMSNorm

40 from vllm.model_executor.layers.linear import (

41 QKVParallelLinear,

39 - ReplicatedLinear,

42 RowParallelLinear,

43 )

44 from vllm.model_executor.layers.logits_processor import LogitsProcessor

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

86 self.hidden_size = hidden_size

87

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

87 - self.gate = ReplicatedLinear(

89 + self.gate = GateLinear(

90 hidden_size,

91 num_experts,

90 - bias=False,

91 - quant_config=None,

92 prefix=f"{prefix}.gate",

93 )

94

vllm/model_executor/models/openpangu.py

+5 -4

@@ -43,7 +43,10 @@ from vllm.model_executor.layers.attention import (

43 Attention,

44 StaticSinkAttention,

45 )

46 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

46 +from vllm.model_executor.layers.fused_moe import (

47 + FusedMoEFactory,

48 + GateLinear,

49 +)

50 from vllm.model_executor.layers.layernorm import RMSNorm

51 from vllm.model_executor.layers.linear import (

52 ColumnParallelLinear,

@@ -149,11 +152,9 @@ class OpenPanguMoE(nn.Module):

152 self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe

153 check_ffn_act_fn(config.hidden_act)

154

152 - self.gate = ReplicatedLinear(

155 + self.gate = GateLinear(

156 config.hidden_size,

157 config.n_routed_experts,

155 - bias=False,

156 - quant_config=None,

158 prefix=f"{prefix}.gate",

159 )

160 if (

vllm/model_executor/models/param2moe.py

+10 -12

@@ -22,7 +22,6 @@ from collections.abc import Iterable, Iterator

22 from itertools import islice

23

24 import torch

25 -import torch.nn.functional as F

25 from torch import nn

26

27 from vllm.config import CacheConfig, VllmConfig

@@ -34,6 +33,7 @@ from vllm.model_executor.layers.activation import SiluAndMul

33 from vllm.model_executor.layers.attention import Attention

34 from vllm.model_executor.layers.fused_moe import (

35 FusedMoEFactory,

36 + GateLinear,

37 MoERunner,

38 )

39 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -289,10 +289,12 @@ class Param2MoEMoEBlock(nn.Module):

289 self.norm_expert_prob: bool = getattr(config, "norm_topk_prob", True)

290 self.score_function: str = getattr(config, "score_function", "sigmoid")

291

292 - self.gate = nn.Linear(

292 + self.gate = GateLinear(

293 self.hidden_size,

294 self.num_experts,

295 - bias=False,

295 + out_dtype=torch.float32,

296 + params_dtype=torch.float32,

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

298 )

299

300 if getattr(config, "moe_router_enable_expert_bias", True):

@@ -350,15 +352,11 @@ class Param2MoEMoEBlock(nn.Module):

352 num_tokens, hidden_dim = hidden_states.shape

353 hidden_states = hidden_states.view(-1, hidden_dim)

354

353 - # Router: both input and weight must be float32 for numerical

354 - # stability (mirrors the original Param2MoEGate behaviour).

355 - # The gate nn.Linear weight lives in the model dtype (bfloat16),

356 - # so we must cast both explicitly via F.linear instead of calling

357 - # self.gate() which would hit a dtype mismatch.

358 - router_logits = F.linear(

359 - hidden_states.float(),

360 - self.gate.weight.float(),

361 - ).to(hidden_states.dtype)

355 + # Router logits are accumulated in float32 for numerical stability

356 + # (mirrors the original Param2MoEGate behaviour) and handed to the

357 + # experts in the activation dtype.

358 + router_logits, _ = self.gate(hidden_states)

359 + router_logits = router_logits.to(hidden_states.dtype)

360

361 expert_output = self.experts(

362 hidden_states=hidden_states,

vllm/model_executor/models/phimoe.py

+5 -5

@@ -35,10 +35,12 @@ from vllm.compilation.decorators import support_torch_compile

35 from vllm.config import CacheConfig, VllmConfig

36 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

37 from vllm.model_executor.layers.attention import Attention

38 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

38 +from vllm.model_executor.layers.fused_moe import (

39 + FusedMoEFactory,

40 + GateLinear,

41 +)

42 from vllm.model_executor.layers.linear import (

43 QKVParallelLinear,

41 - ReplicatedLinear,

44 RowParallelLinear,

45 )

46 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -262,12 +264,10 @@ class PhiMoE(nn.Module):

264 self.hidden_size = hidden_size

265

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

265 - self.gate = ReplicatedLinear(

267 + self.gate = GateLinear(

268 hidden_size,

269 num_experts,

268 - bias=False,

270 params_dtype=params_dtype,

270 - quant_config=None,

271 prefix=f"{prefix}.gate",

272 )

273

vllm/model_executor/models/qwen2_moe.py

+5 -4

@@ -40,7 +40,10 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

40 from vllm.logger import init_logger

41 from vllm.model_executor.layers.activation import SiluAndMul

42 from vllm.model_executor.layers.attention import Attention

43 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

43 +from vllm.model_executor.layers.fused_moe import (

44 + FusedMoEFactory,

45 + GateLinear,

46 +)

47 from vllm.model_executor.layers.layernorm import RMSNorm

48 from vllm.model_executor.layers.linear import (

49 MergedColumnParallelLinear,

@@ -136,11 +139,9 @@ class Qwen2MoeSparseMoeBlock(nn.Module):

139 f"the number of experts {config.num_experts}."

140 )

141

139 - self.gate = ReplicatedLinear(

142 + self.gate = GateLinear(

143 config.hidden_size,

144 config.num_experts,

142 - bias=False,

143 - quant_config=None,

145 prefix=f"{prefix}.gate",

146 )

147

vllm/model_executor/models/qwen3_moe.py

+5 -3

@@ -42,7 +42,10 @@ from vllm.distributed import (

42 from vllm.logger import init_logger

43 from vllm.model_executor.layers.activation import SiluAndMul

44 from vllm.model_executor.layers.attention import Attention

45 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

45 +from vllm.model_executor.layers.fused_moe import (

46 + FusedMoEFactory,

47 + GateLinear,

48 +)

49 from vllm.model_executor.layers.layernorm import RMSNorm

50 from vllm.model_executor.layers.linear import (

51 MergedColumnParallelLinear,

@@ -164,10 +167,9 @@ class Qwen3MoeSparseMoeBlock(nn.Module):

167 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

168 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

169

167 - self.gate = ReplicatedLinear(

170 + self.gate = GateLinear(

171 config.hidden_size,

172 config.num_experts,

170 - bias=False,

173 quant_config=quant_config,

174 prefix=f"{prefix}.gate",

175 )

vllm/model_executor/models/qwen3_next.py

+5 -4

@@ -18,7 +18,10 @@ from vllm.distributed import (

18 tensor_model_parallel_reduce_scatter,

19 )

20 from vllm.model_executor.layers.attention import Attention

21 -from vllm.model_executor.layers.fused_moe import FusedMoEFactory

21 +from vllm.model_executor.layers.fused_moe import (

22 + FusedMoEFactory,

23 + GateLinear,

24 +)

25 from vllm.model_executor.layers.fused_moe.utils import (

26 is_model_fused_shared_expert_compatible,

27 resolve_layer_fused_shared_expert,

@@ -180,11 +183,9 @@ class Qwen3NextSparseMoeBlock(nn.Module):

183 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

184 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

185

183 - self.gate = ReplicatedLinear(

186 + self.gate = GateLinear(

187 config.hidden_size,

188 config.num_experts,

186 - bias=False,

187 - quant_config=None,

189 prefix=f"{prefix}.gate",

190 )

191

vllm/model_executor/models/step3_text.py

+2 -3

@@ -20,6 +20,7 @@ from vllm.model_executor.layers.activation import SiluAndMul

20 from vllm.model_executor.layers.attention import Attention

21 from vllm.model_executor.layers.fused_moe import (

22 FusedMoEFactory,

23 + GateLinear,

24 )

25 from vllm.model_executor.layers.layernorm import RMSNorm

26 from vllm.model_executor.layers.linear import (

@@ -77,11 +78,9 @@ class FusedMoEBlock(nn.Module):

78 quant_config=quant_config,

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

80 )

80 - self.gate = ReplicatedLinear(

81 + self.gate = GateLinear(

82 config.hidden_size,

83 config.moe_num_experts,

83 - bias=False,

84 - quant_config=None,

84 prefix=f"{prefix}.gate",

85 )

86

vllm/model_executor/models/step3p5.py

+5 -17

@@ -8,7 +8,6 @@ from typing import Any

8

9 import torch

10 from torch import nn

11 -from torch.nn.parameter import Parameter

11

12 from vllm.compilation.decorators import support_torch_compile

13 from vllm.config import CacheConfig, VllmConfig

@@ -25,6 +24,7 @@ from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul

24 from vllm.model_executor.layers.attention import Attention

25 from vllm.model_executor.layers.fused_moe import (

26 FusedMoEFactory,

27 + GateLinear,

28 MoERunner,

29 fused_moe_make_expert_params_mapping,

30 )

@@ -33,7 +33,6 @@ from vllm.model_executor.layers.linear import (

33 ColumnParallelLinear,

34 MergedColumnParallelLinear,

35 QKVParallelLinear,

36 - ReplicatedLinear,

36 RowParallelLinear,

37 )

38 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -68,17 +67,6 @@ from .utils import (

67 logger = init_logger(__name__)

68

69

71 -class FP32ReplicatedLinear(ReplicatedLinear):

72 - """Use FP32 for higher precision."""

73 -

74 - def forward(

75 - self,

76 - x: torch.Tensor,

77 - ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:

78 - assert self.params_dtype == torch.float32

79 - return super().forward(x.to(torch.float32))

80 -

81 -

70 class Step3p5MLP(nn.Module):

71 def __init__(

72 self,

@@ -326,12 +314,12 @@ class FusedMoEBlock(nn.Module):

314 f"the number of experts {config.moe_num_experts}."

315 )

316

329 - self.gate = FP32ReplicatedLinear(

317 + # Router logits are accumulated in FP32 for higher precision.

318 + self.gate = GateLinear(

319 config.hidden_size,

320 config.moe_num_experts,

332 - bias=False,

333 - quant_config=None,

334 - params_dtype=torch.float32, # Use FP32 for higher precision.

321 + out_dtype=torch.float32,

322 + params_dtype=torch.float32,

323 prefix=f"{prefix}.gate",

324 )

325 self.use_moe_router_bias = config.use_moe_router_bias

[MoE] Use GateLinear for all MoE models (#58234)