`vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.mla.query_wide` ¶

The MLA query at many local heads (TP1 / DEP: all 96 heads on one GPU), and at TP8 / DCP8 (12 heads) from SPLIT_FROM rows, as three kernels instead of MlaQuery's one:

QueryNorm the q_a RMSNorm -> the q_c rows [B, 1536] (dsa.overlap.QueryNorm, bit-identical to MlaQuery's normalization) Q_B q_c @ Q_B^T -> the BF16 query rows [B, H * 192] (split-K projection: one 128-row weight tile per SM at H = 96, the weight streamed once) MlaAbsorb the W_UK absorption of each head's 128 nope columns, the 64 rope columns and the MLA KV insert -> the query rows [B, H, 576]

MlaQuery gives every (head, Q_B slice) its own CTA of a cluster that stages and normalizes every token row itself: at 96 heads that is 768-1536 CTAs, each re-reading the input rows and running the RMSNorm, several waves of a latency-bound chain (75 us at 5 rows, 306 us at 80 rows on GB300). Split, the Q_B weight (56.6 MB at 96 heads) streams through one GEMM wave and the absorption reads W_UK once per 16 tokens.

Numerics are MlaQuery's: Q_B accumulates in FP32 and rounds to BF16, the absorption takes the BF16 nope rows, accumulates in FP32 (BF16 MMA) and rounds to BF16 before the per-tensor FP8 query scale (fp8_ds_mla: the BF16 query); the KV insert is MlaQuery's own (insert / insert_ds / insert_nvfp4).

Classes:

MlaAbsorb –

W_UK absorption, rope columns and MLA KV insert from the Q_B rows. One
MlaQuant –

The query rows from the W_UK bmm's BF16 latent [H, B, 512] (in MlaAbsorb's

Functions:

absorb_bmm –

Whether batch rows of heads local heads absorb by the W_UK bmm.
absorb_splits –

CTAs per (part, head) of MlaAbsorb: each takes every splits-th token tile.
absorb_tokens –

Tokens per MlaAbsorb tile: 16 up to 16 rows, 32 up to 32, else 64; below
fi_heads –

The FlashInfer sparse MLA query heads of heads local heads: 64 < H < 128
split_query –

Whether the MLA query of batch rows at heads local heads runs split
wide_query –

Whether the MLA query of heads local heads runs split at every batch.

`MlaAbsorb` ¶

Bases: MlaQuery

W_UK absorption, rope columns and MLA KV insert from the Q_B rows. One 128-thread CTA per (128 latent columns, head): 384 CTAs at 96 heads, one wave at three per SM. It stages its 32 KB W_UK slice once, then walks the batch in tiles of absorb_tokens tokens (the next tile's rows load while the current one's output stores); warp w owns latent columns 32w..32w+31 of the part for every 16-token MMA tile (BF16 MMA, FP32 accumulation, K = 128). The four parts of a head share its rope columns, every CTA one KV insert per 4H tokens.

One CTA per (part, head, tile) instead ran 1920 CTAs at 320 rows, each a latency-bound chain (W_UK, PDL wait, rows, MMAs, stores): 36 us against 1.5 us of W_UK bytes.

Small head counts (TP8 / DCP8: 12) split each (part, head)'s token tiles over absorb_splits CTAs (grid z), each staging the W_UK slice.

dcp (with q_heads): each CTA also multicasts the query bytes it stores (latent columns, rope columns) into head block rank of every rank's q mailbox, as MlaQuery.publish; QueryWait polls it (no poll here).

Arguments as MlaQuery's with the Q_B rows [B, H * 192] (BF16) in place of the Q_B weight: (packed rows, q_a norm, kv_a norm, Q_B rows, W_UK, MLA cache, slots, 1 / q_scale, k_scale, eps, query rows[, q mailbox, its multicast address, rank], stream).

Methods:

publish_words –

DCP: words (the query bytes of token's head head from column
rope –

The rope columns of the tile's tokens (8 per 16 B chunk, the PARTS CTAs
stage_rows –

cp.async the nope columns of tokens tok0 .. tok0 + T into qn.