Add files using upload-large-folder tool

all_results.json

@@ -0,0 +1,10 @@
+{
+    "epoch": 0.7839559871158865,
+    "num_tokens": 104891154432,
+    "throughput": 12525.363357673923,
+    "total_flos": 8.619164947133655e+20,
+    "train_loss": 9.173326368447839,
+    "train_runtime": 261696.889,
+    "train_samples_per_second": 195.709,
+    "train_steps_per_second": 0.191
+}

fla3/ops/rwkv7/fused_recurrent.py

@@ -0,0 +1,328 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import warnings
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+
+from fla.ops.generalized_delta_rule import fused_recurrent_dplr_delta_rule
+from fla.ops.utils.op import exp
+from fla.utils import input_guard, use_cuda_graph
+
+
+@triton.heuristics({
+    '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
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages)
+        for BV in [16, 32, 64]
+        for num_warps in [2, 4, 8, 16, 32]
+        for num_stages in [2, 3, 4]
+    ],
+    key=['BK'],
+    use_cuda_graph=use_cuda_graph,
+)
+@triton.jit(do_not_specialize=['T'])
+def fused_recurrent_rwkv7_fwd_kernel(
+    r,
+    w,
+    k,
+    v,
+    kk,
+    a,
+    o,
+    h0,
+    ht,
+    cu_seqlens,
+    scale,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    REVERSE: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    STORE_FINAL_STATE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_DECODE: tl.constexpr,
+):
+    i_v, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+    i_n, i_h = i_nh // H, i_nh % 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_r = r + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k
+    p_w = w + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k
+    p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k
+    p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v
+    p_a = a + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k
+    p_kk = kk + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k
+
+    p_o = o + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v
+
+    mask_k = o_k < K
+    mask_v = o_v < V
+    mask_h = mask_k[None, :] & mask_v[:, None]
+    b_h = tl.zeros([BV, BK], dtype=tl.float32)
+
+    if USE_INITIAL_STATE:
+        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)
+
+    if IS_DECODE:
+        b_r = tl.load(p_r, mask=mask_k, other=0).to(tl.float32) * scale
+        b_w = tl.load(p_w, 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)
+        b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32)
+        b_kk = tl.load(p_kk, mask=mask_k, other=0).to(tl.float32)
+        b_act_a = -b_kk
+        b_b = b_kk * b_a
+
+        tmp = tl.sum(b_h * b_act_a[None, :], axis=1)
+        b_h = exp(b_w)[None, :] * b_h + (tmp[:, None] * b_b[None, :] + b_k[None, :] * b_v[:, None])
+        b_o = tl.sum(b_h * b_r[None, :], axis=1)
+
+        tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
+    else:
+        for _ in range(0, T):
+            b_r = tl.load(p_r, mask=mask_k, other=0).to(tl.float32) * scale
+            b_w = tl.load(p_w, 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)
+            b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32)
+            b_kk = tl.load(p_kk, mask=mask_k, other=0).to(tl.float32)
+            b_act_a = -b_kk
+            b_b = b_kk * b_a
+
+            tmp = tl.sum(b_h * b_act_a[None, :], axis=1)
+            b_h = exp(b_w)[None, :] * b_h + (tmp[:, None] * b_b[None, :] + b_k[None, :] * b_v[:, None])
+            b_o = tl.sum(b_h * b_r[None, :], axis=1)
+
+            tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
+            p_r += (-1 if REVERSE else 1) * H*K
+            p_w += (-1 if REVERSE else 1) * H*K
+            p_k += (-1 if REVERSE else 1) * H*K
+            p_v += (-1 if REVERSE else 1) * H*V
+            p_a += (-1 if REVERSE else 1) * H*K
+            p_kk += (-1 if REVERSE else 1) * H*K
+            p_o += (-1 if REVERSE else 1) * H*V
+
+    if STORE_FINAL_STATE:
+        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)
+
+
+@input_guard
+def fused_recurrent_rwkv7_fwd(
+    r: torch.Tensor,
+    w: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    kk: torch.Tensor,
+    a: torch.Tensor,
+    scale: Optional[float] = 1.0,
+    initial_state: Optional[torch.Tensor] = None,
+    output_final_state: bool = False,
+    reverse: bool = False,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+):
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    N = B if cu_seqlens is None else len(cu_seqlens) - 1
+    BK = triton.next_power_of_2(K)
+    IS_DECODE = (T == 1)
+
+    h0 = initial_state
+    if not output_final_state:
+        ht = None
+    else:
+        ht = r.new_empty(N, H, K, V, dtype=torch.float32)
+    o = torch.empty_like(v)
+
+    def grid(meta): return (triton.cdiv(V, meta['BV']), N * H)
+    fused_recurrent_rwkv7_fwd_kernel[grid](
+        r,
+        w,
+        k,
+        v,
+        kk,
+        a,
+        o,
+        h0,
+        ht,
+        cu_seqlens,
+        scale,
+        T=T,
+        B=B,
+        H=H,
+        K=K,
+        V=V,
+        BK=BK,
+        REVERSE=reverse,
+        IS_DECODE=IS_DECODE
+    )
+    return o, ht
+
+
+def fused_recurrent_rwkv7(
+    r: torch.Tensor,
+    w: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    scale: float = 1.0,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = True,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    head_first: bool = False,
+):
+    """
+    Args:
+        r (torch.Tensor):
+            r of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        w (torch.Tensor):
+            log decay of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        k (torch.Tensor):
+            k of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        v (torch.Tensor):
+            v of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
+        a (torch.Tensor):
+            a of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        b (torch.Tensor):
+            b of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        scale (float):
+            scale of the attention.
+        initial_state (torch.Tensor):
+            initial state of shape `[B, H, K, V]` if cu_seqlens is None else `[N, H, K, V]` where N = len(cu_seqlens) - 1.
+        output_final_state (bool):
+            whether to output the final state.
+        cu_seqlens (torch.LongTensor):
+            Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
+            consistent with the FlashAttention API.
+        head_first (bool):
+            whether to use head first. Recommended to be False to avoid extra transposes.
+            Default: `False`.
+    """
+    return fused_recurrent_dplr_delta_rule(
+        q=r,
+        k=k,
+        v=v,
+        a=a,
+        b=b,
+        gk=w,
+        scale=scale,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+        head_first=head_first,
+    )
+
+
+def fused_mul_recurrent_rwkv7(
+    r: torch.Tensor,
+    w: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    kk: torch.Tensor,
+    a: torch.Tensor,
+    scale: Optional[float] = 1.0,
+    initial_state: Optional[torch.Tensor] = None,
+    output_final_state: bool = False,
+    reverse: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""
+    This function computes the recurrence S_t = S_t @ (I + a_t b_t^T) + v_t k_t^T in a recurrent manner.
+
+    Args:
+        r (torch.Tensor):
+            queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        w (torch.Tensor):
+            keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        k (torch.Tensor):
+            values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
+        v (torch.Tensor):
+            a of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        kk (torch.Tensor):
+            b of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        a (torch.Tensor):
+            gk of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. decay term in log space!
+        scale (Optional[int]):
+            Scale factor for the RetNet attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: 1.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `[N, H, 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, H, K, V]`. Default: `False`.
+        reverse (Optional[bool]):
+            If `True`, process the state passing in reverse order. Default: `False`.
+        cu_seqlens (Optional[torch.Tensor]):
+            Cumulative sequence lengths of shape `[N + 1]` used for variable-length training,
+            consistent with the FlashAttention API.
+        head_first (Optional[bool]):
+            Whether the inputs are in the head-first format, which is not supported for variable-length inputs.
+            Default: `False`.
+    """
+    if head_first:
+        raise DeprecationWarning(
+            "head_first is deprecated and will be removed in a future version. "
+            "Please use head_first=False for now instead."
+        )
+        r, w, k, v, kk, a = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (r, w, k, v, kk, a))
+    if not head_first and r.shape[1] < r.shape[2]:
+        warnings.warn(
+            f"Input tensor shape suggests potential format mismatch: seq_len ({r.shape[1]}) < num_heads ({r.shape[2]}). "
+            "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
+            "when head_first=False was specified. "
+            "Please verify your input tensor format matches the expected shape [B, T, H, ...]."
+        )
+    if cu_seqlens is not None:
+        if r.shape[0] != 1:
+            raise ValueError(
+                f"The batch size is expected to be 1 rather than {r.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 = r.shape[-1] ** -0.5
+    else:
+        assert scale > 0, "scale must be positive"
+    o, final_state = fused_recurrent_rwkv7_fwd(
+        r,
+        w,
+        k,
+        v,
+        kk,
+        a,
+        scale,
+        initial_state,
+        output_final_state,
+        reverse,
+        cu_seqlens,
+    )
+    if head_first:
+        o = rearrange(o, 'b t h ... -> b h t ...')
+    return o, final_state

fla3/ops/simple_gla/parallel.py

@@ -0,0 +1,732 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import warnings
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+
+from fla.ops.utils import prepare_chunk_indices
+from fla.ops.utils.cumsum import chunk_global_cumsum, chunk_local_cumsum
+from fla.ops.utils.op import safe_exp
+from fla.utils import (
+    autocast_custom_bwd,
+    autocast_custom_fwd,
+    check_shared_mem,
+    input_guard,
+    is_intel_alchemist,
+    is_nvidia_hopper
+)
+
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/3449
+triton_config = {'grf_mode': 'large'} if is_intel_alchemist else {}
+NUM_WARPS = [2, 4, 8] if is_nvidia_hopper else [2, 4, 8, 16]
+
+
+@triton.heuristics({
+    'NV': lambda args: triton.cdiv(args['V'], args['BV']),
+    'OUTPUT_ATTENTIONS': lambda args: args['attn'] is not None,
+    'USE_G': lambda args: args['g'] is not None,
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
+        for num_warps in [2, 4, 8, 16]
+        for num_stages in [2, 3, 4]
+    ],
+    key=["BT", "BS", "BK", "BV", "USE_G"],
+)
+@triton.jit
+def parallel_simple_gla_fwd_kernel(
+    q,
+    k,
+    v,
+    g,
+    o,
+    attn,
+    scale,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    NV: tl.constexpr,
+    OUTPUT_ATTENTIONS: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_G: tl.constexpr
+):
+    i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_k, i_v = i_kv // NV, i_kv % NV
+    i_b, i_h = i_bh // H, i_bh % H
+    o += i_k * B * T * H * V
+
+    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
+
+    q += (bos * H + i_h) * K
+    k += (bos * H + i_h) * K
+    v += (bos * H + i_h) * V
+    o += (bos * H + i_h) * V
+    if USE_G:
+        g += bos * H + i_h
+    if OUTPUT_ATTENTIONS:
+        attn += (bos * H + i_h * T) * T + i_k * B * H * T * T
+    stride_qk = H * K
+    stride_vo = H * V
+    stride_g = H
+
+    p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+
+    # the Q block is kept in the shared memory throughout the whole kernel
+    # [BT, BK]
+    b_q = tl.load(p_q, boundary_check=(0, 1))
+    b_q = (b_q * scale).to(b_q.dtype)
+    b_o = tl.zeros([BT, BV], dtype=tl.float32)
+
+    # [BT]
+    o_q = i_t * BT + tl.arange(0, BT)
+    # [BS]
+    o_k = i_t * BT + tl.arange(0, BS)
+    # Q block and K block have overlap.
+    # masks required
+    if USE_G:
+        p_gq = tl.make_block_ptr(g, (T,), (stride_g,), (i_t * BT,), (BT,), (0,))
+        # [BT,]
+        b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32)
+        # rescale interchunk output
+    else:
+        b_gq = None
+
+    for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS):
+        p_k = tl.make_block_ptr(k, (K, T), (1, stride_qk), (i_k * BK, i_s), (BK, BS), (0, 1))
+        p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_s, i_v * BV), (BS, BV), (1, 0))
+        # [BK, BS]
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        # [BS, BV]
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        # [BT, BS]
+        m_s = o_q[:, None] >= o_k[None, :]
+        b_s = tl.dot(b_q, b_k)
+        if USE_G:
+            p_gk = tl.make_block_ptr(g, (T,), (stride_g,), (i_s,), (BS,), (0,))
+            b_gk = tl.load(p_gk, boundary_check=(0,))
+            b_s *= safe_exp(b_gq[:, None] - b_gk[None, :])
+            b_s = tl.where(m_s, b_s, 0)
+        else:
+            b_s = tl.where(m_s, b_s, 0)
+        # [BT, BV]
+        if i_s >= 0:
+            b_o += tl.dot(b_s.to(b_q.dtype), b_v)
+        if OUTPUT_ATTENTIONS:
+            p_a = tl.make_block_ptr(attn, (T, T), (T, 1), (i_t * BT, i_s), (BT, BS), (1, 0))
+            tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1))
+        o_k += BS
+
+    for i_s in range(i_t * BT - BS, -BS, -BS):
+        p_k = tl.make_block_ptr(k, (K, T), (1, stride_qk), (i_k * BK, i_s), (BK, BS), (0, 1))
+        p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_s, i_v * BV), (BS, BV), (1, 0))
+        # [BK, BS]
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        # [BS, BV]
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        b_s = tl.dot(b_q, b_k)
+        if USE_G:
+            p_g = tl.make_block_ptr(g, (T,), (stride_g,), (i_s,), (BS,), (0,))
+            b_g = tl.load(p_g, boundary_check=(0,))
+            b_gn = tl.load(g + (min(i_s + BS, T) - 1) * stride_g)
+            b_gp = tl.load(g + (i_s-1) * stride_g) if i_s % BT > 0 else 0.
+            # No concrete meaning. Just to avoid some layout bugs.
+            b_s *= safe_exp(b_gq[:, None] + (b_gn - b_g)[None, :])
+            b_gq += (b_gn - b_gp)
+        if OUTPUT_ATTENTIONS:
+            p_a = tl.make_block_ptr(attn, (T, T), (T, 1), (i_t * BT, i_s), (BT, BS), (1, 0))
+            tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1))
+        if i_s >= 0:
+            b_o += tl.dot(b_s.to(b_v.dtype), b_v)
+    p_o = tl.make_block_ptr(o, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.jit(do_not_specialize=['T'])
+def parallel_simple_gla_bwd_kernel_dq(
+    i_t,
+    i_k,
+    i_v,
+    q,
+    k,
+    v,
+    g,
+    do,
+    dq,
+    dg,
+    stride_qk,
+    stride_vo,
+    stride_g,
+    scale,
+    T,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr
+):
+    p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+    # [BT, BV]
+    b_do = tl.load(p_do, boundary_check=(0, 1))
+    # [BT, BK]
+    b_dq = tl.zeros([BT, BK], dtype=tl.float32)
+
+    for i_s in range(0, i_t * BT, BS):
+        p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_s, i_k * BK), (BS, BK), (1, 0))
+        p_v = tl.make_block_ptr(v, (V, T), (1, stride_vo), (i_v * BV, i_s), (BV, BS), (0, 1))
+        # [BS, BK]
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        # [BV, BS]
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        # [BT, BV] @ [BV, BS] = [BT, BS]
+        b_ds = tl.dot(b_do, b_v)
+        if USE_G:
+            p_g = tl.make_block_ptr(g, (T,), (stride_g,), (i_s,), (BS,), (0,))
+            b_g = tl.load(p_g, boundary_check=(0,))
+            b_gn = tl.load(g + (min(i_s + BS, T) - 1) * stride_g)
+            b_gp = tl.load(g + (i_s - 1) * stride_g) if i_s % BT > 0 else 0.
+            b_ds *= safe_exp(b_gn - b_g)[None, :]
+            if i_s > 0:
+                b_dq *= safe_exp(b_gn - b_gp)
+        # [BT, BS] @ [BS, BK] = [BT, BK]
+        b_dq += tl.dot(b_ds.to(b_v.dtype), b_k)
+
+    if USE_G:
+        p_gq = tl.make_block_ptr(g, (T,), (stride_g,), (i_t * BT,), (BT,), (0,))
+        # [BT,]
+        b_gq = tl.load(p_gq, boundary_check=(0,))
+        # [BT, BK]
+        b_dq *= safe_exp(b_gq)[:, None]
+
+    # [BT]
+    o_q = i_t * BT + tl.arange(0, BT)
+    # [BS]
+    o_k = i_t * BT + tl.arange(0, BS)
+    # Q block and K block have overlap. masks required
+    for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS):
+        p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_s, i_k * BK), (BS, BK), (1, 0))
+        p_v = tl.make_block_ptr(v, (V, T), (1, stride_vo), (i_v * BV, i_s), (BV, BS), (0, 1))
+        # [BS, BK]
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        # [BV, BS]
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        # [BT, BV] @ [BV, BS] = [BT, BS]
+        b_ds = tl.dot(b_do, b_v)
+        if USE_G:
+            p_gk = tl.make_block_ptr(g, (T,), (stride_g,), (i_s,), (BS,), (0,))
+            b_gk = tl.load(p_gk, boundary_check=(0,))
+            b_ds *= safe_exp(b_gq[:, None] - b_gk[None, :])
+        b_ds = tl.where(o_q[:, None] >= o_k[None, :], b_ds, 0)
+        # [BT, BK]
+        b_dq += tl.dot(b_ds.to(b_k.dtype), b_k)
+        o_k += BS
+
+    b_dq *= scale
+    p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+    if USE_G:
+        p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_dg = tl.sum(b_dq * b_q, 1)
+        p_dg = tl.make_block_ptr(dg, (T,), (stride_g,), (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 parallel_simple_gla_bwd_kernel_dkv(
+    i_t,
+    i_k,
+    i_v,
+    q,
+    k,
+    v,
+    g,
+    do,
+    dk,
+    dv,
+    dg,
+    scale,
+    stride_qk,
+    stride_vo,
+    stride_g,
+    T,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr
+):
+    # [BT, BK]
+    p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    b_k = tl.load(p_k, boundary_check=(0, 1))
+    b_dk = tl.zeros([BT, BK], dtype=tl.float32)
+    # [BT, BV]
+    p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+    b_v = tl.load(p_v, boundary_check=(0, 1))
+    b_dv = tl.zeros([BT, BV], dtype=tl.float32)
+    if USE_G:
+        p_gk = tl.make_block_ptr(g, (T,), (stride_g,), (i_t * BT,), (BT,), (0,))
+        b_gk = tl.load(p_gk, boundary_check=(0,))
+    NTS = tl.cdiv(T, BS)
+    # [BT, BK]
+    for i_s in range(NTS * BS - BS, (i_t + 1) * BT - BS, -BS):
+        p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_s, i_k * BK), (BS, BK), (1, 0))
+        p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_s, i_v * BV), (BS, BV), (1, 0))
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_do = tl.load(p_do, boundary_check=(0, 1))
+        b_ds = tl.dot(b_v, tl.trans(b_do))
+        b_s = tl.dot(b_k, tl.trans(b_q))
+        if USE_G:
+            p_gq = tl.make_block_ptr(g, (T,), (stride_g,), (i_s,), (BS,), (0,))
+            b_gq = tl.load(p_gq, boundary_check=(0,))
+            b_gp = tl.load(g + (min(i_s + BS, T) - 1) * stride_g)
+            b_gn = tl.load(g + (i_s - 1) * stride_g) if i_s % BT > 0 else 0.
+            if i_s >= 0:
+                tmp = safe_exp(b_gp - b_gn)
+                b_dk *= tmp
+                b_dv *= tmp
+                tmp2 = safe_exp(b_gq - b_gn)
+                b_ds *= tmp2[None, :]
+                b_s *= tmp2[None, :]
+        # [BT, BK]
+        b_dk += tl.dot(b_ds.to(b_q.dtype), b_q)
+        # [BT, BV]
+        b_dv += tl.dot(b_s.to(b_do.dtype), b_do)
+
+    if USE_G:
+        b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * stride_g)
+        if i_t >= 0:
+            tmp2 = safe_exp(b_g_last - b_gk)[:, None]
+            b_dk *= tmp2
+            b_dv *= tmp2
+
+    o_q = i_t * BT + tl.arange(0, BS)
+    o_k = i_t * BT + tl.arange(0, BT)
+    for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS):
+        p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_s, i_k * BK), (BS, BK), (1, 0))
+        p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_s, i_v * BV), (BS, BV), (1, 0))
+        # [BS, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        # [BS, BV]
+        b_do = tl.load(p_do, boundary_check=(0, 1))
+        # [BS]
+        b_ds = tl.dot(b_v, tl.trans(b_do))
+        b_s = tl.dot(b_k, tl.trans(b_q))
+        if USE_G:
+            p_gq = tl.make_block_ptr(g, (T,), (stride_g,), (i_s,), (BS,), (0,))
+            b_gq = tl.load(p_gq, boundary_check=(0,))
+            if i_s >= 0:
+                tmp = safe_exp(-b_gk[:, None] + b_gq[None, :])
+                b_ds *= tmp
+                b_s *= tmp
+        m_s = o_k[:, None] <= o_q[None, :]
+        b_s = tl.where(m_s, b_s, 0)
+        b_ds = tl.where(m_s, b_ds, 0)
+        # [BT, BK]
+        b_dk += tl.dot(b_ds.to(b_q.dtype), b_q)
+        b_dv += tl.dot(b_s.to(b_do.dtype), b_do)
+        o_q += BS
+    b_dk *= scale
+    b_dv *= scale
+    p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+    p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+    tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
+    if USE_G:
+        p_dg = tl.make_block_ptr(dg, (T,), (stride_g,), (i_t * BT,), (BT,), (0,))
+        b_dg = tl.load(p_dg, boundary_check=(0,))
+        b_dg -= tl.sum(b_dk * b_k, 1)
+        tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics({
+    'NV': lambda args: triton.cdiv(args['V'], args['BV']),
+    'USE_G': lambda args: args['g'] is not None,
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+})
+@triton.autotune(
+    configs=[
+        triton.Config(triton_config, num_warps=num_warps)
+        for num_warps in NUM_WARPS
+    ],
+    key=['BT', 'BS', 'BK', 'BV', 'USE_G'],
+)
+@triton.jit(do_not_specialize=['T'])
+def parallel_simple_gla_bwd_kernel(
+    q,
+    k,
+    v,
+    g,
+    do,
+    dq,
+    dk,
+    dv,
+    dg,
+    scale,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    NV: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_G: tl.constexpr
+):
+    i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_k, i_v = i_kv // NV, i_kv % NV
+    i_b, i_h = i_bh // H, i_bh % H
+    dq += i_v * B * H * T * K
+    dk += i_v * B * H * T * K
+    dv += i_k * B * H * T * V
+    if USE_G:
+        dg += i_kv * B * H * 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
+
+    q += (bos * H + i_h) * K
+    k += (bos * H + i_h) * K
+    v += (bos * H + i_h) * V
+    do += (bos * H + i_h) * V
+    dq += (bos * H + i_h) * K
+    dk += (bos * H + i_h) * K
+    dv += (bos * H + i_h) * V
+    if USE_G:
+        g += bos * H + i_h
+        dg += bos * H + i_h
+    stride_qk = H * K
+    stride_vo = H * V
+    stride_g = H
+
+    parallel_simple_gla_bwd_kernel_dq(
+        i_t=i_t,
+        i_k=i_k,
+        i_v=i_v,
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        do=do,
+        dq=dq,
+        dg=dg,
+        scale=scale,
+        stride_qk=stride_qk,
+        stride_vo=stride_vo,
+        stride_g=stride_g,
+        T=T,
+        K=K,
+        V=V,
+        BT=BT,
+        BS=BS,
+        BK=BK,
+        BV=BV,
+        USE_G=USE_G
+    )
+    tl.debug_barrier()
+    parallel_simple_gla_bwd_kernel_dkv(
+        i_t=i_t,
+        i_k=i_k,
+        i_v=i_v,
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        do=do,
+        dk=dk,
+        dv=dv,
+        dg=dg,
+        scale=scale,
+        stride_qk=stride_qk,
+        stride_vo=stride_vo,
+        stride_g=stride_g,
+        T=T,
+        K=K,
+        V=V,
+        BT=BT,
+        BS=BS,
+        BK=BK,
+        BV=BV,
+        USE_G=USE_G
+    )
+
+
+def parallel_simple_gla_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    scale: float,
+    output_attentions: bool = False,
+    chunk_size: int = 128,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+):
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    BT, BS = chunk_size, 32
+    if check_shared_mem('hopper', k.device.index):
+        BK = min(256, triton.next_power_of_2(K))
+        BV = min(256, triton.next_power_of_2(V))
+    elif check_shared_mem('ampere', k.device.index):
+        BK = min(128, triton.next_power_of_2(K))
+        BV = min(128, triton.next_power_of_2(V))
+    else:
+        BK = min(64, triton.next_power_of_2(K))
+        BV = min(64, triton.next_power_of_2(V))
+
+    NK = triton.cdiv(K, BK)
+    NV = triton.cdiv(V, BV)
+    assert BT % BS == 0
+
+    chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    # local cumulative decay in log space
+    if g is not None:
+        g = chunk_local_cumsum(g, chunk_size, cu_seqlens=cu_seqlens)
+    grid = (NK * NV, NT, B * H)
+    o = torch.empty(NK, *v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device)
+    attn = q.new_zeros(NK, B, H, T, T) if output_attentions else None
+
+    parallel_simple_gla_fwd_kernel[grid](
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        o=o,
+        attn=attn,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        B=B,
+        H=H,
+        T=T,
+        K=K,
+        V=V,
+        BT=BT,
+        BS=BS,
+        BK=BK,
+        BV=BV,
+    )
+    o = o.sum(0)
+
+    if output_attentions:
+        attn = attn.sum(0)
+    return o, g, attn
+
+
+def parallel_simple_gla_bwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    do: torch.Tensor,
+    scale: float,
+    chunk_size: int = 128,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+):
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    BT, BS = chunk_size, 32
+    if check_shared_mem('hopper', k.device.index):
+        BK = min(256, triton.next_power_of_2(K))
+        BV = min(256, triton.next_power_of_2(V))
+    elif check_shared_mem('ampere', k.device.index):
+        BK = min(128, triton.next_power_of_2(K))
+        BV = min(128, triton.next_power_of_2(V))
+    elif check_shared_mem('ada', k.device.index):
+        BK = min(64, triton.next_power_of_2(K))
+        BV = min(64, triton.next_power_of_2(V))
+    else:
+        BK = min(32, triton.next_power_of_2(K))
+        BV = min(32, triton.next_power_of_2(V))
+
+    NK = triton.cdiv(K, BK)
+    NV = triton.cdiv(V, BV)
+    assert BT % BS == 0
+
+    dq = torch.empty(NV, * q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device)
+    dk = torch.empty(NV, * k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device)
+    dv = torch.empty(NK, * v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device)
+    dg = torch.empty(NK*NV, *g.shape, dtype=torch.float, device=q.device) if g is not None else None
+
+    chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    grid = (NK * NV, NT, B * H)
+    parallel_simple_gla_bwd_kernel[grid](
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        do=do,
+        dq=dq,
+        dk=dk,
+        dv=dv,
+        dg=dg,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        scale=scale,
+        T=T,
+        B=B,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BS=BS,
+        BK=BK,
+        BV=BV,
+    )
+    dq = dq.sum(0)
+    dk = dk.sum(0)
+    dv = dv.sum(0)
+    dg = chunk_global_cumsum(dg.sum(0), reverse=True, cu_seqlens=cu_seqlens) if g is not None else None
+    return dq, dk, dv, dg
+
+
+class ParallelSimpleGLAFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_fwd
+    def forward(ctx, q, k, v, g, scale, output_attentions, cu_seqlens):
+        chunk_size = 128
+        ctx.dtype = q.dtype
+
+        o, g, attn = parallel_simple_gla_fwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            scale=scale,
+            output_attentions=output_attentions,
+            chunk_size=chunk_size,
+            cu_seqlens=cu_seqlens,
+        )
+        ctx.save_for_backward(q, k, v, g, cu_seqlens)
+        ctx.scale = scale
+        ctx.chunk_size = chunk_size
+        return o.to(q.dtype), attn
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_bwd
+    def backward(ctx, do, da=None):
+        q, k, v, g, cu_seqlens = ctx.saved_tensors
+        dq, dk, dv, dg = parallel_simple_gla_bwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            do=do,
+            scale=ctx.scale,
+            chunk_size=ctx.chunk_size,
+            cu_seqlens=cu_seqlens,
+        )
+        return dq.to(q), dk.to(k), dv.to(v), dg.to(ctx.dtype) if dg is not None else None, None, None, None
+
+
+def parallel_simple_gla(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+    output_attentions: bool = False,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    head_first: bool = False
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""
+    Args:
+        q (torch.Tensor):
+            queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`
+        k (torch.Tensor):
+            keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`
+        v (torch.Tensor):
+            values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`
+        g (torch.Tensor):
+            Forget gates of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
+            Compared to GLA, the gating is head-wise instead of elementwise.
+        scale (Optional[int]):
+            Scale factor for attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        output_attentions (bool):
+            Whether to output the materialized attention scores of shape [B, H, T, T]. Default: `False`.
+        head_first (Optional[bool]):
+            Whether the inputs are in the head-first format. 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, H, V]` if `head_first=False` else `[B, H, T, V]`.
+        attn (torch.Tensor):
+            Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`
+    """
+    if head_first:
+        raise DeprecationWarning(
+            "head_first is deprecated and will be removed in a future version. "
+            "Please use head_first=False for now instead."
+        )
+        q, k, v, g = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (q, k, v, g))
+    if not head_first and q.shape[1] < q.shape[2]:
+        warnings.warn(
+            f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
+            "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
+            "when head_first=False was specified. "
+            "Please verify your input tensor format matches the expected shape [B, T, H, ...]."
+        )
+    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 output_attentions:
+        assert cu_seqlens is None, "output_attentions=True is not supported with variable-length sequences"
+
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+    o, attn = ParallelSimpleGLAFunction.apply(
+        q,
+        k,
+        v,
+        g,
+        scale,
+        output_attentions,
+        cu_seqlens
+    )
+    if head_first:
+        o = rearrange(o, 'b t h ... -> b h t ...')
+    return o, attn

