This document is relevant for: Trn2, Trn3

KDA TKG Kernel API Reference#

Kimi Delta Attention (KDA) token generation: output + new state.

Advances the KDA recurrence by exactly one token for each of B requests in a single invocation.

Background#

The kda_tkg kernel implements Kimi Delta Attention (KDA) token generation. It advances the KDA recurrence by one token for each request and returns the attention output and the updated recurrent state.

API Reference#

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

kda_tkg#

nkilib.experimental.attention.kda_tkg(q: nl.NkiTensor, k: nl.NkiTensor, v: nl.NkiTensor, g: nl.NkiTensor, beta: nl.NkiTensor, a_log: nl.NkiTensor, dt_bias: nl.NkiTensor, out: nl.NkiTensor, S: nl.NkiTensor, scale: nl.NkiTensor, lower_bound: nl.NkiTensor)#

Kimi Delta Attention (KDA) token generation: output + new state.

Parameters:

  • q (nl.NkiTensor) – Query tensor in HBM for the current token, raw (pre-L2-norm, pre-scale).

  • Shape – [B, dk, num_head]. Dtype: nl.float32

  • k (nl.NkiTensor) – Key tensor in HBM for the current token, raw (pre-L2-norm).

  • Shape – [B, dk, num_head]. Dtype: nl.float32

  • v (nl.NkiTensor) – Value tensor in HBM for the current token.

  • Shape – [B, num_head, dv]. Dtype: nl.float32

  • g (nl.NkiTensor) – Per-channel forget-gate tensor in HBM, raw (pre-activation; the kernel applies lower_bound * sigmoid(exp(a_log) * (g + dt_bias))).

  • Shape – [B, dk, num_head]. Dtype: nl.float32

  • beta (nl.NkiTensor) – Delta-rule beta LOGITS in HBM (the kernel applies the sigmoid).

  • Shape – [B, 1, num_head]. Dtype: nl.float32

  • a_log (nl.NkiTensor) – Per-head log-gate parameter in HBM; exp(a_log) scales the gate pre-activation. A model parameter, so it carries no batch axis.

  • Shape – [1, num_head]. Dtype: nl.float32

  • dt_bias (nl.NkiTensor) – Per-head, per-CHANNEL gate bias in HBM, added to g before the sigmoid. A model parameter, so it carries no batch axis.

  • Shape – [dk, num_head]. Dtype: nl.float32

  • out (nl.NkiTensor) – Caller-allocated destination for this token’s attention output.

  • Shape – [B, num_head, dv]. Dtype: nl.float32

  • S (nl.NkiTensor) – Recurrent state carried in from the previous token, in HBM. IN/OUT: each item’s [dk, dv] slice is read to seed its step, then overwritten with the advanced state.

  • Shape – [B, num_head, dk, dv]. Dtype: nl.float32

  • scale (nl.NkiTensor) – Query scale, normally 1/sqrt(dk); folded into the L2-normed query so the output needs no post-scale.

  • Shape – [1, 1]. Dtype: nl.float32

  • lower_bound (nl.NkiTensor) – Gate lower bound, expected in [-5, 0]. The activated gate is lower_bound * sigmoid(…), hence <= 0, so exp(gate) is in (0, 1] and the decay can never amplify the state.

  • Shape – [1, 1]. Dtype: nl.float32

Returns:

The caller-provided output tensor, filled with this token’s attention output for every request in the batch.

Return type:

nl.NkiTensor

Returns:

[B, num_head, dv]. Dtype: nl.float32

Return type:

nl.NkiTensor

Returns:

The caller-provided state tensor, advanced in place by one token for every request.

Return type:

nl.NkiTensor

Returns:

[B, num_head, dk, dv]. Dtype: nl.float32

Return type:

nl.NkiTensor

Dimensions:

  • B: Number of requests in the batch (unconstrained)

  • num_head: Number of value heads

  • dk: Key/query head dimension, <= 128

This document is relevant for: Trn2, Trn3