@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 99de48e98f

Fix MLA padding and grouped topk routing in the Transformers modelling backend (#49982)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

Harry Mellor committed Jul 27, 2026 at 19:32 UTC 99de48e98fe9570d52b733a545d27661929c170c

5 files changed +227 -39

tests/models/transformers/fusers/test_moe.py

+80 -3

@@ -29,8 +29,48 @@ class TopKRouter(nn.Module):

29 return logits, value, index

32 +class ScaledRouter(TopKRouter):

33 + """Greedy router scaling its top-k weights (DeepSeek `routed_scaling_factor`)."""

34 +

35 + def forward(self, hidden_states):

36 + logits = F.linear(hidden_states, self.weight)

37 + scores = F.softmax(logits, dim=-1)

38 + value, index = torch.topk(scores, self.top_k, dim=-1)

39 + value = value * 16.0

40 + return logits, value, index

41 +

42 +

43 +class Fp32Router(TopKRouter):

44 + """Router that computes its logits in fp32 (DeepSeek/GLM style)."""

45 +

46 + def forward(self, hidden_states):

47 + logits = F.linear(

48 + hidden_states.type(torch.float32), self.weight.type(torch.float32)

49 + )

50 + scores = F.softmax(logits, dim=-1)

51 + value, index = torch.topk(scores, self.top_k, dim=-1)

52 + return logits, value.to(hidden_states.dtype), index

53 +

54 +

55 +class GroupedRouter(TopKRouter):

56 + """Group-limited router (DeepSeek `group_limited_greedy`), scaled weights."""

57 +

58 + def forward(self, hidden_states):

59 + logits = F.linear(hidden_states, self.weight)

60 + scores = F.softmax(logits, dim=-1)

61 + group_scores = scores.view(-1, 4, 2).max(dim=-1).values

62 + group_idx = torch.topk(group_scores, k=2, dim=-1)[1]

63 + group_mask = torch.zeros_like(group_scores)

64 + group_mask.scatter_(1, group_idx, 1)

65 + score_mask = group_mask.unsqueeze(-1).expand(-1, 4, 2).reshape(-1, 8)

66 + scores = scores.masked_fill(~score_mask.bool(), 0.0)

67 + value, index = torch.topk(scores, self.top_k, dim=-1)

68 + value = value * 16.0

69 + return logits, value, index

70 +

71 +

72 class CorrectionRouter(nn.Module):

33 - """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined."""

73 + """Router with a score-correction bias buffer (DeepSeek-V3 noaux) -> matched."""

75 def __init__(self, num_experts=8, hidden=16):

76 super().__init__()

@@ -228,6 +268,44 @@ def test_moe_fuser_detects_router(sigmoid):

268 assert fuser.shared_name is None and fuser.shared_gate_name is None

269

270

271 +def test_moe_fuser_matches_scaled_router():

272 + """Weight scaling after the top-k (DeepSeek style) does not break matching."""

273 + with torch.device("meta"):

274 + block = MoEBlock(ScaledRouter)

275 + assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser)

276 +

277 +

278 +def test_moe_fuser_matches_grouped_router():

279 + """Group masking between the score and the routing top-k still matches: the

280 + scorer is anchored on the routing (last) top-k, not the group one."""

281 + with torch.device("meta"):

282 + block = MoEBlock(GroupedRouter)

283 + fuser = MoEBlockFuser.match(block, "experts")

284 + assert isinstance(fuser, MoEBlockFuser)

285 + assert fuser.scoring_func == "softmax"

286 +

287 +

288 +def test_moe_fuser_matches_correction_router():

289 + """A score-correction bias buffer (DeepSeek-V3 noaux) is allowed; the

290 + rebuilt gate carries it in fp32 for FusedMoE's biased routers."""

291 + with torch.device("meta"):

292 + block = MoEBlock(CorrectionRouter)

293 + fuser = MoEBlockFuser.match(block, "experts")

294 + assert isinstance(fuser, MoEBlockFuser)

295 + assert fuser.scoring_func == "sigmoid"

296 +

297 +

298 +def test_moe_fuser_reads_router_dtype_from_the_gate():