fla3/ops/ttt/fused_chunk.py

@@ -0,0 +1,835 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang, Yuqi Pan
+
+import warnings
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+
+from fla.modules.layernorm import group_norm
+from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard, is_nvidia_hopper
+
+NUM_WARPS = [1, 2] if is_nvidia_hopper else [1, 2, 4, 8]
+
+
+@triton.heuristics({
+    'USE_INITIAL_STATE': lambda args: args['h0'] is not None,
+    'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None,
+    'STORE_FINAL_STATE': lambda args: args['ht'] is not None,
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4)
+    ],
+    key=['BT', 'BK', 'BV'],
+)
+@triton.jit(do_not_specialize=['T'])
+def fused_chunk_ttt_linear_fwd_kernel(
+    q,
+    k,
+    v,
+    eta,
+    w,
+    b,
+    o,
+    scale,
+    eps,
+    h0,
+    hb0,
+    ht,
+    hbt,
+    cu_seqlens,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    USE_INITIAL_STATE_B: tl.constexpr,
+    STORE_FINAL_STATE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_nh = tl.program_id(0)
+    i_n, i_h = i_nh // H, i_nh % H
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        NT = tl.cdiv(T, BT)
+
+    o_i = tl.arange(0, BT)
+    v_i = tl.arange(0, BV)
+    m_A = o_i[:, None] >= o_i[None, :]
+    b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.)
+    b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.)
+
+    # [BK, BV]
+    b_h = tl.zeros([BK, BV], dtype=tl.float32)
+    # [BV]
+    b_hb = tl.zeros([BV], dtype=tl.float32)
+    if USE_INITIAL_STATE:
+        p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0))
+        b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
+    if USE_INITIAL_STATE_B:
+        p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (0,), (BV,), (0,))
+        b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32)
+
+    for i_t in range(NT):
+        p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0))
+        p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1))
+        p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_o = tl.make_block_ptr(o+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,))
+        p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H
+        # [BK, BT]
+        b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero")
+        # [BT, BV]
+        b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero")
+
+        # [BT, BV]
+        b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :]
+        b_kh = tl.where((v_i < V)[None, :], b_kh, 0.)
+        mean = tl.sum(b_kh, axis=1, keep_dims=True) / V
+        xbar = tl.where((v_i < V)[None, :], b_kh - mean, 0.)
+        var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V
+        rstd = 1 / tl.sqrt(var.to(tl.float32) + eps)
+        b_kh_hat = (b_kh - mean) * rstd
+
+        b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \
+            b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k)
+        b_v = tl.where((v_i < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.)
+        b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype)
+                       * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V
+
+        # [BT, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero")
+        # [BT]
+        b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero")
+        b_q = (b_q * scale).to(b_k.dtype)
+
+        # [BT, BT]
+        b_A = tl.dot(b_q, b_k, allow_tf32=False)
+        b_A = tl.where(m_A, b_A, 0)
+        b_Ae = tl.where(m_A, b_e[:, None], 0.0)
+
+        b_o = - tl.dot(b_e[:, None] * b_A.to(b_v2.dtype), b_v2, allow_tf32=False)
+        b_o += b_hb[None, :] - tl.dot(b_Ae.to(b_v2.dtype), b_v2, allow_tf32=False)
+        b_o += tl.dot(b_q, b_h.to(b_q.dtype), allow_tf32=False)
+        b_e_last = tl.load(p_e_last)
+        b_h = b_h - tl.dot(b_e_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False)
+        b_hb = b_hb - tl.sum(b_e_last * b_v2.to(b_k.dtype), axis=0)
+        b_h = tl.where((v_i < V)[None, :], b_h, 0.)
+        b_hb = tl.where((v_i < V), b_hb, 0.)
+        tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+    if STORE_FINAL_STATE:
+        p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0))
+        p_hbt = tl.make_block_ptr(hbt + i_nh * V, (V,), (1,), (0,), (BV,), (0,))
+        tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
+        tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics({
+    'USE_INITIAL_STATE': lambda args: args['h0'] is not None,
+    'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None,
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4)
+    ],
+    key=['BT', 'BK', 'BV'],
+)
+@triton.jit(do_not_specialize=['T'])
+def fused_chunk_ttt_linear_bwd_kernel_h(
+    k,
+    v,
+    v2,
+    x,
+    y,
+    r,
+    w,
+    b,
+    eta,
+    h0,
+    hb0,
+    h,
+    do,
+    dq,
+    scale,
+    eps,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    USE_INITIAL_STATE_B: tl.constexpr,
+):
+    i_nh = tl.program_id(0)
+    i_n, i_h = i_nh // H, i_nh % H
+    bos, _ = i_n * T, i_n * T + T
+    NT = tl.cdiv(T, BT)
+    boh = i_n * NT
+
+    o_i = tl.arange(0, BT)
+    v_i = tl.arange(0, BV)
+    m_A = o_i[:, None] >= o_i[None, :]
+    b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.)
+    b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.)
+
+    # [BK, BV]
+    b_h = tl.zeros([BK, BV], dtype=tl.float32)
+    # [BV]
+    b_hb = tl.zeros([BV], dtype=tl.float32)
+    if USE_INITIAL_STATE:
+        p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0))
+        b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
+    if USE_INITIAL_STATE_B:
+        p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (0,), (BV,), (0,))
+        b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32)
+
+    for i_t in range(NT):
+        p_h = tl.make_block_ptr(h+((boh+i_t)*H+i_h)*K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0))
+        p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1))
+        p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_v2 = tl.make_block_ptr(v2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0))
+        p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,))
+        p_dq = tl.make_block_ptr(dq+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0))
+        p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H
+        tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1))
+        # [BK, BT]
+        b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero")
+        # [BT, BV]
+        b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero")
+
+        b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :]
+        b_kh = tl.where((v_i < V)[None, :], b_kh, 0.)
+        mean = tl.sum(b_kh, axis=1, keep_dims=True) / V
+        xbar = tl.where((v_i < V)[None, :], b_kh - mean, 0.)
+        var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V
+        rstd = 1 / tl.sqrt(var.to(tl.float32) + eps)
+        b_kh_hat = (b_kh - mean) * rstd
+
+        b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \
+            b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k)
+        b_v = tl.where((v_i < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.)
+        b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype)
+                       * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V
+        tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), boundary_check=(0, 1))
+        tl.store(p_y, b_v.to(p_y.dtype.element_ty), boundary_check=(0, 1))
+        tl.store(p_r, rstd.to(p_r.dtype.element_ty), boundary_check=(0, 1))
+        tl.store(p_v2, b_v2.to(p_v2.dtype.element_ty), boundary_check=(0, 1))
+
+        b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero")
+        b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero")
+
+        b_v2 = tl.where((v_i < V)[None, :], b_v2, 0.)
+        b_ds = tl.dot(b_do, tl.trans(b_v2).to(b_do.dtype))
+        b_ds = tl.where(m_A, b_ds, 0)
+        b_ds = b_ds.to(b_k.dtype)
+        b_dq = tl.dot(b_do, tl.trans(b_h).to(b_do.dtype))
+        b_dq -= tl.dot(b_ds, tl.trans(b_k)) * b_e[:, None]
+        b_dq *= scale
+
+        b_e_last = tl.load(p_e_last)
+        b_h = b_h - tl.dot(b_e_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False)
+        b_hb = b_hb - tl.sum(b_e_last * b_v2.to(b_k.dtype), axis=0)
+        b_h = tl.where((v_i < V)[None, :], b_h, 0.)
+        b_hb = tl.where((v_i < V), b_hb, 0.)
+        tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.heuristics({
+    'USE_INITIAL_STATE': lambda args: args['dh0'] is not None,
+    'USE_INITIAL_STATE_B': lambda args: args['dhb0'] is not None,
+    'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None,
+    'USE_FINAL_STATE_GRADIENT_B': lambda args: args['dhbt'] is not None,
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in NUM_WARPS
+    ],
+    key=['BT', 'BK', 'BV'],
+)
+@triton.jit(do_not_specialize=['T'])
+def fused_chunk_ttt_linear_bwd_kernel_dh(
+    q,
+    k,
+    v,
+    v2,
+    x,
+    y,
+    r,
+    w,
+    b,
+    eta,
+    h,
+    dht,
+    dhbt,
+    dh0,
+    dhb0,
+    do,
+    dk,
+    dv,
+    de,
+    dw,
+    db,
+    scale,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    USE_INITIAL_STATE_B: tl.constexpr,
+    USE_FINAL_STATE_GRADIENT: tl.constexpr,
+    USE_FINAL_STATE_GRADIENT_B: tl.constexpr,
+):
+    i_nh = tl.program_id(0)
+    i_n, i_h = i_nh // H, i_nh % H
+    bos, _ = i_n * T, i_n * T + T
+    NT = tl.cdiv(T, BT)
+    boh = i_n * NT
+
+    # [BK, BV]
+    b_dh = tl.zeros([BK, BV], dtype=tl.float32)
+    # [BV]
+    b_dhb = tl.zeros([BV], dtype=tl.float32)
+    if USE_FINAL_STATE_GRADIENT:
+        p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0))
+        b_dh += tl.load(p_dht, boundary_check=(0, 1), padding_option="zero")
+    if USE_FINAL_STATE_GRADIENT_B:
+        p_dhbt = tl.make_block_ptr(dhbt + i_nh * V, (V,), (1,), (0,), (BV,), (0,))
+        b_dhb += tl.load(p_dhbt, boundary_check=(0,), padding_option="zero")
+
+    # [BV]
+    o_i = tl.arange(0, BT)
+    v_i = tl.arange(0, BV)
+    m_A = o_i[:, None] >= o_i[None, :]
+    m_A_t = o_i[:, None] <= o_i[None, :]
+    b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.)
+    b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.)
+    b_dw = tl.zeros([BV,], dtype=b_w.dtype)
+    b_db = tl.zeros([BV,], dtype=b_b.dtype)
+    p_dw = tl.make_block_ptr(dw + i_nh * V, (V,), (1,), (0,), (BV,), (0,))
+    p_db = tl.make_block_ptr(db + i_nh * V, (V,), (1,), (0,), (BV,), (0,))
+
+    for i_t in range(NT - 1, -1, -1):
+        p_h = tl.make_block_ptr(h+((boh+i_t)*H+i_h)*K*V, (V, K), (1, V), (0, 0), (BV, BK), (0, 1))
+        p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1))
+        p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0))
+        p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_v2 = tl.make_block_ptr(v2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0))
+        p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,))
+        p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_dk = tl.make_block_ptr(dk+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0))
+        p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0))
+        p_de = tl.make_block_ptr(de+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,))
+        p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H
+        b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero")
+        b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero")
+        b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero")
+        b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero")
+        b_e_last = tl.load(p_e_last)
+        b_A = tl.dot(b_k, b_q)
+        b_A = - tl.where(m_A_t, b_A * scale * b_e[None, :], 0).to(do.dtype.element_ty)
+        b_Ae = - tl.where(m_A_t, b_e[None, :], 0).to(do.dtype.element_ty)
+        b_dv_new = tl.dot(b_A.to(b_do.dtype), b_do) + tl.dot(b_Ae.to(b_do.dtype), b_do)
+        b_dv_new -= tl.dot(b_e_last * b_k, b_dh.to(b_k.dtype))
+        b_dv_new -= b_e_last * b_dhb.to(b_k.dtype)[None, :]
+
+        b_v2 = tl.load(p_v2, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype)
+        b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype)
+        b_y = tl.load(p_y, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype)
+        b_rstd = tl.load(p_r, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
+        b_dy = b_rstd * (b_dv_new * V - tl.sum(b_dv_new, axis=1, keep_dims=True) -
+                         b_x * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V
+        b_dx = -b_rstd * (b_dv_new * tl.sum(b_x * b_y, axis=1, keep_dims=True) +
+                          b_y * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V
+        b_drstd = tl.sum(b_dv_new.to(b_rstd.dtype) * b_v2.to(b_rstd.dtype) / b_rstd, axis=1, keep_dims=True)
+
+        b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero")
+        b_w = b_w.to(b_k.dtype)
+        b_b = b_b.to(b_k.dtype)
+        b_dv = -b_w * b_dy.to(b_k.dtype)
+        b_dk = b_w * b_dy.to(b_k.dtype)
+        b_dw += tl.sum(2 * b_w * b_x * b_dy.to(b_k.dtype) +
+                       (b_b - b_v.to(b_k.dtype) + b_k) * b_dy.to(b_k.dtype), axis=0).to(b_dw.dtype)
+        b_db += tl.sum(b_w * b_dy.to(b_k.dtype), axis=0).to(b_db.dtype)
+        b_dx = b_dx.to(b_k.dtype) + b_w * b_w * b_dy.to(b_k.dtype)
+
+        b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero")
+        b_q = (b_q * scale).to(b_q.dtype)
+        b_dkh = b_rstd * (V * b_dx - tl.sum(b_dx, axis=1, keep_dims=True) -
+                          b_x * tl.sum(b_x * b_dx, axis=1, keep_dims=True)) / V
+        b_dkh -= b_rstd * b_rstd * b_drstd * b_x / V
+        b_dkh = tl.where((v_i < V)[None, :] * (o_i < T-i_t*BT)[:, None], b_dkh, 0.)
+        b_dk += tl.dot(b_dkh, b_h.to(b_dkh.dtype)).to(b_k.dtype)
+
+        b_ds = tl.dot(b_do, tl.trans(b_v2))
+        b_ds = tl.where(m_A, b_ds, 0)
+        b_ds = b_ds.to(b_k.dtype)
+        i_last = (BT-1) if (i_t*BT+BT) <= T else (T % BT-1)
+        mask = (o_i == i_last)
+        b_dk -= b_e_last * tl.dot(b_v2, tl.trans(b_dh).to(b_v2.dtype))
+        b_dk -= tl.dot(tl.trans(b_ds), tl.trans(b_q) * b_e[:, None])
+        b_de = mask * tl.sum(- b_dh * tl.trans(tl.dot(tl.trans(b_v2), b_k))).to(b_k.dtype)
+        b_de -= mask * tl.sum(b_dhb * tl.sum(b_v2, axis=0)).to(b_k.dtype)
+        b_de -= tl.sum(tl.dot(b_ds, b_k) * tl.trans(b_q).to(b_k.dtype), axis=1)
+        b_de -= tl.sum(b_ds, axis=1)
+        b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + tl.dot(tl.trans(b_k).to(b_dkh.dtype), b_dkh)
+        b_dhb += tl.sum(b_do + b_dkh, axis=0)
+        b_dh = tl.where((v_i < V)[None, :], b_dh, 0.)
+        b_dhb = tl.where((v_i < V), b_dhb, 0.)
+
+        tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+        tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
+        tl.store(p_de, b_de.to(p_de.dtype.element_ty), boundary_check=(0,))
+    tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0,))
+    tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,))
+
+    if USE_INITIAL_STATE:
+        p_dh0 = tl.make_block_ptr(dh0+i_nh*K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0))
+        tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1))
+    if USE_INITIAL_STATE_B:
+        p_dhb0 = tl.make_block_ptr(dhb0+i_nh*V, (V,), (1,), (0,), (BV,), (0,))
+        tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), boundary_check=(0,))
+
+
+def fused_chunk_ttt_linear_bwd_h(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+    eta: torch.Tensor,
+    scale: float,
+    eps: float,
+    do: torch.Tensor,
+    BT: int = 16,
+    initial_state: torch.Tensor = None,
+    initial_state_bias: torch.Tensor = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+):
+    assert cu_seqlens is None, "bwd of varlen is not implemented yet."
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    # N: the actual number of sequences in the batch with either equal or variable lengths
+    N, NT = B, triton.cdiv(T, BT)
+    BK, BV = triton.next_power_of_2(K), triton.next_power_of_2(V)
+    assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128."
+
+    h = k.new_empty(B, NT, H, K, V)
+    r = v.new_empty(B, T, H, 1, dtype=torch.float32)
+    v2 = torch.empty_like(v)
+    x = torch.empty_like(v)
+    y = torch.empty_like(v)
+    dq = torch.empty_like(q)
+
+    grid = (N * H,)
+    fused_chunk_ttt_linear_bwd_kernel_h[grid](
+        k=k,
+        v=v,
+        v2=v2,
+        x=x,
+        y=y,
+        r=r,
+        w=w,
+        b=b,
+        eta=eta,
+        h0=initial_state,
+        hb0=initial_state_bias,
+        h=h,
+        do=do,
+        dq=dq,
+        scale=scale,
+        eps=eps,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+    )
+    return dq, h, v2, x, y, r
+
+
+def fused_chunk_ttt_linear_bwd_dh(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    v2: torch.Tensor,
+    x: torch.Tensor,
+    y: torch.Tensor,
+    r: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+    eta: torch.Tensor,
+    scale: float,
+    h: torch.Tensor,
+    do: torch.Tensor,
+    dht: torch.Tensor,
+    dhbt: torch.Tensor,
+    BT: int = 16,
+    initial_state: torch.Tensor = None,
+    initial_state_bias: torch.Tensor = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+):
+    assert cu_seqlens is None, "bwd of varlen is not implemented yet."
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    # N: the actual number of sequences in the batch with either equal or variable lengths
+    N = B
+    BK, BV = triton.next_power_of_2(K), triton.next_power_of_2(V)
+    assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128."
+
+    dh0 = torch.empty_like(initial_state, dtype=torch.float32) if initial_state is not None else None
+    dhb0 = torch.empty_like(initial_state_bias, dtype=torch.float32) if initial_state_bias is not None else None
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+    de = torch.empty_like(eta)
+    dw = w.new_empty(B, H, V)
+    db = b.new_empty(B, H, V)
+
+    grid = (N * H,)
+    fused_chunk_ttt_linear_bwd_kernel_dh[grid](
+        q=q,
+        k=k,
+        v=v,
+        v2=v2,
+        x=x,
+        y=y,
+        r=r,
+        w=w,
+        b=b,
+        eta=eta,
+        h=h,
+        dht=dht,
+        dhbt=dhbt,
+        dh0=dh0,
+        dhb0=dhb0,
+        do=do,
+        dk=dk,
+        dv=dv,
+        de=de,
+        dw=dw,
+        db=db,
+        scale=scale,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+    )
+    dw = dw.sum(dim=0)
+    db = db.sum(dim=0)
+    return dk, dv, de, dw, db, dh0, dhb0
+
+
+def fused_chunk_ttt_linear_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+    eta: torch.Tensor,
+    scale: float,
+    eps: float,
+    initial_state: torch.Tensor,
+    initial_state_bias: torch.Tensor,
+    output_final_state: bool,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    BT: int = 16
+):
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    # N: the actual number of sequences in the batch with either equal or variable lengths
+    N = B if cu_seqlens is None else len(cu_seqlens) - 1
+    BK, BV = triton.next_power_of_2(K), triton.next_power_of_2(V)
+    assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128."
+    o = torch.empty_like(v)
+    final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None
+    final_state_bias = k.new_empty(N, H, 1, V, dtype=torch.float32) if output_final_state else None
+
+    grid = (N * H,)
+    fused_chunk_ttt_linear_fwd_kernel[grid](
+        q=q,
+        k=k,
+        v=v,
+        eta=eta,
+        w=w,
+        b=b,
+        o=o,
+        scale=scale,
+        eps=eps,
+        h0=initial_state,
+        hb0=initial_state_bias,
+        ht=final_state,
+        hbt=final_state_bias,
+        cu_seqlens=cu_seqlens,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+    )
+    return o, final_state, final_state_bias
+
+
+def fused_chunk_ttt_linear_bwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+    eta: torch.Tensor,
+    scale: float,
+    eps: float,
+    do: torch.Tensor,
+    dht: torch.Tensor,
+    dhbt: torch.Tensor,
+    BT: int = 16,
+    initial_state: torch.Tensor = None,
+    initial_state_bias: torch.Tensor = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+):
+    assert cu_seqlens is None, "bwd of varlen is not implemented yet."
+    dq, h, v2, x, y, rstd = fused_chunk_ttt_linear_bwd_h(
+        q=q,
+        k=k,
+        v=v,
+        w=w,
+        b=b,
+        eta=eta,
+        scale=scale,
+        eps=eps,
+        do=do,
+        BT=BT,
+        initial_state=initial_state,
+        initial_state_bias=initial_state_bias,
+        cu_seqlens=cu_seqlens,
+    )
+    dk, dv, de, dw, db, dh0, dhb0 = fused_chunk_ttt_linear_bwd_dh(
+        q=q,
+        k=k,
+        v=v,
+        v2=v2,
+        x=x,
+        y=y,
+        r=rstd,
+        w=w,
+        b=b,
+        eta=eta,
+        scale=scale,
+        h=h,
+        do=do,
+        dht=dht,
+        dhbt=dhbt,
+        BT=BT,
+        initial_state=initial_state,
+        initial_state_bias=initial_state_bias,
+        cu_seqlens=cu_seqlens,
+    )
+    return dq, dk, dv, de, dw, db, dh0, dhb0
+
+
+class FusedChunkTTTLinearFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_fwd
+    def forward(ctx, q, k, v, w, b, BT, eta, scale, eps, initial_state,
+                initial_state_bias, output_final_state, cu_seqlens):
+        o, final_state, final_state_bias = fused_chunk_ttt_linear_fwd(
+            q=q,
+            k=k,
+            v=v,
+            w=w,
+            b=b,
+            eta=eta,
+            scale=scale,
+            eps=eps,
+            BT=BT,
+            initial_state=initial_state,
+            initial_state_bias=initial_state_bias,
+            output_final_state=output_final_state,
+            cu_seqlens=cu_seqlens,
+        )
+        ctx.save_for_backward(q, k, v, eta, w, b, initial_state, initial_state_bias)
+        ctx.BT = BT
+        ctx.scale = scale
+        ctx.eps = eps
+        ctx.cu_seqlens = cu_seqlens
+        return o.to(q.dtype), final_state, final_state_bias
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_bwd
+    def backward(ctx, do, dht, dhbt):
+        q, k, v, eta, w, b, initial_state, initial_state_bias = ctx.saved_tensors
+        dq, dk, dv, de, dw, db, dh0, dhb0 = fused_chunk_ttt_linear_bwd(
+            q=q,
+            k=k,
+            v=v,
+            w=w,
+            b=b,
+            eta=eta,
+            scale=ctx.scale,
+            eps=ctx.eps,
+            do=do,
+            dht=dht,
+            dhbt=dhbt,
+            BT=ctx.BT,
+            initial_state=initial_state,
+            initial_state_bias=initial_state_bias,
+            cu_seqlens=ctx.cu_seqlens,
+        )
+        return dq.to(q), dk.to(k), dv.to(v), dw.to(w), db.to(b), None, de.to(eta), None, None, dh0, dhb0, None, None
+
+
+def norm_residual(x, weight, bias, eps):
+    # GroupNorm and Residual
+    B, T, H, D = x.shape
+    x += group_norm(
+        x.reshape(B, T, -1).clone(),
+        weight=weight.reshape(-1).clone(),
+        bias=bias.reshape(-1).clone(),
+        eps=eps,
+        num_groups=H,
+    ).reshape(x.shape)
+    return x
+
+
+def fused_chunk_ttt_linear(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+    eta: torch.Tensor,
+    scale: float = None,
+    eps: float = 1e-6,
+    chunk_size: int = 16,
+    initial_state: torch.Tensor = None,
+    initial_state_bias: torch.Tensor = None,
+    output_final_state: bool = False,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    head_first: bool = False,
+):
+    r"""
+    Args:
+        q (torch.Tensor):
+            queries of shape `(B, H, T, K)`
+        k (torch.Tensor):
+            keys of shape `(B, H, T, K)`
+        v (torch.Tensor):
+            values of shape `(B, H, T, V)`
+        w (torch.Tensor):
+            layer norm weight of shape `(H, V)`
+        b (torch.Tensor):
+            layer norm bias of shape `(H, V)`
+        eta (torch.Tensor):
+            Learning rate for hidden state, of shape `(B, H, T, 1)`.
+        scale (Optional[int]):
+            Scale factor for the RetNet attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        chunk_size (int):
+            chunk size. Default: `16`.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `(B, H, K, V)`. Default: `None`.
+        initial_state_bias (Optional[torch.Tensor]):
+            Initial state bias of shape `(B, H, 1, V)`. Default: `None`.
+        output_final_state (Optional[bool]):
+            Whether to output the final state of shape `(B, H, 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.
+        head_first (Optional[bool]):
+            Whether the inputs are in the head-first format, which is not supported for variable-length inputs.
+            Default: `False`.
+
+    Returns:
+        o (torch.Tensor):
+            Outputs of shape `[B, H, T, V]`
+        final_state (torch.Tensor):
+            Final state of shape `[B, H, K, V]` if `output_final_state=True` else `None`.
+        final_state_bias (torch.Tensor):
+            Final state bias of shape `[B, H, 1, V]` if `output_final_state=True` else `None`.
+    """
+    assert q.dtype == k.dtype == v.dtype
+    assert k.shape[-1] == v.shape[-1], "DK must equal to DV."
+    if isinstance(eta, float):
+        eta = torch.full_like(q[:, :, :, :1], eta)
+    if head_first:
+        raise DeprecationWarning(
+            "head_first is deprecated and will be removed in a future version. "
+            "Please use head_first=False for now instead."
+        )
+        q, k, v, eta = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (q, k, v, eta))
+    if not head_first and q.shape[1] < q.shape[2]:
+        warnings.warn(
+            f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
+            "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
+            "when head_first=False was specified. "
+            "Please verify your input tensor format matches the expected shape [B, T, H, ...]."
+        )
+    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
+    else:
+        assert scale > 0, "Scale must be positive."
+    o, final_state, final_state_bias = FusedChunkTTTLinearFunction.apply(
+        q,
+        k,
+        v,
+        w,
+        b,
+        chunk_size,
+        eta,
+        scale,
+        eps,
+        initial_state,
+        initial_state_bias,
+        output_final_state,
+        cu_seqlens,
+    )
+    o = norm_residual(o, w, b, eps)
+    if head_first:
+        o = rearrange(o, 'b t h ... -> b h t ...')
+    return o, final_state, final_state_bias

