Skip to content

vllm.models.muse_spark.amd.ops.fused_norm_expand

Pre-expert RMSNorm written once per route, in the [T * top_k, H] layout the ROCm MXFP4 (AITER top-1 per route) experts take.

Functions:

fused_rmsnorm_expand(x, weight, eps, gain_center, input_scale, repeats)

RMSNorm of x written repeats times per row.

Source code in vllm/models/muse_spark/amd/ops/fused_norm_expand.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def fused_rmsnorm_expand(
    x: torch.Tensor,
    weight: torch.Tensor | None,
    eps: float,
    gain_center: float,
    input_scale: bool,
    repeats: int,
) -> torch.Tensor:
    """RMSNorm of ``x`` written ``repeats`` times per row."""
    assert x.dim() == 2 and x.is_contiguous()
    T, D = x.shape
    out = torch.empty((T, repeats, D), dtype=x.dtype, device=x.device)
    if T > 0:
        block_d = triton.next_power_of_2(D)
        _rmsnorm_expand_kernel[(T,)](
            x,
            weight if weight is not None else x,
            out,
            D=D,
            eps=eps,
            gain_center=gain_center,
            REPEATS=repeats,
            BLOCK_D=block_d,
            HAS_WEIGHT=weight is not None,
            INPUT_SCALE=input_scale and weight is not None,
            num_warps=8 if block_d >= 4096 else 4,
        )
    first = out[:, 0]
    setattr(first, _EXPANDED_ATTR, out)
    return first

take_expanded(hidden_states, repeats)

The [T * repeats, D] buffer behind a :func:fused_rmsnorm_expand result, or None if hidden_states is not one.

Source code in vllm/models/muse_spark/amd/ops/fused_norm_expand.py 
75
76
77
78
79
80
81
82
83
84
85
86
def take_expanded(hidden_states: torch.Tensor, repeats: int) -> torch.Tensor | None:
    """The ``[T * repeats, D]`` buffer behind a :func:`fused_rmsnorm_expand`
    result, or ``None`` if ``hidden_states`` is not one."""
    out = getattr(hidden_states, _EXPANDED_ATTR, None)
    if (
        out is None
        or out.data_ptr() != hidden_states.data_ptr()
        or out.shape != (hidden_states.shape[0], repeats, hidden_states.shape[1])
        or hidden_states.stride() != (repeats * hidden_states.shape[1], 1)
    ):
        return None
    return out.view(-1, hidden_states.shape[1])