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

feat: import CUDA kernels from xllm/CCCL/FLA upstream repos

Browse Source 
Sources cloned and tree'd (no --depth):
  - jd-opensource/xllm: ILU kernels, CUDA kernels, MoE kernels
  - NVIDIA/cccl: CUB tuning/dispatch headers (block-level primitives)
  - fla-org/flash-linear-attention: Triton GDN kernels
  - NVIDIA/cutlass: grouped GEMM reference (read, not copied)
  - Dao-AILab/flash-attention: attention kernel reference (SM80+, read only)

New CUDA kernels (from xllm, SM-agnostic, portable to BI-V100):
  ex_engine/xllm_kernels/cuda/activation.cu    (188 lines) — silu_and_mul, gelu
  ex_engine/xllm_kernels/cuda/norm.cu          (600 lines) — rms_norm, fused_add_rms_norm
  ex_engine/xllm_kernels/cuda/rope.cu          (258 lines) — rotary_embedding
  ex_engine/xllm_kernels/cuda/block_copy.cu    (209 lines) — copy_blocks, swap_blocks
  ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu (101 lines) — KV cache ops
  ex_engine/xllm_kernels/cuda/headers/         (5 headers for compilation)

ILU bridge kernel sources (from xllm, verified SAME as upstream):
  ex_engine/xllm_kernels/ilu/    (10 files, 925 lines total)
  — activation.cpp, attention.cpp, fused_moe.cpp, group_gemm.cpp,
    matmul.cpp, norm.cpp, rope.cpp, ilu_ops_api.h, ixformer.h, utils.h

FLA Triton GDN kernels (for GatedDeltaNet without SM90+ FlashQLA):
  ex_engine/fla_kernels/gated_delta_rule/  (7 files, 2370 lines)
  — chunk_fwd.py (428), chunk.py (487), wy_fast.py (409),
    fused_recurrent.py (392), naive.py (161), gate.py (380)

CCCL sync (12 tuning + 14 dispatch headers updated from NVIDIA/cccl):
  cccl_upstream/cub/cub/device/dispatch/tuning/ — 12 changed files synced
  cccl_upstream/cub/cub/device/dispatch/ — 14 changed dispatch files synced

Compilation targets for real machine (ivcore10):
  1. CUDA kernels: --cuda-gpu-arch=ivcore10 via corex clang/16
  2. ILU bridges: torch.utils.cpp_extension linking ixformer .so
  3. FLA kernels: Triton JIT (if Triton works on BI-V100)
This commit is contained in:
claude
2026-08-14 07:48:52 +00:00
parent 051b02d3cd
commit 8d75652949
66 changed files with 12834 additions and 5458 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
ex_engine/fla_kernels/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",
]

591
ex_engine/fla_kernels/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

428
ex_engine/fla_kernels/gated_delta_rule/chunk_fwd.py Normal file
View File

@@ -0,0 +1,428 @@
# 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_INTEL, 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)
# The fused kernel keeps ten [BC, BC] fp32 accumulators live across the K loop.
# That fits NVIDIA's register file but spills on Intel GPUs, where the unfused
# two-kernel path measures 2.3-3.0x faster despite the extra HBM round-trip.
if BT == 64 and not IS_INTEL:
# 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
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
ex_engine/fla_kernels/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
ex_engine/fla_kernels/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
ex_engine/fla_kernels/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

351
ex_engine/fla_kernels/gated_delta_rule/wy_fast.py Normal file
View File

@@ -0,0 +1,351 @@
# 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_INTEL, 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]
# Intel keeps scaling past the warp counts NVIDIA prefers: 16 warps is ~1.3x faster
# than 8 for recompute_w_u.
RECOMPUTE_W_U_NUM_WARPS = [2, 4, 8, 16] if IS_INTEL else [2, 4, 8]
@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 RECOMPUTE_W_U_NUM_WARPS
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

65
ex_engine/fla_kernels/utils/__init__.py Normal file
View File