fla3/ops/ttt/naive.py

@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang, Yuqi Pan
+
+import torch
+import torch.nn.functional as F
+
+
+def ttt_linear(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+    eta: torch.Tensor,
+    scale: float,
+    eps: float,
+    mini_batch_size: int,
+    initial_state: torch.Tensor,
+    initial_state_bias: torch.Tensor,
+    output_final_state: bool
+):
+    B, H, T, D = q.shape
+    BT = mini_batch_size
+    NT = T // BT
+    # [NT, B, H, mini_batch_size, D]
+    _q = q.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4)
+    _k = k.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4)
+    _v = v.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4)
+    # [NT, B, H, BT, 1]
+    _eta = eta.reshape(B, H, NT, BT, 1).permute(2, 0, 1, 3, 4)
+    # [H, 1, D]
+    w = w.reshape(H, 1, D).to(torch.float32)
+    b = b.reshape(H, 1, D).to(torch.float32)
+
+    h = torch.zeros((B, H, D, D), device=v.device, dtype=torch.float32) if initial_state is None else initial_state
+    hb = torch.zeros((B, H, 1, D), device=v.device, dtype=torch.float32) if initial_state_bias is None else initial_state_bias
+    q *= scale
+    # [NT, B, H, BT, D]
+    o = torch.empty_like(_v)
+
+    for i in range(NT):
+        q_i, k_i, v_i, eta_i = [x[i] for x in [_q, _k, _v, _eta]]
+        kh = k_i @ h + hb
+        reconstruction_target = v_i - k_i
+
+        mean = kh.mean(-1, True)
+        var = kh.var(-1, unbiased=False, keepdim=True).to(torch.float32)
+        rstd = torch.sqrt(var + eps).to(torch.float32)
+        kh_hat = (kh - mean) / rstd
+
+        g = w * kh_hat + b - reconstruction_target
+        g *= w
+        v_new = (D * g - g.sum(-1, True) - kh_hat * (g * kh_hat).sum(-1, True)) / (rstd * D)
+
+        Attn = torch.tril(q_i @ k_i.transpose(-2, -1))
+        o_i = q_i @ h - (eta_i * Attn) @ v_new + hb - torch.tril(eta_i.expand_as(Attn)) @ v_new
+        h = h - (eta_i[:, :, -1, :, None] * k_i).transpose(-1, -2) @ v_new
+        hb = hb - torch.sum(eta_i[:, :, -1, :, None] * v_new, dim=-2, keepdim=True)
+        # layer norm with residuals
+
+        mean = o_i.mean(dim=-1, keepdim=True)
+        var = o_i.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32)
+        rstd = torch.sqrt(var + eps).to(torch.float32)
+        o[i] = o_i + (o_i - mean) / rstd * w + b
+
+    # [B, H, T, D]
+    o = o.permute(1, 2, 0, 3, 4).reshape(B, H, T, D)
+    h = h if output_final_state else None
+    hb = hb if output_final_state else None
+    return o, h, hb
+
+
+def chunk_ttt_linear_ref(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+    eta: torch.Tensor,
+    scale: float = None,
+    eps: float = 1e-6,
+    mini_batch_size: int = 16,
+    initial_state: torch.Tensor = None,
+    initial_state_bias: torch.Tensor = None,
+    output_final_state: bool = False,
+    head_first: bool = False,
+):
+    assert q.dtype == k.dtype == v.dtype
+    assert k.shape[-1] == v.shape[-1], "The key and value dimension must be the same."
+    if isinstance(eta, float):
+        eta = torch.full_like(q[:, :, :, :1], eta)
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+    if not head_first:
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        eta = eta.transpose(1, 2)
+    T = q.shape[-2]
+    padded = (mini_batch_size - (T % mini_batch_size)) % mini_batch_size
+    if padded > 0:
+        q = F.pad(q, (0, 0, 0, padded))
+        k = F.pad(k, (0, 0, 0, padded))
+        v = F.pad(v, (0, 0, 0, padded))
+        eta = F.pad(eta, (0, 0, 0, padded))
+        eta[:, :, -1, :] = eta[:, :, -(padded+1), :]
+    assert q.shape[-2] % mini_batch_size == 0, "Sequence length should be a multiple of mini_batch_size."
+    q, k, v, eta, w, b = map(lambda x: x.to(torch.float32), [q, k, v, eta, w, b])
+    o, final_state, final_state_bias = ttt_linear(
+        q,
+        k,
+        v,
+        w,
+        b,
+        eta,
+        scale,
+        eps,
+        mini_batch_size,
+        initial_state,
+        initial_state_bias,
+        output_final_state,
+    )
+    o = o[:, :, :T, :].contiguous()
+    if not head_first:
+        o = o.transpose(1, 2)
+    return o, final_state, final_state_bias

