This document is relevant for: Trn2, Trn3

GDN Block TKG Kernel API Reference#

Fused GDN decode block: in_proj + conv1d + GQA + gated delta-rule in ONE kernel.

Runs the full Qwen3.5 GDN decode layer except out_proj / all_reduce: in_proj matmuls -> gdn_conv1d_decode (conv + silu + slide-window state) -> split conv_out + GQA-repeat q,k -> gdn_tkg (l2norm/beta/g/RMSNormGated + single-token recurrence) -> scatter to core_out.

Background#

The gdn_block_tkg kernel fuses a full Qwen3.5 gated delta-net (GDN) decode layer, except the output projection and all-reduce, into one kernel: input projections, causal conv1d, GQA expansion, and the gated delta-rule recurrence.

API Reference#

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

gdn_block_tkg#

nkilib.experimental.gdn.gdn_block_tkg(hidden: nl.NkiTensor, W_in_qkv: nl.NkiTensor, W_in_z: nl.NkiTensor, W_in_a: nl.NkiTensor, W_in_b: nl.NkiTensor, conv_weight: nl.NkiTensor, conv_state: nl.NkiTensor, A_log: nl.NkiTensor, dt_bias: nl.NkiTensor, norm_weight: nl.NkiTensor, recurrent_state: nl.NkiTensor) → Tuple[nl.NkiTensor, nl.NkiTensor, nl.NkiTensor]#

Fused GDN decode block: in_proj + conv1d + GQA + gated delta-rule in ONE kernel.

Parameters:

  • hidden (nl.NkiTensor) – [B, 1, H] @ HBM, bf16 decode-step hidden states.

  • W_in_qkv (nl.NkiTensor) – [H, conv_dim] @ HBM, bf16 q|k|v in_proj weight (== nn.Linear.weight.T).

  • W_in_z (nl.NkiTensor) – [H, value_dim] @ HBM, bf16 gate (z) in_proj weight.

  • W_in_a (nl.NkiTensor) – [H, num_v_heads] @ HBM, bf16 gate-input (a) in_proj weight.

  • W_in_b (nl.NkiTensor) – [H, num_v_heads] @ HBM, bf16 beta-input (b) in_proj weight.

  • conv_weight (nl.NkiTensor) – [conv_dim, K_win] @ HBM, fp32 depthwise conv weight (self.conv1d.weight squeezed).

  • conv_state (nl.NkiTensor) – [B, conv_dim, K_win] @ HBM, bf16 conv window (read at entry).

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

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

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

  • recurrent_state (nl.NkiTensor) – [B, num_v_heads, head_k_dim, head_v_dim] @ HBM, fp32 recurrent state (read at entry).

Returns:

[B, 1, value_dim] @ HBM, RMSNormGated decode output.

Return type:

nl.NkiTensor

Returns:

[B, conv_dim, K_win] @ HBM, updated conv window.

Return type:

nl.NkiTensor

Returns:

[B, num_v_heads, head_k_dim, head_v_dim] @ HBM, updated recurrent state.

Return type:

nl.NkiTensor

Notes:

  • Functional: inputs are not mutated in place; the caller aliases the returned conv_state / recurrent_state back into the persistent caches.

  • gdn_conv1d_decode and gdn_tkg are reused verbatim as nested @nki.jit calls; internal nl.hbm scratch carries the inter-stage activations.

  • All dims are derived from tensor shapes (no python-int args).

Dimensions:

  • B: Batch size

  • H: Hidden dimension (must be a multiple of 128)

  • conv_dim: Conv channel count = 2*key_dim + value_dim (q|k|v packed)

  • key_dim: Query/key projection dim = num_k_heads * head_k_dim

  • value_dim: Value projection dim = num_v_heads * head_v_dim

  • num_k_heads: Number of key heads

  • num_v_heads: Number of value heads (GQA target)

  • head_k_dim: Key/state head dim (<= 128)

  • head_v_dim: Value head dim (<= 128)

  • gqa: Query heads per key head = num_v_heads / num_k_heads

  • BH: Flattened B * num_v_heads

This document is relevant for: Trn2, Trn3