@@ -0,0 +1,65 @@
# 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 .csr import prepare_block_csr
from .cumsum import (
chunk_global_cumsum,
chunk_global_cumsum_scalar,
chunk_global_cumsum_vector,
chunk_local_cumsum,
chunk_local_cumsum_scalar,
chunk_local_cumsum_vector,
)
from .index import (
get_max_num_splits,
prepare_chunk_indices,
prepare_chunk_offsets,
prepare_cu_seqlens_from_lens,
prepare_cu_seqlens_from_mask,
prepare_lens,
prepare_lens_from_mask,
prepare_position_ids,
prepare_sequence_ids,
prepare_token_indices,
)
from .logsumexp import logsumexp_fwd
from .matmul import addmm, matmul
from .pack import pack_sequence, unpack_sequence
from .pooling import mean_pooling
from .softmax import softmax_bwd, softmax_fwd
from .softplus import softplus
from .solve_tril import solve_tril
__all__ = [
"addmm",
"chunk_global_cumsum",
"chunk_global_cumsum_scalar",
"chunk_global_cumsum_vector",
"chunk_local_cumsum",
"chunk_local_cumsum_scalar",
"chunk_local_cumsum_vector",
"get_max_num_splits",
"logsumexp_fwd",
"matmul",
"mean_pooling",
"pack_sequence",
"prepare_block_csr",
"prepare_chunk_indices",
"prepare_chunk_offsets",
"prepare_cu_seqlens_from_lens",
"prepare_cu_seqlens_from_mask",
"prepare_lens",
"prepare_lens_from_mask",
"prepare_position_ids",
"prepare_sequence_ids",
"prepare_token_indices",
"softmax_bwd",
"softmax_fwd",
"softplus",
"solve_tril",
"unpack_sequence",
]

449
ex_engine/fla_kernels/utils/cache.py Normal file
View File