fla3/ops/utils/__init__.py

@@ -0,0 +1,54 @@
+# -*- coding: utf-8 -*-
+
+from .asm import fp32_to_tf32_asm
+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 (
+    prepare_chunk_indices,
+    prepare_chunk_offsets,
+    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 .solve_tril import solve_tril
+
+__all__ = [
+    'chunk_global_cumsum',
+    'chunk_global_cumsum_scalar',
+    'chunk_global_cumsum_vector',
+    'chunk_local_cumsum',
+    'chunk_local_cumsum_scalar',
+    'chunk_local_cumsum_vector',
+    'pack_sequence',
+    'unpack_sequence',
+    'prepare_chunk_indices',
+    'prepare_chunk_offsets',
+    'prepare_cu_seqlens_from_mask',
+    'prepare_lens',
+    'prepare_lens_from_mask',
+    'prepare_position_ids',
+    'prepare_sequence_ids',
+    'prepare_token_indices',
+    'logsumexp_fwd',
+    'addmm',
+    'matmul',
+    'mean_pooling',
+    'softmax_bwd',
+    'softmax_fwd',
+    'fp32_to_tf32_asm',
+    'solve_tril',
+]

fla3/ops/utils/asm.py

@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+
+from ...utils import device_platform
+
+
+def fp32_to_tf32_asm() -> str:
+    """
+    Get the assembly code for converting FP32 to TF32.
+    """
+    ASM_DICT = {
+        'nvidia': 'cvt.rna.tf32.f32 $0, $1;'
+    }
+    if device_platform in ASM_DICT:
+        return ASM_DICT[device_platform]
+    else:
+        # return empty string if the device is not supported
+        return ""

fla3/ops/utils/cumsum.py

@@ -0,0 +1,414 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import warnings
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.index import prepare_chunk_indices
+from ...utils import check_shared_mem, input_guard
+
+BS_LIST = [32, 64] if check_shared_mem() else [16, 32]
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=['B', 'H', 'BT', 'IS_VARLEN', 'REVERSE']
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_local_cumsum_scalar_kernel(
+    s,
+    o,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    REVERSE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    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
+
+    if HEAD_FIRST:
+        p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+        p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+    else:
+        p_s = tl.make_block_ptr(s + 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,))
+    # [BT]
+    b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32)
+    b_o = tl.cumsum(b_s, axis=0)
+    if REVERSE:
+        b_z = tl.sum(b_s, axis=0)
+        b_o = -b_o + b_z[None] + b_s
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BS': BS}, num_warps=num_warps)
+        for BS in BS_LIST
+        for num_warps in [2, 4, 8]
+    ],
+    key=['B', 'H', 'S', 'BT', 'IS_VARLEN', 'REVERSE']
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_local_cumsum_vector_kernel(
+    s,
+    o,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    S: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    REVERSE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+):
+    i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    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
+
+    o_i = tl.arange(0, BT)
+    if REVERSE:
+        m_s = tl.where(o_i[:, None] <= o_i[None, :], 1., 0.)
+    else:
+        m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.)
+
+    if HEAD_FIRST:
+        p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+        p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+    else:
+        p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+        p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+    # [BT, BS]
+    b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
+    b_o = tl.dot(m_s, b_s, allow_tf32=False)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages)
+        for BT in [32, 64, 128, 256]
+        for num_warps in [2, 4, 8]
+        for num_stages in [1, 2, 3, 4]
+    ],
+    key=['B', 'H', 'IS_VARLEN', 'REVERSE']
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_global_cumsum_scalar_kernel(
+    s,
+    o,
+    cu_seqlens,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    REVERSE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+):
+    i_nh = tl.program_id(0)
+    i_n, i_h = i_nh // H, i_nh % H
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+    T = eos - bos
+
+    b_z = tl.zeros([], dtype=tl.float32)
+    NT = tl.cdiv(T, BT)
+    for i_c in range(NT):
+        i_t = NT-1-i_c if REVERSE else i_c
+        if HEAD_FIRST:
+            p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+            p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,))
+        else:
+            p_s = tl.make_block_ptr(s + 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_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32)
+        b_o = tl.cumsum(b_s, axis=0)
+        b_ss = tl.sum(b_s, 0)
+        if REVERSE:
+            b_o = -b_o + b_ss + b_s
+        b_o += b_z
+        if i_c >= 0:
+            b_z += b_ss
+        tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages)
+        for BT in [16, 32, 64, 128]
+        for num_warps in [2, 4, 8]
+        for num_stages in [1, 2, 3, 4]
+    ],
+    key=['B', 'H', 'S', 'IS_VARLEN', 'REVERSE']
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_global_cumsum_vector_kernel(
+    s,
+    z,
+    cu_seqlens,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    S: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    REVERSE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+):
+    i_s, i_nh = tl.program_id(0), tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+    T = eos - bos
+
+    o_i = tl.arange(0, BT)
+    if REVERSE:
+        m_s = tl.where(o_i[:, None] <= o_i[None, :], 1., 0.)
+    else:
+        m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.)
+
+    b_z = tl.zeros([BS], dtype=tl.float32)
+    NT = tl.cdiv(T, BT)
+    for i_c in range(NT):
+        i_t = NT-1-i_c if REVERSE else i_c
+        if HEAD_FIRST:
+            p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+            p_z = tl.make_block_ptr(z + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+        else:
+            p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+            p_z = tl.make_block_ptr(z + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0))
+        # [BT, BS]
+        b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
+        b_c = b_z[None, :] + tl.dot(m_s, b_s, allow_tf32=False)
+        tl.store(p_z, b_c.to(p_z.dtype.element_ty), boundary_check=(0, 1))
+        if i_c >= 0:
+            b_z += tl.sum(b_s, 0)
+
+
+def chunk_local_cumsum_scalar(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if head_first:
+        B, H, T = g.shape
+    else:
+        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
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
+    grid = (NT, B * H)
+    chunk_local_cumsum_scalar_kernel[grid](
+        g_org,
+        g,
+        cu_seqlens,
+        chunk_indices,
+        T=T,
+        B=B,
+        H=H,
+        BT=BT,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse
+    )
+    return g
+
+
+def chunk_local_cumsum_vector(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if head_first:
+        B, H, T, S = g.shape
+    else:
+        B, T, H, S = g.shape
+    BT = chunk_size
+    chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2"
+
+    g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
+    def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H)
+    # keep cummulative normalizer in fp32
+    # this kernel is equivalent to
+    # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1)
+    chunk_local_cumsum_vector_kernel[grid](
+        g_org,
+        g,
+        cu_seqlens,
+        chunk_indices,
+        T=T,
+        B=B,
+        H=H,
+        S=S,
+        BT=BT,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse
+    )
+    return g
+
+
+@input_guard
+def chunk_global_cumsum_scalar(
+    s: torch.Tensor,
+    reverse: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if head_first:
+        B, H, T = s.shape
+    else:
+        B, T, H = s.shape
+    N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
+
+    z = torch.empty_like(s, dtype=output_dtype or s.dtype)
+    grid = (N * H,)
+    chunk_global_cumsum_scalar_kernel[grid](
+        s,
+        z,
+        cu_seqlens,
+        T=T,
+        B=B,
+        H=H,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse
+    )
+    return z
+
+
+@input_guard
+def chunk_global_cumsum_vector(
+    s: torch.Tensor,
+    reverse: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if head_first:
+        B, H, T, S = s.shape
+    else:
+        B, T, H, S = s.shape
+    N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
+    BS = min(32, triton.next_power_of_2(S))
+
+    z = torch.empty_like(s, dtype=output_dtype or s.dtype)
+    grid = (triton.cdiv(S, BS), N * H)
+    chunk_global_cumsum_vector_kernel[grid](
+        s,
+        z,
+        cu_seqlens,
+        T=T,
+        B=B,
+        H=H,
+        S=S,
+        BS=BS,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse
+    )
+    return z
+
+
+@input_guard
+def chunk_global_cumsum(
+    s: torch.Tensor,
+    reverse: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+    output_dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    if cu_seqlens is not None:
+        assert s.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided"
+    if len(s.shape) == 3:
+        return chunk_global_cumsum_scalar(s, reverse, cu_seqlens, head_first, output_dtype)
+    elif len(s.shape) == 4:
+        return chunk_global_cumsum_vector(s, reverse, cu_seqlens, head_first, output_dtype)
+    else:
+        raise ValueError(
+            f"Unsupported input shape {s.shape}. "
+            f"which should be [B, T, H]/[B, T, H, D] if `head_first=False` "
+            f"or [B, H, T]/[B, H, T, D] otherwise"
+        )
+
+
+@input_guard
+def chunk_local_cumsum(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+    output_dtype: Optional[torch.dtype] = torch.float,
+    **kwargs
+) -> torch.Tensor:
+    if not head_first and g.shape[1] < g.shape[2]:
+        warnings.warn(
+            f"Input tensor shape suggests potential format mismatch: seq_len ({g.shape[1]}) < num_heads ({g.shape[2]}). "
+            "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
+            "when head_first=False was specified. "
+            "Please verify your input tensor format matches the expected shape [B, T, H, ...]."
+        )
+    if cu_seqlens is not None:
+        assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided"
+    if len(g.shape) == 3:
+        return chunk_local_cumsum_scalar(g, chunk_size, reverse, cu_seqlens, head_first, output_dtype)
+    elif len(g.shape) == 4:
+        return chunk_local_cumsum_vector(g, chunk_size, reverse, cu_seqlens, head_first, output_dtype)
+    else:
+        raise ValueError(
+            f"Unsupported input shape {g.shape}. "
+            f"which should be (B, T, H, D) if `head_first=False` "
+            f"or (B, H, T, D) otherwise"
+        )

fla3/ops/utils/index.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from ...utils import tensor_cache
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [4, 8, 16, 32]
+    ],
+    key=['B'],
+)
+@triton.jit
+def prepare_position_ids_kernel(
+    y,
+    cu_seqlens,
+    B: tl.constexpr
+):
+    i_n = tl.program_id(0)
+    bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+    T = eos - bos
+
+    o = tl.arange(0, B)
+    for i in range(0, tl.cdiv(T, B) * B, B):
+        o_i = o + i
+        tl.store(y + bos + o_i, o_i, o_i < T)
+
+
+@tensor_cache
+def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor:
+    return cu_seqlens[1:] - cu_seqlens[:-1]
+
+
+@tensor_cache
+def prepare_lens_from_mask(mask: torch.BoolTensor) -> torch.LongTensor:
+    return mask.sum(dim=-1, dtype=torch.int32)
+
+
+@tensor_cache
+def prepare_cu_seqlens_from_mask(mask: torch.BoolTensor, out_dtype: torch.dtype = torch.int32) -> torch.LongTensor:
+    return F.pad(prepare_lens_from_mask(mask).cumsum(dim=0, dtype=out_dtype), (1, 0))
+
+
+@tensor_cache
+def prepare_position_ids(cu_seqlens: torch.LongTensor) -> torch.LongTensor:
+    return torch.cat([
+        torch.arange(n, dtype=cu_seqlens.dtype, device=cu_seqlens.device)
+        for n in prepare_lens(cu_seqlens).unbind()
+    ])
+
+
+@tensor_cache
+def prepare_sequence_ids(cu_seqlens: torch.LongTensor) -> torch.LongTensor:
+    return prepare_position_ids(cu_seqlens).eq(0).cumsum(0) - 1
+
+
+@tensor_cache
+def prepare_token_indices(cu_seqlens: torch.LongTensor) -> torch.LongTensor:
+    position_ids = prepare_position_ids(cu_seqlens)
+    return torch.stack([prepare_sequence_ids(cu_seqlens), position_ids], 1).to(cu_seqlens)
+
+
+@tensor_cache
+def prepare_chunk_indices(
+    cu_seqlens: torch.LongTensor,
+    chunk_size: int
+) -> torch.LongTensor:
+    indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()])
+    return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens)
+
+
+@tensor_cache
+def prepare_chunk_offsets(
+    cu_seqlens: torch.LongTensor,
+    chunk_size: int
+) -> torch.LongTensor:
+    return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1)

