`vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.kda.specialized` ¶

Persistent warp-specialized KDA core for wide steps (KdaWS).

One CTA per SM walks a contiguous range of (request, head) items with two roles:

the scan group (8 warps) holds one item's [128, 128] FP32 state in registers (thread (g, l): rows g + 32 r, r < 4, keys 4 (l + 8 c) + [0, 4), c < 4) and runs the recurrence of the CTA's items one after another;
two prep warpgroups (4 warps each, thread c = channel c; warpgroup w takes items w, w + 2, ...) stream each item in: the state page by bulk copies into the shared-memory page (once the scan holds the previous item in registers), the projection rows and the convolution window by cp.async; they run the convolution, SiLU, q / k normalization, the F_B gate projection (mma.sync) and the decays into one of their two operand blocks; the scan group then runs the item's gated RMSNorm and output from its recurrence outputs.

So the per-item input and output phases (half the instructions of KdaWide) run on their own warps, overlapped with the recurrence, and the scan warps only scan (setmaxnreg: 168 registers for scan threads, 88 for prep).

Per item the math is KdaWide's: the stock Triton formulas and instructions, BF16 roundings of the convolution outputs, raw gates and recurrence outputs, FP32 state; only FP32 reduction order differs (the recurrence's dot products and the q / k / out norms are summed in another order).

Synchronization (mbarriers, phase = use count & 1 of each barrier): state_full prep -> scan (bulk copy bytes) state_empty scan -> scan thread 0 (item j is in registers: it bulk-copies item j + 1's page) ops_full[blk] prep -> scan (operand block blk = 2 w + (j // 2) % 2 ready) done[blk] scan -> prep (item scanned: the operand block is free)

Classes:

KdaWS –

KdaWide's operands and semantics (tokens q, record: RecoverSSM verify,

Functions:

add2 –

Lanewise a + b of FP32 pairs in one packed instruction (IEEE per lane).
fence_async_shared –

Order this thread's shared-memory writes before later bulk copies.
operand_bytes –

One operand block: q, k, v, decay, out (FP32), the output gate (BF16), beta,
pair_total –

(a0 + b0) + (a1 + b1) of a row's two accumulator pairs, one packed add.
store_bulk –

Shared -> global bulk copy (one bulk group).
sync_threads –

bar.sync with a runtime barrier id.
wait_bulk_reads –

The issued bulk stores have read their shared-memory sources.

`KdaWS` ¶

KdaWide's operands and semantics (tokens q, record: RecoverSSM verify, gate_mode the F_B chain order); ctas persistent CTAs (one per SM).

Methods:

block_sums –

Per-value totals over the 128 prep threads (each thread one channel):
decay –

a = exp(A_log) (tl.exp).
epilogue –

Gated RMSNorm and output of the scanned item (scan threads: channel
item_slots –

(head, request, state slot, conv slot, window offset) of an item.
load_page –

Bulk-copy an item's state page (or complete the phase of a padded row).
row_sum –

The four rows' partials of 8 lanes -> the full sum of row lane_in // 2.
silu –

The convolution's SiLU: x / (1 + exp(-x)), stored BF16.
stage_conv –

cp.async the item's convolution window (prep thread p).
stage_rows –

cp.async the item's projection rows: x (q, k, v), output gate, F_A.