vllm.models.muse_spark.amd.mono.kernels.k5m

HM: this rank's routed experts for S = 32 / 64 / 96 / 128 rows as one persistent launch between HA and HB (k5.py); one rank of a TP = EP = 8 group.

Inputs from HA: h [S, 4096] bf16 (pre_expert_norm output), top-8 ids [S, 8] and weights [S, 8]. Every CTA builds the same tables: the picks of each local expert in token order, split into virtual experts of at most 16 picks (one MFMA N tile). Tasks: (virtual expert v, row tile j), 256 a virtual expert (gate_up: 16 gate + 16 up rows; down: 16 output rows); CTA b owns tasks [b V, (b+1) V) (V virtual experts), so it spans at most two virtual experts and keeps their picks' h (gate_up) / a (down) rows in LDS. Weights MXFP4 (w13 / w2 + E8M0 scales) in vLLM's shuffled layouts, streamed with a prefetch queue, decoded in registers (v_cvt_scalef32_pk_bf16_fp4), bf16 MFMA.

gate_up: a = bf16(silu(bf16 g) bf16 u) -> scratch A (tagged pairs, pick rows) down: z = bf16(w_post bf16(y)) -> scratch Z (tagged), per-pick sums of squares a CTA rstd: CTA c: the post-expert-norm rstd of picks c, c + 256, .. (tagged) final: CTA b, columns 16b ..: part[t] = sum over t's local picks (k order) of bf16(z rstd) w -> pushed to the owner rank of the column block (b / 32) = HB's reduce-scatter input (+ flags), and part_out (a local copy). Same tag as HA / HB of the layer (epoch * 64 + layer + 1).

Functions:

• build_vllm –

  The HM build vLLM launches for S rows of a TP = EP = 8 group.

build_vllm(S, W=8)

The HM build vLLM launches for S rows of a TP = EP = 8 group.

Source code in vllm/models/muse_spark/amd/mono/kernels/k5m.py

def build_vllm(S, W=8):
    """The HM build vLLM launches for S rows of a TP = EP = 8 group."""
    assert S in (32, 64, 96, 128) and W == 8, (S, W)
    return BuildM(S=S)