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

K1 kda_pre: one Kimi-K3 KDA layer from the AttnRes seam to the gated core output, for a spec-verify step (R requests x L tokens, S = R L rows).

One launch a layer, BLOCKS CTAs x THREADS threads, in three roles so that no CTA that polls a mailbox has weight loads in flight (a poll's load retires behind every earlier load):

attn CTAs 0 .. 55   slice (128 hidden columns each): prefix (+ delta) ->
                    bf16 prefix out, per-source partial sum of squares and
                    qk dot (PART); gate (token t on CTA t): softmax weights
                    (WGT); mix: the softmax mix, its partial sum of squares
                    (MSQ), every slice's -> output RMSNorm -> bf16 x rows
                    (plain device-scope stores, an XRDY flag a slice)
gemv CTAs 56 ..     in_proj, two 16-row groups each: the first group's
                    weights issued at kernel start, the second's as x lands
                    -> q|k|v|g|f_a|beta (PROJ)
kda tasks           (R x 12 heads x 8 v-chunks) on CTAs 0 ..: inputs that
                    need no hand-off loaded first (metadata, conv window,
                    state rows, f_b rows), then conv update of the owned
                    channels (q / k out through QK), f_b MFMA + lower-bound
                    gate, l2norm, the recurrence on 16 state rows (a state a
                    token written), partial norms (ONRM), gated RMSNorm ->
                    core_out

Rounding follows the original path where it is elementwise (the AttnRes prefix add, every bf16 hand-off between the original kernels); split-K reductions differ only in their order.

Functions:

bf4_words –

Four f32 -> two words of packed bf16.
gemv_loads –

This wave's weight operands of rows 16 rg .. for its K chunks chunks
kda_conv –

Owned channels: the spec conv update of request req's L tokens -> new
kda_gates –

f_a, beta and the head's q / k (one poll batch a thread); g1 = f_b f_a by
kda_norm –

Rows 16 ch ..: gated RMSNorm over the head's 128 (every chunk's partial)
kda_pre –

Host launcher. conv_state [slots, dim, state_len] or a transposed view
kda_recur –

State rows 16 ch .. of head h: the recurrence over the L tokens, a state
kda_stage –

The task's inputs that need no hand-off, issued before its first poll:
ld_bf4 –

bf16 elements i .. i + 4 (i a multiple of 4) as four f32.
load_x –

The bf16 x rows t0 .. t0 + n into LDS rows 0 .. (rows padded).
mix_src –

Token t's source src, slice columns cc .. cc + 4: from vals, or
poll_rows_to_lds –

gemv.rows_to_lds from tagged pairs: columns k0 .. k0 + k of rows
stage_gate –

Token t: the 56 slices' partials -> source weights e_s and their sum.
stage_gates –

The gates of tokens bid, bid + 56, .. (< S) on attn CTA bid.
stage_gemv –

Rows 16 rg .. of in_proj against the LDS x (tokens t0 .. t0 + n) -> bf16
stage_mix –

Columns 128 task ..: the softmax mix of the sources (4 columns a thread,
stage_slice –

Columns 128 task ..: every token's sources into LDS vals (4 columns a
step_tag –

This launch's mailbox tag: the step's epoch (a device counter the runner
vals_all –

Whether every (token, source)'s slice stays in LDS vals (else the
wait_x –

Every slice's x ready.