fla3/ops/utils/logcumsumexp.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import triton
+import triton.language as tl
+
+from ...ops.utils.op import exp, log
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({'BT': BT}, num_warps=num_warps)
+        for BT in [16, 32, 64]
+        for num_warps in [2, 4, 8]
+    ],
+    key=['S']
+)
+@triton.jit(do_not_specialize=['T'])
+def logcumsumexp_fwd_kernel(
+    s,
+    z,
+    T,
+    S: tl.constexpr,
+    BT: tl.constexpr
+):
+    i_bh = tl.program_id(0)
+    o_i = tl.arange(0, BT)
+    m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.)
+
+    b_mp = tl.full([S,], float('-inf'), dtype=tl.float32)
+    b_zp = tl.zeros([S,], dtype=tl.float32)
+    for i_t in range(tl.cdiv(T, BT)):
+        p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0))
+        p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0))
+
+        # [BT, S]
+        b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
+        # [S,]
+        b_mc = tl.max(b_s, 0)
+        b_mc = tl.maximum(b_mp, b_mc)
+        b_zp = b_zp * exp(b_mp - b_mc)
+        # [BT, S]
+        b_s = exp(b_s - b_mc)
+        b_z = tl.dot(m_s, b_s, allow_tf32=False) + b_zp
+        # [S,]
+        b_zc = tl.max(b_z, 0)
+        b_mp = b_mc
+        b_zp = b_zc
+        # [BT, BS]
+        # small eps to prevent underflows
+        b_z = log(tl.where(b_z != 0, b_z, 1e-20)) + b_mc
+        tl.store(p_z, b_z.to(p_z.dtype.element_ty), boundary_check=(0, 1))

fla3/ops/utils/logsumexp.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2024, Songlin Yang, Yu Zhang
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.op import exp, log
+
+
+@triton.heuristics({
+    'HAS_SCALE': lambda args: args['scale'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8, 16, 32]
+    ],
+    key=['D']
+)
+@triton.jit
+def logsumexp_fwd_kernel(
+    x,
+    z,
+    scale,
+    D: tl.constexpr,
+    B: tl.constexpr,
+    HAS_SCALE: tl.constexpr
+):
+    i_n, i_d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+    o_d = i_d * B + tl.arange(0, B)
+    m_d = o_d < D
+
+    b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf'))
+    if HAS_SCALE:
+        b_x = b_x * scale
+    b_m = tl.max(b_x, 0)
+    b_z = log(tl.sum(exp(b_x - b_m), 0)) + b_m
+    tl.store(z + i_n * tl.cdiv(D, B) + i_d, b_z)
+
+
+def logsumexp_fwd(
+    x,
+    scale: Optional[float] = None,
+    dtype: Optional[torch.dtype] = None
+):
+    r"""
+    Compute the logsumexp of the input tensor over the last dimension.
+
+    Args:
+        x (Tensor):
+            The input tensor of any shape.
+        scale (Optional[float]):
+            The scale applied to the input tensor. Default: `None`.
+        dtype (Optional[torch.dtype]):
+            The data type of the output tensor. Default: `None`.
+    Returns:
+        Tensor: The logsumexp of the input tensor.
+    """
+
+    shape = x.shape
+    x = x.view(-1, shape[-1])
+    N, D = x.shape
+    B = min(triton.next_power_of_2(D), 64 * 1024)
+    ND = triton.cdiv(D, B)
+
+    z = x.new_empty(N, ND, dtype=torch.float)
+    logsumexp_fwd_kernel[(N, ND)](
+        x=x,
+        z=z,
+        scale=scale,
+        D=D,
+        B=B
+    )
+    z = z.logsumexp(-1).view(*shape[:-1])
+    if dtype is not None and dtype != torch.float:
+        z = z.to(dtype)
+    return z

fla3/ops/utils/matmul.py

