This document is relevant for: Trn2, Trn3

GDN TKG Kernel API Reference#

Fused single-token gated delta-rule decode step.

Folds l2norm(q,k), beta=sigmoid(b), g=-exp(A_log)*softplus(a+dt_bias), and RMSNormGated(out, z) around the recurrent delta-rule core, then advances the per-head recurrent state by one token.

Background#

The gdn_tkg kernel computes a single-token gated delta-rule (GDN) decode step. It fuses the q/k L2 norm, beta and gate computation, and gated RMSNorm around the recurrence, and advances the per-head recurrent state by one token.

API Reference#

Source code for this kernel API can be found at: gdn_tkg.py

gdn_tkg#

nkilib.experimental.gdn.gdn_tkg(q: nl.NkiTensor, k: nl.NkiTensor, v: nl.NkiTensor, b: nl.NkiTensor, a: nl.NkiTensor, A_log: nl.NkiTensor, dt_bias: nl.NkiTensor, z: nl.NkiTensor, norm_weight: nl.NkiTensor, state_in: nl.NkiTensor) → Tuple[nl.NkiTensor, nl.NkiTensor]#

Fused single-token gated delta-rule decode step.

Parameters:

  • q (nl.NkiTensor) – [BH, K] @ HBM, RAW query (GQA-expanded, not l2-normed).

  • k (nl.NkiTensor) – [BH, K] @ HBM, RAW key (GQA-expanded, not l2-normed).

  • v (nl.NkiTensor) – [BH, V] @ HBM, value.

  • b (nl.NkiTensor) – [BH] @ HBM, RAW beta pre-sigmoid.

  • a (nl.NkiTensor) – [BH] @ HBM, RAW gate input (pre-softplus).

  • A_log (nl.NkiTensor) – [num_v_heads] @ HBM, per-v-head log decay coefficient.

  • dt_bias (nl.NkiTensor) – [num_v_heads] @ HBM, per-v-head softplus bias.

  • z (nl.NkiTensor) – [BH, V] @ HBM, RMSNormGated gate input.

  • norm_weight (nl.NkiTensor) – [V] @ HBM, RMSNormGated weight (gamma).

  • state_in (nl.NkiTensor) – [BH, K, V] @ HBM, incoming recurrent state.

Returns:

[BH, V] @ HBM, RMSNormGated decode output (bf16).

Return type:

nl.NkiTensor

Returns:

[BH, K, V] @ HBM, updated recurrent state (fp32).

Return type:

nl.NkiTensor

Notes:

  • Per-head independent recurrence (no cross-head reduction), so the BH axis is sharded across the SPMD grid; each core writes only its own contiguous head slice. LNC1 is byte-identical to single-core.

  • Do NOT pre-apply exp(g) or the 1/sqrt(K) Q-scale upstream; the kernel applies both internally.

Dimensions:

  • BH: Flattened batch * num_v_heads (heads GQA-expanded, cycling fastest per batch)

  • B: Batch size = BH // num_v_heads

  • num_v_heads: Number of value heads

  • K: Key/state head dimension (<= 128)

This document is relevant for: Trn2, Trn3