lite_boost.ops.chunk_gated_delta_rule

View Source On AtomGit
lite_boost.ops.chunk_gated_delta_rule(query, key, value, beta, initial_state, actual_seq_lengths, g=None, scale_value=1.0)[source]

Chunked (prefill) Gated Delta Rule operator on NPU (ascend_a2 op).

Wraps torch.ops.lite_boost.chunk_gated_delta_rule, which maps to the CANN AscendC operator aclnnChunkGatedDeltaRule on A2. The user-facing BNSD layout is converted to the TND layout expected by the CANN op; query/key/value/beta/initial_state are cast to the low dtype (bfloat16 or float16, following the input dtype; fp32/other inputs default to bfloat16) while the optional gate g stays float32, and actual_seq_lengths is passed through directly (T = sum(actual_seq_lengths)).

At each time step \(t\) the Gated Delta Rule computes the new recurrent state and the attention output as

\[S_t = \alpha_t S_{t-1} + \beta_t (v_t - \alpha_t S_{t-1} k_t) k_t^{\top}\]
\[o_t = S_t q_t \cdot scale\]

where \(\alpha_t = \exp(g_t)\) is the decay factor (with g omitted the decay is disabled, i.e. \(\alpha_t = 1\)), and \(\beta_t\) is the delta update step size. This operator is the chunked (blocked-parallel) implementation of the recurrence above, which is more efficient than the token-by-token form on long sequences and thus suited for the prefill phase; it produces the output at every step as well as the final state.

Supported only on A2; 300I Duo is not supported.

Parameters:

  • query (Tensor) – Query tensor with shape \((B, N_k, S, D_k)\). Cast to the low dtype before computation.

  • key (Tensor) – Key tensor with shape \((B, N_k, S, D_k)\). Cast to the low dtype before computation.

  • value (Tensor) – Value tensor with shape \((B, N_v, S, D_v)\). N_v must be an integer multiple of N_k (GQA: each key head maps to N_v / N_k value heads). Cast to the low dtype before computation.

  • beta (Tensor) – Delta update step size with shape \((B, N_v, S)\), in the range [0, 1]. Cast to the low dtype before computation.

  • initial_state (Tensor) – Incoming recurrent state with shape \((B, N_v, D_k, D_v)\) (transposed to the op's value-first [B, N_v, D_v, D_k] layout).

  • actual_seq_lengths (Tensor) – Per-batch token counts with shape \((B)\), dtype int32. Each element must be within [0, S] — the BNSD S dim is the padded per-batch length and only the first actual_seq_lengths[b] tokens of each batch are valid. Non-uniform lengths are supported: the binding packs the valid tokens into TND tensors with T = sum(actual_seq_lengths), and output positions beyond each batch's length are zero-filled.

  • g (Tensor, optional) – Global decay gate with shape \((B, N_v, S)\), dtype float32, must be negative. This range is not validated: validating it would introduce extra reduction ops (min/max over the tensor) and add noticeable overhead in the inference path. Out-of-range values do not raise but yield meaningless results. None disables the decay gate (hasGamma=0 path). Default: None.

  • scale_value (float, optional) – Attention scale applied to query. Default: 1.0.

Returns:

tuple[Tensor, Tensor]

  • out (Tensor) — Attention output with shape \((B, N_v, S, D_v)\), same dtype as the low-dtype cast of the inputs (bfloat16 by default). Valid for the first actual_seq_lengths[b] positions of each batch; padded positions are zero-filled.

  • final_state (Tensor) — Updated recurrent state with shape \((B, N_v, D_k, D_v)\), same dtype as out.

Raises:

RuntimeError – If the input tensor shapes, dtypes or devices are invalid, or if the CANN operator execution fails.

Note

  • This operator is supported on A2 only; 300I Duo is not supported (there the registered vendor exposes a different aclnn signature, the ascend_300iduo op).

  • All input tensors must reside on the same NPU device.

  • N_v (number of value heads) must be an integer multiple of N_k (number of key/query heads); each key head serves N_v / N_k value heads.

  • The BNSD S dim is the padded per-batch sequence length; only the first actual_seq_lengths[b] tokens of each batch are valid. Non-uniform actual_seq_lengths are packed/unpacked automatically, and padded output positions are zero-filled.

  • The CANN op accepts both bfloat16 and float16 for q/k/v/beta/state via DataTypeList; the input dtype is followed and fp32/other inputs default to bfloat16. The optional gate g is always float32.

Supported Platforms:

Ascend

Examples

>>> import torch
>>> import lite_boost.ops as lite_ops
>>> device = torch.device("npu:0")
>>> B, N, S, Dk, Dv = 1, 8, 16, 32, 64
>>> query = torch.randn(B, N, S, Dk, device=device, dtype=torch.bfloat16)
>>> key = torch.randn(B, N, S, Dk, device=device, dtype=torch.bfloat16)
>>> value = torch.randn(B, N, S, Dv, device=device, dtype=torch.bfloat16)
>>> beta = torch.rand(B, N, S, device=device, dtype=torch.bfloat16) * 0.9 + 0.05
>>> initial_state = torch.zeros(B, N, Dk, Dv, device=device, dtype=torch.bfloat16)
>>> actual_seq_lengths = torch.tensor([S], dtype=torch.int32, device=device)
>>> out, final_state = lite_ops.chunk_gated_delta_rule(
...     query, key, value, beta, initial_state, actual_seq_lengths)
>>> print(out.shape)
torch.Size([1, 8, 16, 64])
>>> print(final_state.shape)
torch.Size([1, 8, 32, 64])
>>> print(out.dtype)
torch.bfloat16