@@ -0,0 +1,245 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# code adapted from
+# https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.op import exp
+from ...utils import input_guard
+
+
+# `triton.jit`'ed functions can be auto-tuned by using the `triton.autotune` decorator, which consumes:
+#   - A list of `triton.Config` objects that define different configurations of
+#       meta-parameters (e.g., `BM`) and compilation options (e.g., `num_warps`) to try
+#   - An auto-tuning *key* whose change in values will trigger evaluation of all the
+#       provided configs
+@triton.heuristics({
+    'HAS_ALPHA': lambda args: args['alpha'] is not None,
+    'HAS_BETA': lambda args: args['beta'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BM': 128, 'BK': 64, 'BN': 256, 'G': 4}, num_stages=3, num_warps=8),
+        triton.Config({'BM': 64, 'BK': 32, 'BN': 256, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 128, 'BK': 32, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 128, 'BK': 32, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 64, 'BK': 32, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 128, 'BK': 32, 'BN': 32, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 64, 'BK': 32, 'BN': 32, 'G': 4}, num_stages=5, num_warps=2),
+        triton.Config({'BM': 32, 'BK': 32, 'BN': 64, 'G': 4}, num_stages=5, num_warps=2),
+        # Good config for fp8 inputs.
+        # triton.Config({'BM': 128, 'BK': 128, 'BN': 256, 'G': 4}, num_stages=3, num_warps=8),
+        # triton.Config({'BM': 256, 'BK': 128, 'BN': 128, 'G': 4}, num_stages=3, num_warps=8),
+        # triton.Config({'BM': 256, 'BK': 128, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 64, 'BK': 128, 'BN': 256, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 128, 'BK': 128, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 128, 'BK': 64, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 64, 'BK': 64, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 128, 'BK': 64, 'BN': 32, 'G': 4}, num_stages=4, num_warps=4)
+    ],
+    key=['M', 'N', 'K']
+)
+@triton.jit
+def matmul_kernel(
+    # Pointers to matrices
+    a,
+    b,
+    c,
+    input,
+    alpha,
+    beta,
+    # Matrix dimensions
+    M,
+    N,
+    K,
+    # The stride variables represent how much to increase the ptr by when moving by 1
+    # element in a particular dimension. E.g. `s_am` is how much to increase `a`
+    # by to get the element one row down (A has M rows).
+    stride_ab, stride_am, stride_ak,  # a: batch, M, K
+    stride_bk, stride_bn,             # b: K, N
+    stride_cb, stride_cm, stride_cn,  # c: batch, M, N
+    # Meta-parameters
+    BM: tl.constexpr,
+    BK: tl.constexpr,
+    BN: tl.constexpr,
+    G: tl.constexpr,
+    ACTIVATION: tl.constexpr,
+    HAS_INPUT: tl.constexpr,
+    HAS_ALPHA: tl.constexpr,
+    HAS_BETA: tl.constexpr,
+    ALLOW_TF32: tl.constexpr,
+    X_DIM: tl.constexpr = 1,
+):
+    """Kernel for computing the matmul C = A x B.
+    A has shape (M, K), B has shape (K, N) and C has shape (M, N)
+    """
+    # -----------------------------------------------------------
+    # Map program ids `pid` to the block of C it should compute.
+    # This is done in a grouped ordering to promote L2 data reuse.
+    # See above `L2 Cache Optimizations` section for details.
+    i_b, i_m, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+
+    NM, NN = tl.num_programs(1), tl.num_programs(2)
+    i_m, i_n = tl.swizzle2d(i_m, i_n, NM, NN, G)
+
+    # ----------------------------------------------------------
+    # Create pointers for the first blocks of A and B.
+    # We will advance this pointer as we move in the K direction
+    # and accumulate
+    # `p_a` is a block of [BM, BK] pointers
+    # `p_b` is a block of [BK, BN] pointers
+    # See above `Pointer Arithmetic` section for details
+    a_batch_ptr = a + i_b * stride_ab
+    o_am = (i_m * BM + tl.arange(0, BM)) % M
+    o_bn = (i_n * BN + tl.arange(0, BN)) % N
+    o_k = tl.arange(0, BK)
+
+    p_a = a_batch_ptr + (o_am[:, None] * stride_am + o_k[None, :] * stride_ak)
+    p_b = b + (o_k[:, None] * stride_bk + o_bn[None, :] * stride_bn)
+
+    b_acc = tl.zeros((BM, BN), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BK)):
+        # Load the next block of A and B, generate a mask by checking the K dimension.
+        # If it is out of bounds, set it to 0.
+        b_a = tl.load(p_a, mask=o_k[None, :] < K - k * BK, other=0.0)
+        b_b = tl.load(p_b, mask=o_k[:, None] < K - k * BK, other=0.0)
+        # We accumulate along the K dimension.
+        b_acc = tl.dot(b_a, b_b, acc=b_acc, allow_tf32=ALLOW_TF32)
+        # Advance the ptrs to the next K block.
+        p_a += BK * stride_ak
+        p_b += BK * stride_bk
+
+    o_cm = i_m * BM + tl.arange(0, BM)
+    o_cn = i_n * BN + tl.arange(0, BN)
+    mask = (o_cm[:, None] < M) & (o_cn[None, :] < N)
+
+    b_c = b_acc
+    # You can fuse arbitrary activation functions here
+    # while the b_acc is still in FP32!
+    if ACTIVATION == "leaky_relu":
+        b_c = leaky_relu(b_c)
+    elif ACTIVATION == "relu":
+        b_c = relu(b_c)
+    elif ACTIVATION == "sigmoid":
+        b_c = sigmoid(b_c)
+    elif ACTIVATION == "tanh":
+        b_c = tanh(b_c)
+
+    if HAS_ALPHA:
+        b_c *= tl.load(alpha)
+
+    if HAS_INPUT:
+        p_i = input + (stride_cm * o_cm[:, None] if X_DIM == 2 else 0) + stride_cn * o_cn[None, :]
+        mask_p = (o_cn[None, :] < N) if X_DIM == 1 else mask
+        b_i = tl.load(p_i, mask=mask_p, other=0.0).to(tl.float32)
+        if HAS_BETA:
+            b_i *= tl.load(beta)
+        b_c += b_i
+
+    # -----------------------------------------------------------
+    # Write back the block of the output matrix C with masks.
+    c_batch_ptr = c + i_b * stride_cb
+    p_c = c_batch_ptr + stride_cm * o_cm[:, None] + stride_cn * o_cn[None, :]
+    tl.store(p_c, b_c.to(c.dtype.element_ty), mask=mask)
+
+
+# We can fuse `leaky_relu` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`.
+@triton.jit
+def leaky_relu(x):
+    return tl.where(x >= 0, x, 0.01 * x)
+
+
+@triton.jit
+def sigmoid(x):
+    # σ(x) = 1 / (1 + exp(-x))
+    return 1.0 / (1.0 + exp(-x))
+
+
+@triton.jit
+def tanh(x):
+    # tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))
+    # 2 * sigmoid(2x) - 1
+    return (exp(x) - exp(-x)) / (exp(x) + exp(-x))
+
+
+@triton.jit
+def relu(x):
+    # ReLU(x) = max(0, x)
+    return tl.maximum(x, 0.0)
+
+
+@input_guard
+def matmul(a, b, activation=''):
+    assert a.dim() in [2, 3], "a must be 2D or 3D"
+    assert b.dim() == 2, "b must be 2D"
+    assert a.shape[-1] == b.shape[0], f"Incompatible dimensions: A {a.shape}, B {b.shape}"
+
+    if a.dim() == 2:
+        a_dim = 2
+        a = a.unsqueeze(0).contiguous()  # (1, M, K)
+    else:
+        a_dim = 3
+    allow_tf32 = False if a.dtype == torch.float32 else True
+
+    B, M, K = a.shape[0], a.shape[1], a.shape[2]
+    K_b, N = b.shape
+    assert K == K_b, f"Incompatible K dimension: A {K} vs B {K_b}"
+    c = a.new_empty(B, M, N)
+
+    def grid(meta): return (B, triton.cdiv(M, meta['BM']), triton.cdiv(N, meta['BN']))
+    matmul_kernel[grid](
+        a, b, c, None, None, None,
+        M, N, K,
+        a.stride(0), a.stride(1), a.stride(2),  # stride_ab, stride_am, stride_ak
+        b.stride(0), b.stride(1),               # stride_bk, stride_bn (b.dim() == 2)
+        c.stride(0), c.stride(1), c.stride(2),  # stride_cb, stride_cm, stride_cn
+        ACTIVATION=activation,
+        ALLOW_TF32=allow_tf32,
+        HAS_INPUT=False,
+    )
+    return c.squeeze(0) if a_dim == 2 else c
+
+
+@input_guard
+def addmm(
+    x: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    alpha: Optional[float] = None,
+    beta: Optional[float] = None,
+) -> torch.Tensor:
+    assert a.dim() in [2, 3], "a must be 2D or 3D"
+    assert b.dim() == 2, "b must be 2D"
+    assert a.shape[-1] == b.shape[0], f"Incompatible dimensions: A {a.shape}, B {b.shape}"
+
+    if a.dim() == 2:
+        a_dim = 2
+        a = a.unsqueeze(0).contiguous()  # (1, M, K)
+    else:
+        a_dim = 3
+    allow_tf32 = False if a.dtype == torch.float32 else True
+
+    B, M, K = a.shape[0], a.shape[1], a.shape[2]
+    K_b, N = b.shape
+    assert K == K_b, f"Incompatible K dimension: A {K} vs B {K_b}"
+    c = a.new_empty(B, M, N)
+
+    def grid(meta): return (B, triton.cdiv(M, meta['BM']), triton.cdiv(N, meta['BN']))
+    matmul_kernel[grid](
+        a, b, c, x, alpha, beta,
+        M, N, K,
+        a.stride(0), a.stride(1), a.stride(2),  # stride_ab, stride_am, stride_ak
+        b.stride(0), b.stride(1),               # stride_bk, stride_bn (b.dim() == 2)
+        c.stride(0), c.stride(1), c.stride(2),  # stride_cb, stride_cm, stride_cn
+        ACTIVATION=None,
+        ALLOW_TF32=allow_tf32,
+        HAS_INPUT=True,
+        X_DIM=x.dim(),
+    )
+    return c.squeeze(0) if a_dim == 2 else c

fla3/ops/utils/pack.py

@@ -0,0 +1,208 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# Code adapted from https://github.com/mayank31398/cute-kernels
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.index import prepare_lens
+from ...utils import input_guard
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [4, 8, 16, 32]
+    ],
+    key=['D', 'PADDING_SIDE', 'PACK']
+)
+@triton.jit
+def packunpack_sequence_kernel(
+    x,
+    y,
+    cu_seqlens,
+    S,
+    D,
+    BD: tl.constexpr,
+    PADDING_SIDE: tl.constexpr,
+    PACK: tl.constexpr,
+):
+    i_d, i_s, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    bos, eos = tl.load(cu_seqlens + i_b), tl.load(cu_seqlens + i_b + 1)
+
+    T = eos - bos
+    if PADDING_SIDE == 'left':
+        NP = S - T
+        if i_s < NP:
+            return
+        i_t = bos + (i_s - NP)
+    else:
+        if i_s >= T:
+            return
+        i_t = bos + i_s
+
+    o_d = i_d * BD + tl.arange(0, BD)
+    mask = o_d < D
+
+    if PACK:
+        b_x = tl.load(x + (i_b * S + i_s) * D + o_d, mask=mask)
+        tl.store(y + i_t * D + o_d, b_x, mask=mask)
+    else:
+        b_x = tl.load(x + i_t * D + o_d, mask=mask)
+        tl.store(y + (i_b * S + i_s) * D + o_d, b_x, mask=mask)
+
+
+def pack_sequence_fwdbwd(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    padding_side: str,
+) -> torch.Tensor:
+    B, S = x.shape[:2]
+    D = x.numel() // (B * S)
+    BD = min(triton.next_power_of_2(D), 4096)
+    ND = triton.cdiv(D, BD)
+
+    y = torch.empty(cu_seqlens[-1].item(), *x.shape[2:], device=x.device, dtype=x.dtype)
+    packunpack_sequence_kernel[ND, S, B](
+        x=x,
+        y=y,
+        cu_seqlens=cu_seqlens,
+        S=S,
+        D=D,
+        BD=BD,
+        PADDING_SIDE=padding_side,
+        PACK=True,
+    )
+    return y
+
+
+def unpack_sequence_fwdbwd(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    padding_side: str,
+    desired_shape: torch.Size,
+) -> torch.Tensor:
+    if desired_shape is None:
+        desired_shape = (len(cu_seqlens) - 1, prepare_lens(cu_seqlens).max().item(), *x.shape[1:])
+    y = torch.zeros(desired_shape, device=x.device, dtype=x.dtype)
+    B, S = y.shape[:2]
+    D = y.numel() // (B * S)
+    BD = min(triton.next_power_of_2(D), 4096)
+    ND = triton.cdiv(D, BD)
+
+    packunpack_sequence_kernel[ND, S, B](
+        x=x,
+        y=y,
+        cu_seqlens=cu_seqlens,
+        S=S,
+        D=D,
+        BD=BD,
+        PADDING_SIDE=padding_side,
+        PACK=False,
+    )
+    return y
+
+
+class PackSequenceFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        padding_side: str,
+    ) -> torch.Tensor:
+        assert padding_side in ['left', 'right']
+        assert x.ndim >= 2
+
+        ctx.cu_seqlens = cu_seqlens
+        ctx.padding_side = padding_side
+        ctx.desired_shape = x.shape
+
+        y = pack_sequence_fwdbwd(
+            x=x,
+            cu_seqlens=cu_seqlens,
+            padding_side=padding_side,
+        )
+        return y
+
+    @staticmethod
+    @input_guard
+    def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None]:
+        dx = unpack_sequence_fwdbwd(
+            x=dy,
+            cu_seqlens=ctx.cu_seqlens,
+            padding_side=ctx.padding_side,
+            desired_shape=ctx.desired_shape,
+        )
+        return dx, *[None] * 10
+
+
+class UnpackSequenceFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        padding_side: str,
+        desired_shape: Optional[torch.Size] = None,
+    ) -> torch.Tensor:
+        assert padding_side in ['left', 'right']
+        assert x.ndim >= 2
+        if desired_shape is not None:
+            assert desired_shape[0] == cu_seqlens.shape[0] - 1
+            assert desired_shape[2:] == x.shape[1:]
+
+        ctx.cu_seqlens = cu_seqlens
+        ctx.padding_side = padding_side
+
+        y = unpack_sequence_fwdbwd(
+            x=x,
+            cu_seqlens=cu_seqlens,
+            padding_side=padding_side,
+            desired_shape=desired_shape,
+        )
+        return y
+
+    @staticmethod
+    @input_guard
+    def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None]:
+        dx = pack_sequence_fwdbwd(
+            x=dy,
+            cu_seqlens=ctx.cu_seqlens,
+            padding_side=ctx.padding_side,
+        )
+        return dx, None, None, None
+
+
+def pack_sequence(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    padding_side: str = 'left'
+) -> torch.Tensor:
+    return PackSequenceFunction.apply(
+        x,
+        cu_seqlens,
+        padding_side,
+    )
+
+
+def unpack_sequence(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    padding_side: str = 'left',
+    desired_shape: Optional[torch.Size] = None,
+) -> torch.Tensor:
+    return UnpackSequenceFunction.apply(
+        x,
+        cu_seqlens,
+        padding_side,
+        desired_shape,
+    )

fla3/ops/utils/pooling.py