299 + """A router that computes its logits in fp32 must keep routing in fp32 when

300 + rebuilt, even though no config field names the dtype. The cast back to the

301 + activation dtype after the top-k must not be mistaken for the routing dtype."""

302 + with torch.device("meta"):

303 + fp32 = MoEBlockFuser.match(MoEBlock(Fp32Router), "experts")

304 + default = MoEBlockFuser.match(MoEBlock(TopKRouter), "experts")

305 + assert fp32.router_dtype == torch.float32

306 + assert default.router_dtype is None

307 +

308 +

309 def test_moe_fuser_detects_shared_experts():

310 with torch.device("meta"):

311 block = MoEBlockShared()

@@ -269,8 +347,7 @@ def test_moe_fuser_detects_non_glu_shared_expert():

347 @pytest.mark.parametrize(

348 "block_cls",

349 [

272 - lambda: MoEBlock(CorrectionRouter), # score-correction buffer (grouped)

273 - lambda: MoEBlock(BiasedRouter), # router not weight-only (extra param)

350 + lambda: MoEBlock(BiasedRouter), # router with an unrecognized extra param

351 MoEBlockTuple, # tuple-returning block (e.g. gpt-oss)

352 MoEBlockTupleVar, # tuple returned via a name binding, not a literal

353 MoEBlockUnaccounted, # weight-bearing child outside the fused dataflow

vllm/model_executor/models/transformers/__init__.py

+12 -1

@@ -18,6 +18,7 @@

19 from typing import TYPE_CHECKING

21 +import torch.nn.functional as F

22 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

24 from vllm.model_executor.models.transformers.base import Base

@@ -59,9 +60,19 @@ def vllm_attention_forward(

60 if scaling is not None:

61 self_attn.impl.scale = float(scaling)

62 hidden = query.shape[-2]

63 + head_dim_qk = query.shape[-1]

64 + head_dim_v = value.shape[-1]

65 query, key, value = (x.transpose(1, 2) for x in (query, key, value))

66 query, key, value = (x.reshape(hidden, -1) for x in (query, key, value))

64 - return self_attn.forward(query, key, value), None

67 + # Pad `value` up to the query/key head size when they differ (expanded MLA).

68 + if head_dim_v != head_dim_qk:

69 + value = F.pad(value.view(-1, head_dim_v), (0, head_dim_qk - head_dim_v))

70 + value = value.reshape(hidden, -1)

71 + attn_output = self_attn.forward(query, key, value)

72 + if head_dim_v != head_dim_qk:

73 + attn_output = attn_output.view(-1, head_dim_qk)[..., :head_dim_v]

74 + attn_output = attn_output.reshape(hidden, -1)

75 + return attn_output, None

78 ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward

vllm/model_executor/models/transformers/fusers/moe.py

+55 -22

@@ -6,14 +6,14 @@ import ast

6 import inspect

7 import textwrap

8 import types

9 -from collections.abc import Iterator

9 +from collections.abc import Iterable, Iterator

10 from dataclasses import dataclass

11 -from itertools import chain

12 import torch

13 from torch import fx, nn

15 from vllm.distributed import tensor_model_parallel_all_gather

16 +from vllm.model_executor.layers.fused_moe import GateLinear

17 from vllm.model_executor.layers.linear import ReplicatedLinear

18 from vllm.model_executor.models.transformers.fx_utils import (

19 find_node,

@@ -21,14 +21,10 @@ from vllm.model_executor.models.transformers.fx_utils import (

21 peel,

22 trace,

23 )

24 +from vllm.model_executor.models.transformers.utils import named_state

25 from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk

27 -def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:

28 - """`module`'s own state (i.e. named parameters and buffers)."""

29 - return chain(module.named_parameters(), module.named_buffers())

30 -

31 -

28 def _own_returns(node: ast.AST) -> Iterator[ast.Return]:

29 """`return` statements in `node`'s own scope, not in nested functions."""

30 stack = list(ast.iter_child_nodes(node))

@@ -79,6 +75,25 @@ def _is_scalar_gate(module: nn.Module) -> bool:

75 )

78 +def _forced_dtype(nodes: Iterable[fx.Node]) -> torch.dtype | None:

79 + """The floating dtype `nodes` cast to, if any.

80 +

81 + Computations that must run in higher precision say so in their forward (e.g.

82 + `hidden_states.type(torch.float32)`), so the dtype is readable from the graph

83 + even when the config does not name it."""

84 + for node in nodes:

85 + if node.op not in ("call_method", "call_function"):

86 + continue

87 + name = str(node.target).rsplit(".", 1)[-1]

88 + if name == "float":

89 + return torch.float32

90 + if name in ("to", "type"):

91 + for arg in (*node.args[1:], *node.kwargs.values()):

92 + if isinstance(arg, torch.dtype) and arg.is_floating_point:

93 + return arg

94 + return None

95 +

96 +

97 def _reaches(node: fx.Node, key: str) -> set[fx.Node]:

98 """Returns the set of nodes reachable from `node` by following `key` edges."""

99 seen: set[fx.Node] = set()

@@ -131,18 +146,24 @@ class MoEBlockFuser:

146 scoring_func: str

147 shared_name: str | None

148 shared_gate_name: str | None

149 + router_dtype: torch.dtype | None = None

150

151 @staticmethod

136 - def _match_router(gate: nn.Module) -> str | None:

137 - """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`."""

138 - if [name for name, _ in named_state(gate)] != ["weight"]:

152 + def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None:

153 + """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`.

154 +

155 + Returns the scoring function and the dtype the router computes in."""

156 + state = {name for name, _ in named_state(gate)}

157 + if "weight" not in state or state - {"weight", "e_score_correction_bias"}:

158 return None

159 graph = trace(gate)

160 if graph is None:

161 return None

143 - topk = find_node(graph, lambda n: is_op(n, "topk"))

144 - if topk is None:

162 + # The routing top-k is the last one; any earlier one scores expert groups.

163 + topks = [node for node in graph.nodes if is_op(node, "topk")]

164 + if not topks:

165 return None

166 + topk = topks[-1]

167 # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear.

168 scorers = [

169 n

@@ -152,9 +173,11 @@ class MoEBlockFuser:

173 if len(scorers) != 1:

174 return None

175 scorer = scorers[0]

155 - if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")):

176 + logits_cone = _reaches(scorer, "all_input_nodes")

177 + if not any(is_op(n, "linear") for n in logits_cone):

178 return None

157 - return "softmax" if is_op(scorer, "softmax") else "sigmoid"

179 + scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid"

180 + return scoring_func, _forced_dtype(logits_cone)

181

182 @staticmethod

183 def _match_shared_experts(

@@ -198,10 +221,14 @@ class MoEBlockFuser:

221 if _returns_tuple(type(moe_block)):

222 return None

223 # Router: the child that scores + top-k selects.

201 - gate_name = scoring_func = None

224 + gate_name = scoring_func = router_dtype = None

225 for name, child in moe_block.named_children():

203 - if name != experts_name and (func := cls._match_router(child)) is not None:

204 - gate_name, scoring_func = name, func

226 + if (

227 + name != experts_name

228 + and (router := cls._match_router(child)) is not None

229 + ):

230 + gate_name = name

231 + scoring_func, router_dtype = router

232 break

233 if gate_name is None or scoring_func is None:

234 return None

@@ -229,17 +256,23 @@ class MoEBlockFuser:

256 for name, child in moe_block.named_children():

257 if name not in accounted and next(named_state(child), None) is not None:

258 return None

232 - return cls(gate_name, scoring_func, shared_name, shared_gate_name)

259 + return cls(gate_name, scoring_func, shared_name, shared_gate_name, router_dtype)

260

234 - def gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear:

235 - """Rebuild the HF gate as a `ReplicatedLinear` for vLLM's fused MoE."""

236 - num_experts, hidden_size = getattr(moe_block, self.gate_name).weight.shape

237 - gate = ReplicatedLinear(

261 + def gate(

262 + self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None

263 + ) -> GateLinear:

264 + """Rebuild the HF gate as a `GateLinear` for vLLM's fused MoE."""

265 + hf_gate = getattr(moe_block, self.gate_name)

266 + num_experts, hidden_size = hf_gate.weight.shape

267 + gate = GateLinear(

268 hidden_size,

269 num_experts,

270 bias=False,

271 + out_dtype=out_dtype or self.router_dtype,

272 prefix=maybe_prefix(prefix, self.gate_name),

273 )

274 + if (bias := getattr(hf_gate, "e_score_correction_bias", None)) is not None:

275 + gate.register_buffer("e_score_correction_bias", bias)

276 setattr(moe_block, self.gate_name, gate)

277 return gate

278

vllm/model_executor/models/transformers/moe.py

+73 -13

@@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any

23 import torch

24 import torch.nn as nn

26 +from vllm._aiter_ops import rocm_aiter_ops

27 from vllm.config.utils import getattr_iter

28 from vllm.distributed import get_dp_group, get_ep_group

29 from vllm.forward_context import ForwardContext, get_forward_context

@@ -30,9 +31,10 @@ from vllm.logger import init_logger

31 from vllm.model_executor.custom_op import PluggableLayer

32 from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts

33 from vllm.model_executor.models.interfaces import MixtureOfExperts

34 +from vllm.model_executor.models.transformers.fuser import get_fuser

35 from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser

36 from vllm.model_executor.models.utils import maybe_prefix

35 -from vllm.utils.torch_utils import direct_register_custom_op

37 +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op

39 from .utils import log_replacement

@@ -176,17 +178,31 @@ class MoEMixin(MixtureOfExperts):

178 0,

179 )

180

179 - # Unused kwargs since we use custom_routing_function:

180 - # - `scoring_func` and `e_score_correction_bias` only used for grouped

181 - # topk routing inside vLLM and are non-trivial to infer

182 - # and hard code `use_grouped_topk=False`

183 - # - `renormalize` passed anyway because it's easy to infer

184 - # - `num_expert_group` and `topk_group` used for inferring expert

185 - # placement strategy in FusedMoE

186 - # - `apply_router_weight_on_input` is already applied in Transformers

181 + # Common kwargs

182 renormalize = getattr(text_config, "norm_topk_prob", top_k > 1)

183 +

184 + # Routed scaling factor kwargs

185 + routed_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0)

186 + # aiter applies routed_scaling_factor internally

187 + apply_routed_scale_to_output = not rocm_aiter_ops.is_fused_moe_enabled()

188 + routed_scaling_factor_kwargs = dict(

189 + routed_scaling_factor=routed_scaling_factor,

190 + apply_routed_scale_to_output=apply_routed_scale_to_output,

191 + )

192 +

193 + # Dtype the router computes in, if it is not the activation dtype.

194 + config_router_dtype = getattr(text_config, "moe_router_dtype", None)

195 + config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype)

196 +

197 + # Grouped topk routing kwargs

198 num_expert_group = getattr(text_config, "n_group", None)

199 topk_group = getattr(text_config, "topk_group", None)

200 + use_grouped_topk = num_expert_group is not None and topk_group is not None

201 + grouped_topk_routing_kwargs = dict(

202 + use_grouped_topk=use_grouped_topk,

203 + num_expert_group=num_expert_group,

204 + topk_group=topk_group,

205 + )

206

207 # MoE activation function

208 activation = "silu"

@@ -211,6 +227,9 @@ class MoEMixin(MixtureOfExperts):

227 self.num_shared_experts = num_shared_experts

228 self.num_redundant_experts = num_redundant_experts

229

230 + # Down projections of shared experts consumed by FusedMoE

231 + shared_down_projs: list[tuple[nn.Module, str]] = []

232 +

233 # Recursively fuse MoE layers

234 def _recursive_replace(module: nn.Module, prefix: str):

235 for child_name, child_module in module.named_children():

@@ -249,7 +268,6 @@ class MoEMixin(MixtureOfExperts):

268 hidden_size=hidden_size,

269 intermediate_size=intermediate_size,

270 renormalize=renormalize,

252 - use_grouped_topk=False,

271 quant_config=self.quant_config,

272 prefix=qual_name,

273 activation=activation,

@@ -259,17 +277,55 @@ class MoEMixin(MixtureOfExperts):

277 routed_experts_cls=TransformersRoutedExperts,

278 )

279 fuser = MoEBlockFuser.match(moe_block, experts_name)

262 - if self.num_expert_groups <= 1 and fuser is not None:

280 + # _maybe_apply_routed_scale_to_output edge case. Transformers

281 + # decoder layers do not compensate for dividing by scaling factor.

282 + reaches_fp16_trick = (

283 + routed_scaling_factor != 1.0

284 + and apply_routed_scale_to_output

285 + and self.model_config.dtype == torch.float16

286 + and fuser is not None

287 + and fuser.shared_name is not None

288 + )

289 + if reaches_fp16_trick:

290 + logger.warning_once(

291 + "%s could be fused but routing it in vLLM would apply "

292 + "`routed_scaling_factor` by dividing the shared expert "

293 + "output, which only fp16 overflow protection expects the "

294 + "decoder layer to compensate for. Falling back to routing "

295 + "in Transformers; run in bfloat16 to fuse it.",

296 + moe_block_cls,

297 + )

298 + if fuser is not None and not reaches_fp16_trick:

299 # MoE block forward is fully replaced.

300 # gate/router and shared expert (if any) runs in FusedMoE.

301 + shared_experts = fuser.shared_experts(moe_block, prefix)

302 + # Store shared experts for later down projection adjustment

303 + if shared_experts is not None:

304 + hf_shared = shared_experts.shared_experts

305 + glu_fuser = get_fuser(hf_shared)

306 + down_name = getattr(glu_fuser, "down_name", None)

307 + if down_name is not None:

308 + shared_down_projs.append((hf_shared, down_name))

309 + # Prefer config, otherwise read it from fuser.

310 + router_dtype = config_router_dtype or fuser.router_dtype

311 + gate = fuser.gate(moe_block, prefix, router_dtype)

312 kwargs |= dict(

313 scoring_func=fuser.scoring_func,

314 is_sequence_parallel=(

315 self.parallel_config.use_sequence_parallel_moe

316 ),

270 - gate=fuser.gate(moe_block, prefix),

271 - shared_experts=fuser.shared_experts(moe_block, prefix),

317 + gate=gate,

318 + shared_experts=shared_experts,

319 )

320 + if router_dtype is not None:

321 + kwargs["router_logits_dtype"] = router_dtype

322 + if use_grouped_topk:

323 + kwargs |= grouped_topk_routing_kwargs

324 + if routed_scaling_factor != 1.0:

325 + kwargs |= routed_scaling_factor_kwargs

326 + bias = getattr(gate, "e_score_correction_bias", None)

327 + if bias is not None:

328 + kwargs["e_score_correction_bias"] = bias

329 fuser.rewrite_forward(moe_block)

330 routed = "gate + experts"

331 if fuser.shared_name:

@@ -333,3 +389,7 @@ class MoEMixin(MixtureOfExperts):

389 self.num_moe_layers = len(self.moe_layers)

390 # Continue with the replacement of layers in Base

391 super().recursive_replace()

392 + # GLUFuser likely fused shared_experts. The down projection after the GLU

393 + # normally immediately reduces but we want FusedMoE to handle the reduction.

394 + for hf_shared, down_name in shared_down_projs:

395 + hf_shared.get_submodule(down_name).reduce_results = False

vllm/model_executor/models/transformers/utils.py

@@ -16,7 +16,9 @@

16 # limitations under the License.

17 """Transformers modeling backend utilities."""

19 +from collections.abc import Iterator

20 from contextlib import contextmanager

21 +from itertools import chain

22 from pathlib import Path

23 from typing import TYPE_CHECKING, Literal

@@ -209,6 +211,11 @@ def recursive_replace_linear(

211 _recursive_replace(model, prefix=prefix)

212

213

214 +def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:

215 + """`module`'s own state (i.e. named parameters and buffers)."""

216 + return chain(module.named_parameters(), module.named_buffers())

217 +

218 +

219 def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module):

220 logger.debug("%s: %s -> %s", name, old_module, new_module)

221