`vllm.models.muse_spark.amd.mono.kernels.k3` ¶

K3: one Muse Spark decoder layer for S = 8 / 16 rows as one persistent launch.

One rank of a TP = EP = 8 group (rank r owns experts [32 r, 32 r + 32)), 256 CTAs x 512 threads, everything between two attention cores:

[attention]  optional: this layer's verify attention (attn_frag.py) at the head
O      o-norm x sigmoid(gate) -> o_proj rows 32 b .. over this rank's K slice; the bf16
       partial goes to the owner rank of the column block (b // 32): reduce-scatter
AB     the owner sums the 8 partials in rank order and all-gathers the reduced o and its
       norm partials; every CTA: post-attention norm, gated residual on the fp32 stream,
       pre-FFN norm -> ffn_in (one token-reduction round for both norms)
A      router logits (fp32) + pre_expert_proj of this rank's shard, K in 8 parts over
       8-CTA groups; logit partials and h_raw go to every rank's peer buffer
route  every CTA routes every token (identical top-8), the local experts' picks
C, D   gate_up / down of the local experts: MXFP4 weights in vLLM's shuffled layouts,
       decoded in registers, bf16 MFMA
E      per-pick post-expert norm and weighted sum -> the EP partial, reduce-scattered to
       the block owner, summed in rank order, all-gathered
F      post_expert_proj (this rank's 1024 rows, K in 4 parts) -> every rank
G      post-FFN norm, gated residual, the next layer's input norm -> stream (fp32) and
       the next attention input (bf16)
QKV    the next layer's qkv + gate GEMV (K in eighths); the owners finish whole heads:
       qk-norm, RoPE (next_rope), q out, bf16 K / V into vLLM's LBHNC KV cache

Peer regions are double buffered by tag parity; tag = epoch * 64 + layer + 1 with the epoch bumped in lockstep by every rank once per forward.

Functions:

bf_mma_operand –

Coalesced bf16 weight registers (4 instructions: rows 8 blk + lane / 8, 16 B at
bf_tile_load –

16 rows (row16 ..) x 128 K (k ..) of a row-major bf16 [*, ldk] weight, coalesced.
bf_tile_load_mma –

16 rows (row16 ..) x 128 K (k ..) of a row-major bf16 [*, ldk] weight, loaded straight in the
build_vllm –

The K3 build vLLM launches for S rows (8 / 16) of a TP = EP = 8 group; next_rope: the next
dn_loads –

Task (u, j)'s w2 operands for this wave (expert etab[u], u clamped below nv).
fp4_mma –

Chunk c (256 K) of a 16-row MXFP4 group: 8 MFMA steps, B column n = LDS row row0 + n (at
mx_chunk –

Rows rbase .. rbase + 16 x K chunk c (256) of a shuffled MXFP4 [*, k] weight w and its e8m0
o_gate –

o-norm (over the head's 8 lanes) x sigmoid(gate) of 8 bf16 (4 pair words xv, gv) -> 4 words.
peer_layout –

Every rank's peer buffer (IPC-shared), [parity] first.
route_token –

Tokens tkns (this wave's, all valid), interleaved for ILP: sigmoid +
scratch_layout –

This rank's scratch regions. ss, ss1, ss2g, ssa1 and ssa2 are unused; they keep the
task_ring –

Tasks task .. end - 1: compute(q[0], task, it, *args) with the next len(q) tasks' loads

`bf_mma_operand(ops, qq, lane)` ¶

Coalesced bf16 weight registers (4 instructions: rows 8 blk + lane / 8, 16 B at 128 h + 16 (lane % 8)) -> the MFMA A operand of step qq (K 32 qq + 8 g of row r).