`vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla_attn` ¶

M1c mla_attn: Kimi-K3's MLA decode attention at TP8 + DCP1, the bf16 query (M1a qbf16) against the fp8 latent cache, split over the KV: one pass a split, every split's partial (bf16 o, f32 LSE) for a combine (M1b's).

One launch, BLOCKS CTAs x THREADS threads. A CTA a (head group, split): 16 query rows (16 / qlen heads x qlen positions) against a chunk of the request's KV. The split's CTAs (every head group) share an XCD, so a chunk streams from HBM once into that XCD's L2.

block of 128 tokens (wave w its tile of 16, the waves in step), its K
rows DMA'd into one of two LDS slots a block ahead (global_load_lds, the
chunks spread over the block before: ``DMA``; this wave's chunks waited
for one at a time, so the top barrier only frees the slot)
    K tile LDS -> registers (a lane 32 B a line of a token: the MFMA's A, the K
    order permuted alike in Q's B), S^T = K Q^T: the fp8 K as is against
    Q's exact fp8 hi + lo (``q_split``), 9 16x16x128 MFMAs a row tile
    the block's row max over the waves (LDS) -> P = exp2(S - m) 2^PLOG,
    fp8 hi + lo -> LDS; the running O, l rescaled (online softmax)
    O^T += V^T P^T: wave w its 64 latents, the fp8 V^T straight from the
    slot (``ds_read_b64_tr_b8``) against P's hi and lo, 4 16x16x128 MFMAs
    a row tile and part
end: l summed over the waves -> o = O / l (bf16), LSE = m + log l

A split with no token of the request (or none before its causal bound) leaves LSE -inf and o 0. Rows are causal: position i of the request's qlen sees its KV up to seq_len - qlen + i.

Functions:

block –

Block b (stage_block) behind the top barrier, its DMA hooks.
chunk_dma –

One chunk DMA (k_dma's): src + imm (an offset off the pool,
cta –

CTA c["bid"]'s (head group, split): its partial stored (cache
dma_lane –

(token, 16 B cc) tail DMA lane lane carries (it lands at position
dma_offs –

k_dma's sources of idx_load's indices: the tail chunk's (16 B
fp8_pack4 –

Four f32 (in range) -> one dword of four E4M3, RNE (ops' but the
fp8_split4 –

Four f32 (in range) -> (hi, lo) dwords of four E4M3: x ~ hi + lo (the
fp8_unpack4 –

One dword of four E4M3 -> four f32 (exact).
host_args –

(qb, kv, pages, indptr, scale, vscale) of the kernel (mla_attn's
idx_dma –

Block b's page indices of this wave's tile (or tiles t0 .. t0 + n) ->
idx_lds –

This lane's DMA tokens' page indices of block b (idx_dma's; as
idx_load –

The page indices of this lane's DMA tokens of block b's tile of this
k_dma –

Block b's K tile of this wave (or tile; dma_offs) -> slot
kpos –

The 16 B position (of 64) in a tile's tail chunk of token t's 16 B cc:
lds_dma –

A global_load_lds_dwordx4: lane i's 16 B at src -> LDS byte
lines_at –

Token t's 16 B col of line s: its byte in the slot (from line s's
mfma_f8 –

16x16x128 E4M3 x E4M3 MFMA, unscaled (scale operands 0: LLVM's
mla_attn –

Host launcher. q [S, 12, 576] bf16; kv [tokens, 576] fp8 rows
q_parts –

Row tile r's B operands from q_split's: QK step s's hi, lo
q_split –

The CTA's bf16 Q rows (slot 1: [row][ENTRY]) -> E4M3 hi, lo of q 2^k
run –

The blocks (block 0's DMA and block 1's indices issued) and the
sig –

The 16 B column swizzle (XOR) of token t (of 16) of a tile of parity
smem –

The LDS of a CTA (cta's ks / xm / xp / xi; xp P's hi | lo, then
stage_block –

One block (its K in slot b % 2) into the running (m, l, O) of each of
tr_read –

ds_read_b64_tr_b8 at LDS byte addr: a lane m of 16 byte m % 8 of
v_reads –

V^T of the slot's block b by transpose reads, the A of a 16x16x128 fp8
vmwait –

s_waitcnt vmcnt(n): at most n of this wave's loads / DMAs in flight
vtr_lane –

This lane's byte in a block's (line wave // 2) of the 8 B it gives a
wait_dma –

This wave's DMAs (and loads) landed; the CTA's after the barrier.
wide_kv –

The build's wide for KV pool kv: a DMA's 32-bit offset off the
ws_dmas –

This wave's DMA hook of block bb's tile's chunks ps (k_dma's),
ws_l –

Row row of row tile r's l: the QK waves' (xm) summed.
ws_offs –

This lane's DMA source of tile t of block b (dma_offs, its index
ws_pv –

The PV waves' steps: step 0 V^T(0) and block 1's DMA, then
ws_pv_block –

PV wave step b: O += V^T P^T of block b - 1 (its V^T va read in
ws_qk –

The QK waves' steps (ws_qk_block), the last block's PV's barriers,
ws_qk_block –

QK wave step b: S of block b's tiles 2 w, 2 w + 1 (stage_block's
ws_run –

Warp-specialized blocks (WSPEC): the QK waves', the PV waves'.
ws_vtile –

V^T latent tile i (of 8: 128 v + 16 i ..) of the slot's block b
xm_row –

The waves' values of row row of row tile r in xm
xp_at –

The P dword of (row tile r, hi / lo h, tiles 4 k ..) of lane.