@@ -0,0 +1,449 @@
# 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 dataclasses
import enum
import json
import logging
import os
import re
from functools import cache, lru_cache
from pathlib import Path
from typing import Any
import torch
import triton
from packaging import version
from triton.runtime.autotuner import Autotuner
TRITON_ABOVE_3_5_1 = version.parse(triton.__version__) >= version.parse("3.5.1")
TRITON_ABOVE_3_4_0 = version.parse(triton.__version__) >= version.parse("3.4.0")
class FlaCacheMode(enum.Enum):
"""Controls how FLA loads kernel configs from its config cache (FLA_CACHE_MODE env var).
DISABLED — skip all cache lookups, always fall back to Triton autotune (default when FLA_CACHE_MODE is unset)
STRICT — exact key match only; falls back to Triton autotune if no match
FUZZY — exact key match → fuzzy key match; falls back to Triton autotune if no match
FULL — exact key match → fuzzy key match → default_config fallback
DEFAULT — use only the top-level default_config field, skip key-based lookup
ALWAYS — like DEFAULT, but re-reads config files on every kernel call;
useful for debugging: edit default_config in a JSON file and the next
kernel call picks it up without restarting the process
"""
DISABLED = "disabled"
STRICT = "strict"
FUZZY = "fuzzy"
FULL = "full"
DEFAULT = "default"
ALWAYS = "always"
def uses_default_config(self) -> bool:
"""Return True for modes that may fall back to default_config (FULL, DEFAULT, ALWAYS)."""
return self in (FlaCacheMode.FULL, FlaCacheMode.DEFAULT, FlaCacheMode.ALWAYS)
@classmethod
def from_env(cls) -> "FlaCacheMode":
mode_str = os.environ.get("FLA_CACHE_MODE", cls.DISABLED.value)
try:
return cls(mode_str)
except ValueError:
valid = [m.value for m in cls]
raise ValueError(
f"Invalid FLA_CACHE_MODE={mode_str!r}. Valid values: {valid}"
) from None
FLA_CACHE_MODE: FlaCacheMode = FlaCacheMode.from_env()
logger = logging.getLogger(__name__)
def sanitize_gpu_name(gpu_name: str) -> str:
sanitized = re.sub(r"[^0-9A-Za-z]+", "_", gpu_name)
sanitized = sanitized.strip("_")
return sanitized or "unknown_gpu"
@lru_cache(maxsize=1)
def get_gpu_info():
"""Get GPU model information.
This function detects the GPU model and returns a sanitized string identifier.
It prioritizes FLA_GPU_NAME environment variable if set, then detects from
available hardware (CUDA, ROCm, Intel GPU, or CPU).
"""
# Check if GPU name is overridden via environment variable
gpu_name = None
# Check if GPU name is overridden via environment variable
if "FLA_GPU_NAME" in os.environ:
gpu_name = os.environ["FLA_GPU_NAME"]
# Try to get device name based on availability
elif torch.cuda.is_available():
# Works for both NVIDIA and AMD GPUs (ROCm)
gpu_name = torch.cuda.get_device_name(0)
elif hasattr(torch, 'xpu') and torch.xpu.is_available():
gpu_name = torch.xpu.get_device_name(0)
if gpu_name:
return sanitize_gpu_name(gpu_name)
# Default to CPU if no GPU available
return "cpu"
def get_fla_config_dir() -> Path:
"""Get FLA's configs directory.
The directory can be overridden by setting the FLA_CONFIG_DIR environment variable.
If set, configs will be loaded directly from $FLA_CONFIG_DIR/. Otherwise FLA
falls back to the default fla/configs/{GPU}/ directory in the project.
"""
# Check if custom config dir is set via environment variable
if "FLA_CONFIG_DIR" in os.environ:
return Path(os.environ["FLA_CONFIG_DIR"])
# Default: project_dir/fla/configs/{GPU}/
project_dir = Path(__file__).parent.parent.parent
return project_dir / "configs" / get_gpu_info()
@dataclasses.dataclass(frozen=True)
class AutotuneKey:
"""Autotune key with exact/fuzzy matching, serialization, and construction helpers."""
autotune_key: tuple[Any, ...]
@staticmethod
def normalize_autotune_key(value: Any) -> Any:
if isinstance(value, (list, tuple)):
return [AutotuneKey.normalize_autotune_key(v) for v in value]
if isinstance(value, dict):
return {k: AutotuneKey.normalize_autotune_key(v) for k, v in value.items()}
return value
@staticmethod
def serialize(key: Any) -> str:
return json.dumps(AutotuneKey.normalize_autotune_key(key), separators=(",", ":"), sort_keys=True)
@staticmethod
def key_hash(key: Any) -> str:
import hashlib
return hashlib.md5(AutotuneKey.serialize(key).encode()).hexdigest()
@staticmethod
def is_numeric(value: Any) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool)
@staticmethod
def keys_fuzzy_match(cached_key: Any, requested_key: Any) -> bool:
# Fuzzy match: numeric leaves are compatible regardless of their actual numeric values
# (e.g. a config tuned for seq_len=1024 can apply to seq_len=2048).
# Structure (type, length, dict keys) must still match exactly.
if AutotuneKey.is_numeric(cached_key) and AutotuneKey.is_numeric(requested_key):
return True
if isinstance(cached_key, (list, tuple)) and isinstance(requested_key, (list, tuple)):
return len(cached_key) == len(requested_key) and all(
AutotuneKey.keys_fuzzy_match(c, r) for c, r in zip(cached_key, requested_key)
)
if isinstance(cached_key, dict) and isinstance(requested_key, dict):
return cached_key.keys() == requested_key.keys() and all(
AutotuneKey.keys_fuzzy_match(cached_key[k], requested_key[k]) for k in cached_key
)
return cached_key == requested_key
@classmethod
def build(
cls,
arg_names: list[str],
key_names: list[str],
positional_args: tuple[Any, ...],
runtime_kwargs: dict[str, Any],
) -> "AutotuneKey":
named_args = dict(zip(arg_names, positional_args))
all_args = {**named_args, **runtime_kwargs}
tracked_args = {k: v for (k, v) in all_args.items() if k in arg_names}
tuning_key = [tracked_args[name] for name in key_names if name in tracked_args]
for arg in tracked_args.values():
if hasattr(arg, "dtype"):
tuning_key.append(str(arg.dtype))
return cls(autotune_key=tuple(tuning_key))
def exact_matches(self, entry_key: Any) -> bool:
return self.serialize(self.autotune_key) == self.serialize(entry_key)
def fuzzy_matches(self, entry_key: Any) -> bool:
self_normalized = self.normalize_autotune_key(self.autotune_key)
entry_normalized = self.normalize_autotune_key(entry_key)
return (
isinstance(self_normalized, list)
and isinstance(entry_normalized, list)
and len(self_normalized) == len(entry_normalized)
and AutotuneKey.keys_fuzzy_match(self_normalized, entry_normalized)
)
@dataclasses.dataclass(frozen=True)
class KernelConfigFile:
"""Validated in-memory representation of a {kernel_name}.json config file."""
kernel_name: str | None
triton_version: str | None
autotune_entries: dict[str, dict[str, Any]] | None
default_config: dict[str, Any] | None
@classmethod
def from_dict(cls, config_file: Path, data: Any) -> "KernelConfigFile | None":
"""Parse and validate a raw JSON dict. Returns None (with a warning) if malformed."""
def fail(msg, *args):
logger.warning(msg, *args)
raise ValueError
try:
if not isinstance(data, dict):
fail("Malformed config %s: root is %s, expected dict", config_file, type(data).__name__)
raw_entries = data.get("autotune_entries")
entries: dict[str, dict[str, Any]] | None = None
if raw_entries is not None:
if not isinstance(raw_entries, dict):
fail("Malformed config %s: 'autotune_entries' is %s, expected dict",
config_file, type(raw_entries).__name__)
for h, entry in raw_entries.items():
if not isinstance(entry, dict):
fail("Malformed config %s: autotune_entries[%r] is %s, expected dict",
config_file, h, type(entry).__name__)
if not isinstance(entry.get("config"), dict):
fail("Malformed config %s: autotune_entries[%r] missing valid 'config' field", config_file, h)
entries = raw_entries
default_config = data.get("default_config")
if default_config is not None and not isinstance(default_config, dict):
fail("Malformed config %s: 'default_config' is %s, expected dict", config_file, type(default_config).__name__)
return cls(
kernel_name=data.get("kernel_name"),
triton_version=data.get("triton_version"),
autotune_entries=entries,
default_config=default_config,
)
except ValueError:
return None
@classmethod
def from_file(cls, config_file: Path) -> "KernelConfigFile | None":
"""Read and validate a config file. Returns None if the file is missing or malformed."""
config_data = read_config_file(config_file)
if config_data is None:
return None
return cls.from_dict(config_file, config_data)
def lookup_exact(self, key: AutotuneKey) -> dict[str, Any] | None:
if self.autotune_entries is None:
return None
return self.autotune_entries.get(AutotuneKey.key_hash(key.autotune_key))
def lookup_fuzzy(self, key: AutotuneKey) -> dict[str, Any] | None:
if self.autotune_entries is None:
return None
for entry in self.autotune_entries.values():
if key.fuzzy_matches(entry.get("autotune_key")):
return entry
return None
@cache
def load_config_file(config_file: Path) -> dict[str, Any] | None:
try:
with open(config_file) as f:
return json.load(f)
except Exception as e:
logger.warning("Error reading config file %s: %s", config_file, e)
return None
def read_config_file(config_file: Path) -> dict[str, Any] | None:
"""Read a config file, bypassing the in-process cache in ALWAYS mode."""
if FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
return load_config_file.__wrapped__(config_file)
return load_config_file(config_file)
def load_cached_config(kernel_name: str, autotune_key: AutotuneKey | None = None) -> dict[str, Any] | None:
"""
Load cached best config for a kernel from FLA configs directory.
This function loads the cached best configuration for a given kernel name
from get_fla_config_dir()/{kernel_name}.json.
Cache files may contain multiple autotune entries keyed by Triton's
runtime tuning key plus a top-level default config.
If the config file is not found or cannot be loaded, a warning is printed
and None is returned, allowing fallback to Triton's autotune.
The lookup mode is controlled by the FLA_CACHE_MODE environment variable (see FlaCacheMode).
Args:
kernel_name: Name of the kernel (e.g., "causal_conv1d_fwd_kernel")
autotune_key: Triton autotune key for the current invocation
Returns:
Best config dictionary or None if not found or disabled
"""
if FLA_CACHE_MODE is FlaCacheMode.DISABLED:
return None
config_dir = get_fla_config_dir()
config_file = config_dir / f"{kernel_name}.json"
if not config_file.exists():
return None
config_data = read_config_file(config_file)
if config_data is None:
return None
config = KernelConfigFile.from_dict(config_file, config_data)
if config is None:
return None
if FLA_CACHE_MODE is FlaCacheMode.DEFAULT or FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
return config.default_config
# STRICT mode: exact match only, no fuzzy fallback
if FLA_CACHE_MODE is FlaCacheMode.STRICT:
if autotune_key is not None:
entry = config.lookup_exact(autotune_key)
if entry is not None:
return entry["config"]
return None
# FULL and FUZZY modes: try exact key match first, then fuzzy match
if autotune_key is not None:
entry = config.lookup_exact(autotune_key) or config.lookup_fuzzy(autotune_key)
if entry is not None:
return entry["config"]
if FLA_CACHE_MODE is FlaCacheMode.FUZZY:
return None
# FULL mode: fall back to default_config, then legacy raw config (no autotune_entries)
if config.default_config is not None:
return config.default_config
if config.autotune_entries is not None:
return None
return config_data
class CachedAutotuner(Autotuner):
"""
A modified autotuner that loads best config from FLA's config directory.
This class extends Triton's Autotuner but overrides the run method to
try loading cached configuration first before falling back to autotune.
"""
def __init__(self, fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs):
super().__init__(fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs)
self.kernel_name = fn.fn.__name__ if hasattr(fn, 'fn') else fn.__name__
# None-safe pre/post hooks: Triton's defaults crash when a restore_value / reset_to_zero arg
# is None (idiomatic for optional pointers gated by a tl.constexpr flag).
# Fixed upstream in triton-lang/triton#10295 — remove this override once FLA's minimum Triton version has it.
if not self.user_defined_pre_hook and (self.reset_to_zero or self.restore_value):
def _pre_hook(kw, reset_only=False):
for n in self.reset_to_zero:
if kw[n] is not None:
kw[n].zero_()
if not reset_only:
self.restore_copies = {n: kw[n].clone() for n in self.restore_value if kw[n] is not None}
self.pre_hook = _pre_hook
if not self.user_defined_post_hook and self.restore_value:
def _post_hook(kw, exception):
for n, copy in self.restore_copies.items():
kw[n].copy_(copy)
self.restore_copies = {}
self.post_hook = _post_hook
def should_check_fla_cache(self, key: AutotuneKey) -> bool:
if FLA_CACHE_MODE is FlaCacheMode.DISABLED:
return False
if FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
return True
return key.autotune_key not in self.cache
def run(self, *args, **kwargs):
key = AutotuneKey.build(self.arg_names, self.keys, args, kwargs)
if self.should_check_fla_cache(key):
self.maybe_load_cached_config(key)
return super().run(*args, **kwargs)
def maybe_load_cached_config(self, key: AutotuneKey):
best_config = load_cached_config(self.kernel_name, key)
if best_config is not None:
kw = best_config["kwargs"]
num_warps = best_config["num_warps"]
num_stages = best_config["num_stages"]
extra = {
"num_ctas": best_config["num_ctas"],
"maxnreg": best_config.get("maxnreg"),
"pre_hook": None,
"ir_override": best_config.get("ir_override"),
} if TRITON_ABOVE_3_5_1 else {}
cfg = triton.Config(kw, num_warps=num_warps, num_stages=num_stages, **extra)
self.cache[key.autotune_key] = cfg
else:
logger.debug(
"No cached config found for kernel %s and key %s; falling back to Triton autotune",
self.kernel_name,
list(key.autotune_key),
)
def fla_cache_autotune(configs, key=None, prune_configs_by=None, reset_to_zero=None, restore_value=None,
pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False,
do_bench=None, cache_results=False):
"""
Decorator for auto-tuning a :code:`triton.jit`'d function with FLA config support.
Extends Triton's autotune to load best configurations from FLA's config directory
(default: fla/configs/{GPU}/, or FLA_CONFIG_DIR/ when overridden), keyed by kernel
name from {kernel_name}.json. Lookup behaviour is controlled by FLA_CACHE_MODE.
Falls back to normal Triton autotuning when no cached config is found.
"""
# key can be None when we want to use cache only (no fallback autotune)
if key is None:
key = []
def decorator(fn):
kwargs = {}
if TRITON_ABOVE_3_4_0:
kwargs = {"cache_results": cache_results}
return CachedAutotuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value,
pre_hook=pre_hook, post_hook=post_hook,
prune_configs_by=prune_configs_by, warmup=warmup, rep=rep,
use_cuda_graph=use_cuda_graph, do_bench=do_bench,
**kwargs,
)
return decorator
def configure_fla_cache_autotune():
triton.autotune = fla_cache_autotune
logger.info(
"configure_fla_cache_autotune() is enabling FLA fla_cache_autotune; "
"triton.autotune will be replaced with fla_cache_autotune."
)
def restore_autotune_backend():
from triton.runtime.autotuner import autotune as original_autotune
triton.autotune = original_autotune
logger.info(
"restore_autotune_backend() is restoring Triton's original autotune; "
"triton.autotune will be replaced with triton.runtime.autotuner.autotune."
)

