chore: import upstream snapshot with attribution

tests/models/transformers/fusers/test_linear.py

@@ -0,0 +1,482 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Transformers modeling backend's linear fusers."""
+
+import inspect
+from types import MethodType, SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm.model_executor.models.transformers.fuser import get_fuser
+from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser
+
+
+class SiluAndMulStub(nn.Module):
+    """Stand-in for vLLM's `SiluAndMul` (no vLLM config required)."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        return F.silu(x[..., :d]) * x[..., d:]
+
+
+class NoDownGLU(nn.Module):
+    """`act(gate(x)) * up(x)` with no output projection -> `down_name` is None."""
+
+    def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden, inter, bias=bias)
+        self.up_proj = nn.Linear(hidden, inter, bias=bias)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        return self.act_fn(self.gate_proj(x)) * self.up_proj(x)
+
+
+class GLUMLP(NoDownGLU):
+    """`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP."""
+
+    def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
+        super().__init__(hidden, inter, bias)
+        self.down_proj = nn.Linear(inter, hidden, bias=bias)
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class ReversedGLUMLP(GLUMLP):
+    """`up(x) * act(gate(x))` — operands swapped (multiply is commutative)."""
+
+    def forward(self, x):
+        return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x)))
+
+
+class NotAnMLP(nn.Module):
+    """Two linears but no activation*linear multiply -> must not match."""
+
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(8, 8)
+        self.fc2 = nn.Linear(8, 8)
+
+    def forward(self, x):
+        return self.fc2(self.fc1(x))
+
+
+class NotAnActGLUMLP(GLUMLP):
+    """GLU-shaped, but the "activation" is not a known activation module."""
+
+    def __init__(self):
+        super().__init__()
+        self.act_fn = nn.Dropout()
+
+
+class UntraceableMLP(GLUMLP):
+    """Data-dependent control flow *before* the GLU -> no match."""
+
+    def forward(self, x):
+        if x.sum() > 0:  # noqa: SIM108 - intentionally untraceable
+            return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return x
+
+
+class UntraceableTailGLUMLP(GLUMLP):
+    """Data-dependent control flow *after* the GLU -> still fusable."""
+
+    def forward(self, x):
+        y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        if y.sum() > torch.inf:  # intentionally untraceable
+            y = y * 0
+        return y
+
+
+class FakeAttention(nn.Module):
+    """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface."""
+
+    is_causal = True
+
+    def __init__(
+        self,
+        hidden: int = 32,
+        head_dim: int = 8,
+        heads: int = 4,
+        kv_heads: int = 4,
+        bias: bool = False,
+        layer_idx: int = 0,
+    ):
+        super().__init__()
+        self.config = SimpleNamespace(_attn_implementation="vllm")
+        self.layer_idx = layer_idx
+        self.head_dim = head_dim
+        self.scaling = head_dim**-0.5
+        self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias)
+        self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
+        self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
+        self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias)
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if past_key_values is not None:
+            k, v = past_key_values.update(k, v, self.layer_idx)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, attn_weights = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        return self.o_proj(attn_output), attn_weights
+
+
+class ReversedFakeAttention(FakeAttention):
+    """Projections computed in (v, k, q) order — q must still be identified."""
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1)), None
+
+
+class ExtraProjAttention(FakeAttention):
+    """A second non-qkv linear of a different width -> `o_proj` still found."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+
+
+class QKNormAttention(FakeAttention):
+    """OLMoE-style: a full-dim norm applied to the whole q/k projection output."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.q_norm = nn.RMSNorm(self.q_proj.out_features)
+        self.k_norm = nn.RMSNorm(self.k_proj.out_features)
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        q = self.q_norm(self.q_proj(hidden_states))
+        k = self.k_norm(self.k_proj(hidden_states))
+        v = self.v_proj(hidden_states)
+        return self.o_proj(q + k + v), None
+
+
+class PerHeadQKNormAttention(FakeAttention):
+    """Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.q_norm = nn.RMSNorm(self.head_dim)
+        self.k_norm = nn.RMSNorm(self.head_dim)
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        shape = (*hidden_states.shape[:-1], -1, self.head_dim)
+        q = self.q_norm(self.q_proj(hidden_states).view(shape))
+        k = self.k_norm(self.k_proj(hidden_states).view(shape))
+        v = self.v_proj(hidden_states).view(shape)
+        return self.o_proj((q + k + v).flatten(-2)), None
+
+
+class FakeSelfAttn(nn.Module):
+    """Stand-in for the vLLM `Attention` looked up in `attention_instances`."""
+
+    def __init__(self):
+        super().__init__()
+        self.impl = SimpleNamespace(scale=None)
+
+    def forward(self, q, k, v):
+        # MHA-shaped stub: any deterministic combination of q/k/v will do
+        return q + 2 * k + 3 * v
+
+
+@pytest.fixture(autouse=True)
+def _clear_fuser_cache():
+    get_fuser.cache_clear()
+    yield
+    get_fuser.cache_clear()
+
+
+def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser):
+    """Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul)."""
+    gate = module.get_submodule(fuser.gate_name)
+    up = module.get_submodule(fuser.up_name)
+    merged = nn.Linear(
+        gate.in_features,
+        gate.out_features + up.out_features,
+        bias=gate.bias is not None,
+    )
+    with torch.no_grad():
+        merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0))
+        if gate.bias is not None:
+            merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0))
+    setattr(module, fuser.merged_name, merged)
+    setattr(module, fuser.act_name, SiluAndMulStub())
+    delattr(module, fuser.gate_name)
+    delattr(module, fuser.up_name)
+    module.forward = MethodType(fuser.fused_forward, module)
+    return module
+
+
+def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
+    """Apply a fuser using a plain merged `nn.Linear` (no TP sharding)."""
+    q, k, v = (
+        module.get_submodule(name)
+        for name in (fuser.q_name, fuser.k_name, fuser.v_name)
+    )
+    merged = nn.Linear(
+        q.in_features,
+        q.out_features + k.out_features + v.out_features,
+        bias=q.bias is not None,
+    )
+    with torch.no_grad():
+        merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
+        if q.bias is not None:
+            merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0))
+    merged.output_sizes = [q.out_features, k.out_features, v.out_features]
+    merged.tp_size = 1
+    setattr(module, fuser.merged_name, merged)
+    for name in (fuser.q_name, fuser.k_name, fuser.v_name):
+        delattr(module, name)
+    module.forward = MethodType(fuser.fused_forward, module)
+    return module
+
+
+@pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP])
+@pytest.mark.parametrize("bias", [False, True])
+def test_detects_and_rewrites_glu(mlp_cls, bias):
+    with torch.device("meta"):
+        meta = mlp_cls(bias=bias)
+    fuser = get_fuser(meta)
+    assert isinstance(fuser, GLUFuser)
+    assert (
+        fuser.gate_name,
+        fuser.up_name,
+        fuser.act_name,
+        fuser.down_name,
+    ) == ("gate_proj", "up_proj", "act_fn", "down_proj")
+
+    # The rewritten forward references the merged projection instead of the
+    # sources; the rest of the forward is untouched.
+    names = fuser.fused_forward.__code__.co_names
+    assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names
+    assert not {"gate_proj", "up_proj"} & set(names)
+
+    # Numerics: the fused forward must match the original on a real instance.
+    real = mlp_cls(bias=bias)
+    for p in real.parameters():
+        nn.init.normal_(p, std=0.05)
+    x = torch.randn(4, 16)
+    expected = real(x)
+    fused = _apply_glu_fuser_with_stubs(real, fuser)
+
+    # Fusion is in place: the module keeps its class and other attributes
+    assert fused is real and type(fused) is mlp_cls
+    torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
+
+
+def test_glu_identifies_down_projection():
+    """The row projection consuming `act(gate(x)) * up(x)` is identified.
+
+    It is forced to `RowParallelLinear` in `update_attrs` so its sharded input
+    matches the column-parallel merged gate/up; `None` when there is no such
+    projection to force (fusion of gate/up still applies)."""
+    with torch.device("meta"):
+        assert get_fuser(GLUMLP()).down_name == "down_proj"
+        assert get_fuser(ReversedGLUMLP()).down_name == "down_proj"
+        assert get_fuser(NoDownGLU()).down_name is None
+
+
+@pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention])
+@pytest.mark.parametrize("kv_heads", [4, 2])
+def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
+    if attn_cls is ReversedFakeAttention and kv_heads == 4:
+        pytest.skip("MHA q/k/v assignment is order-based by design")
+    with torch.device("meta"):
+        meta = attn_cls(kv_heads=kv_heads)
+    fuser = get_fuser(meta)
+    assert isinstance(fuser, QKVFuser)
+    # q (sharded differently under TP) must be identified exactly; k/v may be
+    # swapped for non-canonical compute order, which is numerically consistent
+    # because the weight mapping and the split indices follow the same
+    # assignment.
+    assert fuser.q_name == "q_proj"
+    assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"}
+    assert fuser.o_name == "o_proj"
+
+    # The projections are merged; everything else stays live Python with its
+    # original semantics (branches, kwargs, attribute reads)
+    code = fuser.fused_forward.__code__
+    names = code.co_names
+    assert "qkv_proj" in names and "output_sizes" in names and "o_proj" in names
+    assert "tp_size" in names
+    assert not {"q_proj", "k_proj", "v_proj"} & set(names)
+    if attn_cls is FakeAttention:
+        assert "update" in names  # the cache branch survives
+    assert code.co_flags & inspect.CO_VARKEYWORDS  # **kwargs survives
+
+    # Numerics: the fused forward must match the original on a real instance,
+    # with a different layer_idx than the traced instance (kv_heads == heads so
+    # the q/k/v stub combination is shape-compatible).
+    real = attn_cls(kv_heads=4, layer_idx=3)
+    for p in real.parameters():
+        nn.init.normal_(p, std=0.05)
+    x = torch.randn(1, 5, 32)
+    attention_instances = {3: FakeSelfAttn()}
+    expected, _ = real(x, attention_instances=attention_instances)
+    fused = _apply_qkv_fuser_with_stubs(real, fuser)
+
+    # Fusion is in place: the module keeps its class and other attributes
+    assert fused is real and type(fused) is attn_cls
+    assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None
+    out, _ = fused(x, attention_instances=attention_instances)
+    torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
+
+
+def test_qkv_identifies_output_projection():
+    with torch.device("meta"):
+        assert get_fuser(FakeAttention()).o_name == "o_proj"
+        assert get_fuser(ReversedFakeAttention()).o_name == "o_proj"
+        assert get_fuser(ExtraProjAttention()).o_name == "o_proj"
+        # Norm children (q_norm/k_norm) must not disturb o_proj identification.
+        assert get_fuser(QKNormAttention()).o_name == "o_proj"
+        assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"
+
+
+def test_fuser_is_cached_per_class():
+    with torch.device("meta"):
+        fuser_a = get_fuser(GLUMLP())
+        fuser_b = get_fuser(GLUMLP())
+    assert fuser_a is fuser_b
+    assert GLUMLP in get_fuser.cache
+
+
+@pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])
+def test_non_matching_modules_return_none(cls):
+    with torch.device("meta"):
+        module = cls()
+    assert get_fuser(module) is None
+
+
+def test_untraceable_tail_still_fuses():
+    with torch.device("meta"):
+        meta = UntraceableTailGLUMLP()
+    fuser = get_fuser(meta)
+    assert isinstance(fuser, GLUFuser)
+
+    # Numerics: the live tail must survive the rewrite
+    real = UntraceableTailGLUMLP()
+    for p in real.parameters():
+        nn.init.normal_(p, std=0.05)
+    x = torch.randn(4, 16)
+    expected = real(x)
+    fused = _apply_glu_fuser_with_stubs(real, fuser)
+    torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
+
+
+def test_weight_mappings_are_scoped_to_fused_prefixes():
+    from vllm.model_executor.models.utils import WeightsMapper
+
+    with torch.device("meta"):
+        glu_fuser = get_fuser(GLUMLP())
+        qkv_fuser = get_fuser(FakeAttention())
+
+    mapper = WeightsMapper()
+    for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):
+        mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix))
+    mapper.orig_to_new_stacked.update(
+        qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn")
+    )
+
+    names = [
+        "model.layers.0.mlp.gate_proj.weight",
+        "model.layers.0.mlp.up_proj.weight",
+        "model.layers.1.mlp.gate_proj.weight",
+        "model.layers.0.self_attn.q_proj.weight",
+        "model.layers.0.self_attn.k_proj.weight",
+        "model.layers.0.self_attn.v_proj.weight",
+        # Unfused modules at other prefixes must be left untouched.
+        "model.layers.2.mlp.experts.0.gate_proj.weight",
+        "model.layers.1.self_attn.q_proj.weight",
+    ]
+    # `apply` rewrites the name and stamps the shard id onto each tensor.
+    weights = [(name, torch.empty(0)) for name in names]
+    mapped = list(mapper.apply(weights))
+    mapped_names = [name for name, _ in mapped]
+    shard_ids = [getattr(data, "shard_id", None) for _, data in mapped]
+
+    assert mapped_names == [
+        "model.layers.0.mlp.gate_up_proj.weight",
+        "model.layers.0.mlp.gate_up_proj.weight",
+        "model.layers.1.mlp.gate_up_proj.weight",
+        "model.layers.0.self_attn.qkv_proj.weight",
+        "model.layers.0.self_attn.qkv_proj.weight",
+        "model.layers.0.self_attn.qkv_proj.weight",
+        # Only the exact fused layers are remapped; everything else is untouched.
+        "model.layers.2.mlp.experts.0.gate_proj.weight",
+        "model.layers.1.self_attn.q_proj.weight",
+    ]
+    assert shard_ids == [0, 1, 0, "q", "k", "v", None, None]
+
+    # The fused layers are exposed to the quantization machinery via their
+    # original constituent projection names (what the checkpoint stores).
+    assert glu_fuser.packed_modules_mapping == {
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+    assert qkv_fuser.packed_modules_mapping == {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+    }
+
+
+@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP])
+def test_unfusable_modules_are_not_fused(cls, default_vllm_config):
+    with torch.device("meta"):
+        module = cls()
+    fuser = get_fuser(module)
+    # Either no pattern matches the class, or this instance fails validation
+    # (`recursive_replace` gates fusion and its weight mappings on `validate`)
+    model_config = default_vllm_config.model_config
+    assert fuser is None or not fuser.validate(module, model_config)
+
+
+def test_act_and_mul_derived_from_module(default_vllm_config):
+    from transformers.activations import GELUTanh, SiLUActivation
+
+    from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul
+
+    assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul)
+    assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul)
+    gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh())
+    assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh"
+    gelu = GLUFuser._get_act_and_mul(nn.GELU())
+    assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none"
+    # Not activations at all -> no fusion
+    assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None
+    assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
+    with pytest.raises(ValueError, match="No AndMul equivalent"):
+        GLUFuser._get_act_and_mul(nn.Dropout())

tests/models/transformers/fusers/test_moe.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Transformers modeling backend's MoE fuser."""
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm.model_executor.models.transformers.fusers import MoEBlockFuser
+
+from .test_linear import GLUMLP
+
+
+class TopKRouter(nn.Module):
+    """HF v5 top-k router: `linear -> softmax -> topk (-> renorm)`."""
+
+    def __init__(self, num_experts=8, hidden=16, top_k=2, sigmoid=False):
+        super().__init__()
+        self.top_k = top_k
+        self.sigmoid = sigmoid
+        self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
+
+    def forward(self, hidden_states):
+        logits = F.linear(hidden_states, self.weight)
+        scores = torch.sigmoid(logits) if self.sigmoid else F.softmax(logits, dim=-1)
+        value, index = torch.topk(scores, self.top_k, dim=-1)
+        value = value / value.sum(dim=-1, keepdim=True)
+        return logits, value, index
+
+
+class CorrectionRouter(nn.Module):
+    """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined."""
+
+    def __init__(self, num_experts=8, hidden=16):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
+        self.register_buffer("e_score_correction_bias", torch.zeros(num_experts))
+
+    def forward(self, hidden_states):
+        logits = F.linear(hidden_states, self.weight)
+        scores = torch.sigmoid(logits) + self.e_score_correction_bias
+        _, index = torch.topk(scores, 2, dim=-1)
+        return logits, scores, index
+
+
+class BiasedRouter(TopKRouter):
+    """A valid top-k router but not `weight`-only (extra `bias` param) -> declined."""
+
+    def __init__(self):
+        super().__init__()
+        self.bias = nn.Parameter(torch.zeros(8))
+
+    def forward(self, hidden_states):
+        logits = F.linear(hidden_states, self.weight) + self.bias
+        scores = F.softmax(logits, dim=-1)
+        value, index = torch.topk(scores, self.top_k, dim=-1)
+        return logits, value, index
+
+
+class DisconnectedRouter(TopKRouter):
+    """linear+softmax+top-k present but top-k ignores the logits -> not a router."""
+
+    def forward(self, hidden_states):
+        logits = F.linear(hidden_states, self.weight)
+        _ = F.softmax(logits, dim=-1)  # scored, but not consumed by top-k
+        value, index = torch.topk(hidden_states, self.top_k, dim=-1)
+        return logits, value, index
+
+
+class MoEExperts(nn.Module):
+    """Packed experts (3D weights); only its name (`experts`) matters here."""
+
+    def __init__(self, num_experts=8, hidden=16, inter=32):
+        super().__init__()
+        self.gate_up_proj = nn.Parameter(torch.zeros(num_experts, 2 * inter, hidden))
+        self.down_proj = nn.Parameter(torch.zeros(num_experts, hidden, inter))
+
+    def forward(self, hidden_states, index, weights):
+        return hidden_states
+
+
+class MoEBlock(nn.Module):
+    """Single-tensor MoE block (Qwen3-style); subclasses override `_shared`."""
+
+    def __init__(self, router_cls=TopKRouter):
+        super().__init__()
+        self.experts = MoEExperts()
+        self.gate = router_cls()
+
+    def _shared(self, x, logits):
+        """The term added to the experts' output (none for a plain block)."""
+        return 0
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        logits, weights, index = self.gate(x)
+        out = self.experts(x, index, weights) + self._shared(x, logits)
+        return out.reshape(hidden_states.shape)
+
+
+class MoEBlockNoShared(MoEBlock):
+    """No shared-expert child but a gate-derived add -> trace skipped, still fuses."""
+
+    def _shared(self, x, logits):
+        return logits.sum()
+
+
+class MoEBlockShared(MoEBlock):
+    """A block with a shared expert and its sigmoid gate (Qwen2-style)."""
+
+    def __init__(self):
+        super().__init__()
+        self.shared_expert = GLUMLP()
+        self.shared_expert_gate = nn.Linear(16, 1, bias=False)
+
+    def _shared(self, x, logits):
+        return torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
+
+
+class MoEBlockSharedNoGate(MoEBlock):
+    """A block with an ungated shared expert -> native, shared passed through."""
+
+    def __init__(self):
+        super().__init__()
+        self.shared_expert = GLUMLP()
+
+    def _shared(self, x, logits):
+        return self.shared_expert(x)
+
+
+class MoEBlockTuple(MoEBlock):
+    """A tuple-returning block (gpt-oss-style) -> must decline."""
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        return self.experts(x, index, weights), index
+
+
+class MoEBlockTupleVar(MoEBlock):
+    """Returns a name bound to a tuple, not a literal tuple -> must still decline."""
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        result = self.experts(x, index, weights), index
+        return result
+
+
+class MoEBlockNestedTupleReturn(MoEBlock):
+    """Tuple `return` in a nested helper; block returns one tensor -> still fuses."""
+
+    def forward(self, hidden_states):
+        def keep(a, b):
+            return a, b
+
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        out, _ = keep(self.experts(x, index, weights), index)
+        return out.reshape(hidden_states.shape)
+
+
+class PlainMLP(nn.Module):
+    """A non-GLU FFN: `down(act(up(x)))`, no gating multiply."""
+
+    def __init__(self, hidden: int = 16, inter: int = 32):
+        super().__init__()
+        self.up_proj = nn.Linear(hidden, inter, bias=False)
+        self.down_proj = nn.Linear(inter, hidden, bias=False)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.up_proj(x)))
+
+
+class MoEBlockSharedNonGLU(MoEBlock):
+    """A non-GLU shared expert -> detected by dataflow (no gate/up merge)."""
+
+    def __init__(self):
+        super().__init__()
+        self.shared_expert = PlainMLP()
+
+    def _shared(self, x, logits):
+        return self.shared_expert(x)
+
+
+class MoEBlockUnaccounted(MoEBlock):
+    """A weight-bearing child outside the fused dataflow (pre-router) -> declined."""
+
+    def __init__(self):
+        super().__init__()
+        self.extra = nn.Linear(16, 16, bias=False)
+
+    def forward(self, hidden_states):
+        x = self.extra(hidden_states.reshape(-1, hidden_states.shape[-1]))
+        _, weights, index = self.gate(x)
+        return self.experts(x, index, weights).reshape(hidden_states.shape)
+
+
+class BufferScale(nn.Module):
+    """A stateful child carrying only a buffer (no parameters)."""
+
+    def __init__(self, hidden: int = 16):
+        super().__init__()
+        self.register_buffer("scale", torch.ones(hidden))
+
+    def forward(self, x):
+        return x * self.scale
+
+
+class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted):
+    """Like `MoEBlockUnaccounted`, but the extra child holds only a buffer."""
+
+    def __init__(self):
+        super().__init__()
+        self.extra = BufferScale()
+
+
+@pytest.mark.parametrize("sigmoid", [False, True])
+def test_moe_fuser_detects_router(sigmoid):
+    with torch.device("meta"):
+        block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid))
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    assert fuser.gate_name == "gate"
+    assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax")
+    assert fuser.shared_name is None and fuser.shared_gate_name is None
+
+
+def test_moe_fuser_detects_shared_experts():
+    with torch.device("meta"):
+        block = MoEBlockShared()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    assert fuser.shared_name == "shared_expert"
+    assert fuser.shared_gate_name == "shared_expert_gate"
+
+
+def test_moe_fuser_skips_shared_detection_without_extra_children():
+    """With only experts and gate, shared-expert detection (and its block trace)
+    is skipped, so a gate-derived add is not misread as a shared expert."""
+    with torch.device("meta"):
+        block = MoEBlockNoShared()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    assert fuser.shared_name is None and fuser.shared_gate_name is None
+
+
+def test_moe_fuser_shared_without_gate():
+    with torch.device("meta"):
+        block = MoEBlockSharedNoGate()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    assert fuser.shared_name == "shared_expert"
+    assert fuser.shared_gate_name is None
+
+
+def test_moe_fuser_detects_non_glu_shared_expert():
+    with torch.device("meta"):
+        block = MoEBlockSharedNonGLU()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    # Recognised by dataflow (added to the experts' output), though not a GLU.
+    assert fuser.shared_name == "shared_expert"
+    assert fuser.shared_gate_name is None
+
+
+@pytest.mark.parametrize(
+    "block_cls",
+    [
+        lambda: MoEBlock(CorrectionRouter),  # score-correction buffer (grouped)
+        lambda: MoEBlock(BiasedRouter),  # router not weight-only (extra param)
+        MoEBlockTuple,  # tuple-returning block (e.g. gpt-oss)
+        MoEBlockTupleVar,  # tuple returned via a name binding, not a literal
+        MoEBlockUnaccounted,  # weight-bearing child outside the fused dataflow
+        MoEBlockUnaccountedBuffer,  # buffer-only child outside the fused dataflow
+    ],
+)
+def test_moe_fuser_declines_unsupported(block_cls):
+    with torch.device("meta"):
+        block = block_cls()
+    assert MoEBlockFuser.match(block, "experts") is None
+
+
+def test_moe_fuser_ignores_nested_returns():
+    """A tuple `return` inside a nested helper must not decline a block whose own
+    forward returns a single tensor."""
+    with torch.device("meta"):
+        block = MoEBlockNestedTupleReturn()
+    assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser)
+
+
+def test_moe_fuser_router_requires_connected_dataflow():
+    """A gate with linear + softmax + top-k present but not wired as a router
+    (top-k selects over the input, not the scored logits) is not detected."""
+    with torch.device("meta"):
+        block = MoEBlock(DisconnectedRouter)
+    assert MoEBlockFuser.match(block, "experts") is None

tests/models/transformers/fusers/test_rms_norm.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Transformers modeling backend's RMSNorm fuser."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm.model_executor.models.transformers.fuser import get_fuser
+from vllm.model_executor.models.transformers.fusers import RMSNormFuser
+
+
+class RMSNorm(nn.Module):
+    """The canonical HF RMSNorm: `weight * x * rsqrt(mean(x**2) + eps)`."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-5, weight: bool = True):
+        super().__init__()
+        if weight:
+            self.weight = nn.Parameter(torch.ones(hidden))
+        self.variance_epsilon = eps
+
+    def _rms(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
+
+    def forward(self, x):
+        return self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
+
+
+class GemmaRMSNorm(RMSNorm):
+    """Zero-centered weight: `(1 + weight) * normalized`."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps)
+        self.weight = nn.Parameter(torch.zeros(hidden))
+
+    def forward(self, x):
+        return (1.0 + self.weight) * self._rms(x.to(torch.float32)).to(x.dtype)
+
+
+class WeightlessRMSNorm(RMSNorm):
+    """No scale parameter (e.g. Gemma3n `with_scale=False`)."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps, weight=False)
+
+    def forward(self, x):
+        return self._rms(x.to(torch.float32)).to(x.dtype)
+
+
+class LayerNorm(RMSNorm):
+    """An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast)."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps)
+
+    def forward(self, x):
+        return self.weight * self._rms(x)
+
+
+class NotAnRMSNorm(RMSNorm):
+    """Mean-subtracting LayerNorm-like math -> not an RMSNorm."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps)
+
+    def forward(self, x):
+        x = x - x.mean(-1, keepdim=True)
+        variance = x.var(-1, keepdim=True)
+        return self.weight * x / torch.sqrt(variance + self.variance_epsilon)
+
+
+class GatedRMSNorm(RMSNorm):
+    """Second input and tail compute -> not an RMSNorm."""
+
+    def forward(self, x, gate=None):
+        normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
+        return normed * F.silu(gate)
+
+
+class GatedFusedRMSNorm(nn.Module):
+    """Same as GatedRMSNorm, but built on the fused `rms_norm` op -> not an RMSNorm."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden))
+        self.eps = eps
+
+    def forward(self, x, gate=None):
+        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) * F.silu(gate)
+
+
+class UntraceableGatedRMSNorm(RMSNorm):
+    """Tracer can't see tail compute in forward, but still has a second input (gate)."""
+
+    def forward(self, x, gate=None):
+        normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
+        if gate.sum() > 0:  # untraceable -> partial graph, no visible tail
+            normed = normed * F.silu(gate)
+        return normed
+
+
+@pytest.mark.parametrize(
+    "cls,eps,zero_centered",
+    [
+        (RMSNorm, 1e-5, False),
+        (GemmaRMSNorm, 1e-6, True),
+        (WeightlessRMSNorm, 1e-6, False),
+        (LayerNorm, 1e-6, False),
+        (torch.nn.RMSNorm, 1e-5, False),  # fused `F.rms_norm` op
+    ],
+)
+def test_detects_rms_norm_variants(cls, eps, zero_centered):
+    with torch.device("meta"):
+        fuser = get_fuser(cls(16, eps=eps))
+    assert isinstance(fuser, RMSNormFuser)
+    assert fuser.zero_centered == zero_centered
+
+
+@pytest.mark.parametrize("cls", [NotAnRMSNorm, nn.LayerNorm, nn.SiLU])
+def test_non_rms_norms_are_not_matched(cls):
+    with torch.device("meta"):
+        module = cls(16) if cls is nn.LayerNorm else cls()
+    assert not isinstance(get_fuser(module), RMSNormFuser)
+
+
+@pytest.mark.parametrize(
+    "cls", [GatedRMSNorm, GatedFusedRMSNorm, UntraceableGatedRMSNorm]
+)
+def test_gated_rms_norm_is_not_fused(cls):
+    with torch.device("meta"):
+        assert not isinstance(get_fuser(cls()), RMSNormFuser)
+
+
+@pytest.mark.parametrize(
+    "cls,expected,zero_centered",
+    [
+        (RMSNorm, "RMSNorm", False),
+        (GemmaRMSNorm, "GemmaRMSNorm", True),
+        (WeightlessRMSNorm, "RMSNorm", False),
+    ],
+)
+def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_config):
+    from vllm.model_executor.layers.layernorm import GemmaRMSNorm as VLLMGemmaRMSNorm
+    from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
+
+    # `default_vllm_config` supplies the config context the CustomOp needs; the
+    # weightless path reads hidden size from the model config, so stub it.
+    model_config = SimpleNamespace(get_hidden_size=lambda: 16)
+    with torch.device("meta"):
+        module = cls()
+        fuser = get_fuser(module)
+        built = fuser.fuse(module, "norm", model_config, None)
+    from vllm.model_executor.models.transformers.fusers.rms_norm import (
+        TPAwareNormMixin,
+    )
+
+    types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm}
+    assert isinstance(built, types_by_name[expected])
+    assert isinstance(built, TPAwareNormMixin)  # fused norms self-correct under TP
+    assert built.variance_epsilon == module.variance_epsilon
+    assert isinstance(built.weight, nn.Parameter) == (
+        getattr(module, "weight", None) is not None
+    )
+
+
+def test_fused_rms_norm_op_default_eps(default_vllm_config):
+    """`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path;
+    its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`."""
+    from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
+
+    with torch.device("meta"):
+        module = torch.nn.RMSNorm(16)  # forward is a single `F.rms_norm` call
+        fuser = get_fuser(module)
+        assert isinstance(fuser, RMSNormFuser)
+        assert not fuser.zero_centered
+        model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)
+        built = fuser.fuse(module, "norm", model_config, None)
+    assert isinstance(built, VLLMRMSNorm)
+    assert built.variance_epsilon == torch.finfo(torch.float32).eps
+
+
+def test_eps_is_derived_per_instance(default_vllm_config):
+    """Two instances of the same norm class with different eps must fuse to their
+    own eps: the type-cached fuser holds only structure, not this value."""
+    model_config = SimpleNamespace(get_hidden_size=lambda: 16)
+    with torch.device("meta"):
+        for eps in (1e-5, 1e-6):
+            module = RMSNorm(16, eps=eps)
+            built = get_fuser(module).fuse(module, "norm", model_config, None)
+            assert built.variance_epsilon == eps
+
+
+def test_fused_norm_is_gather_capable(default_vllm_config):
+    """Every fused norm is emitted gather-capable, so a norm on a head-sharded
+    projection (OLMoE-style) self-corrects at runtime with no QKV-specific
+    plumbing. A full-width input skips the gather and equals a plain norm."""
+    from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
+    from vllm.model_executor.models.transformers.fusers import rms_norm
+
+    torch.manual_seed(0)
+    x = torch.randn(4, 16)
+    for gathered_cls, plain_cls in [
+        (rms_norm.TPAwareRMSNorm, RMSNorm),
+        (rms_norm.TPAwareGemmaRMSNorm, GemmaRMSNorm),
+    ]:
+        gathered = gathered_cls(hidden_size=16, eps=1e-6)
+        assert isinstance(gathered, rms_norm.TPAwareNormMixin)
+        plain = plain_cls(hidden_size=16, eps=1e-6)
+        with torch.no_grad():
+            weight = torch.randn(16)
+            gathered.weight.copy_(weight)
+            plain.weight.copy_(weight)
+        torch.testing.assert_close(gathered(x), plain(x))
+
+
+def test_gathered_norm_rejects_uneven_sharding(default_vllm_config):
+    """A sharded input (narrower than the full-width weight) that does not tile
+    the weight evenly across ranks is rejected before any collective."""
+    from vllm.model_executor.models.transformers.fusers import rms_norm
+
+    norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6)
+    norm.tp_size = 2  # emulate TP=2 without a real process group
+    with pytest.raises(ValueError, match="does not tile it evenly"):
+        norm(torch.randn(2, 3))  # 3 * 2 != 8