Module kda 

Kimi Delta Attention (KDA) — the glm5_next (GLM-5.3-Flash) linear-attention mixer.

Arithmetic contract: memra_reference::kimi_delta_net, pinned by kimi_delta_net_matches_hand_derived_three_token_recurrence. Every step below cites the reference stage it reproduces; the GPU-vs-reference gate is crates/memra-engine/tests/kda_fixture_gpu.rs.

Geometry (research/glm53-flash-bringup-20260827/CENSUS.md): 64 heads x 128, q/k/v all the same width, short conv kernel 4, forget-gate lower bound -5.0. Symmetric widths and no GQA repeat mean channel c == h*head_dim + i IS the (head, dim) pair, so every per-token tensor stays token-major end to end — there is no analogue of GDN’s qkv_to_gdn_repack scatter here.

PREFILL DISPATCH — SEQUENTIAL SCAN, not the chunked UT transform (deliberate). memra_kda_scan_s128 runs prefill and decode alike, which is exactly the shipped GDN arrangement next door: gdn_scan_s128 IS the default prefill path and the chunked WY kernels sit behind MEMRA_GDN_CHUNKED. One kernel for both also keeps the decode==verify dispatch identity that cu/hybrid.cu’s headers require. A chunked twin exists but is SHELVED, ATTRIBUTED-NEGATIVE — it is not a pending tuning follow-up. It was built as L3 of the prefill-gap plan (MEMRA_KDA_CHUNKED, unmerged branch lane/glm5-kda-chunk-scan), and the box prefill census then attributed the wall elsewhere: on a cold 4626-token prime the whole kda family is 221.6 GPU ms of 6598 (3.4%, “confirms L3’s ATTRIBUTED-NEGATIVE: scan ~2.4%”) while mla-prefill-attn owns 75.8% — receipts research/glm53-flash-bringup-20260827/launch-diet-20260830/WINDOW-20260830.md §4 and box-receipts-20260830/census-analysis.txt. No A/B is owed on the scan; a revival needs a new attribution first. The algebra stays banked for that day: it is NOT a transcription of the GDN K1-K5 chain — KDA’s decay is per channel, so the chunk form needs a per-channel cumulative log gate Gcum[t][i] with k scaled by exp(-Gcum) and q by exp(+Gcum) (banked chunk_kimi_delta_attention in research/glm53-flash-bringup-20260827/modular_glm5_next-ref.py), where GDN gets away with one scalar G per (token, head).

CONV FUSION — fused WEIGHTS and a fused RING, per-plane launches. The checkpoint ships three per-plane conv weights; they are concatenated once at load into one [3*qkv, kernel] f32 buffer, because the plan already declares the state carrier fused (StatePlan::Recurrent conv_width = 3*qkv) and that makes a plane’s weight offset and its ring offset the same plane*qkv arithmetic. The three PROJECTIONS stay separate: they are independently quantized tensors, and concatenating them would mean dequantizing to build one matmul. Applying each plane’s taps to its own plane is the fused grouped conv exactly (the reference says so in-line), so nothing is approximated by the split.

Structs

KdaAttnLayer

One loaded KDA mixer. Field names follow the reference’s tensor roles, not the HF spellings.

KdaRowsStash

The rollback stash of one BATCHED verify-rows KDA call (lane/glm5-verify-batch): the per-layer t=K+1 twin of the per-row KdaScanInputs steal. Everything here is either stolen from buffers the call allocated anyway (raws, scan — zero copies) or one small clone per layer per round (ring_snap, 3*qkv*(kernel-1) floats ~ 96 KiB).

KdaScanInputs

The scan-input buffers of one KDA step, STOLEN from the step instead of dropped (lane/glm5-loop-port, port 3 — the module doc’s named GdnStash/ReplaySSM diet): the glm5 verify walk’s rollback checkpoint keeps these ~160 KB of already-allocated buffers per row per layer and retires the per-row 4 MiB recurrent-state clones (~0.95 GiB transient at K=7). Replaying kda_scan over them from a pre-round state snapshot rebuilds the post-row state EXACTLY: each replay is the ORIGINAL t=1 launch re-issued — same kernel, same inputs, same shape — so the rebuilt state is byte-identical to the clone it replaces by construction, not by a numeric argument.