@@ -0,0 +1,207 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.index import prepare_chunk_indices
+from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BD': BD}, num_warps=num_warps)
+        for BD in [16, 32, 64, 128]
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=['BT']
+)
+@triton.jit(do_not_specialize=['T'])
+def mean_pooling_fwd_kernel(
+    x,
+    o,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    BT: tl.constexpr,
+    BD: tl.constexpr,
+    IS_VARLEN: tl.constexpr
+):
+    i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    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
+        NT = tl.cdiv(T, BT)
+    else:
+        NT = tl.cdiv(T, BT)
+        i_tg = i_b * NT + i_t
+        bos, eos = i_b * T, i_b * T + T
+
+    p_x = tl.make_block_ptr(x + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0))
+    p_o = tl.make_block_ptr(o + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,))
+    # [BT, BD]
+    b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
+    # [BD]
+    b_o = tl.sum(b_x, axis=0) / min(BT, T - i_t * BT)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BD': BD}, num_warps=num_warps)
+        for BD in [16, 32, 64, 128]
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=['BT']
+)
+@triton.jit(do_not_specialize=['T'])
+def mean_pooling_bwd_kernel(
+    do,
+    dx,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    BT: tl.constexpr,
+    BD: tl.constexpr,
+    IS_VARLEN: tl.constexpr
+):
+    i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    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
+        NT = tl.cdiv(T, BT)
+    else:
+        NT = tl.cdiv(T, BT)
+        i_tg = i_b * NT + i_t
+        bos, eos = i_b * T, i_b * T + T
+
+    p_dx = tl.make_block_ptr(dx + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0))
+    p_do = tl.make_block_ptr(do + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,))
+    # [BD]
+    b_do = tl.load(p_do, boundary_check=(0,)).to(tl.float32)
+    # [BT, BD]
+    b_dx = b_do / tl.full((BT,), min(BT, T - i_t * BT), dtype=tl.float32)[:, None]
+    tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1))
+
+
+def mean_pooling_fwd(
+    x: torch.Tensor,
+    chunk_size: int,
+    cu_seqlens: Optional[torch.LongTensor] = None
+) -> torch.Tensor:
+    B, T, H, D = x.shape
+    BT = chunk_size
+    chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    o = x.new_empty(B, NT, H, D)
+    def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H)
+    mean_pooling_fwd_kernel[grid](
+        x,
+        o,
+        cu_seqlens,
+        chunk_indices,
+        T=T,
+        H=H,
+        D=D,
+        BT=BT,
+    )
+    return o
+
+
+def mean_pooling_bwd(
+    do: torch.Tensor,
+    batch_size: int,
+    seq_len: int,
+    chunk_size: int,
+    cu_seqlens: Optional[torch.LongTensor] = None
+) -> torch.Tensor:
+    B, T, H, D = batch_size, seq_len, *do.shape[-2:]
+    BT = chunk_size
+    chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    dx = do.new_empty(B, T, H, D)
+    def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H)
+    mean_pooling_bwd_kernel[grid](
+        do,
+        dx,
+        cu_seqlens,
+        chunk_indices,
+        T=T,
+        H=H,
+        D=D,
+        BT=BT,
+    )
+    return dx
+
+
+class MeanPoolingFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_fwd
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        chunk_size: int,
+        cu_seqlens: Optional[torch.LongTensor] = None
+    ) -> torch.Tensor:
+        o = mean_pooling_fwd(x, chunk_size, cu_seqlens)
+        ctx.batch_size = x.shape[0]
+        ctx.seq_len = x.shape[1]
+        ctx.chunk_size = chunk_size
+        ctx.cu_seqlens = cu_seqlens
+        return o
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_bwd
+    def backward(
+        ctx, do
+    ) -> Tuple[torch.Tensor, None, None]:
+        batch_size = ctx.batch_size
+        seq_len = ctx.seq_len
+        chunk_size = ctx.chunk_size
+        cu_seqlens = ctx.cu_seqlens
+        dx = mean_pooling_bwd(do, batch_size, seq_len, chunk_size, cu_seqlens)
+        return dx, None, None
+
+
+def mean_pooling(
+    x: torch.Tensor,
+    chunk_size: int,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    head_first: bool = False
+) -> torch.Tensor:
+    if head_first:
+        x = x.transpose(1, 2)
+    if cu_seqlens is not None:
+        if x.shape[0] != 1:
+            raise ValueError(
+                f"The batch size is expected to be 1 rather than {x.shape[0]} when using `cu_seqlens`."
+                f"Please ..tten variable-length inputs before processing."
+            )
+    o = MeanPoolingFunction.apply(x, chunk_size, cu_seqlens)
+    if head_first:
+        o = o.transpose(1, 2)
+    return o

fla3/ops/utils/softmax.py

