Skip to content

vllm.models.minimax_m3.amd.mono.weights

Zero-copy views (and two converted copies) of one vLLM MiniMaxM3DecoderLayer for the mono kernels.

The kernels read tensors through raw pointers in fixed layouts: a differently laid-out weight would not fault, it would produce plausible garbage. So every layout assumption is asserted once and a model that does not match is refused (MonoUnsupported). The experts already have the kernels' layout (MXFP4, aiter shuffle_weight(16, 16) + e8m0_shuffle scales, fused shared expert in slot 128). The attention linears are MXFP8 in vLLM and the router gate is fp32, so those are converted once when the runner is built: MXFP8 -> bf16 -> per-channel fp8, preshuffled (16, 16); gate -> bf16.

Without shared-expert fusion the routed table holds the 128 routed experts and the shared expert stays the separate MXFP8 MLP; its weights are copied once into the f8f6f4 MFMA's fp8 A-operand order (mxfp8_a_tiles, the same bytes) with e8m0_shuffle-ordered scales, so the kernels run it in MXFP8.

Classes:

  • SparseMoeLayer –

    Everything the layer kernel reads for one sparse-attention MoE layer.

Functions:

  • mxfp8_a_tiles –

    MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> its bytes in the

SparseMoeLayer dataclass

Everything the layer kernel reads for one sparse-attention MoE layer.

Methods:

  • shared_bytes –

    Bytes of the kernels' copy of the separate shared expert.

Source code in vllm/models/minimax_m3/amd/mono/weights.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
@dataclass(frozen=True)
class SparseMoeLayer:
    """Everything the layer kernel reads for one sparse-attention MoE layer."""

    layer_id: int
    attn: object  # MiniMaxM3SparseAttention: layer name, KV cache views, scales
    index_cache: torch.Tensor  # [blocks, block, head_dim] fp8, shared block ids
    g_in: torch.Tensor
    w_qkv: torch.Tensor
    s_qkv: torch.Tensor
    g_q: torch.Tensor
    g_k: torch.Tensor
    g_iq: torch.Tensor
    g_ik: torch.Tensor
    cos_sin: torch.Tensor
    w_o: torch.Tensor
    s_o: torch.Tensor
    g_post: torch.Tensor
    gate: torch.Tensor
    bias: torch.Tensor
    w13: torch.Tensor
    s13: torch.Tensor
    w2: torch.Tensor
    s2: torch.Tensor
    # the separate MXFP8 shared expert (``mxfp8_a_tiles``) and the int64 table of
    # their addresses the kernels take (``SH_ARGS`` order); None when it is fused
    sh_w13: torch.Tensor | None = None
    sh_s13: torch.Tensor | None = None
    sh_w2: torch.Tensor | None = None
    sh_s2: torch.Tensor | None = None
    sh_args: torch.Tensor | None = None

    @property
    def shared_fp8(self) -> bool:
        return self.sh_args is not None

    def shared_bytes(self) -> int:
        """Bytes of the kernels' copy of the separate shared expert."""
        if not self.shared_fp8:
            return 0
        nbytes = 0
        for t in (self.sh_w13, self.sh_s13, self.sh_w2, self.sh_s2):
            assert t is not None
            nbytes += t.numel()
        return nbytes

shared_bytes()

Bytes of the kernels' copy of the separate shared expert.

Source code in vllm/models/minimax_m3/amd/mono/weights.py 
 99
100
101
102
103
104
105
106
107
def shared_bytes(self) -> int:
    """Bytes of the kernels' copy of the separate shared expert."""
    if not self.shared_fp8:
        return 0
    nbytes = 0
    for t in (self.sh_w13, self.sh_s13, self.sh_w2, self.sh_s2):
        assert t is not None
        nbytes += t.numel()
    return nbytes

_build_cp_qkv(linear, rank, tp_group, layer_id)

The wider fused projection of indexer context parallelism: every rank's index q head, all-gathered once at build time (IndexHeads(TP, rank)).

One rank's rows (2560): q(2048) | k(128) | v(128) | index_q(128) | index_k(128) CP rows (2944): q(2048) | k(128) | v(128) | iq_rank0..3 | index_k(128)

The gather's list order is the TP group's rank order, which IndexHeads.own indexes; the rank's own slot is checked against its local rows.

