dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

ref(upstream): 搬运 3 大 GDN 上游仓库 — FLA naive ops + vllm GDN 子树 + xllm C++ 参考

Browse Source 
来源:
  1. fla-org/flash-linear-attention (5538 stars)
     → upstream_ref/fla/ops/gated_delta_rule/naive.py (正确的纯 PyTorch GDN)
     → upstream_ref/fla/ops/gated_delta_rule/chunk.py (Triton chunk kernel)
     → upstream_ref/fla/layers/gated_deltanet.py (层集成)

  2. vllm-project/vllm main (88717 stars)
     → upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py (1751行, Qwen3.5 原生 GDN)
     → upstream_ref/vllm_gdn/ops/causal_conv1d.py (1289行, 正确的 Conv1d)
     → upstream_ref/vllm_gdn/third_party/ops/ (FLA Triton ops vendored)
     → upstream_ref/vllm_gdn/models/qwen3_5.py (vllm 最新 Qwen3.5 模型)

  3. Deep-Spark/xllm (BI-V100 硬件厂商)
     → upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp (576行)
     → upstream_ref/xllm_latest/core/kernels/npu/npu_causal_conv1d.cpp
     → upstream_ref/xllm_latest/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp

目的: 修复 corex_gdn.py Conv1d groups 接口不匹配问题
  错误: conv1d_weight shape (2560,1,4) 被当成 (num_k_heads,1,4) 索引
  conv_dim = key_dim*2 + value_dim = 10240, TP=4 后 2560
  FLA naive.py 和 vllm qwen_gdn_linear_attn.py 有正确的实现可直接对接
This commit is contained in:
Claude
2026-08-11 03:55:50 +00:00
parent 5862708b32
commit 6cdf2ec87b
46 changed files with 13857 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files

364
upstream_ref/fla/layers/gated_deltanet.py Normal file
View File