@@ -0,0 +1,111 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2024, Songlin Yang, Yu Zhang
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.op import exp
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16),
+        triton.Config({}, num_warps=32)
+    ],
+    key=['D']
+)
+@triton.jit
+def softmax_fwd_kernel(
+    x,
+    p,
+    D: tl.constexpr,
+    B: tl.constexpr
+):
+    i_n = tl.program_id(0)
+    o_d = tl.arange(0, B)
+    m_d = o_d < D
+
+    b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf'))
+    b_m = tl.max(b_x, 0)
+    b_x = exp(b_x - b_m)
+    b_p = b_x / tl.sum(b_x, 0)
+
+    tl.store(p + i_n * D + o_d, b_p.to(p.dtype.element_ty), mask=m_d)
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16),
+        triton.Config({}, num_warps=32)
+    ],
+    key=['D']
+)
+@triton.jit
+def softmax_bwd_kernel(
+    p,
+    dp,
+    ds,
+    D: tl.constexpr,
+    B: tl.constexpr
+):
+    i_n = tl.program_id(0)
+    o_d = tl.arange(0, B)
+    m_d = o_d < D
+
+    b_p = tl.load(p + i_n * D + o_d, mask=m_d, other=0.)
+    b_dp = tl.load(dp + i_n * D + o_d, mask=m_d, other=0.)
+    b_pp = tl.sum(b_p * b_dp, 0)
+    b_ds = b_p * b_dp - b_p * b_pp
+    tl.store(ds + i_n * D + o_d, b_ds.to(ds.dtype.element_ty), mask=m_d)
+
+
+def softmax_fwd(
+    x: torch.Tensor,
+    dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    shape = x.shape
+    x = x.view(-1, x.shape[-1])
+
+    N, D = x.shape
+    B = triton.next_power_of_2(D)
+
+    p = torch.empty_like(x, dtype=dtype)
+    softmax_fwd_kernel[(N,)](
+        x=x,
+        p=p,
+        D=D,
+        B=B
+    )
+    return p.view(*shape)
+
+
+def softmax_bwd(
+    p: torch.Tensor,
+    dp: torch.Tensor,
+    dtype: Optional[torch.dtype] = torch.float
+) -> torch.Tensor:
+    shape = p.shape
+    p = p.view(-1, p.shape[-1])
+    ds = torch.empty_like(p, dtype=dtype)
+
+    N, D = p.shape
+    B = triton.next_power_of_2(D)
+    softmax_bwd_kernel[(N,)](
+        p=p,
+        dp=dp,
+        ds=ds,
+        D=D,
+        B=B
+    )
+    return ds.view(*shape)

fla3/ops/utils/solve_tril.py

@@ -0,0 +1,276 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.index import prepare_chunk_indices
+from ...utils import input_guard
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
+        for num_warps in [1, 2, 4, 8]
+        for num_stages in [2, 3, 4, 5]
+    ],
+    key=['BT'],
+)
+@triton.jit(do_not_specialize=['T'])
+def solve_tril_16x16_kernel(
+    A,
+    Ad,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    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
+
+    A = A + (bos*H + i_h) * BT
+    Ad = Ad + (bos*H + i_h) * 16
+
+    offset = (i_t * 16) % BT
+    p_A = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * 16, offset), (16, 16), (1, 0))
+    p_Ai = tl.make_block_ptr(Ad, (T, 16), (H*16, 1), (i_t * 16, 0), (16, 16), (1, 0))
+    b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
+    b_A = -tl.where(tl.arange(0, 16)[:, None] > tl.arange(0, 16)[None, :], b_A, 0)
+
+    o_i = tl.arange(0, 16)
+    for i in range(1, min(16, T-i_t*16)):
+        b_a = -tl.load(A + (i_t * 16 + i) * H*BT + o_i + offset)
+        b_a = b_a + tl.sum(b_a[:, None] * b_A, 0)
+        mask = o_i == i
+        b_A = tl.where(mask[:, None], b_a, b_A)
+    b_A += o_i[:, None] == o_i[None, :]
+    tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
+        for num_warps in [1, 2, 4, 8]
+        for num_stages in [2, 3, 4, 5]
+    ],
+    key=['H', 'BT', 'IS_VARLEN'],
+)
+@triton.jit(do_not_specialize=['T'])
+def merge_16x16_to_32x32_inverse_kernel(
+    A,
+    Ad,
+    Ai,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    IS_VARLEN: tl.constexpr
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    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
+
+    A += (bos*H + i_h) * 32
+    Ad += (bos*H + i_h) * 16
+    Ai += (bos*H + i_h) * 32
+
+    p_A_21 = tl.make_block_ptr(A, (T, 32), (H*32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0))
+    p_Ad_11 = tl.make_block_ptr(Ad, (T, 16), (H*16, 1), (i_t * 32, 0), (16, 16), (1, 0))
+    p_Ad_22 = tl.make_block_ptr(Ad, (T, 16), (H*16, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0))
+    p_Ai_11 = tl.make_block_ptr(Ai, (T, 32), (H*32, 1), (i_t * 32, 0), (16, 16), (1, 0))
+    p_Ai_22 = tl.make_block_ptr(Ai, (T, 32), (H*32, 1), (i_t * 32 + 16, 16), (16, 16), (1, 0))
+    p_Ai_21 = tl.make_block_ptr(Ai, (T, 32), (H*32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0))
+
+    A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
+    Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)
+    Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)
+    Ai_21 = -tl.dot(tl.dot(Ai_22, A_21, input_precision='ieee'), Ai_11, input_precision='ieee')
+    tl.store(p_Ai_11, Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_22, Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_21, Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.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, 5]
+    ],
+    key=['H', 'BT', 'IS_VARLEN'],
+)
+@triton.jit(do_not_specialize=['T'])
+def merge_16x16_to_64x64_inverse_kernel(
+    A,
+    Ad,
+    Ai,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    IS_VARLEN: tl.constexpr
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    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
+
+    A += (bos*H + i_h) * 64
+    Ad += (bos*H + i_h) * 16
+    Ai += (bos*H + i_h) * 64
+
+    p_A_21 = tl.make_block_ptr(A, (T, 64), (H*64, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0))
+    p_A_32 = tl.make_block_ptr(A, (T, 64), (H*64, 1), (i_t * 64 + 32, 16), (16, 16), (1, 0))
+    p_A_31 = tl.make_block_ptr(A, (T, 64), (H*64, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0))
+    p_A_43 = tl.make_block_ptr(A, (T, 64), (H*64, 1), (i_t * 64 + 48, 32), (16, 16), (1, 0))
+    p_A_42 = tl.make_block_ptr(A, (T, 64), (H*64, 1), (i_t * 64 + 48, 16), (16, 16), (1, 0))
+    p_A_41 = tl.make_block_ptr(A, (T, 64), (H*64, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0))
+    p_Ad_11 = tl.make_block_ptr(Ad, (T, 16), (H*16, 1), (i_t * 64, 0), (16, 16), (1, 0))
+    p_Ad_22 = tl.make_block_ptr(Ad, (T, 16), (H*16, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0))
+    p_Ad_33 = tl.make_block_ptr(Ad, (T, 16), (H*16, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0))
+    p_Ad_44 = tl.make_block_ptr(Ad, (T, 16), (H*16, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0))
+
+    A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
+    A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32)
+    A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32)
+    A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32)
+    A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32)
+    A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32)
+
+    Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)
+    Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)
+    Ai_33 = tl.load(p_Ad_33, boundary_check=(0, 1)).to(tl.float32)
+    Ai_44 = tl.load(p_Ad_44, boundary_check=(0, 1)).to(tl.float32)
+
+    Ai_21 = -tl.dot(tl.dot(Ai_22, A_21, input_precision='ieee'), Ai_11, input_precision='ieee')
+    Ai_32 = -tl.dot(tl.dot(Ai_33, A_32, input_precision='ieee'), Ai_22, input_precision='ieee')
+    Ai_43 = -tl.dot(tl.dot(Ai_44, A_43, input_precision='ieee'), Ai_33, input_precision='ieee')
+
+    Ai_31 = -tl.dot(
+        Ai_33,
+        tl.dot(A_31, Ai_11, input_precision='ieee') +
+        tl.dot(A_32, Ai_21, input_precision='ieee'),
+        input_precision='ieee'
+    )
+    Ai_42 = -tl.dot(
+        Ai_44,
+        tl.dot(A_42, Ai_22, input_precision='ieee') +
+        tl.dot(A_43, Ai_32, input_precision='ieee'),
+        input_precision='ieee'
+    )
+    Ai_41 = -tl.dot(
+        Ai_44,
+        tl.dot(A_41, Ai_11, input_precision='ieee') +
+        tl.dot(A_42, Ai_21, input_precision='ieee') +
+        tl.dot(A_43, Ai_31, input_precision='ieee'),
+        input_precision='ieee'
+    )
+
+    p_Ai_11 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64, 0), (16, 16), (1, 0))
+    p_Ai_22 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 16, 16), (16, 16), (1, 0))
+    p_Ai_33 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 32, 32), (16, 16), (1, 0))
+    p_Ai_44 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 48, 48), (16, 16), (1, 0))
+    p_Ai_21 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0))
+    p_Ai_31 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0))
+    p_Ai_32 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 32, 16), (16, 16), (1, 0))
+    p_Ai_41 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0))
+    p_Ai_42 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 48, 16), (16, 16), (1, 0))
+    p_Ai_43 = tl.make_block_ptr(Ai, (T, 64), (H*64, 1), (i_t * 64 + 48, 32), (16, 16), (1, 0))
+    tl.store(p_Ai_11, Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_22, Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_33, Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_44, Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_21, Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_31, Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_32, Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_41, Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_42, Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    tl.store(p_Ai_43, Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+
+
+@input_guard
+def solve_tril(
+    A: torch.Tensor,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    output_dtype: torch.dtype = torch.float
+) -> torch.Tensor:
+    """
+    Compute the inverse of the lower triangular matrix
+    A should be strictly lower triangular, i.e., A.triu() == 0.
+
+    Args:
+        A (torch.Tensor):
+            [B, T, H, K]
+        cu_seqlens (torch.Tensor):
+            The cumulative sequence lengths of the input tensor.
+            Default: None.
+        output_dtype (torch.dtype):
+            The dtype of the output tensor. Default: `torch.float`
+
+    Returns:
+        (I + A)^-1 with the same shape as A
+    """
+    assert A.shape[-1] in [16, 32, 64]
+
+    B, T, H, BT = A.shape
+    Ad = torch.empty(B, T, H, 16, device=A.device, dtype=torch.float if BT != 16 else output_dtype)
+
+    chunk_indices = prepare_chunk_indices(cu_seqlens, 16) if cu_seqlens is not None else None
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 16)
+    solve_tril_16x16_kernel[NT, B * H](
+        A=A,
+        Ad=Ad,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        BT=BT,
+    )
+    if BT == 16:
+        return Ad
+
+    Ai = torch.zeros(B, T, H, BT, device=A.device, dtype=output_dtype)
+    merge_fn = merge_16x16_to_32x32_inverse_kernel if BT == 32 else merge_16x16_to_64x64_inverse_kernel
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
+    merge_fn[NT, B * H](
+        A=A,
+        Ad=Ad,
+        Ai=Ai,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        BT=BT,
+    )
+    return Ai

flame/data.py

@@ -0,0 +1,246 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, List, Union
+
+import numpy as np
+import torch
+from datasets import Dataset, IterableDataset
+from flame.logging import get_logger
+from transformers import PreTrainedTokenizer
+
+logger = get_logger(__name__)
+
+
+class HuggingfaceDataset(IterableDataset):
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        tokenizer: PreTrainedTokenizer,
+        context_len: int = 2048,
+        rank: int = 0,
+        world_size: int = 1,
+        buffer_size: int = 1024
+    ) -> HuggingfaceDataset:
+
+        self.dataset = dataset
+        self.tokenizer = tokenizer
+
+        self.data = dataset.shard(world_size, rank)
+        self.context_len = context_len
+        self.rank = rank
+        self.world_size = world_size
+        self.buffer_size = buffer_size
+
+        if tokenizer.vocab_size < torch.iinfo(torch.int16).max:
+            self.dtype = torch.int16
+        elif tokenizer.vocab_size < torch.iinfo(torch.int32).max:
+            self.dtype = torch.int32
+        else:
+            self.dtype = torch.int64
+        self.states = None
+        self.buffer = torch.tensor([], dtype=self.dtype)
+        self.tokens = []
+        self.rand_id = 0
+        self.token_id = 0
+        self.rng_state = None
+        self._epoch = 0
+
+    def __iter__(self):
+        g = torch.Generator()
+        g.manual_seed(self._epoch + self.rank)
+        if self.rng_state is not None:
+            g.set_state(self.rng_state)
+
+        rand_it = self.randint(0, self.buffer_size, g=g)
+        if self.states is not None:
+            self.data.load_state_dict(self.states)
+
+        # max number of tokens allowed in the chunk buffer
+        n_tokens = self.buffer_size * self.context_len
+
+        while True:
+            for sample in self.tokenize(self.data):
+                # keep appending the samples to the token buffer
+                self.tokens += sample
+                # if the token buffer is full, start sampling
+                # NOTE: we first convert the token ids to a tensor of shape [n_chunks, context_len] for efficiency
+                if len(self.buffer) == 0 and len(self.tokens) >= n_tokens:
+                    self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=self.dtype).view(self.buffer_size, -1)
+                    self.tokens = self.tokens[n_tokens:]
+                if len(self.buffer) == self.buffer_size:
+                    yield from self.sample(rand_it)
+
+            n_chunks = len(self.tokens) // self.context_len
+            # handle the left tokens in the buffer
+            if n_chunks > 0:
+                n_tokens = n_chunks * self.context_len
+                indices = torch.randperm(n_chunks, generator=g).tolist()
+                self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=torch.long).view(n_chunks, -1)
+                self.tokens = self.tokens[n_tokens:]
+                for i in indices:
+                    yield {'input_ids': self.buffer[i]}
+
+    def tokenize(self, data, batch_size: int = 64):
+        texts, states = [], []
+        for sample in data:
+            texts.append(sample['text'])
+            states.append(self.data.state_dict())
+            if len(texts) == batch_size:
+                for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']):
+                    self.states = s
+                    yield tokenized
+                texts, states = [], []
+        if len(texts) > 0:
+            for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']):
+                self.states = s
+                yield tokenized
+
+    def sample(self, indices):
+        n_tokens = (len(self.tokens) // self.context_len) * self.context_len
+        while self.token_id < n_tokens:
+            i = next(indices)
+            start, end = self.token_id, self.token_id + self.context_len
+            self.token_id += self.context_len
+            yield {'input_ids': self.buffer[i].to(torch.long)}
+            self.buffer[i] = torch.tensor(self.tokens[start:end], dtype=self.dtype)
+        self.token_id = 0
+        self.tokens = self.tokens[n_tokens:]
+
+    def randint(
+        self,
+        low: int,
+        high: int,
+        batch_size: int = 1024,
+        g: torch.Generator = torch.Generator()
+    ) -> Iterable[int]:
+        indices = torch.empty(batch_size, dtype=torch.long)
+        while True:
+            # record the generator states before sampling
+            self.rng_state = g.get_state()
+            indices = torch.randint(low, high, (batch_size,), out=indices, generator=g)
+            for i in indices[self.rand_id:].tolist():
+                self.rand_id += 1
+                yield i
+            self.rand_id = 0
+
+    def set_epoch(self, epoch):
+        self._epoch = epoch
+        if hasattr(self.dataset, "set_epoch"):
+            self.dataset.set_epoch(epoch)
+
+    def state_dict(self):
+        return {
+            'states': self.states,
+            'buffer': self.buffer.clone(),
+            'tokens': deepcopy(self.tokens),
+            'rand_id': self.rand_id,
+            'token_id': self.token_id,
+            'rng_state': self.rng_state,
+            'epoch': self._epoch
+        }
+
+    def load_state_dict(self, state_dict):
+        self.states = state_dict['states']
+        self.buffer = state_dict['buffer'].clone()
+        self.tokens = deepcopy(state_dict['tokens'])
+        self.rand_id = state_dict['rand_id']
+        self.token_id = state_dict['token_id']
+        self.rng_state = state_dict['rng_state'].clone() if state_dict['rng_state'] is not None else None
+        self._epoch = state_dict['epoch']
+
+
+@dataclass
+class DataCollatorForLanguageModeling:
+    """
+    Data collator used for language modeling.
+
+    Args:
+        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
+            The tokenizer used for encoding the data.
+        varlen (`bool`):
+            Whether to return sequences with variable lengths.
+            If `True`, the offsets indicating the start and end of each sequence will be returned.
+            For example, if the sequence lengths are `[4, 8, 12]`,
+            the returned `input_ids` will be a long flattened tensor of shape `[1, 24]`, with `offsets` being `[0, 4, 12, 24]`.
+            If `False`, the `input_ids` with shape `[batch_size, seq_len]` will be returned directly.
+        return_tensors (`str`):
+            The type of Tensor to return. Allowable values are "pt".
+    """
+
+    tokenizer: PreTrainedTokenizer
+    varlen: bool = False
+    return_tensors: str = "pt"
+
+    def __call__(
+        self,
+        examples: List[Union[List[int], Dict[str, Any]]]
+    ) -> Dict[str, Any]:
+        if not isinstance(examples[0], Dict):
+            examples = [{'input_ids': example} for example in examples]
+
+        def tensorize(example: Dict[str, Any]) -> Dict[str, Any]:
+            tensorized = {}
+            for key in ['input_ids', 'offsets']:
+                if key not in example:
+                    continue
+                if isinstance(example[key], List):
+                    tensorized[key] = torch.tensor(example[key], dtype=torch.long)
+                elif isinstance(example[key], np.ndarray):
+                    tensorized[key] = torch.from_numpy(example[key])
+                else:
+                    tensorized[key] = example[key]
+            return tensorized
+
+        examples = list(map(tensorize, examples))
+
+        if not self.varlen:
+            length_of_first = examples[0]['input_ids'].size(0)
+            # Check if padding is necessary.
+            if all(example['input_ids'].size(0) == length_of_first for example in examples):
+                batch = {
+                    'input_ids': torch.stack([example['input_ids'] for example in examples], dim=0),
+                }
+            else:
+                # If yes, check if we have a `pad_token`.
+                if self.tokenizer._pad_token is None:
+                    raise ValueError(
+                        f"You are attempting to pad samples but the tokenizer you are using "
+                        f"({self.tokenizer.__class__.__name__}) does not have a pad token."
+                    )
+                batch = self.tokenizer.pad(examples, return_tensors=self.return_tensors, return_attention_mask=False)
+        else:
+            if len(examples) > 1:
+                raise ValueError("The batch size must be 1 for variable length inputs.")
+            batch = {
+                'input_ids': torch.cat([example['input_ids'] for example in examples], dim=0).unsqueeze(0)
+            }
+            if 'offsets' in examples[0]:
+                batch['offsets'] = torch.cat([example['offsets'] for example in examples], dim=0).unsqueeze(0)
+            else:
+                # determine boundaries by bos/eos positions
+                if self.tokenizer.add_bos_token:
+                    offsets = []
+                    if batch['input_ids'][0, 0] != self.tokenizer.bos_token_id:
+                        offsets.append(torch.tensor([0], dtype=torch.long))
+                    offsets.append(torch.where(batch['input_ids'].eq(self.tokenizer.bos_token_id))[1])
+                    offsets.append(torch.tensor([len(batch['input_ids'][0])], dtype=torch.long))
+                    batch['offsets'] = torch.cat(offsets, dim=0)
+                elif self.tokenizer.add_eos_token:
+                    offsets = [torch.tensor([0], dtype=torch.long)]
+                    offsets.append(torch.where(batch['input_ids'].eq(self.tokenizer.eos_token_id))[1] + 1)
+                    if batch['input_ids'][0, -1] != self.tokenizer.eos_token_id:
+                        offsets.append(torch.tensor([len(batch['input_ids'][0])], dtype=torch.long))
+                    batch['offsets'] = torch.cat(offsets, dim=0)
+                else:
+                    raise ValueError("You must allow the tokenizer to add either a bos or eos token as separators.")
+
+        labels = batch['input_ids'].clone()
+        if self.tokenizer.pad_token_id is not None:
+            labels[labels == self.tokenizer.pad_token_id] = -100
+        batch["labels"] = labels
+        return batch

flame/logging.py

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+
+import json
+import logging
+import os
+import sys
+import time
+
+from transformers.trainer_callback import (ExportableState, TrainerCallback,
+                                           TrainerControl, TrainerState)
+from transformers.training_args import TrainingArguments
+
+
+def get_logger(name: str = None) -> logging.Logger:
+    formatter = logging.Formatter(
+        fmt="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S"
+    )
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setFormatter(formatter)
+
+    logger = logging.getLogger(name)
+    if 'RANK' in os.environ and int(os.environ['RANK']) == 0:
+        logger.setLevel(logging.INFO)
+        logger.addHandler(handler)
+
+    return logger
+
+
+logger = get_logger(__name__)
+
+LOG_FILE_NAME = "trainer_log.jsonl"
+
+
+class LogCallback(TrainerCallback, ExportableState):
+    def __init__(self, start_time: float = None, elapsed_time: float = None):
+
+        self.start_time = time.time() if start_time is None else start_time
+        self.elapsed_time = 0 if elapsed_time is None else elapsed_time
+        self.last_time = self.start_time
+
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs
+    ):
+        r"""
+        Event called at the beginning of training.
+        """
+        if state.is_local_process_zero:
+            if not args.resume_from_checkpoint:
+                self.start_time = time.time()
+                self.elapsed_time = 0
+            else:
+                self.start_time = state.stateful_callbacks['LogCallback']['start_time']
+                self.elapsed_time = state.stateful_callbacks['LogCallback']['elapsed_time']
+
+        if args.save_on_each_node:
+            if not state.is_local_process_zero:
+                return
+        else:
+            if not state.is_world_process_zero:
+                return
+
+        self.last_time = time.time()
+        if os.path.exists(os.path.join(args.output_dir, LOG_FILE_NAME)) and args.overwrite_output_dir:
+            logger.warning("Previous log file in this folder will be deleted.")
+            os.remove(os.path.join(args.output_dir, LOG_FILE_NAME))
+
+    def on_log(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        logs,
+        **kwargs
+    ):
+        if args.save_on_each_node:
+            if not state.is_local_process_zero:
+                return
+        else:
+            if not state.is_world_process_zero:
+                return
+
+        self.elapsed_time += time.time() - self.last_time
+        self.last_time = time.time()
+        if 'num_input_tokens_seen' in logs:
+            logs['num_tokens'] = logs.pop('num_input_tokens_seen')
+            state.log_history[-1].pop('num_input_tokens_seen')
+            throughput = logs['num_tokens'] / args.world_size / self.elapsed_time
+            state.log_history[-1]['throughput'] = logs['throughput'] = throughput
+        state.stateful_callbacks["LogCallback"] = self.state()
+
+        logs = dict(
+            current_steps=state.global_step,
+            total_steps=state.max_steps,
+            loss=state.log_history[-1].get("loss", None),
+            eval_loss=state.log_history[-1].get("eval_loss", None),
+            predict_loss=state.log_history[-1].get("predict_loss", None),
+            learning_rate=state.log_history[-1].get("learning_rate", None),
+            epoch=state.log_history[-1].get("epoch", None),
+            percentage=round(state.global_step / state.max_steps * 100, 2) if state.max_steps != 0 else 100,
+        )
+
+        os.makedirs(args.output_dir, exist_ok=True)
+        with open(os.path.join(args.output_dir, "trainer_log.jsonl"), "a", encoding="utf-8") as f:
+            f.write(json.dumps(logs) + "\n")
+
+    def state(self) -> dict:
+        return {
+            'start_time': self.start_time,
+            'elapsed_time': self.elapsed_time
+        }
+
+    @classmethod
+    def from_state(cls, state):
+        return cls(state['start_time'], state['elapsed_time'])

flame/parser.py

@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+import transformers
+from transformers import HfArgumentParser, TrainingArguments
+
+from flame.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class TrainingArguments(TrainingArguments):
+
+    model_name_or_path: str = field(
+        default=None,
+        metadata={
+            "help": "Path to the model weight or identifier from huggingface.co/models or modelscope.cn/models."
+        },
+    )
+    tokenizer: str = field(
+        default="fla-hub/gla-1.3B-100B",
+        metadata={"help": "Name of the tokenizer to use."}
+    )
+    use_fast_tokenizer: bool = field(
+        default=False,
+        metadata={"help": "Whether or not to use one of the fast tokenizer (backed by the tokenizers library)."},
+    )
+    from_config: bool = field(
+        default=True,
+        metadata={"help": "Whether to initialize models from scratch."},
+    )
+    dataset: Optional[str] = field(
+        default=None,
+        metadata={"help": "The dataset(s) to use. Use commas to separate multiple datasets."},
+    )
+    dataset_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "The name of provided dataset(s) to use."},
+    )
+    cache_dir: str = field(
+        default=None,
+        metadata={"help": "Path to the cached tokenized dataset."},
+    )
+    split: str = field(
+        default="train",
+        metadata={"help": "Which dataset split to use for training and evaluation."},
+    )
+    streaming: bool = field(
+        default=False,
+        metadata={"help": "Enable dataset streaming."},
+    )
+    hf_hub_token: Optional[str] = field(
+        default=None,
+        metadata={"help": "Auth token to log in with Hugging Face Hub."},
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the pre-processing."},
+    )
+    buffer_size: int = field(
+        default=2048,
+        metadata={"help": "Size of the buffer to randomly sample examples from in dataset streaming."},
+    )
+    context_length: int = field(
+        default=2048,
+        metadata={"help": "The context length of the tokenized inputs in the dataset."},
+    )
+    varlen: bool = field(
+        default=False,
+        metadata={"help": "Enable training with variable length inputs."},
+    )
+
+
+def get_train_args():
+    parser = HfArgumentParser(TrainingArguments)
+    args, unknown_args = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+
+    if unknown_args:
+        print(parser.format_help())
+        print("Got unknown args, potentially deprecated arguments: {}".format(unknown_args))
+        raise ValueError("Some specified arguments are not used by the HfArgumentParser: {}".format(unknown_args))
+
+    if args.should_log:
+        transformers.utils.logging.set_verbosity(args.get_process_log_level())
+        transformers.utils.logging.enable_default_handler()
+        transformers.utils.logging.enable_explicit_format()
+    # set seeds manually
+    transformers.set_seed(args.seed)
+    return args