Source code in vllm/models/minimax_m3/amd/mono/weights.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def _build_cp_qkv(
    linear, rank: int, tp_group, layer_id: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """The wider fused projection of indexer context parallelism: every rank's
    index q head, all-gathered once at build time (``IndexHeads(TP, rank)``).

    One rank's rows (2560): q(2048) | k(128) | v(128) | index_q(128) | index_k(128)
    CP rows (2944):         q(2048) | k(128) | v(128) | iq_rank0..3 | index_k(128)

    The gather's list order is the TP group's rank order, which ``IndexHeads.own``
    indexes; the rank's own slot is checked against its local rows."""
    one_rows = (LOCAL_Q_HEADS + 3 + 1) * HEAD_DIM  # 2560
    w = getattr(linear, "weight", None)
    s = getattr(linear, "weight_scale", None)
    _need(
        w is not None
        and w.dtype == dtypes.fp8
        and tuple(w.shape) == (one_rows, HIDDEN),
        f"layer {layer_id} qkv_proj CP: weight "
        f"{None if w is None else (w.dtype, tuple(w.shape))}",
    )
    _need(
        s is not None and tuple(s.shape) == (one_rows, HIDDEN // 32),
        f"layer {layer_id} qkv_proj CP: scale {None if s is None else tuple(s.shape)}",
    )
    assert w is not None and s is not None
    _need(
        dist.get_rank(tp_group) == rank and dist.get_world_size(tp_group) == TP,
        f"layer {layer_id} qkv_proj CP: TP group rank {dist.get_rank(tp_group)}"
        f" of {dist.get_world_size(tp_group)}, not {rank} of {TP}",
    )
    bf = dequant_mxfp8_to_bf16(w.data, s.data.view(torch.uint8))

    iq_start = (LOCAL_Q_HEADS + 2) * HEAD_DIM  # 2304
    iq_end = iq_start + HEAD_DIM  # 2432
    my_iq = bf[iq_start:iq_end].contiguous()

    all_iq = [torch.empty_like(my_iq) for _ in range(TP)]
    dist.all_gather(all_iq, my_iq, group=tp_group)
    _need(
        torch.equal(all_iq[rank], my_iq),
        f"layer {layer_id} qkv_proj CP: gathered index q heads not in rank order",
    )

    cp_bf = torch.cat([bf[:iq_start], *all_iq, bf[iq_end:]], dim=0)
    cp_rows = (LOCAL_Q_HEADS + 3 + TP) * HEAD_DIM  # 2944
    assert cp_bf.shape[0] == cp_rows, f"CP QKV rows {cp_bf.shape[0]} != {cp_rows}"
    return _ptpc_fp8_from_bf16(cp_bf, cp_rows)

_ptpc_fp8_from_bf16(bf, rows)

bf16 [rows, cols] -> per-channel fp8 preshuffled (16, 16) + fp32 scale [rows].

Source code in vllm/models/minimax_m3/amd/mono/weights.py 
179
180
181
182
183
184
185
186
187
def _ptpc_fp8_from_bf16(
    bf: torch.Tensor, rows: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """bf16 [rows, cols] -> per-channel fp8 preshuffled (16, 16) + fp32 scale [rows]."""
    q, scale = pertoken_quant(bf, quant_dtype=dtypes.fp8)
    return (
        shuffle_weight(q, layout=(16, 16)).contiguous(),
        scale.view(rows).to(torch.float32).contiguous(),
    )

_ptpc_fp8_from_mxfp8(linear, rows, cols, name)

MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> per-channel fp8 preshuffled (16, 16) + fp32 scale [rows], as the kernels' PTPC GEMV reads.

Source code in vllm/models/minimax_m3/amd/mono/weights.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def _ptpc_fp8_from_mxfp8(
    linear, rows: int, cols: int, name: str
) -> tuple[torch.Tensor, torch.Tensor]:
    """MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> per-channel fp8
    preshuffled (16, 16) + fp32 scale [rows], as the kernels' PTPC GEMV reads."""
    w = getattr(linear, "weight", None)
    s = getattr(linear, "weight_scale", None)
    _need(
        w is not None and w.dtype == dtypes.fp8 and tuple(w.shape) == (rows, cols),
        f"{name}: weight {None if w is None else (w.dtype, tuple(w.shape))}",
    )
    _need(
        s is not None and tuple(s.shape) == (rows, cols // 32),
        f"{name}: scale {None if s is None else tuple(s.shape)}",
    )
    assert w is not None and s is not None
    bf = dequant_mxfp8_to_bf16(w.data, s.data.view(torch.uint8))
    q, scale = pertoken_quant(bf, quant_dtype=dtypes.fp8)
    return (
        shuffle_weight(q, layout=(16, 16)).contiguous(),
        scale.view(rows).to(torch.float32).contiguous(),
    )

mxfp8_a_tiles(linear, rows, cols, name)

MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> its bytes in the f8f6f4 MFMA's fp8 A order and its scales in e8m0_shuffle order.

A 16-row x 128-k tile is 2 KiB; lane g * 16 + r of a tile holds row r's k 16 g .. + 16 then 64 + 16 g .. + 16 (the B operand's order), 32 bytes.

Source code in vllm/models/minimax_m3/amd/mono/weights.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def mxfp8_a_tiles(
    linear, rows: int, cols: int, name: str
) -> tuple[torch.Tensor, torch.Tensor]:
    """MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> its bytes in the
    f8f6f4 MFMA's fp8 A order and its scales in ``e8m0_shuffle`` order.

    A 16-row x 128-k tile is 2 KiB; lane g * 16 + r of a tile holds row r's k
    16 g .. + 16 then 64 + 16 g .. + 16 (the B operand's order), 32 bytes."""
    w = getattr(linear, "weight", None)
    s = getattr(linear, "weight_scale", None)
    _need(
        w is not None and w.dtype == dtypes.fp8 and tuple(w.shape) == (rows, cols),
        f"{name}: weight {None if w is None else (w.dtype, tuple(w.shape))}",
    )
    _need(
        s is not None and tuple(s.shape) == (rows, cols // 32),
        f"{name}: scale {None if s is None else tuple(s.shape)}",
    )
    assert w is not None and s is not None
    # the row-major layout the native MXFP8 linear kernel keeps; a kernel that
    # preshuffles or pads would make the tiles silently wrong
    _need(
        not getattr(w, "is_shuffled", False) and not getattr(s, "is_shuffled", False),
        f"{name}: weight preshuffled",
    )
    _need(rows % 32 == 0 and cols % 256 == 0, f"{name}: shape {rows} x {cols}")
    tiles = (
        w.data.contiguous()
        .view(torch.uint8)
        .view(rows // 16, 16, cols // 128, 2, 4, 16)
        .permute(0, 2, 4, 1, 3, 5)
        .contiguous()
        .view(-1)
    )
    scales = (
        s.data.contiguous()
        .view(torch.uint8)
        .view(rows // 32, 2, 16, cols // 256, 2, 4)
        .permute(0, 3, 5, 2, 4, 1)
        .contiguous()
        .view(-1)
    )
    return tiles, scales