101
ex_engine/fla_kernels/utils/op.py Normal file
View File

@@ -0,0 +1,101 @@
# 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 os
import triton
import triton.language as tl
import triton.language.extra.libdevice as tldevice
from fla.utils import IS_GATHER_SUPPORTED, IS_NVIDIA_BLACKWELL
if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
@triton.jit
def exp(x): return tldevice.fast_expf(x.to(tl.float32))
@triton.jit
def exp2(x): return tldevice.exp2(x.to(tl.float32))
@triton.jit
def log(x): return tldevice.fast_logf(x.to(tl.float32))
@triton.jit
def log2(x): return tldevice.fast_log2f(x.to(tl.float32))
@triton.jit
def tanh(x): return tldevice.fast_tanhf(x.to(tl.float32))
else:
@triton.jit
def exp(x): return tl.exp(x.to(tl.float32))
@triton.jit
def exp2(x): return tl.math.exp2(x.to(tl.float32))
@triton.jit
def log(x): return tl.log(x.to(tl.float32))
@triton.jit
def log2(x): return tl.log2(x.to(tl.float32))
@triton.jit
def tanh(x): return tldevice.tanh(x.to(tl.float32))
if IS_NVIDIA_BLACKWELL:
"""
Compute tl.dot with Blackwell workaround.
On SM100 datacenter and SM120 consumer Blackwell GPUs, wraps the result in
inline assembly to prevent the TritonGPUHoistTMEMAlloc pass from incorrectly
fusing add and dot operations.
See: https://github.com/fla-org/flash-linear-attention/issues/638
TODO: Remove this workaround once the Triton compiler bug is fixed.
Track upstream issue at: https://github.com/triton-lang/triton/issues/8695
"""
@triton.jit
def safe_dot(a, b, allow_tf32: tl.constexpr = None):
return tl.inline_asm_elementwise(
asm="mov.f32 $0, $1;",
constraints="=r,r",
args=[tl.dot(a, b, allow_tf32=allow_tf32)],
dtype=tl.float32,
is_pure=True,
pack=1,
)
else:
@triton.jit
def safe_dot(a, b, allow_tf32: tl.constexpr = None):
return tl.dot(a, b, allow_tf32=allow_tf32)
if not IS_GATHER_SUPPORTED:
@triton.jit
def gather(src, index, axis, _builder=None):
"""
Gather operation that works when tl.gather is not supported.
This is a fallback implementation that returns None.
Just to make triton compiler happy.
"""
return None
else:
gather = tl.gather
if hasattr(triton.language, '_experimental_make_tensor_descriptor'):
# For Triton 3.3.x
make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor
elif hasattr(triton.language, 'make_tensor_descriptor'):
# For Triton 3.4.x and later
make_tensor_descriptor = triton.language.make_tensor_descriptor
else:
"""
Fallback implementation when TMA is not supported.
Returns None to indicate TMA descriptors are unavailable.
Just make triton compiler happy.
"""
@triton.jit
def make_tensor_descriptor(
base,
shape,
strides,
block_shape,
_builder=None,
):
return None