Constants

KDA_HEAD_DIM

The only head width memra_kda_scan_s128 is instantiated for, and the only one glm5_next ships (linear_attn_config.head_dim = 128).

Statics

KDA_FUSED6_BF16_DISPATCHES

Same door, BF16 operand arm (qmatvec_kda6_bf16f32, lane/glm5-decode-diet lever 3). Counted separately so a box A/B on the serving recipe (MEMRA_BF16_MMV=1, where the q8 arm refuses by design) can attribute engagement to the arm that actually ran.

KDA_FUSED6_DISPATCHES

Engagement counter for the fused 6-way projection door (MEMRA_KDA_FUSED_PROJ), the grouped-prefill moe_grouped_prefill_dispatches precedent: gates and box A/B arms count dispatches at the arm’s own call site instead of inferring engagement from a 200.

Functions

kda_attn

STATELESS prefill from a zero conv ring and a zero recurrent state — the arm the logits-only forward paths take. Allocates and discards both state buffers.

kda_attn_decode

T=1 decode step. Same math as a one-token prime; separate conv arm so the fused assemble+conv+roll kernel keeps decode and the spec verify on one dispatch class.

kda_attn_prime

STATEFUL prefill: carries the ring forward and advances the recurrent state from state_in into state_out. Callers own the ping-pong; the two state buffers must be distinct.

kda_decode_cached

One decode step through the cache’s KDA state for layer il.

kda_decode_cached_stash

kda_decode_cached with the step’s scan inputs STOLEN for a rollback stash (loop-port 3; doc on KdaScanInputs). Identical launches — the steal is a move of buffers the step allocated either way.

kda_prime_cached

Stateful prefill of t tokens through the cache’s KDA state for layer il.

kda_scan_replay

Rebuild layer il’s recurrent state to “after row inputs.len()-1” by REPLAYING the stashed scan inputs from the pre-round snapshot snap (loop-port 3, the module-doc diet made concrete): each replay is the original t=1 memra_kda_scan_s128 launch re-issued over the very buffers that step consumed, so the rebuilt state is byte-identical to the per-row clone it replaces BY CONSTRUCTION. The readout is discarded; the conv ring is not touched (the walk still clones it per row — 288 KiB against the 4 MiB ssm plane this retires). The ping-pong rides the resident pair and ends with the rebuilt state under the ssm_state name, matching kda_cached’s own swap discipline.

kda_verify_rollback_rows

Roll layer il back to “after row keep-1” from a BATCHED verify-rows round (lane/glm5-verify-batch; the KdaRowsStash doc states the two-plane contract): restore the pre-round conv ring and re-roll keep raw rows (pure placement), then replay the scan ONCE at T=keep from the pre-round ssm snapshot over the batched inputs. Full accept (keep == rows) never calls this — the resident state IS the state after the last kept row.

kda_verify_rollback_rows_on

kda_verify_rollback_rows over a CALLER-OWNED state plane — the glm5 spec x TP seam (lane/glm5-composition): under MEMRA_GLM5_TP each rank’s shard-geometry conv ring + ssm ping-pong lives in cache.glm5_tp_recur[il][rank] on that rank’s engine, so the rollback restores per rank through this entry with the rank’s own (engine, shard, snapshot, stash) tuple. The cache wrapper above delegates here — one body, byte-for-byte the pre-refactor walk on the plain path.

kda_verify_rows_cached

THE BATCHED VERIFY-ROWS KDA CALL (lane/glm5-verify-batch): one t=K+1 kda_core pass per layer per round, replacing t per-row kda_decode_cached_stash calls. Projections, gates and norms batch m=t through the decode-exact matmul classes (matmul_rows_exact); the conv takes the prefill dispatch (per-token bit-identical to the decode arm’s taps); the recurrence stays SEQUENTIAL inside one memra_kda_scan_s128 launch (the in-kernel T-loop over register-resident state == the chained t=1 program). Per-row bit-identity vs the t=1 chain is held by the walk gates (glm5_tparallel_verify_gpu) and the kernel bit-gates (glm5_verify_batch_gpu).