Skip to content

vllm.models.minimax_m3.amd.mono.layout

Where the mono layer kernel (K4) puts its tasks and data, per decode token count.

Plain Python, evaluated at build time: which CTA runs each stage's tasks, the scratch and peer-symmetric regions, and the LDS row pitches. Kept free of FlyDSL so it is testable without a GPU.

Functions:

  • pool_words –

    LDS words of the activation pool x. Narrow: every token's bf16 xn (the

  • shared_after_routing –

    Every CTA carries the same number of routed up / gate tasks (S = 4): a

  • stage_bases –

    Task t of a stage runs on CTA (base + t) % BLOCKS; the merge (S > 1), o and

  • sym_layout –

    Per-rank symmetric buffer: attn / ffn partial regions, src rank[HIDDEN],

  • ug_ctas_per_token –

    CTAs per token of the routed up / gate tasks, token k on CTAs k C .. (k + 1)

  • ug_tasks_of –

    Routed up / gate tasks CTA cta runs.

  • wide_moe –

    The MoE stage groups its work by expert (every distinct expert's weights

_layout(items, align=256)

Byte offsets of (value, tag) pair regions, 8 bytes per pair.

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
56
57
58
59
60
61
62
63
def _layout(items, align=256):
    """Byte offsets of (value, tag) pair regions, 8 bytes per pair."""
    off, out = 0, {}
    for name, pairs in items:
        out[name] = off
        off += (pairs * 8 + align - 1) // align * align
    out["_bytes"] = off
    return out

_shared_base(tokens, router)

First of the SH_TASKS consecutive CTAs that run the shared expert's up / gate before the routing lands. Ranked by the most routed up / gate tasks a window CTA carries (they queue behind it); then, only if that is the most any CTA carries, by its router CTAs (they start it late and would delay the last routed tasks: S = 4, 4 us); then by its down CTAs (S = 2: 43.1 vs 43.7 us); then by its router CTAs anyway.

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
def _shared_base(tokens: int, router: int) -> int:
    """First of the SH_TASKS consecutive CTAs that run the shared expert's up /
    gate before the routing lands. Ranked by the most routed up / gate tasks a
    window CTA carries (they queue behind it); then, only if that is the most any
    CTA carries, by its router CTAs (they start it late and would delay the last
    routed tasks: S = 4, 4 us); then by its down CTAs (S = 2: 43.1 vs 43.7 us);
    then by its router CTAs anyway."""
    n_router = N_ROUTER * tokens
    n_split = N_SPLIT * tokens
    most = max(ug_tasks_of(c, tokens) for c in range(BLOCKS))

    def cost(base):
        ctas = [(base + j) % BLOCKS for j in range(SH_TASKS)]
        busiest = max(ug_tasks_of(c, tokens) for c in ctas)
        in_router = sum((c - router) % BLOCKS < n_router for c in ctas)
        in_down = sum(n_split <= c < n_split + N_DN for c in ctas)
        return busiest, in_router if busiest == most else 0, in_down, in_router

    return min(range(BLOCKS), key=cost)

pool_words(tokens)

LDS words of the activation pool x. Narrow: every token's bf16 xn (the largest stage operand). Wide: the largest of the split's PV partials, the up / gate stage's MXFP8 xn + scales + up partials + router sigmoids, and the down stage's MXFP8 mids + scales + a zero row.

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
def pool_words(tokens: int) -> int:
    """LDS words of the activation pool ``x``. Narrow: every token's bf16 xn (the
    largest stage operand). Wide: the largest of the split's PV partials, the up /
    gate stage's MXFP8 xn + scales + up partials + router sigmoids, and the down
    stage's MXFP8 mids + scales + a zero row."""
    if not wide_moe(tokens):
        return HIDDEN // 2 * tokens
    rows = MOE_SLOTS * tokens + 1
    return max(
        WAVES * O_K,
        tokens * (XN8_ROW + XSC_ROW + N_ROUTED) + WAVES * 64 * 4,
        rows * (MID_ROW + MIDSC_ROW),
    )

shared_after_routing(tokens)

Every CTA carries the same number of routed up / gate tasks (S = 4): a shared-expert CTA routes first and runs its shared half task while its first routed task's weights are in flight, instead of holding its routing back (S = 4 54.51 -> 54.09 us; S = 2, where the shared window sits on CTAs with one task fewer, +0.1).

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
187
188
189
190
191
192
193
194
def shared_after_routing(tokens: int) -> bool:
    """Every CTA carries the same number of routed up / gate tasks (S = 4): a
    shared-expert CTA routes first and runs its shared half task while its first
    routed task's weights are in flight, instead of holding its routing back
    (S = 4 54.51 -> 54.09 us; S = 2, where the shared window sits on CTAs with one
    task fewer, +0.1)."""
    loads = [ug_tasks_of(c, tokens) for c in range(BLOCKS)]
    return min(loads) == max(loads)

stage_bases(tokens)

Task t of a stage runs on CTA (base + t) % BLOCKS; the merge (S > 1), o and down tasks start past the tokens' split CTAs.

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def stage_bases(tokens: int):
    """Task t of a stage runs on CTA (base + t) % BLOCKS; the merge (S > 1), o and
    down tasks start past the tokens' split CTAs."""
    n_split = N_SPLIT * tokens
    n_merge = tokens if tokens > 1 else 0
    router = (n_split + n_merge + N_O) % BLOCKS
    return {
        "split": 0,
        "merge": n_split,
        "o": n_split + n_merge,
        "router": router,
        "shared": _shared_base(tokens, router),
        "ug": 0,
        "down": n_split,
    }

sym_layout(npes)

Per-rank symmetric buffer: attn / ffn partial regions, src rank[HIDDEN], and the peers' index-selection candidates (sel_cp).

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def sym_layout(npes: int):
    """Per-rank symmetric buffer: attn / ffn partial regions, [src rank][token][HIDDEN],
    and the peers' index-selection candidates (``sel_cp``)."""
    rows = npes * MAX_TOKENS * HIDDEN // 2
    # a_ag / ffn_ag: the finished attention / FFN all-reduce rows, each written by
    # its owner rank (the reduce-scatter + all-gather form, S >= AG_RS_FROM)
    ag = MAX_TOKENS * HIDDEN // 2
    return _layout(
        [
            ("attn", rows),
            ("ffn", rows),
            ("a_ag", ag),
            ("ffn_ag", ag),
            # indexer context parallelism: token k's candidates for this rank's
            # index head, source rank s's half h's r-th best at pairs
            # 2 ((k N_SPLIT + 2 s + h) TOPK_BLOCKS + r) + (0 key, 1 block)
            ("sel_cp", SEL_LONG_PAIRS * MAX_TOKENS),
        ]
    )

ug_ctas_per_token(tokens)

CTAs per token of the routed up / gate tasks, token k on CTAs k C .. (k + 1) C - 1: a CTA stages one token's xn, not all (fused layer, S = 4: 57.9 -> 57.6 us). None: round-robin -- S < 4 (S = 2: +0.4 us) or S not dividing the grid.

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
150
151
152
153
154
def ug_ctas_per_token(tokens: int):
    """CTAs per token of the routed up / gate tasks, token k on CTAs k C .. (k + 1)
    C - 1: a CTA stages one token's xn, not all (fused layer, S = 4: 57.9 -> 57.6
    us). None: round-robin -- S < 4 (S = 2: +0.4 us) or S not dividing the grid."""
    return BLOCKS // tokens if tokens >= 4 and BLOCKS % tokens == 0 else None

ug_tasks_of(cta, tokens)

Routed up / gate tasks CTA cta runs.

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
157
158
159
160
161
162
163
def ug_tasks_of(cta: int, tokens: int) -> int:
    """Routed up / gate tasks CTA ``cta`` runs."""
    per_token = TOP_K * UG_PER_SLOT
    cpt = ug_ctas_per_token(tokens)
    if cpt is None:
        return (per_token * tokens - cta + BLOCKS - 1) // BLOCKS
    return (per_token - cta % cpt + cpt - 1) // cpt

wide_moe(tokens)

The MoE stage groups its work by expert (every distinct expert's weights read once, a B column per token) instead of by (token, slot): past four tokens the per-token tasks re-read the experts the tokens share, and their operands no longer fit in LDS.

Source code in vllm/models/minimax_m3/amd/mono/layout.py 
200
201
202
203
204
205
def wide_moe(tokens: int) -> bool:
    """The MoE stage groups its work by expert (every distinct expert's weights
    read once, a B column per token) instead of by (token, slot): past four tokens
    the per-token tasks re-read the experts the tokens share, and their operands
    no longer fit in LDS."""
    return tokens >= WIDE_FROM