@@ -0,0 +1,364 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
from __future__ import annotations
import math
import warnings
from typing import TYPE_CHECKING
import torch
import torch.nn as nn
from einops import rearrange
from torch.nn import functional as F
from fla.layers.utils import get_layer_cache, repad_hidden_states, unpad_hidden_states, update_layer_cache
from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
from fla.modules.convolution import causal_conv1d
from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
if TYPE_CHECKING:
from transformers.processing_utils import Unpack
from fla.models.utils import Cache
class GatedDeltaNet(nn.Module):
"""
Gated Delta Networks (GDN) layer implementation.
Reference: `Gated Delta Networks: Improving Mamba2 with Delta Rule <https://arxiv.org/abs/2412.06464>`_
Similar to Mamba2, each layer contains around 6*hidden_size*hidden_size parameters.
Parameter allocation when use_gate=True:
- 0.75 * hidden_size * hidden_size for the q_proj and k_proj each
- 1.5 * hidden_size * hidden_size for the v_proj, g_proj and o_proj each
- Others are ignorably small.
- In total = 0.75 * 2 + 1.5 * 3 = 6 * hidden_size * hidden_size
NOTE: num_heads * head_dim = 0.75 * hidden_size, please make sure to set the correct num_heads and head_dim.
Parameter allocation when use_gate=False:
- 1 * hidden_size * hidden_size for the q_proj and k_proj each
- 2 * hidden_size * hidden_size for the v_proj and o_proj each
- Others are ignorably small.
- In total = 1 * 2 + 2 * 2 = 6 * hidden_size * hidden_size
Args:
hidden_size (int, Optional):
The hidden size of the input. Default: 2048.
expand_v (float, Optional):
The expansion ratio for the value dimension. Default: 2.0.
head_dim (int, Optional):
The dimension of each head. Default: 256.
num_heads (int, Optional):
The number of heads. Default: 6.
num_v_heads (int, Optional):
The number of heads for the value projection, equal to `num_heads` if `None`.
GVA (Grouped Value Attention) is applied if `num_v_heads` > `num_heads`,
where `num_v_heads` must be divisible by `num_heads`.
The kernels natively support GVA by mapping multiple value heads to each query/key head.
Default: `None`.
mode (str, Optional):
Which Gated DeltaNet kernel to use.
Currently available: `chunk` and `fused_recurrent`.
Default: `chunk`.
use_gate (bool, Optional):
Whether to use output gate. Default: `True`.
use_short_conv (bool, Optional):
Whether to use short convolutions. Default: `True`.
allow_neg_eigval (bool, Optional):
Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2.
See reference:
`Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues <https://arxiv.org/abs/2411.12537>`_
conv_size (int, Optional):
The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
conv_bias (bool, Optional):
Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
layer_idx (int, Optional):
The index of the layer. Default: None.
norm_eps (float, Optional):
The epsilon value for the normalization layer. Default: 1e-5.
"""
def __init__(
self,
hidden_size: int = 2048,
expand_v: float = 2,
head_dim: int = 256,
num_heads: int = 6,
num_v_heads: int = None,
mode: str = 'chunk',
use_gate: bool = True,
use_short_conv: bool = True,
allow_neg_eigval: bool = False,
conv_size: int = 4,
conv_bias: bool = False,
layer_idx: int = None,
norm_eps: float = 1e-5,
**kwargs,
) -> GatedDeltaNet:
super().__init__()
self.mode = mode
self.allow_neg_eigval = allow_neg_eigval
self.hidden_size = hidden_size
self.expand_v = expand_v
self.use_gate = use_gate
self.use_short_conv = use_short_conv
self.conv_size = conv_size
self.conv_bias = conv_bias
self.head_dim = head_dim
self.num_heads = num_heads
self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads
self.head_k_dim = head_dim
self.head_v_dim = int(self.head_dim * self.expand_v)
self.key_dim = int(self.num_heads * self.head_k_dim)
self.value_dim = int(self.num_v_heads * self.head_v_dim)
self.layer_idx = layer_idx
# Consistency check: Ensure expand_v produces integer values
if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5):
raise ValueError(
f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. "
f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.",
)
if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0:
raise ValueError(
f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.",
)
if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5):
raise ValueError(
f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. "
f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.",
)
assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False)
self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False)
A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16)
self.A_log = nn.Parameter(torch.log(A))
self.A_log._no_weight_decay = True
# hard coded for now
dt_min = 0.001
dt_max = 0.1
dt_init_floor = 1e-4
dt = torch.exp(
torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min))
+ math.log(dt_min),
)
dt = torch.clamp(dt, min=dt_init_floor)
# Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
inv_dt = dt + torch.log(-torch.expm1(-dt))
self.dt_bias = nn.Parameter(inv_dt)
# Just to be explicit. Without this we already don't put wd on dt_bias because of the check
# name.endswith("bias") in param_grouping.py
self.dt_bias._no_weight_decay = True
if use_short_conv:
self.conv_size = conv_size
self.q_conv1d = ShortConvolution(
hidden_size=self.key_dim,
kernel_size=conv_size,
bias=conv_bias,
activation='silu',
)
self.k_conv1d = ShortConvolution(
hidden_size=self.key_dim,
kernel_size=conv_size,
bias=conv_bias,
activation='silu',
)
self.v_conv1d = ShortConvolution(
hidden_size=self.value_dim,
kernel_size=conv_size,
bias=conv_bias,
activation='silu',
)
else:
warnings.warn(
"ShortConvolution is crucial to the performance. "
"Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.",
)
if use_gate:
self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps)
else:
self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32)
self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
def _use_fused_qkv_conv(
self,
last_state: dict | None,
use_cache: bool | None,
cu_seqlens: torch.Tensor | None,
) -> bool:
# The dense no-cache q/k/v short convolutions can collapse into a single
# causal_conv1d only when there is no cache/varlen state and the three convs
# share the same backend, activation, and kernel size.
if not (self.use_short_conv and last_state is None and not use_cache and cu_seqlens is None):
return False
return (
self.q_conv1d.backend == self.k_conv1d.backend == self.v_conv1d.backend and
self.q_conv1d.activation == self.k_conv1d.activation == self.v_conv1d.activation and
self.q_conv1d.kernel_size == self.k_conv1d.kernel_size == self.v_conv1d.kernel_size
)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
past_key_values: Cache | None = None,
use_cache: bool | None = False,
output_attentions: bool | None = False,
**kwargs: Unpack[dict],
) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
if attention_mask is not None:
assert len(attention_mask.shape) == 2, (
"Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
"for padding purposes (0 indicating padding). "
"Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
)
batch_size, q_len, _ = hidden_states.shape
# change to inference mode.
mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
if self.training:
assert mode == 'chunk', "Only chunk mode is supported in training."
last_state = get_layer_cache(self, past_key_values)
cu_seqlens = kwargs.get('cu_seqlens')
hidden_states, indices, cu_seqlens = unpad_hidden_states(hidden_states, cu_seqlens, attention_mask, q_len)
conv_state_q, conv_state_k, conv_state_v = None, None, None
if self._use_fused_qkv_conv(last_state, use_cache, cu_seqlens):
qkv = torch.cat(
[
self.q_proj(hidden_states),
self.k_proj(hidden_states),
self.v_proj(hidden_states),
],
dim=-1,
)
qkv_weight = torch.cat(
[
self.q_conv1d.weight.squeeze(1),
self.k_conv1d.weight.squeeze(1),
self.v_conv1d.weight.squeeze(1),
],
dim=0,
)
if self.conv_bias:
qkv_bias = torch.cat([self.q_conv1d.bias, self.k_conv1d.bias, self.v_conv1d.bias], dim=0)
else:
qkv_bias = None
qkv, _ = causal_conv1d(
x=qkv,
weight=qkv_weight,
bias=qkv_bias,
activation=self.q_conv1d.activation,
backend=self.q_conv1d.backend,
)
q, k, v = torch.split(qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
elif self.use_short_conv:
if last_state is not None:
conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
q, conv_state_q = self.q_conv1d(
x=self.q_proj(hidden_states),
cache=conv_state_q,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
k, conv_state_k = self.k_conv1d(
x=self.k_proj(hidden_states),
cache=conv_state_k,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
v, conv_state_v = self.v_conv1d(
x=self.v_proj(hidden_states),
cache=conv_state_v,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
else:
q = F.silu(self.q_proj(hidden_states))
k = F.silu(self.k_proj(hidden_states))
v = F.silu(self.v_proj(hidden_states))
q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
beta = self.b_proj(hidden_states)
recurrent_state = last_state['recurrent_state'] if last_state is not None else None
if mode == 'chunk':
o, recurrent_state = chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=self.a_proj(hidden_states),
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=use_cache,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
allow_neg_eigval=self.allow_neg_eigval,
state_v_first=True,
cu_seqlens=cu_seqlens,
)
elif mode == 'fused_recurrent':
o, recurrent_state = fused_recurrent_gated_delta_rule(
q=q,
k=k,
v=v,
g=self.a_proj(hidden_states),
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=use_cache,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
allow_neg_eigval=self.allow_neg_eigval,
state_v_first=True,
cu_seqlens=cu_seqlens,
)
else:
raise NotImplementedError(f"Not supported mode `{mode}`.")
update_layer_cache(
self,
past_key_values,
recurrent_state=recurrent_state,
conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
offset=q_len,
)
if self.use_gate:
g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
o = self.o_norm(o, g)
else:
o = self.o_norm(o)
o = rearrange(o, 'b t h d -> b t (h d)')
o = self.o_proj(o)
o = repad_hidden_states(o, indices, batch_size, q_len)
return o, None, past_key_values

52
upstream_ref/fla/modules/__init__.py Normal file
View File

@@ -0,0 +1,52 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
from fla.modules.convolution import ImplicitLongConvolution, LongConvolution, ShortConvolution
from fla.modules.fused_bitlinear import BitLinear, FusedBitLinear
from fla.modules.fused_cross_entropy import FusedCrossEntropyLoss
from fla.modules.fused_kl_div import FusedKLDivLoss
from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
from fla.modules.fused_norm_gate import (
FusedLayerNormGated,
FusedLayerNormSwishGate,
FusedLayerNormSwishGateLinear,
FusedRMSNormGated,
FusedRMSNormSwishGate,
FusedRMSNormSwishGateLinear,
)
from fla.modules.l2norm import L2Norm
from fla.modules.layernorm import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear
from fla.modules.mlp import GatedMLP
from fla.modules.rotary import RotaryEmbedding
from fla.modules.token_shift import TokenShift
__all__ = [
'BitLinear',
'FusedBitLinear',
'FusedCrossEntropyLoss',
'FusedKLDivLoss',
'FusedLayerNormGated',
'FusedLayerNormSwishGate',
'FusedLayerNormSwishGateLinear',
'FusedLinearCrossEntropyLoss',
'FusedRMSNormGated',
'FusedRMSNormSwishGate',
'FusedRMSNormSwishGateLinear',
'GatedMLP',
'GroupNorm',
'GroupNormLinear',
'ImplicitLongConvolution',
'L2Norm',
'LayerNorm',
'LayerNormLinear',
'LongConvolution',
'RMSNorm',
'RMSNormLinear',
'RotaryEmbedding',
'ShortConvolution',
'TokenShift',
]

42
upstream_ref/fla/modules/convolution.py Normal file
View File

@@ -0,0 +1,42 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
from fla.modules.conv import (
ImplicitLongConvolution,
LongConvolution,
PositionalEmbedding,
ShortConvolution,
causal_conv1d,
fft_conv,
)
from fla.modules.conv.cp import CausalConv1dFunctionCP, causal_conv1d_cp
from fla.modules.conv.cuda import FastCausalConv1dFn, fast_causal_conv1d_fn
from fla.modules.conv.triton import (
CausalConv1dFunction,
causal_conv1d_bwd,
causal_conv1d_fwd,
causal_conv1d_update,
causal_conv1d_update_states,
)
__all__ = [
'CausalConv1dFunction',
'CausalConv1dFunctionCP',
'FastCausalConv1dFn',
'ImplicitLongConvolution',
'LongConvolution',
'PositionalEmbedding',
'ShortConvolution',
'causal_conv1d',
'causal_conv1d_bwd',
'causal_conv1d_cp',
'causal_conv1d_fwd',
'causal_conv1d_update',
'causal_conv1d_update_states',
'fast_causal_conv1d_fn',
'fft_conv',
]

91
upstream_ref/fla/ops/__init__.py Normal file
View File

@@ -0,0 +1,91 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
from .abc import chunk_abc
from .attn import parallel_attn
from .attnres import fused_attnres
from .based import fused_chunk_based, parallel_based
from .comba import chunk_comba, fused_recurrent_comba
from .delta_rule import chunk_delta_rule, fused_chunk_delta_rule, fused_recurrent_delta_rule
from .forgetting_attn import parallel_forgetting_attn
from .gated_delta_rule import chunk_gated_delta_rule, chunk_gdn, fused_recurrent_gated_delta_rule, fused_recurrent_gdn
from .generalized_delta_rule import (
chunk_dplr_delta_rule,
chunk_iplr_delta_rule,
fused_recurrent_dplr_delta_rule,
fused_recurrent_iplr_delta_rule,
)
from .gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla
from .gsa import chunk_gsa, fused_recurrent_gsa
from .hgrn import fused_recurrent_hgrn
from .kda import chunk_kda, fused_recurrent_kda
from .lightning_attn import chunk_lightning_attn, fused_recurrent_lightning_attn
from .linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn
from .log_linear_attn import chunk_log_linear_attn
from .mesa_net import chunk_mesa_net
from .nsa import parallel_nsa
from .parallax import parallel_parallax
from .path_attn import parallel_path_attn
from .retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention
from .rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6
from .rwkv7 import chunk_rwkv7, fused_recurrent_rwkv7
from .simple_gla import chunk_simple_gla, fused_chunk_simple_gla, fused_recurrent_simple_gla, parallel_simple_gla
from .wall_attn import parallel_wall_attn, parallel_wall_attn_decode
__all__ = [
'chunk_abc',
'chunk_comba',
'chunk_delta_rule',
'chunk_dplr_delta_rule',
'chunk_gated_delta_rule',
'chunk_gdn',
'chunk_gla',
'chunk_gsa',
'chunk_iplr_delta_rule',
'chunk_kda',
'chunk_lightning_attn',
'chunk_linear_attn',
'chunk_log_linear_attn',
'chunk_mesa_net',
'chunk_retention',
'chunk_rwkv6',
'chunk_rwkv7',
'chunk_simple_gla',
'fused_attnres',
'fused_chunk_based',
'fused_chunk_delta_rule',
'fused_chunk_gla',
'fused_chunk_linear_attn',
'fused_chunk_retention',
'fused_chunk_simple_gla',
'fused_recurrent_comba',
'fused_recurrent_delta_rule',
'fused_recurrent_dplr_delta_rule',
'fused_recurrent_gated_delta_rule',
'fused_recurrent_gdn',
'fused_recurrent_gla',
'fused_recurrent_gsa',
'fused_recurrent_hgrn',
'fused_recurrent_iplr_delta_rule',
'fused_recurrent_kda',
'fused_recurrent_lightning_attn',
'fused_recurrent_linear_attn',
'fused_recurrent_retention',
'fused_recurrent_rwkv6',
'fused_recurrent_rwkv7',
'fused_recurrent_simple_gla',
'parallel_attn',
'parallel_based',
'parallel_forgetting_attn',
'parallel_nsa',
'parallel_parallax',
'parallel_path_attn',
'parallel_retention',
'parallel_simple_gla',
'parallel_wall_attn',
'parallel_wall_attn_decode',
]

17
upstream_ref/fla/ops/gated_delta_rule/__init__.py Normal file
View File

@@ -0,0 +1,17 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
from .chunk import chunk_gated_delta_rule, chunk_gdn
from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gdn
from .naive import naive_chunk_gated_delta_rule, naive_recurrent_gated_delta_rule
__all__ = [
"chunk_gated_delta_rule", "chunk_gdn",
"fused_recurrent_gated_delta_rule", "fused_recurrent_gdn",
"naive_chunk_gated_delta_rule",
"naive_recurrent_gated_delta_rule",
]

20
upstream_ref/fla/ops/gated_delta_rule/backends/__init__.py Normal file
View File

@@ -0,0 +1,20 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
"""GDR backends."""
from fla.ops.backends import BackendRegistry, dispatch
from fla.ops.gated_delta_rule.backends.flash_qla import FlashQLABackend
from fla.ops.gated_delta_rule.backends.triton_ascend import TritonAscendGDNBackend
gdr_registry = BackendRegistry("gated_delta_rule")
gdr_registry.register(TritonAscendGDNBackend())
gdr_registry.register(FlashQLABackend())
__all__ = ['dispatch', 'gdr_registry']

132
upstream_ref/fla/ops/gated_delta_rule/backends/flash_qla.py Normal file
View File

@@ -0,0 +1,132 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
#
# Copyright (c) 2026 Qwen Team, Alibaba Cloud
"""FlashQLA backend for chunk_gated_delta_rule."""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from fla.ops.backends import BaseBackend
from fla.utils import IS_NVIDIA_HOPPER, IS_NVIDIA_SM100, IS_NVIDIA_SM120
if TYPE_CHECKING:
from fla.ops.cp import FLACPContext
def _needs_backward(*tensors: torch.Tensor | None) -> bool:
"""Whether autograd would later require a backward pass over these inputs."""
if not torch.is_grad_enabled():
return False
return any(isinstance(t, torch.Tensor) and t.requires_grad for t in tensors)
class FlashQLABackend(BaseBackend):
"""Copyright (c) 2026 Qwen Team, Alibaba Cloud
Fused TileLang forward and backward with intra-card CP (replaces the multi-kernel Triton path).
https://github.com/QwenLM/FlashQLA
SM90/SM100/SM103 run both directions. SM120 (consumer/workstation Blackwell) ships a
bfloat16 forward kernel only, so it is dispatched exclusively for grad-free bf16 calls
(inference, frozen weights) and falls back to Triton otherwise.
Disable with ``FLA_FLASH_QLA=0``.
"""
backend_type = "flash_qla"
package_name = "flash_qla"
env_var = "FLA_FLASH_QLA"
default_enable = True
priority = 3
def chunk_gated_delta_rule_verifier(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cu_seqlens_cpu: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
**kwargs,
) -> tuple[bool, str | None]:
if not (IS_NVIDIA_HOPPER or IS_NVIDIA_SM100 or IS_NVIDIA_SM120):
return False, "FlashQLA requires NVIDIA SM90, SM100/SM103 or SM120"
if IS_NVIDIA_SM120 and _needs_backward(q, k, v, g, beta, initial_state):
return False, "FlashQLA on SM120 implements the forward pass only, but an input requires grad"
if q.dtype != torch.float16 and q.dtype != torch.bfloat16:
return False, f"FlashQLA requires dtype float16 or bfloat16, got {q.dtype}"
if not (q.dtype == k.dtype == v.dtype):
return False, f"FlashQLA requires q, k, v to have the same dtype, got {q.dtype}, {k.dtype}, {v.dtype}"
# NOTE: the masked tail-store in FlashQLA's blackwell_sm120 forward kernel emits
# tl::pack_float16x4 on cutlass::half_t, which fails to compile under nvcc.
if IS_NVIDIA_SM120 and q.dtype == torch.float16:
return False, "FlashQLA's SM120 forward kernel does not compile for float16"
if q.shape[-1] != 128:
return False, f"FlashQLA requires K=128, got {q.shape[-1]}"
if v.shape[-1] != 128:
return False, f"FlashQLA requires V=128, got {v.shape[-1]}"
if kwargs.get('use_gate_in_kernel'):
return False, "FlashQLA does not support use_gate_in_kernel"
if use_beta_sigmoid_in_kernel:
return False, "FlashQLA does not support use_beta_sigmoid_in_kernel"
if allow_neg_eigval:
return False, "FlashQLA does not support allow_neg_eigval"
if 'transpose_state_layout' in kwargs:
return False, "FlashQLA does not support the deprecated transpose_state_layout"
if cp_context is not None:
return False, "FlashQLA does not support inter-card context parallel"
return True, None
def chunk_gated_delta_rule(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cu_seqlens_cpu: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
**kwargs,
):
import flash_qla
return flash_qla.chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=g,
beta=beta,
scale=scale,
initial_state=initial_state,
output_final_state=output_final_state,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
state_v_first=state_v_first,
cu_seqlens=cu_seqlens,
auto_cp=True,
)

96
upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/__init__.py Normal file
View File

@@ -0,0 +1,96 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
"""Triton-Ascend Ascend NPU backend for GDN gated_delta_rule ops."""
from __future__ import annotations
import torch
from fla.ops.backends import BaseBackend
class TritonAscendGDNBackend(BaseBackend):
"""Ascend NPU backend for GDN gate and WY-representation kernels."""
backend_type = "triton_ascend"
package_name = None
env_var = None
priority = 0
@classmethod
def is_available(cls) -> bool:
from fla.utils import IS_NPU
return IS_NPU
def gdn_gate_fwd_verifier(self, *args, **kwargs):
return True, None
def gdn_gate_fwd(self, *args, **kwargs):
from fla.ops.gated_delta_rule.backends.triton_ascend.gate import gdn_gate_fwd_npu
return gdn_gate_fwd_npu(*args, **kwargs)
def gdn_gate_chunk_cumsum_verifier(self, *args, **kwargs):
return True, None
def gdn_gate_chunk_cumsum(self, *args, **kwargs):
from fla.ops.gated_delta_rule.backends.triton_ascend.gate import gdn_gate_chunk_cumsum_npu
return gdn_gate_chunk_cumsum_npu(*args, **kwargs)
def gdn_gate_bwd_verifier(self, *args, **kwargs):
return True, None
def gdn_gate_bwd(self, *args, **kwargs):
from fla.ops.gated_delta_rule.backends.triton_ascend.gate import gdn_gate_bwd_npu
return gdn_gate_bwd_npu(*args, **kwargs)
def recompute_w_u_fwd_verifier(
self,
k,
v,
beta,
A,
g=None,
cu_seqlens=None,
chunk_indices=None,
) -> tuple[bool, str | None]:
from fla.utils import IS_NPU
if not IS_NPU:
return False, "not running on NPU"
if k.device.type != "npu":
return False, "input device is not NPU"
if all(t.dtype in (torch.float32, torch.float16, torch.bfloat16)
for t in (k, v, beta, A)):
return True, None
return False, "unsupported dtype for NPU recompute_w_u_fwd"
def recompute_w_u_fwd(
self,
k,
v,
beta,
A,
g=None,
cu_seqlens=None,
chunk_indices=None,
):
from fla.ops.gated_delta_rule.backends.triton_ascend.wy_fast import recompute_w_u_fwd_npu
return recompute_w_u_fwd_npu(k, v, beta, A, g, cu_seqlens, chunk_indices)
def prepare_wy_repr_bwd_verifier(self, *args, **kwargs):
return True, None
def prepare_wy_repr_bwd(self, *args, **kwargs):
from fla.ops.gated_delta_rule.backends.triton_ascend.wy_fast import prepare_wy_repr_bwd_npu
return prepare_wy_repr_bwd_npu(*args, **kwargs)
def chunk_gated_delta_rule_fwd_intra_verifier(self, *args, **kwargs):
return True, None
def chunk_gated_delta_rule_fwd_intra(self, *args, **kwargs):
from fla.ops.gated_delta_rule.backends.triton_ascend.chunk_fwd import chunk_gated_delta_rule_fwd_intra_npu
return chunk_gated_delta_rule_fwd_intra_npu(*args, **kwargs)

62
upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/chunk_fwd.py Normal file
View File

@@ -0,0 +1,62 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
"""chunk_gated_delta_rule_fwd_intra adapted for triton-ascend on Ascend NPU."""
from __future__ import annotations
import torch
from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd
from fla.ops.utils import prepare_chunk_indices, solve_tril
from fla.utils import input_guard
@input_guard
def chunk_gated_delta_rule_fwd_intra_npu(
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_size: int = 64,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
if chunk_size not in (16, 32, 64):
raise ValueError(f'`chunk_size` must be 16, 32, or 64, got {chunk_size}.')
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
# Unfused kkt + solve_tril path to stay within UB budget.
A = chunk_scaled_dot_kkt_fwd(
k=k,
g=g,
beta=beta,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_size=BT,
output_dtype=torch.float32,
)
A = solve_tril(
A=A,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
output_dtype=k.dtype,
)
w, u = recompute_w_u_fwd(
k=k,
v=v,
beta=beta,
A=A,
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
return w, u, A

375
upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/gate.py Normal file
View File

@@ -0,0 +1,375 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
"""GDN gate kernels adapted for triton-ascend on Ascend NPU."""
from __future__ import annotations
import torch
import triton
import triton.language as tl
from fla.ops.utils.index import prepare_chunk_indices
from fla.ops.utils.op import exp
from fla.ops.utils.softplus import softplus
from fla.utils import input_guard
from fla.utils.ascend_ub_manager import (
ASCEND_MAX_GRID_DIM,
compute_ub_block_size,
max_grid_axis_chunks,
)
_NUM_WARPS = 4
# Peak live fp32 vectors: input + output (+ bias path).
_GATE_FWD_MEM_MULT = 3.0
_GATE_BWD_MEM_MULT = 5.0
_SAFETY_MARGIN = 0.85
_FALLBACK_BT = 32
_FALLBACK_BT_FWD = 64
def _get_gate_fwd_bt(T: int) -> int:
return compute_ub_block_size(
T,
_GATE_FWD_MEM_MULT,
safety_margin=_SAFETY_MARGIN,
dtype_size=4,
fallback=_FALLBACK_BT_FWD,
desired=min(triton.next_power_of_2(T), _FALLBACK_BT_FWD),
)
def _get_gate_bwd_bt(T: int) -> int:
return compute_ub_block_size(
T,
_GATE_BWD_MEM_MULT,
safety_margin=_SAFETY_MARGIN,
dtype_size=4,
fallback=_FALLBACK_BT,
desired=min(triton.next_power_of_2(T), _FALLBACK_BT),
)
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
})
@triton.jit(do_not_specialize=['T'])
def gdn_gate_fwd_kernel_npu(
g,
A_log,
dt_bias,
yg,
T,
H: tl.constexpr,
BT: tl.constexpr,
HAS_BIAS: tl.constexpr,
NT_OFFSET: tl.constexpr,
H_OFFSET: tl.constexpr,
):
i_t = tl.program_id(0) + NT_OFFSET
i_h = tl.program_id(1) + H_OFFSET
b_A = tl.load(A_log + i_h).to(tl.float32)
p_g = tl.make_block_ptr(g + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
p_yg = tl.make_block_ptr(yg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
b_yg = -exp(b_A) * softplus(b_g)
tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), boundary_check=(0,))
def _launch_gate_fwd(
*,
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None,
yg: torch.Tensor,
T: int,
H: int,
BT: int,
) -> None:
NT = triton.cdiv(T, BT)
kernel_kwargs = dict(
g=g,
A_log=A_log,
dt_bias=dt_bias,
yg=yg,
T=T,
H=H,
BT=BT,
num_warps=_NUM_WARPS,
)
max_nt = max_grid_axis_chunks(NT, H, max_grid=ASCEND_MAX_GRID_DIM)
for nt_off in range(0, NT, max_nt):
nt_len = min(max_nt, NT - nt_off)
max_h = max_grid_axis_chunks(H, nt_len, max_grid=ASCEND_MAX_GRID_DIM)
for h_off in range(0, H, max_h):
h_len = min(max_h, H - h_off)
gdn_gate_fwd_kernel_npu[(nt_len, h_len)](
**kernel_kwargs,
NT_OFFSET=nt_off,
H_OFFSET=h_off,
)
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
'HAS_SCALE': lambda args: args['scale'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@triton.jit(do_not_specialize=['T'])
def gdn_gate_chunk_cumsum_scalar_kernel_npu(
g,
A_log,
dt_bias,
o,
scale,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
BT: tl.constexpr,
REVERSE: tl.constexpr,
HAS_BIAS: tl.constexpr,
HAS_SCALE: tl.constexpr,
IS_VARLEN: tl.constexpr,
NT_OFFSET: tl.constexpr,
BH_OFFSET: tl.constexpr,
):
i_t = tl.program_id(0) + NT_OFFSET
i_bh = tl.program_id(1) + BH_OFFSET
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
b_A = tl.load(A_log + i_h).to(tl.float32)
b_gate = -exp(b_A) * softplus(b_g)
b_o = tl.cumsum(b_gate, axis=0)
if REVERSE:
b_z = tl.sum(b_gate, axis=0)
b_o = -b_o + b_z[None] + b_gate
if HAS_SCALE:
b_o *= scale
tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
def _launch_gate_chunk_cumsum(
*,
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None,
o: torch.Tensor,
scale: float | None,
cu_seqlens: torch.LongTensor | None,
chunk_indices: torch.LongTensor | None,
T: int,
B: int,
H: int,
BT: int,
NT: int,
reverse: bool,
) -> None:
bh_total = B * H
kernel_kwargs = dict(
g=g,
A_log=A_log,
dt_bias=dt_bias,
o=o,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
BT=BT,
REVERSE=reverse,
num_warps=_NUM_WARPS,
)
max_nt = max_grid_axis_chunks(NT, bh_total, max_grid=ASCEND_MAX_GRID_DIM)
for nt_off in range(0, NT, max_nt):
nt_len = min(max_nt, NT - nt_off)
max_bh = max_grid_axis_chunks(bh_total, nt_len, max_grid=ASCEND_MAX_GRID_DIM)
for bh_off in range(0, bh_total, max_bh):
bh_len = min(max_bh, bh_total - bh_off)
gdn_gate_chunk_cumsum_scalar_kernel_npu[(nt_len, bh_len)](
**kernel_kwargs,
NT_OFFSET=nt_off,
BH_OFFSET=bh_off,
)
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
})
@triton.jit(do_not_specialize=['T'])
def gdn_gate_bwd_kernel_npu(
g,
A_log,
dt_bias,
dyg,
dg,
dA,
T,
H: tl.constexpr,
BT: tl.constexpr,
HAS_BIAS: tl.constexpr,
NT_OFFSET: tl.constexpr,
H_OFFSET: tl.constexpr,
):
i_t = tl.program_id(0) + NT_OFFSET
i_h = tl.program_id(1) + H_OFFSET
b_A = tl.load(A_log + i_h).to(tl.float32)
p_g = tl.make_block_ptr(g + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
p_dg = tl.make_block_ptr(dg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
p_dyg = tl.make_block_ptr(dyg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32)
b_dyg = tl.load(p_dyg, boundary_check=(0,)).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
b_neg_expA = -exp(b_A)
b_yg = b_neg_expA * softplus(b_g)
b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g))
b_dA = tl.sum(b_dyg * b_yg, 0)
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))
tl.store(dA + i_t * H + i_h, b_dA)
def _launch_gate_bwd(
*,
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None,
dyg: torch.Tensor,
dg: torch.Tensor,
dA: torch.Tensor,
T: int,
H: int,
BT: int,
) -> None:
NT = triton.cdiv(T, BT)
kernel_kwargs = dict(
g=g,
A_log=A_log,
dt_bias=dt_bias,
dyg=dyg,
dg=dg,
dA=dA,
T=T,
H=H,
BT=BT,
num_warps=_NUM_WARPS,
)
max_nt = max_grid_axis_chunks(NT, H, max_grid=ASCEND_MAX_GRID_DIM)
for nt_off in range(0, NT, max_nt):
nt_len = min(max_nt, NT - nt_off)
max_h = max_grid_axis_chunks(H, nt_len, max_grid=ASCEND_MAX_GRID_DIM)
for h_off in range(0, H, max_h):
h_len = min(max_h, H - h_off)
gdn_gate_bwd_kernel_npu[(nt_len, h_len)](
**kernel_kwargs,
NT_OFFSET=nt_off,
H_OFFSET=h_off,
)
@input_guard
def gdn_gate_fwd_npu(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
H = g.shape[-1]
T = g.numel() // H
BT = _get_gate_fwd_bt(T)
yg = torch.empty_like(g, dtype=output_dtype)
_launch_gate_fwd(g=g, A_log=A_log, dt_bias=dt_bias, yg=yg, T=T, H=H, BT=BT)
return yg
@input_guard
def gdn_gate_chunk_cumsum_npu(
g: torch.Tensor,
A_log: torch.Tensor,
chunk_size: int,
scale: float = None,
dt_bias: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
output_dtype: torch.dtype | None = torch.float,
) -> torch.Tensor:
B, T, H = g.shape
assert chunk_size == 2 ** (chunk_size.bit_length() - 1), "chunk_size must be a power of 2"
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
o = torch.empty_like(g, dtype=output_dtype or g.dtype)
_launch_gate_chunk_cumsum(
g=g,
A_log=A_log,
dt_bias=dt_bias,
o=o,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
B=B,
H=H,
BT=BT,
NT=NT,
reverse=False,
)
return o
def gdn_gate_bwd_npu(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None,
dyg: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
H = g.shape[-1]
T = g.numel() // H
BT = _get_gate_bwd_bt(T)
dg = torch.empty_like(g, dtype=torch.float32)
NT = triton.cdiv(T, BT)
dA = A_log.new_empty(NT, H, dtype=torch.float32)
_launch_gate_bwd(
g=g,
A_log=A_log,
dt_bias=dt_bias,
dyg=dyg,
dg=dg,
dA=dA,
T=T,
H=H,
BT=BT,
)
dg = dg.view_as(g).type_as(g)
dA = dA.sum(0).view_as(A_log).type_as(A_log)
dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None
return dg, dA, dbias

823
upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py Normal file
View File

@@ -0,0 +1,823 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
"""WY-representation kernels adapted for triton-ascend on Ascend NPU."""
from __future__ import annotations
import torch
import triton
import triton.language as tl
import triton.runtime.driver as driver
from fla.ops.utils import prepare_chunk_indices
from fla.ops.utils.op import exp2
from fla.utils import input_guard
from fla.utils.ascend_ub_manager import (
ASCEND_MAX_GRID_DIM,
compute_row_tile_block_size,
max_grid_axis_chunks,
)
def get_npu_properties():
device = torch.npu.current_device()
return driver.active.utils.get_device_properties(device)
# prepare_wy_repr_bwd stage-specific UB models
_PREPARE_BWD_K_MEM_MULT = 4.5
_PREPARE_BWD_V_MEM_MULT = 8.0
_SAFETY_MARGIN = 0.75
_FALLBACK_TILE = 8
_MAX_TILE_BWD = 128
def _g_npu_arg(g: torch.Tensor | None, HV: int) -> tuple[torch.Tensor | None, bool]:
if g is None or HV == 1:
return g, False
return g.transpose(1, 2).contiguous(), True
def _beta_npu_arg(beta: torch.Tensor, HV: int) -> tuple[torch.Tensor, bool]:
if HV == 1:
return beta, False
return beta.transpose(1, 2).contiguous(), True
def _t_npu_buf(
B: int, T: int, HV: int, *, dtype: torch.dtype, device: torch.device,
) -> tuple[torch.Tensor, bool]:
if HV == 1:
return torch.empty(B, T, HV, dtype=dtype, device=device), False
return torch.empty(B, HV, T, dtype=dtype, device=device), True
def _get_bwd_k_tile(BT: int, K: int) -> int:
return compute_row_tile_block_size(
BT, K, _PREPARE_BWD_K_MEM_MULT,
tiling_row=False,
safety_margin=_SAFETY_MARGIN,
fallback=_FALLBACK_TILE,
min_block=8,
max_block=min(_MAX_TILE_BWD, triton.next_power_of_2(K)),
)
def _get_bwd_v_tile(BT: int, V: int) -> int:
return compute_row_tile_block_size(
BT, V, _PREPARE_BWD_V_MEM_MULT,
tiling_row=False,
safety_margin=_SAFETY_MARGIN,
fallback=_FALLBACK_TILE,
min_block=8,
max_block=min(_MAX_TILE_BWD, triton.next_power_of_2(V)),
)
def _get_bwd_tiles(BT: int, K: int, V: int) -> tuple[int, int]:
return _get_bwd_k_tile(BT, K), _get_bwd_v_tile(BT, V)
@triton.jit
def _g_contig_base(g, bos, i_b, i_h, T_seq, HV, IS_VARLEN: tl.constexpr):
if IS_VARLEN:
return g + bos + i_h * T_seq
return g + i_b * HV * T_seq + i_h * T_seq
@triton.jit
def _t_block_ptr(base, T, offset, BLK, CONTIG: tl.constexpr, HV: tl.constexpr):
if CONTIG:
return tl.make_block_ptr(base, (T,), (1,), (offset,), (BLK,), (0,))
return tl.make_block_ptr(base, (T,), (HV,), (offset,), (BLK,), (0,))
def _launch_wy_kernel(kernel, *, NT: int, bh_total: int, kernel_kwargs: dict) -> None:
max_nt = max_grid_axis_chunks(NT, bh_total, max_grid=ASCEND_MAX_GRID_DIM)
chunk_indices = kernel_kwargs.get('chunk_indices')
cu_seqlens = kernel_kwargs.get('cu_seqlens')
for nt_off in range(0, NT, max_nt):
nt_len = min(max_nt, NT - nt_off)
if cu_seqlens is not None and chunk_indices is not None:
kernel_kwargs['chunk_indices'] = chunk_indices[nt_off:nt_off + nt_len]
kernel_kwargs['NT_OFFSET'] = 0
else:
kernel_kwargs['NT_OFFSET'] = nt_off
max_bh = max_grid_axis_chunks(bh_total, nt_len, max_grid=ASCEND_MAX_GRID_DIM)
for bh_off in range(0, bh_total, max_bh):
bh_len = min(max_bh, bh_total - bh_off)
kernel_kwargs['BH_OFFSET'] = bh_off
kernel[(nt_len, bh_len)](**kernel_kwargs)
def _launch_wy_core_grid(kernel, *, task_num: int, kernel_kwargs: dict) -> None:
num_core = get_npu_properties()["num_aicore"]
kernel[(num_core,)](task_num=task_num, num_core=num_core, **kernel_kwargs)
@triton.heuristics({
"IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
"USE_G": lambda args: args["g"] is not None,
})
@triton.jit(do_not_specialize=["T", "B", "task_num", "num_core"])
def recompute_w_u_fwd_kernel_npu(
k,
v,
beta,
w,
u,
A,
g,
cu_seqlens,
chunk_indices,
T,
B,
task_num,
num_core,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
IS_VARLEN: tl.constexpr,
USE_G: tl.constexpr,
):
T_max = T
core_id = tl.program_id(0)
for task_id in tl.range(core_id, task_num, num_core):
i_t_o = task_id // (B * HV)
i_bh = task_id % (B * HV)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t_o * 2).to(tl.int32), tl.load(
chunk_indices + i_t_o * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
T = eos - bos
bos_bh = bos
else:
i_t = i_t_o
bos, eos = i_b * T, i_b * T + T
bos_bh = i_b * HV * T_max
offs_t = tl.arange(0, BT)
global_offs_t = i_t * BT + offs_t
mask_t = global_offs_t < T
offs_t_2d = global_offs_t[:, None]
offs_bt = tl.arange(0, BT)[None, :]
ptr_A = A + (bos * HV + i_h) * BT + offs_t_2d * (HV * BT) + offs_bt * 1
mask_A = mask_t[:, None]
b_A = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32)
ptr_beta = beta + bos_bh + i_h * T_max + global_offs_t
b_beta = tl.load(ptr_beta, mask=mask_t, other=0.0).to(tl.float32)
for i_v in range(tl.cdiv(V, BV)):
offs_v = i_v * BV + tl.arange(0, BV)[None, :]
mask_v = (mask_t[:, None]) & (offs_v < V)
ptr_v = v + (bos * HV + i_h) * V + offs_t_2d * (HV * V) + offs_v * 1
b_v = tl.load(ptr_v, mask=mask_v, other=0.0).to(tl.float32)
b_vb = b_v * b_beta[:, None]
b_u = tl.dot(b_A, b_vb, allow_tf32=False)
ptr_u = u + (bos * HV + i_h) * V + offs_t_2d * (HV * V) + offs_v * 1
tl.store(ptr_u, b_u.to(ptr_u.dtype.element_ty), mask=mask_v)
if USE_G:
ptr_g = g + bos_bh + i_h * T_max + global_offs_t
b_g = exp2(tl.load(ptr_g, mask=mask_t, other=0.0)).to(tl.float32)
for i_k in range(tl.cdiv(K, BK)):
offs_k = i_k * BK + tl.arange(0, BK)[None, :]
mask_k = (mask_t[:, None]) & (offs_k < K)
ptr_k = (
k
+ (bos * H + i_h // (HV // H)) * K
+ offs_t_2d * (H * K)
+ offs_k * 1
)
b_k = tl.load(ptr_k, mask=mask_k, other=0.0).to(tl.float32)
b_kb = b_k * b_beta[:, None]
if USE_G:
b_kb = b_kb * b_g[:, None]
b_w = tl.dot(b_A, b_kb, allow_tf32=False)
ptr_w = w + (bos * HV + i_h) * K + offs_t_2d * (HV * K) + offs_k * 1
tl.store(ptr_w, b_w.to(ptr_w.dtype.element_ty), mask=mask_k)
@triton.heuristics({
"IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
})
@triton.jit(do_not_specialize=["T", "B", "task_num", "num_core"])
def prepare_wy_repr_bwd_kv_npu(
k, v, beta, g, A, dw, du, dk, dv, dA_scr, db, dg,
cu_seqlens, chunk_indices, T, B,
task_num, num_core,
H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr,
BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
USE_G: tl.constexpr, IS_VARLEN: tl.constexpr,
G_T_CONTIG: tl.constexpr, BETA_T_CONTIG: tl.constexpr,
DG_T_CONTIG: tl.constexpr, DB_T_CONTIG: tl.constexpr,
G_EXP_PRECOMP: tl.constexpr,
):
T_seq = T
core_id = tl.program_id(0)
for task_id in tl.range(core_id, task_num, num_core):
i_t_o = task_id // (B * HV)
i_bh = task_id % (B * HV)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t_o * 2).to(tl.int32), tl.load(
chunk_indices + i_t_o * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
T = eos - bos
else:
i_t = i_t_o
bos, eos = i_b * T_seq, i_b * T_seq + T_seq
if BETA_T_CONTIG:
beta_ptr = _g_contig_base(beta, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
p_b = _t_block_ptr(beta_ptr, T, i_t * BT, BT, True, HV)
else:
p_b = tl.make_block_ptr(beta + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,))
if DB_T_CONTIG:
db_ptr = _g_contig_base(db, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
p_db = _t_block_ptr(db_ptr, T, i_t * BT, BT, True, HV)
else:
p_db = tl.make_block_ptr(db + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,))
p_A = tl.make_block_ptr(
A + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1),
)
p_dA = tl.make_block_ptr(
dA_scr + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1),
)
b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)
b_db = tl.zeros([BT], dtype=tl.float32)
b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
b_dA = tl.zeros([BT, BT], dtype=tl.float32)
if USE_G:
if G_T_CONTIG:
g_ptr = _g_contig_base(g, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
p_g = _t_block_ptr(g_ptr, T, i_t * BT, BT, True, HV)
else:
p_g = tl.make_block_ptr(g + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,))
b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32)
b_g_exp = b_g if G_EXP_PRECOMP else exp2(b_g)
b_dg = tl.zeros([BT], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
p_k = tl.make_block_ptr(
k + (bos * H + i_h // (HV // H)) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0),
)
p_dk = tl.make_block_ptr(
dk + (bos * HV + i_h) * K, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0),
)
p_dw = tl.make_block_ptr(
dw + (bos * HV + i_h) * K, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0),
)
b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32)
if USE_G:
b_kbg = b_k * (b_b * b_g_exp)[:, None]
else:
b_kbg = b_k * b_b[:, None]
b_dw = tl.load(p_dw, boundary_check=(0, 1)).to(tl.float32)
b_dA += tl.dot(b_dw, tl.trans(b_kbg), allow_tf32=False)
b_dkbg = tl.dot(b_A, b_dw, allow_tf32=False)
if USE_G:
b_dk = b_dkbg * (b_g_exp * b_b)[:, None]
b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1)
b_dg += tl.sum(b_dkbg * b_kbg, 1)
else:
b_dk = b_dkbg * b_b[:, None]
b_db += tl.sum(b_dkbg * b_k, 1)
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
for i_v in range(tl.cdiv(V, BV)):
p_v = tl.make_block_ptr(
v + (bos * HV + i_h) * V, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0),
)
p_dv = tl.make_block_ptr(
dv + (bos * HV + i_h) * V, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0),
)
p_du = tl.make_block_ptr(
du + (bos * HV + i_h) * V, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0),
)
b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32)
b_du = tl.load(p_du, boundary_check=(0, 1)).to(tl.float32)
b_vb = b_v * b_b[:, None]
b_dA += tl.dot(b_du, tl.trans(b_vb), allow_tf32=False)
b_dvb = tl.dot(b_A, b_du, allow_tf32=False)
b_dv = b_dvb * b_b[:, None]
b_db += tl.sum(b_dvb * b_v, 1)
tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,))
if USE_G:
if DG_T_CONTIG:
dg_ptr = _g_contig_base(dg, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
p_dg = _t_block_ptr(dg_ptr, T, i_t * BT, BT, True, HV)
else:
p_dg = tl.make_block_ptr(dg + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,))
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))
@triton.jit(do_not_specialize=['T'])
def prepare_wy_repr_bwd_da_mask_dot1_npu(
A, dA_scr, dA_mid,
cu_seqlens, chunk_indices, T,
HV: tl.constexpr, BT: tl.constexpr,
IS_VARLEN: tl.constexpr,
NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr,
):
i_t = tl.program_id(0) + NT_OFFSET
i_bh = tl.program_id(1) + BH_OFFSET
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
p_A = tl.make_block_ptr(
A + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1),
)
p_in = tl.make_block_ptr(
dA_scr + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1),
)
p_out = tl.make_block_ptr(
dA_mid + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1),
)
b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
b_dA = tl.load(p_in, boundary_check=(0, 1)).to(tl.float32)
o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
b_dA = tl.where(m_A, b_dA, 0)
b_out = tl.dot(b_dA, b_A, allow_tf32=False)
tl.store(p_out, b_out.to(p_out.dtype.element_ty), boundary_check=(0, 1))
@triton.jit(do_not_specialize=['T'])
def prepare_wy_repr_bwd_da_dot2_npu(
A, dA_mid, dA_out,
cu_seqlens, chunk_indices, T,
HV: tl.constexpr, BT: tl.constexpr,
IS_VARLEN: tl.constexpr,
NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr,
):
i_t = tl.program_id(0) + NT_OFFSET
i_bh = tl.program_id(1) + BH_OFFSET
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
p_A = tl.make_block_ptr(A + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1))
p_in = tl.make_block_ptr(dA_mid + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1))
p_out = tl.make_block_ptr(dA_out + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1))
b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
b_dA = tl.load(p_in, boundary_check=(0, 1)).to(tl.float32)
b_dA = tl.dot(b_A, b_dA, allow_tf32=False)
o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
b_dA = tl.where(m_A, -b_dA, 0)
tl.store(p_out, b_dA.to(p_out.dtype.element_ty), boundary_check=(0, 1))
_DG_BLK = 16
@triton.jit(do_not_specialize=['T'])
def prepare_wy_repr_bwd_da_gate_npu(
g, dA_out,
cu_seqlens, chunk_indices, T,
HV: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr,
IS_VARLEN: tl.constexpr, G_T_CONTIG: tl.constexpr,
NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr,
):
i_t = tl.program_id(0) + NT_OFFSET
i_bh = tl.program_id(1) + BH_OFFSET
i_b, i_h = i_bh // HV, i_bh % HV
T_seq = T
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
n_sub = BT // BC
if G_T_CONTIG:
g_ptr = _g_contig_base(g, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
else:
g_ptr = g + (bos * HV + i_h)
for r in range(n_sub):
i_tr = i_t * BT + r * BC
p_gr = _t_block_ptr(g_ptr, T, i_tr, BC, G_T_CONTIG, HV)
b_gr = tl.load(p_gr, boundary_check=(0,)).to(tl.float32)
for c in range(n_sub):
i_tc = i_t * BT + c * BC
p_dA = tl.make_block_ptr(
dA_out + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT),
(r * BC, i_t * BT + c * BC), (BC, BC), (0, 1),
)
b_dA = tl.load(p_dA, boundary_check=(0, 1)).to(tl.float32)
p_gc = _t_block_ptr(g_ptr, T, i_tc, BC, G_T_CONTIG, HV)
b_gc = tl.load(p_gc, boundary_check=(0,)).to(tl.float32)
b_gate = exp2(b_gr[:, None] - b_gc[None, :])
b_prod = b_dA * b_gate
b_dA = tl.where(b_prod == b_prod, b_prod, 0.0)
tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1))
@triton.heuristics({
"IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
})
@triton.jit(do_not_specialize=["T", "B", "task_num", "num_core"])
def prepare_wy_repr_bwd_finalize_k_npu(
k, beta, dA_out, dk, db,
cu_seqlens, chunk_indices, T, B,
task_num, num_core,
H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr,
BT: tl.constexpr, BK: tl.constexpr,
IS_VARLEN: tl.constexpr, BETA_T_CONTIG: tl.constexpr, DB_T_CONTIG: tl.constexpr,
):
T_seq = T
core_id = tl.program_id(0)
for task_id in tl.range(core_id, task_num, num_core):
i_t_o = task_id // (B * HV)
i_bh = task_id % (B * HV)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t_o * 2).to(tl.int32), tl.load(
chunk_indices + i_t_o * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
T = eos - bos
else:
i_t = i_t_o
bos, eos = i_b * T_seq, i_b * T_seq + T_seq
if BETA_T_CONTIG:
beta_ptr = _g_contig_base(beta, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
p_b = _t_block_ptr(beta_ptr, T, i_t * BT, BT, True, HV)
else:
p_b = tl.make_block_ptr(beta + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,))
if DB_T_CONTIG:
db_ptr = _g_contig_base(db, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
p_db = _t_block_ptr(db_ptr, T, i_t * BT, BT, True, HV)
else:
p_db = tl.make_block_ptr(db + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,))
p_dA = tl.make_block_ptr(
dA_out + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1),
)
b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)
b_db = tl.load(p_db, boundary_check=(0,)).to(tl.float32)
b_dA = tl.load(p_dA, boundary_check=(0, 1)).to(tl.float32)
for i_k in range(tl.cdiv(K, BK)):
p_k = tl.make_block_ptr(
k + (bos * H + i_h // (HV // H)) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0),
)
p_dk = tl.make_block_ptr(
dk + (bos * HV + i_h) * K, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0),
)
b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32)
b_kb = b_k * b_b[:, None]
b_dkb = tl.dot(b_dA, b_k, allow_tf32=False)
b_db += tl.sum(b_dkb * b_k, 1)
b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb), b_dA, allow_tf32=False))
b_dk += tl.load(p_dk, boundary_check=(0, 1)).to(tl.float32)
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,))
@triton.jit(do_not_specialize=['T'])
def prepare_wy_repr_bwd_finalize_a2_npu(
k, beta, a2_scr,
cu_seqlens, chunk_indices, T,
H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr,
BT: tl.constexpr, BK: tl.constexpr,
IS_VARLEN: tl.constexpr, BETA_T_CONTIG: tl.constexpr,
NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr,
):
i_t = tl.program_id(0) + NT_OFFSET
i_bh = tl.program_id(1) + BH_OFFSET
i_b, i_h = i_bh // HV, i_bh % HV
T_seq = T
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
if BETA_T_CONTIG:
beta_ptr = _g_contig_base(beta, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
p_b = _t_block_ptr(beta_ptr, T, i_t * BT, BT, True, HV)
else:
p_b = tl.make_block_ptr(beta + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,))
p_a2 = tl.make_block_ptr(a2_scr + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1))
b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)
b_A2 = tl.zeros([BT, BT], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
p_k = tl.make_block_ptr(
k + (bos * H + i_h // (HV // H)) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0),
)
b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32)
b_A2 += tl.dot(b_k, tl.trans(b_k), allow_tf32=False)
b_A2 *= b_b[:, None]
tl.store(p_a2, b_A2.to(p_a2.dtype.element_ty), boundary_check=(0, 1))
@triton.jit(do_not_specialize=['T'])
def prepare_wy_repr_bwd_finalize_dg_npu(
dA_out, a2_scr, dg, col_acc_scr,
cu_seqlens, chunk_indices, T,
HV: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr,
IS_VARLEN: tl.constexpr, DG_T_CONTIG: tl.constexpr,
NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr,
):
i_t = tl.program_id(0) + NT_OFFSET
i_bh = tl.program_id(1) + BH_OFFSET
i_b, i_h = i_bh // HV, i_bh % HV
T_seq = T
if IS_VARLEN:
i_tg = i_t
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
T = eos - bos
else:
NT = tl.cdiv(T, BT)
i_tg = i_b * NT + i_t
bos, eos = i_b * T, i_b * T + T
n_sub = BT // BC
col_off = (i_tg * HV + i_h) * BT
p_col0 = tl.make_block_ptr(col_acc_scr + col_off, (BT,), (1,), (0,), (BT,), (0,))
tl.store(p_col0, tl.zeros([BT], dtype=tl.float32), boundary_check=(0,))
if DG_T_CONTIG:
dg_ptr = _g_contig_base(dg, bos, i_b, i_h, T_seq, HV, IS_VARLEN)
else:
dg_ptr = dg + (bos * HV + i_h)
for r in range(n_sub):
i_tr = i_t * BT + r * BC
p_dg_r = _t_block_ptr(dg_ptr, T, i_tr, BC, DG_T_CONTIG, HV)
b_dg_r = tl.load(p_dg_r, boundary_check=(0,)).to(tl.float32)
for c in range(n_sub):
p_dA = tl.make_block_ptr(
dA_out + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT),
(r * BC, i_t * BT + c * BC), (BC, BC), (0, 1),
)
p_a2 = tl.make_block_ptr(
a2_scr + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT),
(r * BC, i_t * BT + c * BC), (BC, BC), (0, 1),
)
b_dA = tl.load(p_dA, boundary_check=(0, 1)).to(tl.float32)
b_a2 = tl.load(p_a2, boundary_check=(0, 1)).to(tl.float32)
prod = b_dA * b_a2
b_dg_r += tl.sum(prod, axis=1)
p_col = tl.make_block_ptr(
col_acc_scr + col_off, (BT,), (1,), (c * BC,), (BC,), (0,),
)
b_col = tl.load(p_col, boundary_check=(0,)).to(tl.float32)
b_col += tl.sum(prod, axis=0)
tl.store(p_col, b_col.to(p_col.dtype.element_ty), boundary_check=(0,))
tl.store(p_dg_r, b_dg_r.to(p_dg_r.dtype.element_ty), boundary_check=(0,))
p_dg = _t_block_ptr(dg_ptr, T, i_t * BT, BT, DG_T_CONTIG, HV)
p_col = tl.make_block_ptr(col_acc_scr + col_off, (BT,), (1,), (0,), (BT,), (0,))
b_dg = tl.load(p_dg, boundary_check=(0,)).to(tl.float32)
b_col = tl.load(p_col, boundary_check=(0,)).to(tl.float32)
b_dg -= b_col
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))
@input_guard
def recompute_w_u_fwd_npu(
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
g: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
B, T, H, K, V = *k.shape, v.shape[-1]
HV = v.shape[2]
BT = A.shape[-1]
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
BK = 64
BV = 64
u = torch.empty_like(v)
w = k.new_empty(B, T, HV, K)
beta = beta.transpose(1, 2).contiguous()
if g is not None:
g = g.transpose(1, 2).contiguous()
num_core = get_npu_properties()["num_aicore"]
task_num = NT * B * HV
recompute_w_u_fwd_kernel_npu[(num_core,)](
k=k,
v=v,
beta=beta,
w=w,
u=u,
A=A,
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
B=B,
task_num=task_num,
num_core=num_core,
H=H,
HV=HV,
K=K,
V=V,
BT=BT,
BK=BK,
BV=BV,
)
return w, u
def prepare_wy_repr_bwd_npu(
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
dw: torch.Tensor,
du: torch.Tensor,
g: torch.Tensor = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2]
BT = A.shape[-1]
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
BK, BV = _get_bwd_tiles(BT, K, V)
use_g = g is not None
is_varlen = cu_seqlens is not None
dk = k.new_empty(B, T, HV, K)
dv = torch.empty_like(v)
dg, dg_t_contig = None, False
if use_g:
dg, dg_t_contig = _t_npu_buf(B, T, HV, dtype=g.dtype, device=k.device)
db, db_t_contig = _t_npu_buf(B, T, HV, dtype=beta.dtype, device=k.device)
beta_arg, beta_t_contig = _beta_npu_arg(beta, HV)
g_gate, g_t_contig = None, False
g_exp_precomp = False
if use_g:
g_gate, g_t_contig = _g_npu_arg(g, HV)
g_k_arg = g_gate
if not is_varlen:
g_k_arg = torch.exp2(g_gate.float()).to(g_gate.dtype)
g_exp_precomp = True
dg_arg = dg if use_g else beta
dA_scr = torch.zeros_like(A, dtype=torch.float32)
dA_mid = torch.zeros_like(A, dtype=torch.float32)
dA_out = torch.zeros_like(A, dtype=torch.float32)
a2_scr = torch.zeros_like(A, dtype=torch.float32)
col_acc_scr = torch.zeros(B, triton.cdiv(T, BT) if cu_seqlens is None else len(
chunk_indices), HV, BT, dtype=torch.float32, device=k.device)
base = dict(
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
BT=BT,
IS_VARLEN=is_varlen,
)
core_base = dict(B=B, **base)
task_num = NT * B * HV
_launch_wy_core_grid(
prepare_wy_repr_bwd_kv_npu,
task_num=task_num,
kernel_kwargs=dict(
k=k, v=v, beta=beta_arg, g=g_k_arg if use_g else k, A=A, dw=dw, du=du,
dk=dk, dv=dv, dA_scr=dA_scr, db=db, dg=dg_arg,
H=H, HV=HV, K=K, V=V, BK=BK, BV=BV, USE_G=use_g,
G_T_CONTIG=g_t_contig, BETA_T_CONTIG=beta_t_contig,
DG_T_CONTIG=dg_t_contig, DB_T_CONTIG=db_t_contig,
G_EXP_PRECOMP=g_exp_precomp,
**core_base,
),
)
_launch_wy_kernel(
prepare_wy_repr_bwd_da_mask_dot1_npu,
NT=NT,
bh_total=B * HV,
kernel_kwargs=dict(
A=A, dA_scr=dA_scr, dA_mid=dA_mid,
HV=HV,
**base,
),
)
_launch_wy_kernel(
prepare_wy_repr_bwd_da_dot2_npu,
NT=NT,
bh_total=B * HV,
kernel_kwargs=dict(
A=A, dA_mid=dA_mid, dA_out=dA_out,
HV=HV,
**base,
),
)
if use_g:
_launch_wy_kernel(
prepare_wy_repr_bwd_da_gate_npu,
NT=NT,
bh_total=B * HV,
kernel_kwargs=dict(
g=g_gate, dA_out=dA_out,
HV=HV, BC=_DG_BLK, G_T_CONTIG=g_t_contig,
**base,
),
)
_launch_wy_core_grid(
prepare_wy_repr_bwd_finalize_k_npu,
task_num=task_num,
kernel_kwargs=dict(
k=k, beta=beta_arg, dA_out=dA_out, dk=dk, db=db,
H=H, HV=HV, K=K, BK=BK, BETA_T_CONTIG=beta_t_contig, DB_T_CONTIG=db_t_contig,
**core_base,
),
)
if use_g:
_launch_wy_kernel(
prepare_wy_repr_bwd_finalize_a2_npu,
NT=NT,
bh_total=B * HV,
kernel_kwargs=dict(
k=k, beta=beta_arg, a2_scr=a2_scr,
H=H, HV=HV, K=K, BK=BK, BETA_T_CONTIG=beta_t_contig,
**base,
),
)
_launch_wy_kernel(
prepare_wy_repr_bwd_finalize_dg_npu,
NT=NT,
bh_total=B * HV,
kernel_kwargs=dict(
dA_out=dA_out, a2_scr=a2_scr, dg=dg_arg, col_acc_scr=col_acc_scr,
HV=HV, BC=_DG_BLK, DG_T_CONTIG=dg_t_contig,
**base,
),
)
if H != HV:
dk = dk.view(B, T, H, HV // H, K).sum(3)
if db_t_contig:
db = db.transpose(1, 2).contiguous()
if use_g and dg_t_contig:
dg = dg.transpose(1, 2).contiguous()
return dk, dv, db, dg

591
upstream_ref/fla/ops/gated_delta_rule/chunk.py Normal file
View File

@@ -0,0 +1,591 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import warnings
import torch
from fla.modules.l2norm import l2norm_bwd, l2norm_fwd
from fla.ops.backends import dispatch
from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h
from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o
from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd
from fla.ops.cp import FLACPContext
from fla.ops.cp.chunk_delta_h import (
chunk_gated_delta_rule_bwd_dhu_pre_process,
chunk_gated_delta_rule_fwd_h_pre_process,
compress_h0,
expand_h0,
)
from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra
from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum
from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd
from fla.ops.utils import chunk_local_cumsum
from fla.ops.utils.constant import RCP_LN2
from fla.ops.utils.index import prepare_chunk_indices
from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
def chunk_gated_delta_rule_fwd(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
output_final_state: bool,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
chunk_indices: torch.LongTensor | None = None,
use_gate_in_kernel: bool = False,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
chunk_size: int = 64,
):
g_input = g if use_gate_in_kernel else None
if use_gate_in_kernel:
g = gdn_gate_chunk_cumsum(
g=g,
A_log=A_log,
chunk_size=chunk_size,
scale=RCP_LN2,
dt_bias=dt_bias,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
else:
g = chunk_local_cumsum(
g,
chunk_size=chunk_size,
scale=RCP_LN2,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
# obtain WY representation. u is actually the new v.
# fused kkt + solve_tril + recompute_w_u
w, u, A = chunk_gated_delta_rule_fwd_intra(
k=k,
v=v,
g=g,
beta=beta,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_size=chunk_size,
)
if cp_context is not None:
initial_state = chunk_gated_delta_rule_fwd_h_pre_process(
k=k,
w=w,
u=u,
g=g,
cu_seqlens=cu_seqlens,
initial_state=initial_state,
context=cp_context,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
k=k,
w=w,
u=u,
g=g,
initial_state=initial_state,
output_final_state=output_final_state,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
if cp_context is not None:
initial_state = compress_h0(initial_state, context=cp_context)
o = chunk_fwd_o(
q=q,
k=k,
v=v_new,
h=h,
g=g,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
return g, o, A, final_state, initial_state, g_input
def chunk_gated_delta_rule_bwd(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
do: torch.Tensor,
dht: torch.Tensor,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
chunk_indices: torch.LongTensor | None = None,
use_gate_in_kernel: bool = False,
g_input: torch.Tensor | None = None,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
chunk_size: int = 64,
):
w, u = recompute_w_u_fwd(
k=k,
v=v,
beta=beta,
A=A,
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
if cp_context is not None:
initial_state = expand_h0(initial_state, context=cp_context)
h, v_new, _ = chunk_gated_delta_rule_fwd_h(
k=k,
w=w,
u=u,
g=g,
initial_state=initial_state,
output_final_state=False,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
dv = chunk_bwd_dv_local(
q=q,
k=k,
g=g,
do=do,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_size=chunk_size,
)
if cp_context is not None:
# initial_state is None in the CP mode
# We only need to compute dht of current rank and pass it to the backward kernel
dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process(
q=q,
k=k,
w=w,
do=do,
dv=dv,
g=g,
scale=scale,
cu_seqlens=cu_seqlens,
dht=dht,
initial_state=initial_state,
context=cp_context,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu(
q=q,
k=k,
w=w,
g=g,
h0=initial_state,
dht=dht,
do=do,
dv=dv,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
dq, dk, dw, dg = chunk_bwd_dqkwg(
q=q,
k=k,
v=v_new,
w=w,
g=g,
h=h,
dv=dv,
do=do,
dh=dh,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
dk2, dv, db, dg2 = prepare_wy_repr_bwd(
k=k,
v=v,
beta=beta,
g=g,
A=A,
dw=dw,
du=dv,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
dk.add_(dk2)
dg.add_(dg2)
dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices)
dA_log, ddt_bias = None, None
if use_gate_in_kernel:
dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg)
return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias
class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
@staticmethod
@input_guard
@autocast_custom_fwd
def forward(
ctx,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
output_final_state: bool,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cu_seqlens_cpu: torch.LongTensor | None = None,
use_qk_l2norm_in_kernel: bool = False,
use_gate_in_kernel: bool = False,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
cp_context: FLACPContext | None = None,
chunk_size: int = 64,
):
q_rstd, k_rstd = None, None
if use_qk_l2norm_in_kernel:
q, q_rstd = l2norm_fwd(q)
k, k_rstd = l2norm_fwd(k)
beta_raw = beta
if use_beta_sigmoid_in_kernel:
beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0)
chunk_indices = None
if cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu)
g, o, A, final_state, initial_state, g_input = chunk_gated_delta_rule_fwd(
q=q,
k=k,
v=v,
g=g,
beta=beta,
scale=scale,
initial_state=initial_state,
output_final_state=output_final_state,
cu_seqlens=cu_seqlens,
cp_context=cp_context,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
use_gate_in_kernel=use_gate_in_kernel,
A_log=A_log,
dt_bias=dt_bias,
chunk_size=chunk_size,
)
ctx.save_for_backward(
q,
q_rstd,
k,
k_rstd,
v,
g,
beta_raw,
beta,
A,
initial_state,
cu_seqlens,
chunk_indices,
g_input,
A_log,
dt_bias,
)
ctx.scale = scale
ctx.chunk_size = chunk_size
ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel
ctx.allow_neg_eigval = allow_neg_eigval
ctx.cp_context = cp_context
ctx.state_v_first = state_v_first
ctx.use_gate_in_kernel = use_gate_in_kernel
return o.to(q.dtype), final_state
@staticmethod
@input_guard
@autocast_custom_bwd
def backward(
ctx,
do: torch.Tensor,
dht: torch.Tensor,
):
(
q,
q_rstd,
k,
k_rstd,
v,
g,
beta_raw,
beta,
A,
initial_state,
cu_seqlens,
chunk_indices,
g_input,
A_log,
dt_bias,
) = ctx.saved_tensors
dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A=A,
scale=ctx.scale,
initial_state=initial_state,
do=do,
dht=dht,
cu_seqlens=cu_seqlens,
cp_context=ctx.cp_context,
chunk_indices=chunk_indices,
state_v_first=ctx.state_v_first,
use_gate_in_kernel=ctx.use_gate_in_kernel,
g_input=g_input,
A_log=A_log,
dt_bias=dt_bias,
chunk_size=ctx.chunk_size,
)
if ctx.use_qk_l2norm_in_kernel:
dq = l2norm_bwd(q, q_rstd, dq)
dk = l2norm_bwd(k, k_rstd, dk)
if ctx.use_beta_sigmoid_in_kernel:
db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0)
return (
dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta_raw),
None, dh0, None, None, None, None, None, None, dA_log, ddt_bias,
None, None, None, None,
)
@dispatch('gated_delta_rule')
@torch.compiler.disable
def chunk_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cu_seqlens_cpu: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
**kwargs,
):
r"""
Args:
q (torch.Tensor):
queries of shape `[B, T, H, K]`.
k (torch.Tensor):
keys of shape `[B, T, H, K]`.
v (torch.Tensor):
values of shape `[B, T, HV, V]`.
GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.
g (torch.Tensor):
(forget) gating tensor of shape `[B, T, HV]`.
When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay).
When `use_gate_in_kernel=True`, `g` is the raw input before gate activation;
the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally.
beta (torch.Tensor):
betas of shape `[B, T, HV]`.
scale (Optional[float]):
Scale factor for the RetNet attention scores.
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
initial_state (Optional[torch.Tensor]):
Initial state of shape `[N, HV, K, V]` for `N` input sequences.
For equal-length input sequences, `N` equals the batch size `B`.
Default: `None`.
output_final_state (Optional[bool]):
Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.
use_qk_l2norm_in_kernel (bool):
Whether to apply L2norm to the q/k tensor internally. Default: `False`.
use_gate_in_kernel (bool):
Whether to compute the log-space GDN decay internally.
When `True`, the passed `g` is the raw input, and `A_log` must be provided.
The kernel fuses gate activation + chunk cumsum in a single pass.
Default: `False`.
A_log (Optional[torch.Tensor]):
Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.
dt_bias (Optional[torch.Tensor]):
Bias added to `g` before activation, of shape `[HV]`.
Only used when `use_gate_in_kernel=True`.
use_beta_sigmoid_in_kernel (bool):
Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel.
- If `True`, the passed `beta` acts as the raw beta logits.
- If `False`, `beta` is expected to already be in post-sigmoid space.
Default: `False`.
allow_neg_eigval (bool):
Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.
Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case
the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.
state_v_first (Optional[bool]):
Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.
cu_seqlens (torch.LongTensor):
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
consistent with the FlashAttention API.
cp_context (Optional[FLACPContext]):
Context parallel context for distributed training across multiple devices.
When provided, `initial_state` and `output_final_state` are not supported,
and `cu_seqlens` will be overridden by the context. Default: `None`.
Returns:
o (torch.Tensor):
Outputs of shape `[B, T, HV, V]`.
final_state (torch.Tensor):
Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.
Examples::
>>> import torch
>>> import torch.nn.functional as F
>>> from einops import rearrange
>>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
# inputs with equal lengths
>>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
>>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
>>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
>>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')
>>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid()
>>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda'))
>>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')
>>> o, ht = chunk_gated_delta_rule(
q, k, v, g, beta,
initial_state=h0,
output_final_state=True
)
# for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
>>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
# for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
>>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
>>> o, ht = chunk_gated_delta_rule(
q, k, v, g, beta,
initial_state=h0,
output_final_state=True,
cu_seqlens=cu_seqlens
)
"""
if 'transpose_state_layout' in kwargs:
if state_v_first:
raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.")
warnings.warn(
"`transpose_state_layout` is deprecated and renamed to `state_v_first`.",
DeprecationWarning,
stacklevel=2,
)
state_v_first = kwargs.pop('transpose_state_layout')
# Validate head dimensions
if q.shape[2] != k.shape[2]:
raise ValueError(
f"q and k must have the same number of heads, "
f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}"
)
H, HV = q.shape[2], v.shape[2]
if HV % H != 0:
raise ValueError(
f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by "
f"num_heads (H={H}), but got HV % H = {HV % H}"
)
if 'head_first' in kwargs:
raise DeprecationWarning(
"head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
)
chunk_size = kwargs.pop('chunk_size', 64)
if chunk_size not in (16, 32, 64):
raise ValueError(f"`chunk_size` must be 16, 32, or 64 for Gated Delta Rule, got {chunk_size}.")
if cp_context is not None:
assert initial_state is None, "Initial state is not supported for CP"
assert output_final_state is False, "Output final state is not supported for CP"
assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP"
cu_seqlens = cp_context.cu_seqlens
if cp_context.cu_seqlens_cpu is not None:
cu_seqlens_cpu = cp_context.cu_seqlens_cpu
if cu_seqlens is not None:
if q.shape[0] != 1:
raise ValueError(
f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
f"Please flatten variable-length inputs before processing.",
)
if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
raise ValueError(
f"The number of initial states is expected to be equal to the number of input sequences, "
f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.",
)
use_gate_in_kernel = kwargs.get('use_gate_in_kernel', False)
A_log = kwargs.get('A_log')
dt_bias = kwargs.get('dt_bias')
if use_gate_in_kernel:
assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True."
if allow_neg_eigval and not use_beta_sigmoid_in_kernel:
raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.")
if scale is None:
scale = k.shape[-1] ** -0.5
o, final_state = ChunkGatedDeltaRuleFunction.apply(
q,
k,
v,
g,
beta,
scale,
initial_state,
output_final_state,
state_v_first,
cu_seqlens,
cu_seqlens_cpu,
use_qk_l2norm_in_kernel,
use_gate_in_kernel,
A_log,
dt_bias,
use_beta_sigmoid_in_kernel,
allow_neg_eigval,
cp_context,
chunk_size,
)
return o, final_state
chunk_gdn = chunk_gated_delta_rule

425
upstream_ref/fla/ops/gated_delta_rule/chunk_fwd.py Normal file
View File

@@ -0,0 +1,425 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import triton
import triton.language as tl
from fla.ops.backends import dispatch
from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd
from fla.ops.utils import prepare_chunk_indices, solve_tril
from fla.ops.utils.cache import fla_cache_autotune
from fla.ops.utils.op import exp2
from fla.utils import IS_TF32_SUPPORTED, autotune_cache_kwargs
if IS_TF32_SUPPORTED:
SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32')
else:
SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee')
@triton.heuristics({
'USE_G': lambda args: args['g'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({'BK': BK}, num_warps=num_warps)
for BK in [32, 64]
for num_warps in [1, 2, 4]
],
key=['H', 'HV', 'K', 'BC'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def chunk_gated_delta_rule_fwd_kkt_solve_kernel(
k,
g,
beta,
A,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
BT: tl.constexpr,
BC: tl.constexpr,
BK: tl.constexpr,
USE_G: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
"""
Fused kernel: compute beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass.
This kernel fuses chunk_scaled_dot_kkt_fwd and solve_tril into a single kernel,
avoiding the HBM round-trip for the intermediate A matrix.
Steps:
1. Compute all 10 lower-triangular [BC, BC] blocks of beta * K @ K^T in registers
2. Apply gate and beta scaling
3. Forward substitution on diagonal blocks
4. Block merge to get full (I+A)^{-1}
5. Write result to A (output)
"""
i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
if i_t * BT >= T:
return
i_tc0 = i_t * BT
i_tc1 = i_t * BT + BC
i_tc2 = i_t * BT + 2 * BC
i_tc3 = i_t * BT + 3 * BC
k += (bos * H + i_h // (HV // H)) * K
A += (bos * HV + i_h) * BT
o_i = tl.arange(0, BC)
m_tc0 = (i_tc0 + o_i) < T
m_tc1 = (i_tc1 + o_i) < T
m_tc2 = (i_tc2 + o_i) < T
m_tc3 = (i_tc3 + o_i) < T
# load beta for each sub-chunk
p_b0 = beta + bos * HV + i_h + (i_tc0 + o_i) * HV
p_b1 = beta + bos * HV + i_h + (i_tc1 + o_i) * HV
p_b2 = beta + bos * HV + i_h + (i_tc2 + o_i) * HV
p_b3 = beta + bos * HV + i_h + (i_tc3 + o_i) * HV
b_b0 = tl.load(p_b0, mask=m_tc0, other=0.0).to(tl.float32)
b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32)
b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32)
b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32)
# load gate if used
if USE_G:
p_g0 = g + bos * HV + i_h + (i_tc0 + o_i) * HV
p_g1 = g + bos * HV + i_h + (i_tc1 + o_i) * HV
p_g2 = g + bos * HV + i_h + (i_tc2 + o_i) * HV
p_g3 = g + bos * HV + i_h + (i_tc3 + o_i) * HV
b_g0 = tl.load(p_g0, mask=m_tc0, other=0.0).to(tl.float32)
b_g1 = tl.load(p_g1, mask=m_tc1, other=0.0).to(tl.float32)
b_g2 = tl.load(p_g2, mask=m_tc2, other=0.0).to(tl.float32)
b_g3 = tl.load(p_g3, mask=m_tc3, other=0.0).to(tl.float32)
############################################################################
# Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T
############################################################################
# 4 diagonal blocks
b_A00 = tl.zeros([BC, BC], dtype=tl.float32)
b_A11 = tl.zeros([BC, BC], dtype=tl.float32)
b_A22 = tl.zeros([BC, BC], dtype=tl.float32)
b_A33 = tl.zeros([BC, BC], dtype=tl.float32)
# 6 off-diagonal blocks
b_A10 = tl.zeros([BC, BC], dtype=tl.float32)
b_A20 = tl.zeros([BC, BC], dtype=tl.float32)
b_A21 = tl.zeros([BC, BC], dtype=tl.float32)
b_A30 = tl.zeros([BC, BC], dtype=tl.float32)
b_A31 = tl.zeros([BC, BC], dtype=tl.float32)
b_A32 = tl.zeros([BC, BC], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
p_k0 = k + (i_tc0 + o_i)[:, None] * (H*K) + o_k[None, :]
b_k0 = tl.load(p_k0, mask=m_tc0[:, None] & (o_k[None, :] < K), other=0.0)
# diagonal block 0
b_A00 += tl.dot(b_k0, tl.trans(b_k0))
if i_tc1 < T:
p_k1 = k + (i_tc1 + o_i)[:, None] * (H*K) + o_k[None, :]
b_k1 = tl.load(p_k1, mask=m_tc1[:, None] & (o_k[None, :] < K), other=0.0)
# diagonal block 1
b_A11 += tl.dot(b_k1, tl.trans(b_k1))
# off-diagonal (1,0)
b_A10 += tl.dot(b_k1, tl.trans(b_k0))
if i_tc2 < T:
p_k2 = k + (i_tc2 + o_i)[:, None] * (H*K) + o_k[None, :]
b_k2 = tl.load(p_k2, mask=m_tc2[:, None] & (o_k[None, :] < K), other=0.0)
# diagonal block 2
b_A22 += tl.dot(b_k2, tl.trans(b_k2))
# off-diagonal (2,0), (2,1)
b_A20 += tl.dot(b_k2, tl.trans(b_k0))
b_A21 += tl.dot(b_k2, tl.trans(b_k1))
if i_tc3 < T:
p_k3 = k + (i_tc3 + o_i)[:, None] * (H*K) + o_k[None, :]
b_k3 = tl.load(p_k3, mask=m_tc3[:, None] & (o_k[None, :] < K), other=0.0)
# diagonal block 3
b_A33 += tl.dot(b_k3, tl.trans(b_k3))
# off-diagonal (3,0), (3,1), (3,2)
b_A30 += tl.dot(b_k3, tl.trans(b_k0))
b_A31 += tl.dot(b_k3, tl.trans(b_k1))
b_A32 += tl.dot(b_k3, tl.trans(b_k2))
############################################################################
# Step 2: apply gate and beta scaling
############################################################################
# apply gate, beta scaling, and masking
# m_d: strictly lower triangular mask for diagonal blocks
# m_tc: boundary mask to prevent NaN from 0 * inf (IEEE 754) when
# out-of-bounds g loads as 0 via boundary_check and exp2(0 - g_inbounds) overflows
m_d = o_i[:, None] > o_i[None, :]
m_I = o_i[:, None] == o_i[None, :]
if USE_G:
b_A00 *= tl.where(m_d & m_tc0[:, None] & m_tc0[None, :], exp2(b_g0[:, None] - b_g0[None, :]), 0.)
b_A11 *= tl.where(m_d & m_tc1[:, None] & m_tc1[None, :], exp2(b_g1[:, None] - b_g1[None, :]), 0.)
b_A22 *= tl.where(m_d & m_tc2[:, None] & m_tc2[None, :], exp2(b_g2[:, None] - b_g2[None, :]), 0.)
b_A33 *= tl.where(m_d & m_tc3[:, None] & m_tc3[None, :], exp2(b_g3[:, None] - b_g3[None, :]), 0.)
b_A10 *= tl.where(m_tc1[:, None] & m_tc0[None, :], exp2(b_g1[:, None] - b_g0[None, :]), 0.)
b_A20 *= tl.where(m_tc2[:, None] & m_tc0[None, :], exp2(b_g2[:, None] - b_g0[None, :]), 0.)
b_A21 *= tl.where(m_tc2[:, None] & m_tc1[None, :], exp2(b_g2[:, None] - b_g1[None, :]), 0.)
b_A30 *= tl.where(m_tc3[:, None] & m_tc0[None, :], exp2(b_g3[:, None] - b_g0[None, :]), 0.)
b_A31 *= tl.where(m_tc3[:, None] & m_tc1[None, :], exp2(b_g3[:, None] - b_g1[None, :]), 0.)
b_A32 *= tl.where(m_tc3[:, None] & m_tc2[None, :], exp2(b_g3[:, None] - b_g2[None, :]), 0.)
else:
b_A00 = tl.where(m_d, b_A00, 0.)
b_A11 = tl.where(m_d, b_A11, 0.)
b_A22 = tl.where(m_d, b_A22, 0.)
b_A33 = tl.where(m_d, b_A33, 0.)
# diagonal blocks: scaled by beta
b_A00 = b_A00 * b_b0[:, None]
b_A11 = b_A11 * b_b1[:, None]
b_A22 = b_A22 * b_b2[:, None]
b_A33 = b_A33 * b_b3[:, None]
# off-diagonal blocks: full block, scaled by beta
b_A10 = b_A10 * b_b1[:, None]
b_A20 = b_A20 * b_b2[:, None]
b_A21 = b_A21 * b_b2[:, None]
b_A30 = b_A30 * b_b3[:, None]
b_A31 = b_A31 * b_b3[:, None]
b_A32 = b_A32 * b_b3[:, None]
############################################################################
# Step 3: forward substitution on diagonal blocks -> (I + A_diag)^{-1}
#
# Same algorithm as solve_tril, but rows are extracted from in-register
# [BC, BC] tensor via tl.sum(tl.where(mask, tensor, 0), 0) instead of
# tl.load from HBM.
############################################################################
b_Ai00 = -b_A00
b_Ai11 = -b_A11
b_Ai22 = -b_A22
b_Ai33 = -b_A33
for i in range(2, min(BC, T - i_tc0)):
b_a00 = tl.sum(tl.where((o_i == i)[:, None], -b_A00, 0.), 0)
b_a00 = tl.where(o_i < i, b_a00, 0.)
b_a00 = b_a00 + tl.sum(b_a00[:, None] * b_Ai00, 0)
b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
for i in range(2, min(BC, T - i_tc1)):
b_a11 = tl.sum(tl.where((o_i == i)[:, None], -b_A11, 0.), 0)
b_a11 = tl.where(o_i < i, b_a11, 0.)
b_a11 = b_a11 + tl.sum(b_a11[:, None] * b_Ai11, 0)
b_Ai11 = tl.where((o_i == i)[:, None], b_a11, b_Ai11)
for i in range(2, min(BC, T - i_tc2)):
b_a22 = tl.sum(tl.where((o_i == i)[:, None], -b_A22, 0.), 0)
b_a22 = tl.where(o_i < i, b_a22, 0.)
b_a22 = b_a22 + tl.sum(b_a22[:, None] * b_Ai22, 0)
b_Ai22 = tl.where((o_i == i)[:, None], b_a22, b_Ai22)
for i in range(2, min(BC, T - i_tc3)):
b_a33 = tl.sum(tl.where((o_i == i)[:, None], -b_A33, 0.), 0)
b_a33 = tl.where(o_i < i, b_a33, 0.)
b_a33 = b_a33 + tl.sum(b_a33[:, None] * b_Ai33, 0)
b_Ai33 = tl.where((o_i == i)[:, None], b_a33, b_Ai33)
b_Ai00 += m_I
b_Ai11 += m_I
b_Ai22 += m_I
b_Ai33 += m_I
############################################################################
# Step 4: block merge -> full (I + A)^{-1}
############################################################################
b_Ai10 = -tl.dot(
tl.dot(b_Ai11, b_A10, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai00,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai21 = -tl.dot(
tl.dot(b_Ai22, b_A21, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai11,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai32 = -tl.dot(
tl.dot(b_Ai33, b_A32, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai22,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai20 = -tl.dot(
b_Ai22,
tl.dot(b_A20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_A21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION,
)
b_Ai31 = -tl.dot(
b_Ai33,
tl.dot(b_A31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_A32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION,
)
b_Ai30 = -tl.dot(
b_Ai33,
tl.dot(b_A30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_A31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_A32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION,
)
############################################################################
# Step 5: store full (I + A)^{-1} to output A
############################################################################
p_A00 = A + (i_tc0 + o_i)[:, None] * (HV*BT) + o_i[None, :]
p_A10 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + o_i[None, :]
p_A11 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
p_A20 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + o_i[None, :]
p_A21 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
p_A22 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :]
p_A30 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + o_i[None, :]
p_A31 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
p_A32 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :]
p_A33 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (3*BC + o_i)[None, :]
m_A0 = m_tc0[:, None] & (o_i[None, :] < BT)
m_A1 = m_tc1[:, None] & (o_i[None, :] < BT)
m_A2 = m_tc2[:, None] & (o_i[None, :] < BT)
m_A3 = m_tc3[:, None] & (o_i[None, :] < BT)
m_A11 = m_tc1[:, None] & ((BC + o_i)[None, :] < BT)
m_A21 = m_tc2[:, None] & ((BC + o_i)[None, :] < BT)
m_A22 = m_tc2[:, None] & ((2*BC + o_i)[None, :] < BT)
m_A31 = m_tc3[:, None] & ((BC + o_i)[None, :] < BT)
m_A32 = m_tc3[:, None] & ((2*BC + o_i)[None, :] < BT)
m_A33 = m_tc3[:, None] & ((3*BC + o_i)[None, :] < BT)
tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), mask=m_A0)
tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), mask=m_A1)
tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), mask=m_A11)
tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), mask=m_A2)
tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), mask=m_A21)
tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), mask=m_A22)
tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), mask=m_A3)
tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), mask=m_A31)
tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), mask=m_A32)
tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), mask=m_A33)
@dispatch('gated_delta_rule')
def chunk_gated_delta_rule_fwd_intra(
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_size: int = 64,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
r"""
GDN intra-chunk forward: fused or unfused kkt + solve_tril + recompute_w_u.
For ``chunk_size == 64``, this uses the fused kkt + solve_tril path. For
other supported chunk sizes, it computes the mathematically equivalent
representation with ``chunk_scaled_dot_kkt_fwd`` followed by ``solve_tril``.
Args:
k (torch.Tensor):
The key tensor of shape `[B, T, H, K]`.
v (torch.Tensor):
The value tensor of shape `[B, T, HV, V]`.
g (torch.Tensor):
The cumulative sum of the gate tensor of shape `[B, T, HV]`. Default: `None`.
beta (torch.Tensor):
The beta tensor of shape `[B, T, HV]`.
cu_seqlens (torch.LongTensor):
The cumulative sequence lengths. Default: `None`.
chunk_size (int):
The chunk size. Default: 64.
chunk_indices (torch.LongTensor):
Precomputed chunk indices. Default: `None`.
Returns:
w (torch.Tensor): shape `[B, T, HV, K]`
u (torch.Tensor): shape `[B, T, HV, V]`
A (torch.Tensor): shape `[B, T, HV, BT]`, the solved (I+A)^{-1} matrix
"""
if chunk_size not in (16, 32, 64):
raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.")
B, T, H, K, HV = *k.shape, beta.shape[2]
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
if BT == 64:
# Step 1: fused kkt + solve_tril
BC = 16
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
A = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype)
chunk_gated_delta_rule_fwd_kkt_solve_kernel[(NT, B * HV)](
k=k,
g=g,
beta=beta,
A=A,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
HV=HV,
K=K,
BT=BT,
BC=BC,
)
else:
# Step 1: mathematically equivalent unfused kkt + solve_tril for non-64 chunks
A = chunk_scaled_dot_kkt_fwd(
k=k,
g=g,
beta=beta,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_size=BT,
output_dtype=torch.float32,
)
A = solve_tril(
A=A,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
output_dtype=k.dtype,
)
# Step 2: recompute_w_u
w, u = recompute_w_u_fwd(
k=k,
v=v,
beta=beta,
A=A,
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
return w, u, A

478
upstream_ref/fla/ops/gated_delta_rule/fused_recurrent.py Normal file
View File

@@ -0,0 +1,478 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import warnings
import torch
import triton
import triton.language as tl
from fla.ops.utils.op import exp
from fla.ops.utils.softplus import softplus
from fla.utils import input_guard
@triton.heuristics({
'USE_G': lambda args: args['g'] is not None,
'USE_GK': lambda args: args['gk'] is not None,
'USE_GV': lambda args: args['gv'] is not None,
'USE_INITIAL_STATE': lambda args: args['h0'] is not None,
'STORE_FINAL_STATE': lambda args: args['ht'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
'USE_GATE_IN_KERNEL': lambda args: args['A_log'] is not None,
'HAS_DT_BIAS': lambda args: args['dt_bias'] is not None,
})
@triton.jit(do_not_specialize=['T'])
def fused_recurrent_gated_delta_rule_fwd_kernel(
q,
k,
v,
g,
gk,
gv,
beta,
A_log,
dt_bias,
o,
h0,
ht,
cu_seqlens,
scale,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
USE_G: tl.constexpr,
USE_GK: tl.constexpr,
USE_GV: tl.constexpr,
USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
IS_BETA_HEADWISE: tl.constexpr,
USE_INITIAL_STATE: tl.constexpr,
STORE_FINAL_STATE: tl.constexpr,
STATE_V_FIRST: tl.constexpr,
IS_VARLEN: tl.constexpr,
USE_GATE_IN_KERNEL: tl.constexpr,
HAS_DT_BIAS: tl.constexpr,
APPLY_BETA_SIGMOID: tl.constexpr,
ALLOW_NEG_EIGVAL: tl.constexpr,
):
pid = tl.program_id(0)
NV = tl.cdiv(V, BV)
i_v, i_nh = pid % NV, (pid // NV).to(tl.int64)
i_n, i_hv = i_nh // HV, i_nh % HV
i_h = i_hv // (HV // H)
if IS_VARLEN:
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_n * T, i_n * T + T
o_k = tl.arange(0, BK)
o_v = i_v * BV + tl.arange(0, BV)
p_q = q + (bos * H + i_h) * K + o_k
p_k = k + (bos * H + i_h) * K + o_k
p_v = v + (bos * HV + i_hv) * V + o_v
if USE_G:
p_g = g + bos * HV + i_hv
if USE_GK:
p_gk = gk + (bos * HV + i_hv) * K + o_k
if USE_GV:
p_gv = gv + (bos * HV + i_hv) * V + o_v
if IS_BETA_HEADWISE:
p_beta = beta + bos * HV + i_hv
else:
p_beta = beta + (bos * HV + i_hv) * V + o_v
p_o = o + (bos * HV + i_hv) * V + o_v
mask_k = o_k < K
mask_v = o_v < V
if STATE_V_FIRST:
mask_h = mask_v[:, None] & mask_k[None, :]
else:
mask_h = mask_k[:, None] & mask_v[None, :]
if STATE_V_FIRST:
b_h = tl.zeros([BV, BK], dtype=tl.float32)
else:
b_h = tl.zeros([BK, BV], dtype=tl.float32)
if USE_INITIAL_STATE:
if STATE_V_FIRST:
p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :]
else:
p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :]
b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
for _ in tl.range(0, T):
b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
if USE_QK_L2NORM_IN_KERNEL:
b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
b_q = b_q * scale
if IS_BETA_HEADWISE:
b_beta = tl.load(p_beta).to(tl.float32)
else:
b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
if APPLY_BETA_SIGMOID:
b_beta = tl.sigmoid(b_beta)
if ALLOW_NEG_EIGVAL:
b_beta = b_beta * 2
if USE_G:
b_g = tl.load(p_g).to(tl.float32)
if USE_GATE_IN_KERNEL:
b_A = tl.load(A_log + i_hv).to(tl.float32)
if HAS_DT_BIAS:
b_g = b_g + tl.load(dt_bias + i_hv).to(tl.float32)
b_g = -exp(b_A) * softplus(b_g)
b_h *= exp(b_g)
if USE_GK:
b_gk = tl.load(p_gk).to(tl.float32)
if STATE_V_FIRST:
b_h *= exp(b_gk[None, :])
else:
b_h *= exp(b_gk[:, None])
if USE_GV:
b_gv = tl.load(p_gv).to(tl.float32)
if STATE_V_FIRST:
b_h *= exp(b_gv[:, None])
else:
b_h *= exp(b_gv[None, :])
if STATE_V_FIRST:
b_v = b_beta * (b_v - tl.sum(b_h * b_k[None, :], 1))
b_h += b_v[:, None] * b_k[None, :]
b_o = tl.sum(b_h * b_q[None, :], 1)
else:
b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0))
b_h += b_k[:, None] * b_v
b_o = tl.sum(b_h * b_q[:, None], 0)
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
p_q += H*K
p_k += H*K
p_v += HV*V
if USE_G:
p_g += HV
if USE_GK:
p_gk += HV*K
if USE_GV:
p_gv += HV*V
p_beta += HV * (1 if IS_BETA_HEADWISE else V)
p_o += HV*V
if STORE_FINAL_STATE:
if STATE_V_FIRST:
p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :]
else:
p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :]
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
def fused_recurrent_gated_delta_rule_fwd(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
gk: torch.Tensor | None = None,
gv: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
B, T, H, K, V = *k.shape, v.shape[-1]
HV = v.shape[2]
N = B if cu_seqlens is None else len(cu_seqlens) - 1
BK = triton.next_power_of_2(K)
BV = min(8, triton.next_power_of_2(V)) if gv is None else triton.next_power_of_2(V)
NV = triton.cdiv(V, BV)
o = torch.empty_like(v)
if output_final_state:
if state_v_first:
final_state = q.new_empty(N, HV, V, K, dtype=torch.float32)
else:
final_state = q.new_empty(N, HV, K, V, dtype=torch.float32)
else:
final_state = None
grid = (NV * N * HV,)
fused_recurrent_gated_delta_rule_fwd_kernel[grid](
q=q,
k=k,
v=v,
g=g,
gk=gk,
gv=gv,
beta=beta,
A_log=A_log,
dt_bias=dt_bias,
o=o,
h0=initial_state,
ht=final_state,
cu_seqlens=cu_seqlens,
scale=scale,
T=T,
H=H,
HV=HV,
K=K,
V=V,
BK=BK,
BV=BV,
IS_BETA_HEADWISE=beta.ndim != v.ndim,
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel,
ALLOW_NEG_EIGVAL=allow_neg_eigval,
STATE_V_FIRST=state_v_first,
num_warps=1,
num_stages=3,
)
return o, final_state
class FusedRecurrentFunction(torch.autograd.Function):
@staticmethod
@input_guard
def forward(
ctx,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
gk: torch.Tensor | None = None,
gv: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
):
o, final_state = fused_recurrent_gated_delta_rule_fwd(
q=q,
k=k,
v=v,
g=g,
gk=gk,
gv=gv,
beta=beta,
A_log=A_log,
dt_bias=dt_bias,
scale=scale,
initial_state=initial_state,
output_final_state=output_final_state,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
allow_neg_eigval=allow_neg_eigval,
state_v_first=state_v_first,
cu_seqlens=cu_seqlens,
)
return o, final_state
@staticmethod
@input_guard
def backward(ctx, do, dht):
raise NotImplementedError(
"Backward pass is not implemented yet and we do not have plans to implement it "
"because we haven't figured out how to compute dg without materializing the full "
"hidden states for all time steps.",
)
def fused_recurrent_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
gk: torch.Tensor | None = None,
gv: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_gate_in_kernel: bool = False,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
**kwargs,
) -> tuple[torch.Tensor, torch.Tensor]:
r"""
Args:
q (torch.Tensor):
queries of shape `[B, T, H, K]`.
k (torch.Tensor):
keys of shape `[B, T, H, K]`.
v (torch.Tensor):
values of shape `[B, T, HV, V]`.
GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.
g (torch.Tensor):
g (decays) of shape `[B, T, HV]`. Default: `None`.
When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay).
When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses
`-exp(A_log) * softplus(g + dt_bias)` internally per step.
gk (torch.Tensor):
gk (decays) of shape `[B, T, HV, K]`. Default: `None`.
gv (torch.Tensor):
gv (decays) of shape `[B, T, HV, V]`. Default: `None`.
beta (torch.Tensor):
betas of shape `[B, T, HV]`.
scale (Optional[float]):
Scale factor for the RetNet attention scores.
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
initial_state (Optional[torch.Tensor]):
Initial state of shape `[N, HV, K, V]` for `N` input sequences.
For equal-length input sequences, `N` equals the batch size `B`.
Default: `None`.
output_final_state (Optional[bool]):
Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.
use_qk_l2norm_in_kernel (Optional[bool]):
Whether to use L2 normalization in the kernel. Default: `False`.
use_gate_in_kernel (bool):
Whether to compute the log-space GDN decay internally.
When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses
gate activation into the recurrence. Default: `False`.
A_log (Optional[torch.Tensor]):
Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.
dt_bias (Optional[torch.Tensor]):
Bias added to `g` before activation, of shape `[HV]`.
Only used when `use_gate_in_kernel=True`.
use_beta_sigmoid_in_kernel (Optional[bool]):
Whether to apply `torch.sigmoid(beta)` inside the kernel.
- If `True`, the passed `beta` acts as the raw beta logits.
- If `False`, `beta` is expected to already be in post-sigmoid space.
Default: `False`.
allow_neg_eigval (Optional[bool]):
Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.
Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case
the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.
state_v_first (Optional[bool]):
Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.
cu_seqlens (torch.LongTensor):
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
consistent with the FlashAttention API.
Returns:
o (torch.Tensor):
Outputs of shape `[B, T, HV, V]`.
final_state (torch.Tensor):
Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.
Examples::
>>> import torch
>>> import torch.nn.functional as F
>>> from einops import rearrange
>>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
# inputs with equal lengths
>>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
>>> q = torch.randn(B, T, H, K, device='cuda')
>>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)
>>> v = torch.randn(B, T, HV, V, device='cuda')
>>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
>>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
>>> h0 = torch.randn(B, HV, K, V, device='cuda')
>>> o, ht = fused_gated_recurrent_delta_rule(
q, k, v, g, beta,
initial_state=h0,
output_final_state=True
)
# for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
>>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))
# for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
>>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
>>> o, ht = fused_gated_recurrent_delta_rule(
q, k, v, g, beta,
initial_state=h0,
output_final_state=True,
cu_seqlens=cu_seqlens
)
"""
if 'transpose_state_layout' in kwargs:
if state_v_first:
raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.")
warnings.warn(
"`transpose_state_layout` is deprecated and renamed to `state_v_first`.",
DeprecationWarning,
stacklevel=2,
)
state_v_first = kwargs.pop('transpose_state_layout')
if cu_seqlens is not None:
if q.shape[0] != 1:
raise ValueError(
f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
f"Please flatten variable-length inputs before processing.",
)
if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
raise ValueError(
f"The number of initial states is expected to be equal to the number of input sequences, "
f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.",
)
if scale is None:
scale = k.shape[-1] ** -0.5
if beta is None:
beta = torch.ones_like(q[..., 0])
if use_gate_in_kernel:
if A_log is None:
raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True`.")
if g is None:
raise ValueError("`g` (raw pre-activation) must be provided when `use_gate_in_kernel=True`.")
else:
A_log = None
dt_bias = None
if allow_neg_eigval and not use_beta_sigmoid_in_kernel:
raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.")
o, final_state = FusedRecurrentFunction.apply(
q,
k,
v,
g,
gk,
gv,
beta,
A_log,
dt_bias,
scale,
initial_state,
output_final_state,
use_qk_l2norm_in_kernel,
use_beta_sigmoid_in_kernel,
allow_neg_eigval,
state_v_first,
cu_seqlens,
)
return o, final_state
fused_recurrent_gdn = fused_recurrent_gated_delta_rule

344
upstream_ref/fla/ops/gated_delta_rule/gate.py Normal file
View File

@@ -0,0 +1,344 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from fla.ops.backends import dispatch
from fla.ops.utils.cache import fla_cache_autotune
from fla.ops.utils.index import prepare_chunk_indices
from fla.ops.utils.op import exp
from fla.ops.utils.softplus import softplus
from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard
def naive_gdn_gate(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""
Torch reference implementation for GDN gate computation.
Computes: ``g = -A_log.exp() * softplus(g + dt_bias)``
Args:
g (torch.Tensor):
Input tensor of shape `[..., HV]`.
A_log (torch.Tensor):
Decay parameter tensor with `HV` elements.
dt_bias (torch.Tensor | None):
Optional bias tensor added to `g` before activation, shape `[HV]`.
Returns:
Output tensor of shape `[..., HV]`.
"""
g = g.float()
if dt_bias is not None:
g = g + dt_bias.float()
return (-A_log.float().exp() * F.softplus(g)).to(output_dtype)
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
'HAS_SCALE': lambda args: args['scale'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({}, num_warps=num_warps)
for num_warps in [1, 2, 4, 8]
],
key=['H', 'BT', 'IS_VARLEN', 'REVERSE'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def gdn_gate_chunk_cumsum_scalar_kernel(
g,
A_log,
dt_bias,
o,
scale,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
BT: tl.constexpr,
REVERSE: tl.constexpr,
HAS_BIAS: tl.constexpr,
HAS_SCALE: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
p_g = g + bos * H + i_h + o_t * H
p_o = o + bos * H + i_h + o_t * H
b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
b_A = tl.load(A_log + i_h).to(tl.float32)
b_gate = -exp(b_A) * softplus(b_g)
b_o = tl.cumsum(b_gate, axis=0)
if REVERSE:
b_z = tl.sum(b_gate, axis=0)
b_o = -b_o + b_z[None] + b_gate
if HAS_SCALE:
b_o *= scale
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t)
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({}, num_warps=num_warps)
for num_warps in [1, 2, 4, 8]
],
key=['H', 'BT'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def gdn_gate_bwd_kernel(
g,
A_log,
dt_bias,
dyg,
dg,
dA,
T,
H: tl.constexpr,
BT: tl.constexpr,
HAS_BIAS: tl.constexpr,
):
i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1)
b_A = tl.load(A_log + i_h).to(tl.float32)
o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
p_g = g + i_h + o_t * H
p_dg = dg + i_h + o_t * H
p_dyg = dyg + i_h + o_t * H
b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
b_dyg = tl.load(p_dyg, mask=m_t, other=0.0).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
# gate = -exp(A_log) * softplus(g + bias)
# d(gate)/d(g) = -exp(A_log) * sigmoid(g + bias) (softplus' = sigmoid)
# d(gate)/d(A_log) = -exp(A_log) * softplus(g + bias) = gate
b_neg_expA = -exp(b_A)
b_yg = b_neg_expA * softplus(b_g)
b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g))
b_dA = tl.sum(b_dyg * b_yg, 0)
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t)
tl.store(dA + i_t * H + i_h, b_dA)
@input_guard
@dispatch('gated_delta_rule')
def gdn_gate_chunk_cumsum(
g: torch.Tensor,
A_log: torch.Tensor,
chunk_size: int,
scale: float = None,
dt_bias: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
output_dtype: torch.dtype | None = torch.float,
) -> torch.Tensor:
B, T, H = g.shape
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
o = torch.empty_like(g, dtype=output_dtype or g.dtype)
gdn_gate_chunk_cumsum_scalar_kernel[(NT, B * H)](
g=g,
A_log=A_log,
dt_bias=dt_bias,
o=o,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
BT=BT,
REVERSE=False,
)
return o
@dispatch('gated_delta_rule')
def gdn_gate_bwd(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None,
dyg: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
H = g.shape[-1]
T = g.numel() // H
BT = 32
NT = triton.cdiv(T, BT)
dg = torch.empty_like(g, dtype=torch.float32)
dA = A_log.new_empty(NT, H, dtype=torch.float32)
gdn_gate_bwd_kernel[(NT, H)](
g=g,
A_log=A_log,
dt_bias=dt_bias,
dyg=dyg,
dg=dg,
dA=dA,
T=T,
H=H,
BT=BT,
)
dg = dg.view_as(g).type_as(g)
dA = dA.sum(0).view_as(A_log).type_as(A_log)
dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None
return dg, dA, dbias
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages)
for BT in [32, 64, 128]
for num_warps in [1, 2, 4, 8]
for num_stages in [2, 3]
],
key=['H'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def gdn_gate_fwd_kernel(
g,
A_log,
dt_bias,
yg,
T,
H: tl.constexpr,
BT: tl.constexpr,
HAS_BIAS: tl.constexpr,
):
i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1)
b_A = tl.load(A_log + i_h).to(tl.float32)
o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
p_g = g + i_h + o_t * H
p_yg = yg + i_h + o_t * H
b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
b_yg = -exp(b_A) * softplus(b_g)
tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), mask=m_t)
@dispatch('gated_delta_rule')
def gdn_gate_fwd(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
H = g.shape[-1]
T = g.numel() // H
yg = torch.empty_like(g, dtype=output_dtype)
def grid(meta):
return (triton.cdiv(T, meta['BT']), H)
gdn_gate_fwd_kernel[grid](
g=g,
A_log=A_log,
dt_bias=dt_bias,
yg=yg,
T=T,
H=H,
)
return yg
class GDNGateFunction(torch.autograd.Function):
@staticmethod
@input_guard
@autocast_custom_fwd
def forward(
ctx,
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
yg = gdn_gate_fwd(g=g, A_log=A_log, dt_bias=dt_bias, output_dtype=output_dtype)
ctx.save_for_backward(g, A_log, dt_bias)
return yg
@staticmethod
@input_guard
@autocast_custom_bwd
def backward(ctx, dyg: torch.Tensor):
g, A_log, dt_bias = ctx.saved_tensors
dg, dA, dbias = gdn_gate_bwd(g=g, A_log=A_log, dt_bias=dt_bias, dyg=dyg)
return dg, dA, dbias, None
@torch.compiler.disable
def fused_gdn_gate(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
r"""
Fused GDN gate computation with autograd support.
Computes: ``g = -A_log.exp() * softplus(g + dt_bias)``
Args:
g (torch.Tensor):
Input tensor of shape `[..., HV]`.
A_log (torch.Tensor):
Decay parameter tensor with `HV` elements.
dt_bias (torch.Tensor | None):
Optional bias tensor added to `g` before activation, shape `[HV]`.
output_dtype (torch.dtype):
The dtype of the output tensor. Default: `torch.float32`.
Returns:
Output tensor of shape `[..., HV]`.
"""
return GDNGateFunction.apply(g, A_log, dt_bias, output_dtype)

161
upstream_ref/fla/ops/gated_delta_rule/naive.py Normal file
View File

@@ -0,0 +1,161 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import torch.nn.functional as F
from einops import rearrange
def naive_recurrent_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
g: torch.Tensor,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
):
"""
Reference PyTorch implementation of recurrent gated delta rule.
Args:
q: [B, T, H, K]
k: [B, T, H, K]
v: [B, T, H, V]
beta: [B, T, H]
g: [B, T, H]
scale: float, optional
initial_state: [B, H, K, V], optional
output_final_state: bool
Returns:
o: [B, T, H, V]
final_state: [B, H, K, V] if output_final_state else None
"""
q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
B, H, T, K, V = *k.shape, v.shape[-1]
o = torch.zeros(B, H, T, V).to(v)
h = torch.zeros(B, H, K, V).to(v)
if initial_state is not None:
h = initial_state.to(torch.float32)
if scale is None:
scale = 1 / (q.shape[-1] ** 0.5)
q = q * scale
for i in range(T):
b_q = q[:, :, i]
b_k = k[:, :, i]
b_v = v[:, :, i].clone()
h = h.clone() * g[:, :, i].exp()[..., None, None]
b_beta = beta[:, :, i]
b_v = b_v - (h.clone() * b_k[..., None]).sum(-2)
b_v = b_v * b_beta[..., None]
h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2)
o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h)
if not output_final_state:
h = None
o = o.transpose(1, 2).contiguous()
return o, h
def naive_chunk_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
chunk_size: int = 64,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
):
"""
Reference PyTorch implementation of chunk gated delta rule.
Args:
q: [B, T, H, K]
k: [B, T, H, K]
v: [B, T, H, V]
g: [B, T, H]
beta: [B, T, H]
chunk_size: int
scale: float, optional
initial_state: [B, H, K, V], optional
output_final_state: bool
Returns:
o: [B, T, H, V]
final_state: [B, H, K, V] if output_final_state else None
"""
BT = chunk_size
if scale is None:
scale = 1 / (q.shape[-1] ** 0.5)
q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
T = q.shape[-2]
pad_len = (BT - (T % BT)) % BT
if pad_len > 0:
q = F.pad(q, (0, 0, 0, pad_len))
k = F.pad(k, (0, 0, 0, pad_len))
v = F.pad(v, (0, 0, 0, pad_len))
beta = F.pad(beta, (0, pad_len))
g = F.pad(g, (0, pad_len))
q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g])
decay = g
chunk_size = BT
b, h, l, d_k = q.shape
d_v = v.shape[-1]
q = q * scale
v = v * beta[..., None]
k_beta = k * beta[..., None]
assert l % chunk_size == 0
# note that diagonal is masked.
mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)
q, k, v, k_beta, decay = map(
lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size),
[q, k, v, k_beta, decay.unsqueeze(-1)],
)
decay = decay.squeeze(-1).cumsum(-1)
decay_exp = decay.exp()[..., None]
L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0)
for i in range(1, chunk_size):
attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2)
attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
attn = attn
k_cumsum = attn @ v
k_cumdecay = attn @ (k_beta * decay_exp)
v = k_cumsum
S = k.new_zeros(b, h, d_k, d_v)
if initial_state is not None:
S = initial_state.to(torch.float32)
o = torch.zeros_like(v)
mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)
for i in range(0, l // chunk_size):
q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i]
attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0)
v_prime = (k_cumdecay[:, :, i]) @ S
v_new = v_i - v_prime
o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S
o[:, :, i] = o_inter + attn @ v_new
S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()
[..., None]).transpose(-1, -2) @ v_new
if not output_final_state:
S = None
# unpad
o = rearrange(o, 'b h n c d -> b h (n c) d')
o = o[:, :, :T]
o = o.transpose(1, 2)
return o, S

347
upstream_ref/fla/ops/gated_delta_rule/wy_fast.py Normal file
View File

@@ -0,0 +1,347 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import triton
import triton.language as tl
from fla.ops.backends import dispatch
from fla.ops.utils import prepare_chunk_indices
from fla.ops.utils.cache import fla_cache_autotune
from fla.ops.utils.op import exp2
from fla.utils import IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem
# Blackwell can select unstable Triton configs for prepare_wy_repr_bwd_kernel
# during autotuning (see #913). Restrict it to the config that has been
# validated on B200 until the wider config space is re-validated.
PREPARE_WY_REPR_BWD_NUM_WARPS = [2] if IS_NVIDIA_BLACKWELL else [2, 4]
PREPARE_WY_REPR_BWD_NUM_STAGES = [4] if IS_NVIDIA_BLACKWELL else [2, 3, 4]
@triton.heuristics({
'USE_G': lambda args: args['g'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in [2, 4, 8]
for num_stages in [2, 3, 4]
],
key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def recompute_w_u_fwd_kernel(
k,
v,
beta,
w,
u,
A,
g,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
USE_G: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
o_t = i_t * BT + tl.arange(0, BT)
o_A = tl.arange(0, BT)
m_t = o_t < T
m_A = m_t[:, None] & (o_A[None, :] < BT)
p_b = beta + bos*HV + i_h + o_t * HV
b_b = tl.load(p_b, mask=m_t, other=0.0)
p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (HV*BT) + o_A[None, :]
b_A = tl.load(p_A, mask=m_A, other=0.0)
for i_v in range(tl.cdiv(V, BV)):
o_v = i_v * BV + tl.arange(0, BV)
m_v = m_t[:, None] & (o_v[None, :] < V)
p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
p_u = u + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
b_v = tl.load(p_v, mask=m_v, other=0.0)
b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
b_u = tl.dot(b_A, b_vb, allow_tf32=False)
tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v)
if USE_G:
p_g = g + (bos*HV + i_h) + o_t * HV
b_g = exp2(tl.load(p_g, mask=m_t, other=0.0))
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
m_k = m_t[:, None] & (o_k[None, :] < K)
p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
p_w = w + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
b_k = tl.load(p_k, mask=m_k, other=0.0)
b_kb = b_k * b_b[:, None]
if USE_G:
b_kb *= b_g[:, None]
b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k)
@triton.heuristics({
'USE_G': lambda args: args['g'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in PREPARE_WY_REPR_BWD_NUM_WARPS
for num_stages in PREPARE_WY_REPR_BWD_NUM_STAGES
],
key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def prepare_wy_repr_bwd_kernel(
k,
v,
beta,
g,
A,
dw,
du,
dk,
dv,
db,
dg,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
USE_G: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
o_t = i_t * BT + tl.arange(0, BT)
o_A = tl.arange(0, BT)
m_t = o_t < T
m_AT = (o_A[:, None] < BT) & m_t[None, :]
p_b = beta + (bos*HV + i_h) + o_t * HV
p_db = db + (bos*HV + i_h) + o_t * HV
p_A = A + (bos*HV + i_h) * BT + o_A[:, None] + o_t[None, :] * (HV*BT)
b_b = tl.load(p_b, mask=m_t, other=0.0)
b_db = tl.zeros([BT], dtype=tl.float32)
b_A = tl.load(p_A, mask=m_AT, other=0.0)
b_dA = tl.zeros([BT, BT], dtype=tl.float32)
if USE_G:
p_g = g + (bos*HV + i_h) + o_t * HV
b_g = tl.load(p_g, mask=m_t, other=0.0)
b_g_exp = exp2(b_g)
b_dg = tl.zeros([BT], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
m_k = m_t[:, None] & (o_k[None, :] < K)
p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
p_dw = dw + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
# [BT, BK]
b_k = tl.load(p_k, mask=m_k, other=0.0)
if USE_G:
b_kbg = b_k * (b_b * b_g_exp)[:, None]
else:
b_kbg = b_k * b_b[:, None]
b_dw = tl.load(p_dw, mask=m_k, other=0.0)
b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype))
b_dkbg = tl.dot(b_A, b_dw)
if USE_G:
b_dk = b_dkbg * (b_g_exp * b_b)[:, None]
b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1)
b_dg += tl.sum(b_dkbg * b_kbg, 1)
else:
b_dk = b_dkbg * b_b[:, None]
b_db += tl.sum(b_dkbg * b_k, 1)
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k)
for i_v in range(tl.cdiv(V, BV)):
o_v = i_v * BV + tl.arange(0, BV)
m_v = m_t[:, None] & (o_v[None, :] < V)
p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
p_dv = dv + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
p_du = du + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
b_v = tl.load(p_v, mask=m_v, other=0.0)
b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
b_du = tl.load(p_du, mask=m_v, other=0.0)
b_dA += tl.dot(b_du, tl.trans(b_vb))
b_dvb = tl.dot(b_A, b_du)
b_dv = b_dvb * b_b[:, None]
b_db += tl.sum(b_dvb * b_v, 1)
tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v)
m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
b_dA = tl.where(m_A, b_dA, 0)
b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)
b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))
if USE_G:
b_dA *= exp2(b_g[:, None] - b_g[None, :])
b_A = tl.zeros([BT, BT], dtype=tl.float32)
b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty)
tl.debug_barrier()
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
m_k = m_t[:, None] & (o_k[None, :] < K)
p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
b_k = tl.load(p_k, mask=m_k, other=0.0)
b_kt = tl.trans(b_k)
b_kb = b_k * b_b[:, None]
b_A += tl.dot(b_k, b_kt)
b_dkb = tl.dot(b_dA, b_k)
b_db += tl.sum(b_dkb * b_k, 1)
b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb).to(b_dA.dtype), b_dA))
b_dk += tl.load(p_dk, mask=m_k, other=0.0)
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k)
tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t)
b_A *= b_b[:, None]
if USE_G:
b_AdA = b_dA * b_A
p_dg = dg + (bos*HV + i_h) + o_t * HV
b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0)
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t)
@dispatch('gated_delta_rule')
def recompute_w_u_fwd(
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
g: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2]
BT = A.shape[-1]
BK = 64
BV = 64
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
w = k.new_empty(B, T, HV, K)
u = torch.empty_like(v)
recompute_w_u_fwd_kernel[(NT, B*HV)](
k=k,
v=v,
beta=beta,
w=w,
u=u,
A=A,
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
HV=HV,
K=K,
V=V,
BT=BT,
BK=BK,
BV=BV,
)
return w, u
@dispatch('gated_delta_rule')
def prepare_wy_repr_bwd(
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
dw: torch.Tensor,
du: torch.Tensor,
g: torch.Tensor = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2]
BT = A.shape[-1]
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
CONST_TILING = 64 if check_shared_mem() else 32
BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING)
BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING)
dk = k.new_empty(B, T, HV, K)
dv = torch.empty_like(v)
dg = torch.empty_like(g) if g is not None else None
db = torch.empty_like(beta)
prepare_wy_repr_bwd_kernel[(NT, B * HV)](
k=k,
v=v,
beta=beta,
g=g,
A=A,
dw=dw,
du=du,
dk=dk,
dv=dv,
db=db,
dg=dg,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
HV=HV,
K=K,
V=V,
BT=BT,
BK=BK,
BV=BV,
)
if H != HV:
dk = dk.view(B, T, H, HV // H, K).sum(3)
return dk, dv, db, dg
fwd_recompute_w_u = recompute_w_u_fwd
bwd_prepare_wy_repr = prepare_wy_repr_bwd