lite_boost.ops.recurrent_gated_delta_rule

View Source On AtomGit
lite_boost.ops.recurrent_gated_delta_rule(query, key, value, beta, state, actual_seq_lengths, ssm_state_indices, g, gk, num_accepted_tokens, scale_value=1.0)[source]

Recurrent GatedDeltaRule operator — CANN aclnn-backed recurrent linear attention decode.

Implements the token-by-token recurrent forward pass of the Gated Delta Rule, updating the recurrent state matrix and producing the attention output. Primarily used for decode-phase inference acceleration in hybrid linear attention models such as Qwen3.5.

Algorithm flow (executed sequentially for each token in each batch). The state decay is

\[S = S * \exp(g) * \exp(gk)\]

The memory retrieval is

\[kv\_mem = S^{\top} k\]

The delta update is

\[S = S + k^{\top} ((v - kv\_mem) * \beta)\]

The output is computed as

\[o = S^{\top} q\]

where \(S\) is the recurrent state matrix of shape \((N_v, D_k, D_v)\), storing the key-value associations of linear attention.

Parameters:

  • query (Tensor) – Query tensor of shape \((B, N_k, S, D_k)\), dtype=bfloat16. Must be L2-normalized (L2 norm of each head vector is 1, value range [0, 1]). B=batch_size, N_k=num_key_heads, S=seq_len, D_k=key_dim.

  • key (Tensor) – Key tensor of shape \((B, N_k, S, D_k)\), dtype=bfloat16. Must be L2-normalized (same as query).

  • value (Tensor) – Value tensor of shape \((B, N_v, S, D_v)\), dtype=bfloat16. N_v=num_value_heads, D_v=value_dim. N_v must be divisible by N_k.

  • beta (Tensor) – Delta update step size of shape \((B, N_v, S)\), dtype=bfloat16. Value range [0, 1]. Controls the magnitude of each delta update: a larger beta causes new information to overwrite old memory more aggressively; a smaller beta tends to preserve existing memory.

  • state (Tensor) – Recurrent state pool of shape \((state\_slots, N_v, D_k, D_v)\), dtype=bfloat16. state_slots is the number of state slots in the pool; each slot independently stores the cumulative key-value associations of one sequence, and each token selects its slot via ssm_state_indices (for standard inference each batch occupies one slot, so state_slots usually equals B). D_k is the key dimension (rows), D_v is the value dimension (columns). Can be initialized to zeros for the first call.

  • actual_seq_lengths (Tensor) – Actual sequence lengths of shape \((B)\), dtype=int32. Used for variable-length sequence inference. Each element represents the number of valid tokens in the corresponding batch. E.g., [4, 3, 5] means 3 batches with sequence lengths 4, 3, and 5.

  • ssm_state_indices (Tensor) – State-slot indices of shape \((T)\), dtype=int32, where T = B * S (one entry per flattened token). Each token selects one state slot in the global state pool (dim 0 of state, state_slots slots in total).

  • g (Tensor) – Global decay gate of 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. exp(g) serves as the state decay factor with range (0, 1). The more negative g is, the faster historical information is forgotten. E.g., when g=-1, approximately 37% of the historical state is retained per step.

  • gk (Tensor) – Key-dimension gate of shape \((B, N_v, S, D_k)\), dtype=float32. Must be negative (same no-validation rationale as g). exp(gk) applies per-dimension decay independently along the key dimension, enabling finer-grained memory control. Unlike the global gate g, gk operates element-wise along the D_k dimension.

  • num_accepted_tokens (Tensor) – Number of accepted tokens of shape \((B)\), dtype=int32. Used in speculative decoding and similar scenarios to mark the number of actually accepted (non-rejected) tokens. For standard inference, this is the same as actual_seq_lengths.

  • scale_value (float, optional) – Attention scale factor, default 1.0. Typically set to 1.0 / sqrt(D_k), consistent with standard attention scaling. The query is multiplied by this scale factor before computation.

Returns:

tuple[Tensor, Tensor]

  • out (Tensor) — Attention output of shape \((B, N_v, S, D_v)\), dtype=bfloat16. The linear attention result at each token position.

  • state_out (Tensor) — Updated recurrent state pool with the same shape as state, dtype=bfloat16. Must be passed as state input in the next recurrent step to form a state-passing chain.

Raises:

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

  • ValueError – If input tensor shapes are invalid, i.e. N_v is not an integer multiple of N_k, or ssm_state_indices does not contain exactly one entry per flattened token (B * S entries).

Note

  • This operator only supports the decode phase (token-by-token inference), with sequence length S not exceeding 8. For parallel prefill computation, use the chunk-level operator.

  • Supports grouped recurrent heads where N_v is an integer multiple of N_k.

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

  • The CANN operator stores state internally as \((state\_slots, N_v, D_v, D_k)\) layout (value dimension first). This function automatically performs the layout conversion.

Supported Platforms:

Ascend

Examples

>>> import torch
>>> import lite_boost.ops as lite_ops
>>> device = torch.device("npu:0")
>>> B, N, S, Dk, Dv = 1, 64, 1, 64, 512
>>> 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
>>> state  = torch.zeros(B, N, Dk, Dv, device=device, dtype=torch.bfloat16)
>>> g      = -(torch.rand(B, N, S, device=device) + 0.01)
>>> gk     = -(torch.rand(B, N, S, Dk, device=device) + 0.01)
>>> actual_seq_lengths  = torch.tensor([S], dtype=torch.int32, device=device)
>>> ssm_state_indices   = torch.tensor([0], dtype=torch.int32, device=device)
>>> num_accepted_tokens = torch.tensor([S], dtype=torch.int32, device=device)
>>> output, state_out = lite_ops.recurrent_gated_delta_rule(
...     query, key, value, beta, state,
...     actual_seq_lengths, ssm_state_indices,
...     g, gk, num_accepted_tokens,
...     scale_value=1.0 / (Dk ** 0.5))