Skip to content

vllm.models.muse_spark.amd.mono.kernels.k5

HA and HB: one Muse Spark decoder layer for S = 32 / 64 / 96 / 128 rows as two persistent launches around HM (k5m.py, the local experts); one rank of a TP = EP = 8 group.

[attention] -> HA -> HM -> HB
this layer's verify attention (attn_frag.py) at the head; o-norm x gate

(distributed) -> o_proj (rows 32 b .., K = this rank's 1024) -> the bf16 partials go to the owner rank -> attention boundary (post_attn_norm, gated residual, pre-FFN norm) -> ffn_in -> router (fp32, 32 experts a rank) and pre_expert_proj (512 rows a rank), K in 8 parts over 8-CTA groups -> logits and h_raw all-gathered -> top-8 (CTA t routes token t) -> h = pre_expert_norm(h_raw) (CTA t: token t) -> h, top-k ids / weights for HM.

HB EP all-reduce of HM's partial (reduce-scatter to the column-block owner + all-gather) -> post_expert_proj (this rank's 1024 rows, K in 4 parts) -> post-FFN norm, gated residual, next input norm -> xn -> qkv + gate (K eighths) -> head finishers (qk-norm, RoPE, q out, bf16 K / V into vLLM's LBHNC KV cache).

The boundaries (Build.bnd). S > 32, token owners: rank r owns tokens t = r (mod 8); the o_proj partials and the post_expert_proj outputs go to the owner rank of each token, one CTA a token does the whole row (all-reduce sum, both norms, gated residual) and the ffn_in / xn rows are all-gathered. The residual stream lives token-sharded (rank r: its S / 8 rows, fp32, all 8192 columns); HA reads it as rows so the first layer can read a full replicated stream (s_full); HB writes it the same way and (gather) all-gathers it. S = 32, column blocks: the o_proj partials go to the owner of the column block, which reduces and all-gathers them; both boundaries run on all 256 CTAs (token groups x column groups, one-hop token norms) and the stream is replicated.

Every row tile is 16 tokens (MFMA N); dense stages keep their weights in VGPRs and loop over the S / 16 row tiles with the activations read straight from L2 / peer memory in MFMA layout. Tags: tag = epoch * 64 + layer + 1 (HA and HB of a layer share it, regions are disjoint); peer regions double buffered by tag parity.

Classes:

Functions:

  • build_vllm –

    The HA / HB build vLLM launches for S rows (32 / 64 / 96 / 128) of a TP = EP = 8 group;

Build dataclass

Attributes:

  • bnd –

    The boundaries: 3 at S = 32 (o reduce-scattered to the column-block owner and

Source code in vllm/models/muse_spark/amd/mono/kernels/k5.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
@dataclass(frozen=True)
class Build:
    S: int = 32  # rows (32 / 64 / 96 / 128)
    next_rope: bool = True  # the next layer has RoPE (SWA) or not (full, NoPE)
    # this layer's verify attention fused at the HA head ("swa" / "full")
    attn: str = ""

    @property
    def bnd(self):
        """The boundaries: 3 at S = 32 (o reduce-scattered to the column-block owner and
        all-gathered, both boundaries on all 256 CTAs with one-hop token norms over column
        groups, the stream replicated); 0 (S > 32): the token-owner boundaries (one CTA a
        token, the stream token-sharded and all-gathered by HB)."""
        return 3 if self.S == 32 else 0

bnd property

The boundaries: 3 at S = 32 (o reduce-scattered to the column-block owner and all-gathered, both boundaries on all 256 CTAs with one-hop token norms over column groups, the stream replicated); 0 (S > 32): the token-owner boundaries (one CTA a token, the stream token-sharded and all-gathered by HB).

build_vllm(S=32, W=8, next_rope=True, attn='')

The HA / HB build vLLM launches for S rows (32 / 64 / 96 / 128) of a TP = EP = 8 group; next_rope: the next layer has RoPE (HB); attn = "swa" / "full": this layer's verify attention fused at the HA head.

Source code in vllm/models/muse_spark/amd/mono/kernels/k5.py 
2384
2385
2386
2387
2388
2389
def build_vllm(S=32, W=8, next_rope=True, attn=""):
    """The HA / HB build vLLM launches for S rows (32 / 64 / 96 / 128) of a TP = EP = 8 group;
    next_rope: the next layer has RoPE (HB); attn = "swa" / "full": this layer's verify attention
    fused at the HA head."""
    assert W == 8 and S in (32, 64, 96, 128), (S, W)
    return Build(S=S, next_rope=bool(next_rope), attn=attn)