Add Echo-Memory codebase used for this run (CC BY 4.0, JD Echo Team) (part 4)

code/flash-linear-attention/tests/models/test_modeling_rwkv7.py

@@ -0,0 +1,57 @@
+
+import pytest
+import torch
+
+from fla.models import RWKV7Config
+
+from .test_modeling_base import run_test_generation, run_test_model_forward_backward
+
+
+# ===================================================================================
+# Test for Modeling (Forward/Backward Pass)
+# ===================================================================================
+@pytest.mark.parametrize(
+    ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'],
+    [
+        pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test))
+        for test in [
+            (4, 4, 1024, 4, 64, True, torch.bfloat16),
+            (4, 4, 1024, 4, 64, False, torch.bfloat16),
+            (4, 4, 1024, 4, 128, False, torch.bfloat16),
+        ]
+    ],
+)
+def test_modeling(
+    L: int,
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    use_l2warp: bool,
+    dtype: torch.dtype,
+):
+    run_test_model_forward_backward(L, B, T, H, D, RWKV7Config, use_l2warp=use_l2warp, dtype=dtype)
+
+# ===================================================================================
+# Test for Generation
+# ===================================================================================
+
+
+@pytest.mark.parametrize(
+    ['L', 'B', 'T', 'H', 'D', 'dtype'],
+    [
+        pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (2, 4, 2000, 8, 64, torch.float16),
+        ]
+    ],
+)
+def test_generation(
+    L: int,
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    run_test_generation(L, B, T, H, D, RWKV7Config, dtype)

code/flash-linear-attention/tests/models/test_modeling_samba.py

@@ -0,0 +1,56 @@
+
+import pytest
+import torch
+
+from fla.models import SambaConfig
+
+from .test_modeling_base import run_test_generation, run_test_model_forward_backward
+
+
+# ===================================================================================
+# Test for Modeling (Forward/Backward Pass)
+# ===================================================================================
+@pytest.mark.parametrize(
+    ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'],
+    [
+        pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test))
+        for test in [
+            (4, 4, 1024, 18, 64, True, torch.bfloat16),
+            (4, 4, 1024, 18, 64, False, torch.bfloat16),
+            (4, 4, 1024, 9, 128, False, torch.bfloat16),
+        ]
+    ],
+)
+def test_modeling(
+    L: int,
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    use_l2warp: bool,
+    dtype: torch.dtype,
+):
+    run_test_model_forward_backward(L, B, T, H, D, SambaConfig, use_l2warp=use_l2warp, dtype=dtype)
+
+
+# ===================================================================================
+# Test for Generation
+# ===================================================================================
+@pytest.mark.parametrize(
+    ['L', 'B', 'T', 'H', 'D', 'dtype'],
+    [
+        pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (2, 4, 2000, 18, 64, torch.float16),
+        ]
+    ],
+)
+def test_generation(
+    L: int,
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    run_test_generation(L, B, T, H, D, SambaConfig, dtype)

code/flash-linear-attention/tests/models/test_modeling_transformer.py

@@ -0,0 +1,56 @@
+
+import pytest
+import torch
+
+from fla.models import TransformerConfig
+
+from .test_modeling_base import run_test_generation, run_test_model_forward_backward
+
+
+# ===================================================================================
+# Test for Modeling (Forward/Backward Pass)
+# ===================================================================================
+@pytest.mark.parametrize(
+    ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'],
+    [
+        pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test))
+        for test in [
+            (4, 4, 1024, 4, 64, True, torch.bfloat16),
+            (4, 4, 1024, 4, 64, False, torch.bfloat16),
+            (4, 4, 1024, 4, 128, False, torch.bfloat16),
+        ]
+    ],
+)
+def test_modeling(
+    L: int,
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    use_l2warp: bool,
+    dtype: torch.dtype,
+):
+    run_test_model_forward_backward(L, B, T, H, D, TransformerConfig, use_l2warp=use_l2warp, dtype=dtype)
+
+
+# ===================================================================================
+# Test for Generation
+# ===================================================================================
+@pytest.mark.parametrize(
+    ['L', 'B', 'T', 'H', 'D', 'dtype'],
+    [
+        pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (2, 4, 2000, 8, 64, torch.float16),
+        ]
+    ],
+)
+def test_generation(
+    L: int,
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    run_test_generation(L, B, T, H, D, TransformerConfig, dtype)

code/flash-linear-attention/tests/models/test_modeling_utils.py

@@ -0,0 +1,84 @@
+
+import math
+
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM
+
+from fla.utils import device
+
+# Models that do not yet support variable sequence lengths (for modeling tests)
+MODELING_UNSUPPORTED_VARLEN = [
+    "ABCConfig", "ForgettingTransformerConfig", "LinearAttentionConfig", "LightNetConfig",
+    "Mamba2Config", "MambaConfig", "MesaNetConfig", "SambaConfig",
+    "RodimusConfig",
+]
+
+# Models not yet ready for basic testing
+NOT_READY_FOR_TESTING = ['RodimusConfig']
+
+# Models requiring specific hardware (e.g., NVIDIA Hopper)
+HOPPER_EXCLUSIVE = []
+
+GENERATION_UNSUPPORTED = [
+    "ABCConfig", "LinearAttentionConfig", "LightNetConfig",
+    "Mamba2Config", "MambaConfig", "NSAConfig", "SambaConfig", "RWKV6Config", "RWKV7Config",
+    "DeltaFormerConfig",
+]
+
+
+def create_model_and_config(config_class, L, H, D, dtype, **kwargs):
+    """
+    A helper function to create a model and its configuration.
+    """
+    config_params = {
+        'hidden_size': H * D,
+        'num_hidden_layers': L,
+        **({'num_heads': H} if config_class.__name__ != 'NSAConfig' else {}),
+        **kwargs,
+    }
+    config = config_class(**config_params)
+    model = AutoModelForCausalLM.from_config(config)
+    model.apply(init_weights_recursively)
+    model.to(dtype).to(device)
+    return model, config
+
+
+def init_weights_with_asymmetric_pattern(module):
+    """Initialize weights with asymmetric patterns for debugging.
+
+    Args:
+        module: The module to initialize weights for.
+    """
+    if isinstance(module, (nn.Linear, nn.Conv1d)):
+        nn.init.kaiming_normal_(module.weight, a=math.sqrt(5))
+        with torch.no_grad():
+            shape = module.weight.shape
+            if len(shape) > 1:
+                quarter_size = shape[0] // 4
+                module.weight[:quarter_size] *= 1.2
+                module.weight[-quarter_size:] *= 0.8
+                if shape[0] == shape[1]:
+                    idx = torch.arange(min(shape[0], shape[1]))
+                    module.weight[idx, idx] += 0.05
+        if module.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(module.bias, -bound, bound)
+            with torch.no_grad():
+                module.bias[::3] *= 1.1
+                module.bias[1::3] *= 0.9
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        with torch.no_grad():
+            vocab_size, dim = module.weight.shape
+            pattern = 0.01 * torch.sin(torch.arange(dim) * (6.28 / dim))
+            for i in range(min(100, vocab_size)):
+                module.weight[i] += pattern * (1 + i % 5) * 0.2
+
+
+def init_weights_recursively(module):
+    if hasattr(module, 'weight'):
+        init_weights_with_asymmetric_pattern(module)
+    for submodule in module.children():
+        init_weights_recursively(submodule)

code/flash-linear-attention/tests/modules/test_activation.py

@@ -0,0 +1,133 @@
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.modules.activations import logsigmoid, sigmoid, swiglu, swiglu_linear, swish
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'compile'),
+    [
+        (1, 1, 64, False),
+        (2, 500, 128, False),
+        (2, 512, 128, True),
+        (3, 2048, 1200, True),
+    ],
+)
+def test_sigmoid(B: int, T: int, D: int, compile: bool):
+    torch.manual_seed(42)
+    x = torch.randn(B, T, D, device=device, requires_grad=True)
+    y_ref = torch.sigmoid(x)
+    y_tri = sigmoid(x) if not compile else torch.compile(sigmoid)(x)
+
+    g = torch.randn_like(y_ref)
+    dx_ref = torch.autograd.grad(y_ref, x, g)[0]
+    dx_tri = torch.autograd.grad(y_tri, x, g)[0]
+
+    assert_close('sigmoid fwd', y_ref, y_tri, 1e-3)
+    assert_close('sigmoid bwd', dx_ref, dx_tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'temperature', 'compile'),
+    [
+        (1, 1, 64, 1.0, False),
+        (2, 500, 128, 0.5, False),
+        (2, 512, 128, 0.5, True),
+        (3, 2048, 1200, 2.0, True),
+    ],
+)
+def test_logsigmoid(B: int, T: int, D: int, temperature: float, compile: bool):
+    torch.manual_seed(42)
+    x = torch.randn(B, T, D, device=device, requires_grad=True)
+    y_ref = F.logsigmoid(x) / temperature
+    y_tri = logsigmoid(x, temperature) if not compile else torch.compile(logsigmoid)(x, temperature)
+
+    g = torch.randn_like(y_ref)
+    dx_ref = torch.autograd.grad(y_ref, x, g)[0]
+    dx_tri = torch.autograd.grad(y_tri, x, g)[0]
+
+    assert_close('logsigmoid fwd', y_ref, y_tri, 1e-3)
+    assert_close('logsigmoid bwd', dx_ref, dx_tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'compile'),
+    [
+        (1, 1, 64, True),
+        (2, 500, 128, True),
+        (2, 512, 128, False),
+        (3, 2048, 1200, False),
+    ],
+)
+def test_swish(B: int, T: int, D: int, compile: bool):
+    torch.manual_seed(42)
+    x = torch.randn(B, T, D, device=device, requires_grad=True)
+    y_ref = F.silu(x)
+    y_tri = swish(x) if not compile else torch.compile(swish)(x)
+
+    g = torch.randn_like(y_ref)
+    dx_ref = torch.autograd.grad(y_ref, x, g)[0]
+    dx_tri = torch.autograd.grad(y_tri, x, g)[0]
+
+    assert_close('swish fwd', y_ref, y_tri, 1e-3)
+    assert_close('swish bwd', dx_ref, dx_tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'compile'),
+    [
+        (1, 1, 64, True),
+        (2, 500, 128, True),
+        (2, 512, 128, False),
+        (3, 2048, 1200, False),
+    ],
+)
+def test_swiglu(B: int, T: int, D: int, compile: bool):
+    torch.manual_seed(42)
+    x = torch.randn(B, T, D, device=device, requires_grad=True)
+    y = torch.randn(B, T, D, device=device, requires_grad=True)
+
+    y_ref = F.silu(x) * y
+    y_tri = swiglu(x, y) if not compile else torch.compile(swiglu)(x, y)
+
+    g = torch.randn_like(y_ref)
+    dx_ref, dy_ref = torch.autograd.grad(y_ref, (x, y), g)
+    dx_tri, dy_tri = torch.autograd.grad(y_tri, (x, y), g)
+
+    assert_close('swiglu fwd', y_ref, y_tri, 1e-3)
+    assert_close('swiglu dx',  dx_ref, dx_tri, 1e-3)
+    assert_close('swiglu dy',  dy_ref, dy_tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'O', 'compile'),
+    [
+        (2, 512, 128, 256, True),
+        (1, 1, 64, 32, False),
+        (2, 500, 128, 64, True),
+        (3, 2048, 1200, 600, False),
+    ],
+)
+def test_swiglu_linear(B: int, T: int, D: int, O: int, compile: bool):  # noqa: E741
+    torch.manual_seed(42)
+    x = torch.randn(B, T, D, device=device, requires_grad=True)
+    y = torch.randn(B, T, D, device=device, requires_grad=True)
+    w = torch.randn(O, D, device=device, requires_grad=True)
+    b = torch.randn(O, device=device, requires_grad=True)
+
+    z_ref = F.silu(x) * y
+    out_ref = F.linear(z_ref, w, b)
+    out_tri = swiglu_linear(x, y, w, b) if not compile else torch.compile(swiglu_linear)(x, y, w, b)
+
+    g = torch.randn_like(out_ref)
+    dx_ref, dy_ref, dw_ref, db_ref = torch.autograd.grad(out_ref, (x, y, w, b), g)
+    dx_tri, dy_tri, dw_tri, db_tri = torch.autograd.grad(out_tri, (x, y, w, b), g)
+
+    assert_close('swiglu_linear out', out_ref, out_tri, 1e-3)
+    assert_close('swiglu_linear dx',  dx_ref,  dx_tri,  1e-3)
+    assert_close('swiglu_linear dy',  dy_ref,  dy_tri,  1e-3)
+    assert_close('swiglu_linear dw',  dw_ref,  dw_tri,  1e-3)
+    assert_close('swiglu_linear db',  db_ref,  db_tri,  1e-3)

code/flash-linear-attention/tests/modules/test_conv.py

@@ -0,0 +1,713 @@
+
+import pytest
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.modules.convolution import ShortConvolution, causal_conv1d, causal_conv1d_update
+from fla.utils import assert_close, device
+
+try:
+    from causal_conv1d import causal_conv1d_fn
+except ImportError:
+    causal_conv1d_fn = None
+
+
+def causal_conv1d_ref_torch(
+    x,
+    weight,
+    bias=None,
+    initial_state=None,
+    output_final_state=False,
+    final_states_out=None,
+    activation=None,
+):
+    """
+    x: (batch, dim, seqlen)
+    weight: (dim, width)
+    bias: (dim,)
+    initial_state: (batch, dim, width - 1)
+    final_states_out: (batch, dim, width - 1)
+
+    out: (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+    dtype_in = x.dtype
+    x = x.to(weight.dtype)
+    seqlen = x.shape[-1]
+    dim, width = weight.shape
+    if initial_state is None:
+        out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim)
+    else:
+        x = torch.cat([initial_state, x], dim=-1)
+        out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim)
+    out = out[..., :seqlen]
+    if output_final_state:
+        final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to(
+            dtype_in,
+        )  # (batch, dim, width - 1)
+        if final_states_out is not None:
+            final_states_out.copy_(final_states)
+        else:
+            final_states_out = final_states
+    out = (out if activation is None else F.silu(out)).to(dtype=dtype_in)
+    return out if not output_final_state else (out, final_states_out)
+
+
+def causal_conv1d_update_ref_torch(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None):
+    """
+    x: (batch, dim) or (batch, dim, seqlen)
+    conv_state: (batch, dim, state_len), where state_len >= width - 1
+    weight: (dim, width)
+    bias: (dim,)
+    cache_seqlens: (batch,), dtype int32.
+        If not None, the conv_state is treated as a circular buffer.
+        The conv_state will be updated by copying x to the conv_state starting at the index
+        @cache_seqlens % state_len before performing the convolution.
+
+    out: (batch, dim) or (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+    dtype_in = x.dtype
+    unsqueeze = x.dim() == 2
+    if unsqueeze:
+        x = x.unsqueeze(-1)
+    batch, dim, seqlen = x.shape
+    width = weight.shape[1]
+    state_len = conv_state.shape[-1]
+    assert conv_state.shape == (batch, dim, state_len)
+    assert weight.shape == (dim, width)
+    if cache_seqlens is None:
+        x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype)  # (batch, dim, state_len + seqlen)
+        conv_state.copy_(x_new[:, :, -state_len:])
+    else:
+        width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1)
+        width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1)
+        x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype)
+        copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1)
+        copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1)
+        conv_state.scatter_(2, copy_idx, x)
+    out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:]
+    if unsqueeze:
+        out = out.squeeze(-1)
+    return (out if activation is None else F.silu(out)).to(dtype=dtype_in)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype'),
+    [
+        pytest.param(*test, id="B{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}".format(*test))
+        for test in [
+            (2, 64, 128, 3, "swish", True, True, torch.float32),
+            (2, 128, 128, 4, "swish", False, True, torch.float32),
+            (2, 64, 128, 3, "swish", True, False, torch.float32),
+            (2, 128, 128, 4, "swish", False, False, torch.float32),
+            (2, 500, 1024, 3, None, True, True, torch.float32),
+            (2, 1024, 1024, 4, None, False, True, torch.float32),
+            (2, 64, 128, 3, None, True, False, torch.float16),
+            (2, 128, 128, 4, None, False, False, torch.float16),
+        ]
+    ],
+)
+def test_conv(
+    B: int,
+    T: int,
+    D: int,
+    W: int,
+    activation: str,
+    has_bias: bool,
+    has_residual: bool,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    x = torch.randn(B, T, D).to(device, dtype).requires_grad_(True)
+    weight = torch.randn(D, W).to(device, dtype).requires_grad_(True)
+    bias = torch.randn(D).to(device, dtype).requires_grad_(True) if has_bias else None
+    residual = x.detach().clone().requires_grad_(True) if has_residual else None
+    dy = torch.randn(B, T, D).to(device, dtype)
+
+    ref = causal_conv1d_ref_torch(
+        x=rearrange(x, "b t d -> b d t"),
+        weight=weight,
+        bias=bias,
+        activation=activation,
+    )
+    ref = rearrange(ref, "b d t -> b t d")
+    if has_residual:
+        ref += residual
+    ref.backward(dy)
+    ref_dx, x.grad = x.grad, None
+    ref_dw, weight.grad = weight.grad, None
+    if has_bias:
+        ref_db, bias.grad = bias.grad, None
+    if has_residual:
+        ref_dr, residual.grad = residual.grad, None
+
+    tri, _ = causal_conv1d(x, weight, bias, residual=residual, activation=activation)
+    tri.backward(dy)
+    tri_dx, x.grad = x.grad, None
+    tri_dw, weight.grad = weight.grad, None
+    if has_bias:
+        tri_db, bias.grad = bias.grad, None
+    if has_residual:
+        tri_dr, residual.grad = residual.grad, None
+
+    assert_close(" y", ref, tri, 1e-3)
+    assert_close("dx", ref_dx, tri_dx, 1e-3)
+    assert_close("dw", ref_dw, tri_dw, 1e-3)
+    if has_bias:
+        assert_close("db", ref_db, tri_db, 1e-3)
+    if has_residual:
+        assert_close("dr", ref_dr, tri_dr, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('N', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype'),
+    [
+        pytest.param(*test, id="N{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}".format(*test))
+        for test in [
+            (4, 500, 128, 3, "swish", True, True, torch.float32),
+            (4, 1024, 200, 4, "swish", False, True, torch.float32),
+            (4, 500, 128, 3, None, True, False, torch.float16),
+            (4, 1024, 1024, 4, None, False, False, torch.float16),
+        ]
+    ],
+)
+def test_conv_varlen(
+    N: int,
+    T: int,
+    D: int,
+    W: int,
+    activation: str,
+    has_bias: bool,
+    has_residual: bool,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    cu_seqlens = torch.cat([
+        torch.tensor([0], dtype=torch.long),
+        torch.arange(16, T)[torch.randperm(T - 16)[:N-1]],
+        torch.tensor([T], dtype=torch.long),
+    ], 0).to(device).sort()[0]
+
+    x = torch.randn(1, T, D).to(device, dtype).requires_grad_(True)
+    weight = torch.randn(D, W).to(device, dtype).requires_grad_(True)
+    bias = torch.randn(D).to(device, dtype).requires_grad_(True) if has_bias else None
+    residual = x.detach().clone().requires_grad_(True) if has_residual else None
+    dy = torch.randn(1, T, D).to(device, dtype)
+
+    ref = torch.cat([
+        rearrange(
+            causal_conv1d_ref_torch(
+                x=rearrange(x[:, bos:eos].contiguous(), "b t d -> b d t"),
+                weight=weight,
+                bias=bias,
+                activation=activation,
+            ),
+            "b t d -> b d t",
+        ) + (residual[:, bos:eos] if has_residual else torch.zeros_like(x[:, bos:eos]))
+        for bos, eos in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)
+    ], 1)
+    ref.backward(dy)
+    ref_dx, x.grad = x.grad, None
+    ref_dw, weight.grad = weight.grad, None
+    if has_bias:
+        ref_db, bias.grad = bias.grad, None
+    if has_residual:
+        ref_dr, residual.grad = residual.grad, None
+
+    tri, _ = causal_conv1d(x, weight, bias, residual=residual, activation=activation, cu_seqlens=cu_seqlens)
+    tri.backward(dy)
+    tri_dx, x.grad = x.grad, None
+    tri_dw, weight.grad = weight.grad, None
+    if has_bias:
+        tri_db, bias.grad = bias.grad, None
+    if has_residual:
+        tri_dr, residual.grad = residual.grad, None
+
+    assert_close(" y", ref, tri, 1e-3)
+    assert_close("dx", ref_dx, tri_dx, 1e-3)
+    assert_close("dw", ref_dw, tri_dw, 1e-3)
+    if has_bias:
+        assert_close("db", ref_db, tri_db, 1e-3)
+    if has_residual:
+        assert_close("dr", ref_dr, tri_dr, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype'),
+    [
+        pytest.param(*test, id="B{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}".format(*test))
+        for test in [
+            (2, 64, 128, 3, "swish", True, True, torch.float32),
+            (2, 128, 128, 4, "swish", False, True, torch.float32),
+            (2, 64, 128, 3, "swish", True, False, torch.float32),
+            (2, 128, 128, 4, "swish", False, False, torch.float32),
+            (2, 500, 1024, 3, None, True, True, torch.float32),
+            (2, 1024, 1024, 4, None, False, True, torch.float32),
+            (2, 64, 128, 3, None, True, False, torch.float16),
+            (2, 128, 128, 4, None, False, False, torch.float16),
+        ]
+    ],
+)
+@torch.no_grad
+def test_conv_decoding(
+        B: int,
+        T: int,
+        D: int,
+        W: int,
+        activation: str,
+        has_bias: bool,
+        has_residual: bool,
+        dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    x = torch.randn(B, T, D).to(device, dtype)
+    weight = torch.randn(D, W).to(device, dtype) * 0
+    bias = torch.randn(D).to(device, dtype) if has_bias else None
+    residual = x.clone() if has_residual else None
+
+    ref = causal_conv1d_ref_torch(
+        x=rearrange(x, "b t d -> b d t"),
+        weight=weight,
+        bias=bias,
+        activation=activation,
+    )
+    ref = rearrange(ref, "b d t -> b t d")
+    if has_residual:
+        ref += residual
+    ref_cache = x.new_zeros(B, D, W)
+    ref_cache[:, :, -min(W, T):].copy_(rearrange(x[..., -min(W, T):, :], 'n w d -> n d w'))
+
+    tri = torch.zeros_like(x)
+    tri_cache = x.new_zeros(B, D, W)
+    for i in range(T):
+        y, tri_cache = causal_conv1d_update(
+            x=x[:, i:i+1, :],
+            cache=tri_cache,
+            residual=residual[:, i:i+1, :] if has_residual else None,
+            weight=weight,
+            bias=bias,
+            activation=activation,
+        )
+        tri[:, i:i+1, :] = y
+
+    assert_close("    y", ref, tri, 1e-3)
+    assert_close("cache", ref_cache, tri_cache, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype', 'backend'),
+    [
+        pytest.param(
+            *test, id="B{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}_{8}".format(*test))
+        for test in [
+            (2, 64, 128, 3, "swish", True, True, torch.float32, 'triton'),
+            (2, 128, 128, 4, "swish", False, True, torch.float32, 'triton'),
+            (2, 64, 128, 3, "swish", True, False, torch.float32, 'triton'),
+            (2, 128, 128, 4, "swish", False, False, torch.float32, 'triton'),
+            (2, 500, 1024, 3, None, True, True, torch.float32, 'triton'),
+            (2, 1024, 1024, 4, None, False, True, torch.float32, 'triton'),
+            (2, 64, 128, 3, None, True, False, torch.float16, 'triton'),
+            (2, 128, 128, 4, None, False, False, torch.float16, 'triton'),
+            (2, 64, 128, 3, "swish", True, True, torch.float32, 'cuda'),
+            (2, 128, 128, 4, "swish", False, True, torch.float32, 'cuda'),
+            (2, 64, 128, 3, "swish", True, False, torch.float32, 'cuda'),
+            (2, 128, 128, 4, "swish", False, False, torch.float32, 'cuda'),
+            (2, 2, 128, 4, "swish", True, True, torch.float32, 'cuda'),  # T_prefill < W
+            (2, 2, 128, 4, "swish", True, True, torch.float32, 'triton'),
+            (2, 3, 128, 4, "swish", True, True, torch.float32, 'triton'),
+            (2, 4, 128, 4, "swish", True, True, torch.float32, 'triton'),
+            (2, 2, 128, 3, "swish", True, True, torch.float32, 'triton'),
+        ]
+    ],
+)
+@torch.no_grad
+def test_conv_with_cache_prefill_fwd(
+    B: int,
+    T: int,
+    D: int,
+    W: int,
+    activation: str,
+    has_bias: bool,
+    has_residual: bool,
+    dtype: torch.dtype,
+    backend: str,
+):
+    if causal_conv1d_fn is None and backend == 'cuda':
+        pytest.skip("causal_conv1d is not installed for CUDA backend")
+    torch.manual_seed(42)
+
+    x = torch.randn(B, T, D).to(device, dtype)
+    residual = torch.randn(B, T, D).to(device, dtype) if has_residual else None
+
+    conv = ShortConvolution(
+        hidden_size=D,
+        kernel_size=W,
+        bias=has_bias,
+        activation=activation,
+        backend=backend,
+        device=device,
+        dtype=dtype,
+    )
+
+    cache = torch.randn(B, D, W - 1).to(device, dtype)
+
+    ref = causal_conv1d_ref_torch(
+        x=x.transpose(1, 2),                    # (B, D, T)
+        weight=rearrange(conv.weight, "d 1 w -> d w"),
+        bias=conv.bias,
+        initial_state=cache,                    # (B, D, W-1)
+        activation=activation,
+    ).transpose(1, 2)                           # (B, T, D)
+    if has_residual:
+        ref += residual
+
+    zero_padding = torch.zeros(B, D, 1).to(device, dtype)
+    tri_cache = torch.cat([zero_padding, cache], dim=-1)  # (B, D, W)
+    tri, cache_out = conv(x, residual=residual, cache=tri_cache.clone(), output_final_state=True)
+
+    assert_close("y", ref, tri, 1e-3)
+    for p in range(1, W):
+        if p <= T:
+            expected = x[:, -p, :]
+        else:
+            expected = tri_cache[:, :, -(p - T)]
+        torch.testing.assert_close(
+            cache_out[:, :, -p],
+            expected,
+            atol=1e-3, rtol=1e-3,
+        )
+
+
+@pytest.mark.parametrize(
+    ('N', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype', 'backend'),
+    [
+        pytest.param(
+            *test,
+            id="N{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}_{8}".format(*test),
+        )
+        for test in [
+            (3, 128, 64, 4, "swish", True, True, torch.float32, 'triton'),
+            (4, 256, 128, 3, None,  False, True, torch.float32, 'triton'),
+            (2,  64, 128, 4, "swish", True, False, torch.float16, 'cuda'),
+            (3, 200,  64, 3, None,  False, False, torch.float16, 'cuda'),
+            (2,   3,  64, 4, "swish", True, True, torch.float32, 'triton'),  # T < W
+            (2,   3,  64, 3, None,  False, True, torch.float32, 'cuda'),     # T < W
+        ]
+    ],
+)
+@torch.no_grad
+def test_conv_varlen_with_cache_prefill_fwd(
+    N: int,
+    T: int,
+    D: int,
+    W: int,
+    activation: str,
+    has_bias: bool,
+    has_residual: bool,
+    dtype: torch.dtype,
+    backend: str,
+):
+    if causal_conv1d_fn is None and backend == 'cuda':
+        pytest.skip("causal_conv1d is not installed for CUDA backend")
+    torch.manual_seed(42)
+
+    min_len_each = max(1, T // N)
+    lengths = [min_len_each] * N
+    lengths[-1] += T % N
+    assert all(length >= 1 for length in lengths), "all lengths must >= 1"
+    cu_seqlens = torch.tensor([0] + torch.cumsum(torch.tensor(lengths), 0).tolist(),
+                              device=device, dtype=torch.int32)
+
+    x = torch.randn(1, T, D).to(device, dtype)
+    residual = torch.randn(1, T, D).to(device, dtype) if has_residual else None
+
+    conv = ShortConvolution(
+        hidden_size=D,
+        kernel_size=W,
+        bias=has_bias,
+        activation=activation,
+        backend=backend,
+        device=device,
+        dtype=dtype,
+    )
+
+    cache = torch.randn(N, D, W - 1).to(device, dtype)
+    ref_list = []
+    for i, (bos, eos) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)):
+        xi = x[:, bos:eos, :].transpose(1, 2)  # (1, D, l)
+        ci = cache[i:i + 1]                    # (1, D, W-1)
+        refi = causal_conv1d_ref_torch(
+            x=xi,
+            weight=rearrange(conv.weight, "d 1 w -> d w"),
+            bias=conv.bias,
+            initial_state=ci,
+            activation=activation,
+        ).transpose(1, 2)                      # (1, l, D)
+        if has_residual:
+            refi += residual[:, bos:eos, :]
+        ref_list.append(refi)
+    ref = torch.cat(ref_list, dim=1)           # (1, T, D)
+
+    zero_pad = torch.zeros(N, D, 1, device=device, dtype=dtype)
+    tri_cache = torch.cat([zero_pad, cache], dim=-1)  # (N, D, W)
+    tri, cache_out = conv(x,
+                          residual=residual,
+                          cache=tri_cache.clone(),
+                          cu_seqlens=cu_seqlens,
+                          output_final_state=True)
+
+    assert_close("varlen y", ref, tri, 1e-3)
+
+    for i, (bos, eos) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)):
+        length = eos - bos
+        for p in range(1, W):
+            if p <= length:
+                expected = x[0, eos - p, :]
+            else:
+                expected = tri_cache[i, :, -(p - length)]
+            torch.testing.assert_close(
+                cache_out[i, :, -p],
+                expected,
+                atol=1e-3,
+                rtol=1e-3,
+            )
+
+
+@pytest.mark.parametrize(
+    ('B', 'D', 'W', 'has_bias', 'has_residual', 'activation', 'dtype', 'backend'),
+    [
+        pytest.param(*test, id="B{0}_D{1}_W{2}_has_bias{3}_has_residual{4}_activation{5}_{6}_{7}".format(*test))
+        for test in [
+            (2, 128, 3, True, True, "swish", torch.float32, 'triton'),
+            (2, 128, 4, False, True, "swish", torch.float32, 'triton'),
+            (2, 128, 3, True, False, "swish", torch.float32, 'triton'),
+            (2, 128, 4, False, False, "swish", torch.float32, 'triton'),
+            (2, 128, 3, True, True, "swish", torch.float32, 'cuda'),
+            (2, 128, 4, False, True, "swish", torch.float32, 'cuda'),
+            (2, 128, 3, True, False, "swish", torch.float32, 'cuda'),
+            (2, 128, 4, False, False, "swish", torch.float32, 'cuda'),
+            (2, 128, 4, False, False, None, torch.float32, 'cuda'),
+            (2, 128, 4, False, False, None, torch.float32, 'triton'),
+        ]
+    ],
+)
+@torch.no_grad
+def test_conv_decoding_with_cache(
+    B: int,
+    D: int,
+    W: int,
+    activation: str,
+    has_bias: bool,
+    has_residual: bool,
+    dtype: torch.dtype,
+    backend: str,
+):
+    if causal_conv1d_fn is None and backend == 'cuda':
+        pytest.skip("causal_conv1d is not installed for CUDA backend")
+    torch.manual_seed(42)
+
+    x = torch.randn(B, 1, D).to(device, dtype)        # (B, 1, D)
+    residual = x.clone() if has_residual else None
+
+    conv = ShortConvolution(
+        hidden_size=D,
+        kernel_size=W,
+        bias=has_bias,
+        activation=activation,
+        backend=backend,
+        device=device,
+        dtype=dtype,
+    )
+
+    state = torch.randn(B, D, W).to(device, dtype)
+
+    # reference
+    ref = causal_conv1d_update_ref_torch(
+        x.squeeze(1),                           # (B, D)
+        conv_state=state.clone(),
+        weight=rearrange(conv.weight, "d 1 w -> d w"),
+        bias=conv.bias,
+        activation=activation,
+    ).unsqueeze(1)                             # (B, 1, D)
+    if has_residual:
+        ref += residual
+
+    # ShortConvolution step
+    with torch.no_grad():
+        y, _ = conv.step(x, residual, state.clone())
+
+    assert_close("y", ref, y, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'W', 'has_bias', 'has_residual', 'activation', 'dtype'),
+    [
+        pytest.param(*test, id="B{0}_T{1}_D{2}_W{3}_has_bias{4}_has_residual{5}_activation{6}_{7}".format(*test))
+        for test in [
+            (2, 64, 128, 3, True, True, "swish", torch.float32),
+            (2, 128, 128, 4, False, True, "swish", torch.float32),
+            (2, 64, 128, 3, True, False, "swish", torch.float32),
+            (2, 128, 128, 4, False, False, "swish", torch.float32),
+        ]
+    ],
+)
+@torch.no_grad
+def test_mixed_backend(
+    B: int,
+    T: int,
+    D: int,
+    W: int,
+    has_bias: bool,
+    has_residual: bool,
+    activation: str,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    T_decode = 1
+    x = torch.randn(B, T + T_decode, D, device=device, dtype=dtype)
+    residual = torch.randn_like(x) if has_residual else None
+
+    conv = ShortConvolution(
+        hidden_size=D,
+        kernel_size=W,
+        bias=has_bias,
+        activation=activation,
+        backend="cuda",
+        device=device,
+        dtype=dtype,
+    )
+
+    cache = torch.randn(B, D, W-1, device=device, dtype=dtype)
+    y_cuda_prefill, final_state = conv(
+        x[:, :T],
+        residual=residual[:, :T] if has_residual else None,
+        cache=cache,
+        output_final_state=True,
+    )
+
+    conv.backend = "triton"
+    y_triton_decode, _ = conv(
+        x[:, T:],
+        residual=residual[:, T:] if has_residual else None,
+        cache=final_state,
+        output_final_state=True,
+    )
+
+    conv.backend = "triton"
+    cache = torch.cat((torch.zeros_like(cache[..., :1]), cache), -1)
+    y_triton_full, _ = conv(x, residual=residual, cache=cache)
+
+    y_mixed = torch.cat([y_cuda_prefill, y_triton_decode], dim=1)
+    assert_close("cuda→triton vs triton", y_mixed, y_triton_full, 1e-3)
+
+    conv.backend = "triton"
+    y_triton_prefill, final_state = conv(
+        x[:, :T],
+        residual=residual[:, :T] if has_residual else None,
+        cache=cache,
+        output_final_state=True,
+    )
+
+    conv.backend = "cuda"
+    y_cuda_decode, _ = conv(
+        x[:, T:],
+        residual=residual[:, T:] if has_residual else None,
+        cache=final_state,
+        output_final_state=True,
+    )
+
+    y_mixed2 = torch.cat([y_triton_prefill, y_cuda_decode], dim=1)
+    assert_close("triton→cuda vs triton", y_mixed2, y_triton_full,  1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'W', 'has_bias', 'has_residual', 'activation', 'dtype'),
+    [
+        pytest.param(*test, id="B{0}_T{1}_D{2}_W{3}_has_bias{4}_has_residual{5}_activation{6}_{7}".format(*test))
+        for test in [
+            (2, 64, 100, 3, True, True, "swish", torch.float32),
+            (2, 128, 128, 4, True, True, "swish", torch.float32),
+            (3, 128, 128, 4, True, True, "swish", torch.float32),
+            (3, 128, 256, 4, True, True, "swish", torch.float32),
+            (3, 128, 512, 4, True, True, "swish", torch.float32),
+            (2, 128, 1024, 4, True, True, "swish", torch.float32),
+            (2, 128, 2048, 3, True, True, "swish", torch.float32),
+            (2, 128, 4096, 4, True, True, "swish", torch.float32),
+            (2, 128, 8192, 4, True, True, "swish", torch.float32),
+        ]
+    ],
+)
+def test_conv_cache_backward(
+    B: int,
+    T: int,
+    D: int,
+    W: int,
+    has_bias: bool,
+    has_residual: bool,
+    activation: str,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    x = torch.randn(B, T, D, device=device, dtype=dtype, requires_grad=True)
+    weight = torch.randn(D, W, device=device, dtype=dtype, requires_grad=True)
+    bias = torch.randn(D, device=device, dtype=dtype, requires_grad=True) if has_bias else None
+    residual = torch.randn(B, T, D, device=device, dtype=dtype, requires_grad=True) if has_residual else None
+    cache = torch.randn(B, D, W - 1, device=device, dtype=dtype, requires_grad=True)
+
+    def ref_func(x, weight, bias, residual, cache):
+        out, cache_out = causal_conv1d_ref_torch(
+            x.transpose(1, 2),
+            weight,
+            bias,
+            initial_state=cache,
+            output_final_state=True,
+            activation=activation,
+        )
+        out = out.transpose(1, 2)
+        if residual is not None:
+            out += residual
+        return out, cache_out
+
+    def triton_func(x, weight, bias, residual, cache):
+        zero_padding = torch.zeros(B, D, 1, device=device, dtype=dtype)
+        triton_cache = torch.cat([zero_padding, cache], dim=-1).contiguous()
+        tri, cache_out_triton = causal_conv1d(
+            x,
+            weight=weight,
+            bias=bias,
+            residual=residual,
+            initial_state=triton_cache,
+            output_final_state=True,
+            activation=activation,
+        )
+        cache_out_triton = cache_out_triton[..., 1:].clone()  # [B, D, W-1]
+        return tri, cache_out_triton
+
+    d_tri = torch.randn_like(x)
+    d_cache_out = torch.randn_like(cache)
+
+    def get_grads(func, *inputs):
+        out, cache_out = func(*inputs)
+        loss = (out * d_tri).sum() + (cache_out * d_cache_out).sum()
+        grads = torch.autograd.grad(
+            loss,
+            inputs,
+            retain_graph=True,
+            create_graph=False,
+        )
+        return grads
+
+    inputs = (x, weight, bias, residual, cache)
+    grads_ref = get_grads(ref_func, *inputs)
+    grads_tri = get_grads(triton_func, *inputs)
+
+    names = ["x", "weight", "bias", "residual", "cache"]
+    for name, g_ref, g_tri in zip(names, grads_ref, grads_tri, strict=False):
+        assert_close(name, g_ref, g_tri, ratio=1e-3)

code/flash-linear-attention/tests/modules/test_cross_entropy.py

@@ -0,0 +1,81 @@
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss
+from fla.utils import assert_close, device, device_platform
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("T", [512, 1024])
+@pytest.mark.parametrize("D", [1024, 2048])
+@pytest.mark.parametrize("V", [32000, 100000])
+@pytest.mark.parametrize("reduction", ['mean'])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason="Intel Triton Failure",
+)
+def test_fused_cross_entropy(B: int, T: int, D: int, V: int, reduction: str, dtype: torch.dtype):
+    torch.manual_seed(42)
+    logits = torch.randn(B * T, V).to(device).to(dtype=dtype).requires_grad_()
+    target = torch.randint(0, V, (B, T)).to(device)
+    target = torch.cat((target[..., 1:], torch.full_like(target[..., :1], -100)), -1)
+    target = target.flatten()
+
+    ref = nn.CrossEntropyLoss(reduction=reduction)(logits, target).to(dtype=dtype)
+    do = torch.randn_like(ref).to(device).to(dtype=dtype)
+
+    ref.backward(do)
+    ref_d, logits.grad = logits.grad.clone(), None
+
+    tri = FusedCrossEntropyLoss(reduction=reduction)(logits, target).to(dtype=dtype)
+    tri.backward(do)
+    tri_d, logits.grad = logits.grad.clone(), None
+
+    assert_close(" o", ref, tri, ratio=1e-2)
+    assert_close("dl", ref_d, tri_d, ratio=1e-2)
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("T", [512, 1024])
+@pytest.mark.parametrize("D", [1024, 2048])
+@pytest.mark.parametrize("V", [32000, 100000])
+@pytest.mark.parametrize("scale", [1., 0.5])
+@pytest.mark.parametrize("reduction", ['mean'])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason="Intel Triton Failure",
+)
+def test_fused_linear_cross_entropy(B: int, T: int, D: int, V: int, scale: float, reduction: str, dtype: torch.dtype):
+    torch.manual_seed(42)
+
+    x = torch.randn(B * T, D).to(device).to(dtype=dtype).requires_grad_()
+    target = torch.randint(0, V, (B, T)).to(device)
+    target = torch.cat((target[..., 1:], torch.full_like(target[..., :1], -100)), -1)
+    target = target.flatten()
+    weight = torch.randn(V, D).to(device).to(dtype=dtype).requires_grad_()
+    bias = torch.randn(V).to(device).to(dtype=dtype).requires_grad_()
+
+    logits = F.linear(x, weight, bias)
+    ref = FusedCrossEntropyLoss(logit_scale=scale, reduction=reduction)(logits, target)
+    do = torch.randn_like(ref).to(device).to(dtype=dtype)
+
+    ref.backward(do)
+    ref_dx, x.grad = x.grad.clone(), None
+    ref_dw, weight.grad = weight.grad.clone(), None
+    ref_db, bias.grad = bias.grad.clone(), None
+
+    tri = FusedLinearCrossEntropyLoss(logit_scale=scale, reduction=reduction)(x, target, weight, bias)
+    tri.backward(do)
+    tri_dx, x.grad = x.grad.clone(), None
+    tri_dw, weight.grad = weight.grad.clone(), None
+    tri_db, bias.grad = bias.grad.clone(), None
+
+    assert_close(" o", ref, tri, ratio=1e-2)
+    assert_close("dx", ref_dx, tri_dx, ratio=1e-2)
+    assert_close("dw", ref_dw, tri_dw, ratio=1e-2)
+    assert_close("db", ref_db, tri_db, ratio=1e-2)

code/flash-linear-attention/tests/modules/test_grpo.py

@@ -0,0 +1,56 @@
+
+import pytest
+import torch
+
+from fla.modules.grpo import fused_grpo_loss, grpo_loss_torch
+from fla.utils import assert_close, device, device_torch_lib, is_nvidia_hopper
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("T", [16, 1024, 4096])
+@pytest.mark.parametrize("V", [32000, 65536, 131072])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("inplace", [True, False])
+@pytest.mark.parametrize("repeat", [100])
+def test_fused_grpos(B: int, T: int, V: int, dtype: torch.dtype, inplace: bool, repeat: int):
+    device_torch_lib.manual_seed(42)
+    for i in range(repeat):
+        if not is_nvidia_hopper and T == 4096:
+            pytest.skip("Skip test for T=4096 on Intel Alchemist")
+
+        def get_random_ref_log_probs(logits, input_ids):
+            with torch.inference_mode():
+                logits = logits[:, :-1]
+                per_token_logps = []
+                for logits_row, input_ids_row in zip(logits, input_ids[:, -logits.size(1):], strict=False):
+                    log_probs = torch.randn_like(logits_row).log_softmax(dim=-1)
+                    token_log_prob = torch.gather(log_probs, dim=1, index=input_ids_row.unsqueeze(1)).squeeze(1)
+                    per_token_logps.append(token_log_prob)
+                device_torch_lib.empty_cache()
+                return torch.stack(per_token_logps)
+
+        logits = torch.randn(B, T + 1, V, device=device, dtype=dtype)
+        logits.requires_grad_(True)
+        advantages = torch.randn(B, device=device, dtype=torch.float32)
+        input_ids = torch.randint(0, V-1, (B, T + 64), device=device)
+        ref_logp = get_random_ref_log_probs(logits, input_ids)
+        beta = 0.04
+        completion_mask = torch.ones(B, T, dtype=torch.int32, device=device)
+        completion_mask[::2, T//3: T//2] = 0
+        save_kl = True
+
+        gold_logits = logits.detach().clone().float()
+        gold_logits.requires_grad_(True)
+        gold_ref_logp = ref_logp.clone().float()
+        device_torch_lib.empty_cache()
+        y1 = fused_grpo_loss(logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl=save_kl, inplace=inplace)
+        y2 = grpo_loss_torch(gold_logits, gold_ref_logp, input_ids, advantages, beta, completion_mask, save_kl)
+        if save_kl:
+            y1, kl2 = y1
+            y2, kl3 = y2
+            assert (kl2-kl3).abs().max() < 1e-3
+        dy = torch.randn_like(y1) * 10
+        y1.backward(dy)
+        y2.backward(dy.float())
+        assert (y1-y2).abs().max() < 1e-3
+        assert_close(" dlogits", gold_logits.grad, logits.grad, 3e-3)

code/flash-linear-attention/tests/modules/test_kl_div.py

@@ -0,0 +1,44 @@
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.modules import FusedKLDivLoss
+from fla.utils import assert_close, device, device_platform
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("T", [16, 32])
+@pytest.mark.parametrize("D", [1024, 2048])
+@pytest.mark.parametrize("V", [32000, 100000])
+@pytest.mark.parametrize("reduction", ["batchmean"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason="Intel Triton Failure",
+)
+def test_fused(B: int, T: int, D: int, V: int, reduction: str, dtype: torch.dtype):
+    torch.manual_seed(42)
+    x = torch.randn(B * T, D).to(device).to(dtype=dtype).requires_grad_()
+    x_weight = torch.randn(V, D).to(device).to(dtype=dtype).requires_grad_()
+    target_x = torch.randn(B * T, D).to(device).to(dtype=dtype)
+    target_weight = torch.randn(V, D).to(device).to(dtype=dtype)
+
+    ref = F.kl_div(
+        F.linear(x, x_weight).log_softmax(-1),
+        F.linear(target_x, target_weight).softmax(-1),
+        reduction=reduction,
+    ).to(dtype)
+    do = torch.randn_like(ref).to(device)
+    ref.backward(do)
+    ref_dx, x.grad = x.grad.clone(), None
+    ref_dw, x_weight.grad = x_weight.grad.clone(), None
+
+    tri = FusedKLDivLoss(reduction)(x, target_x, x_weight, target_weight).to(dtype=dtype)
+    tri.backward(do)
+    tri_dx, x.grad = x.grad.clone(), None
+    tri_dw, x_weight.grad = x_weight.grad.clone(), None
+
+    assert_close("  o", ref, tri, 1e-2)
+    assert_close(" dx", ref_dx, tri_dx, 1e-2)
+    assert_close(" dw", ref_dw, tri_dw, 1e-2)

code/flash-linear-attention/tests/modules/test_l2norm.py

@@ -0,0 +1,37 @@
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.modules.l2norm import l2_norm
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 60, torch.float),
+            (2, 500, 4, 64, torch.float),
+            (2, 1000, 2, 100, torch.float),
+            (3, 1024, 4, 128, torch.float),
+            (4, 1024, 5, 1024, torch.float16),
+            (4, 1024, 5, 1024, torch.bfloat16),
+            (5, 1024, 6, 2048, torch.float16),
+            (5, 1024, 6, 2048, torch.bfloat16),
+        ]
+    ],
+)
+def test_l2norm(B: int, T: int, H: int, D: int, dtype: torch.dtype):
+    torch.manual_seed(42)
+    x = torch.randn(B, T, H, D, dtype=dtype).to(device).requires_grad_(True)
+    x = x * 0.5 + 0.3
+
+    ref = F.normalize(x, dim=-1, p=2)
+    tri = l2_norm(x)
+    ref_dx = torch.autograd.grad(ref.sum(), x)[0]
+    tri_dx = torch.autograd.grad(tri.sum(), x)[0]
+
+    assert_close('y', ref, tri, 0.005)
+    assert_close('dx', ref_dx, tri_dx, 0.005)

code/flash-linear-attention/tests/modules/test_l2warp.py

@@ -0,0 +1,70 @@
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fla.modules import FusedLinearCrossEntropyLoss
+from fla.modules.l2warp import l2_warp as standalone_l2_warp
+from fla.utils import assert_close, device, is_intel_alchemist
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("B", [4, 8])
+@pytest.mark.parametrize("T", [1024])
+@pytest.mark.parametrize("H", [256])
+@pytest.mark.parametrize("V", [2000])
+@pytest.mark.parametrize("l2_penalty_factor", [1e-4, 1])
+@pytest.mark.skipif(
+    is_intel_alchemist is True,
+    reason="Intel Triton Failure",
+)
+def test_fused_linear_cross_entropy_l2_warp(
+    B: int,
+    T: int,
+    H: int,
+    V: int,
+    l2_penalty_factor: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    lm_head = nn.Linear(H, V, bias=True, device=device, dtype=dtype)
+    x = torch.randn(B, T, H, device=device, dtype=dtype, requires_grad=True)
+    labels = torch.randint(0, V, (B, T), device=device)
+
+    ignore_index = -100
+    shift_labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], ignore_index)), 1)
+
+    ref_criterion = nn.CrossEntropyLoss()
+
+    ref_logits = F.linear(x.view(-1, H), lm_head.weight, lm_head.bias)
+    ref_loss_ce = ref_criterion(ref_logits.view(B * T, V), shift_labels.view(-1))
+    ref_loss = standalone_l2_warp(ref_loss_ce, ref_logits.view(B, T, V), l2_penalty_factor)
+
+    ref_loss.backward()
+    ref_x_grad = x.grad.clone()
+    ref_w_grad = lm_head.weight.grad.clone()
+    ref_b_grad = lm_head.bias.grad.clone()
+
+    x.grad = None
+    lm_head.zero_grad()
+
+    fused_criterion = FusedLinearCrossEntropyLoss(
+        l2_penalty_factor=l2_penalty_factor,
+        use_l2warp=True,  # Make sure to enable it
+    )
+
+    fused_loss = fused_criterion(x, shift_labels, lm_head.weight, lm_head.bias)
+
+    fused_loss.backward()
+    fused_x_grad = x.grad.clone()
+    fused_w_grad = lm_head.weight.grad.clone()
+    fused_b_grad = lm_head.bias.grad.clone()
+
+    ratio = 4e-3 if dtype == torch.bfloat16 else 1e-3
+
+    assert_close("Loss", ref_loss, fused_loss, ratio)
+    assert_close("dx", ref_x_grad, fused_x_grad, ratio)
+    assert_close("dw", ref_w_grad, fused_w_grad, ratio)
+    assert_close("db", ref_b_grad, fused_b_grad, ratio)

code/flash-linear-attention/tests/modules/test_layernorm.py

@@ -0,0 +1,217 @@
+
+import pytest
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+
+from fla.modules import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear
+from fla.modules.layernorm import GroupNormRef
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("H", [2])
+@pytest.mark.parametrize("T", [512])
+@pytest.mark.parametrize("D", [50, 64, 128])
+@pytest.mark.parametrize("elementwise_affine", [False, True])
+@pytest.mark.parametrize("bias", [False, True])
+def test_layernorm(B: int, H: int, T: int, D: int, elementwise_affine: bool, bias: bool):
+    x = torch.randn(B, H, T, D).to(device).requires_grad_(True)
+    ref = nn.LayerNorm(D, elementwise_affine=elementwise_affine, bias=bias).to(device)
+    tri = LayerNorm(D, elementwise_affine=elementwise_affine, bias=bias).to(device)
+    if ref.weight is not None:
+        nn.init.normal_(ref.weight)
+        tri.weight.data.copy_(ref.weight.data)
+    if ref.bias is not None:
+        nn.init.normal_(ref.bias)
+        tri.bias.data.copy_(ref.bias.data)
+
+    ref_y = ref(x)
+    tri_y = tri(x)
+    ref_dx = torch.autograd.grad(ref(x).sum(), x)[0]
+    tri_dx = torch.autograd.grad(tri(x).sum(), x)[0]
+
+    if ref.weight is not None:
+        ref_dw = torch.autograd.grad(ref(x).sum(), ref.weight)[0]
+        tri_dw = torch.autograd.grad(tri(x).sum(), tri.weight)[0]
+    if ref.bias is not None:
+        ref_db = torch.autograd.grad(ref(x).sum(), ref.bias)[0]
+        tri_db = torch.autograd.grad(tri(x).sum(), tri.bias)[0]
+
+    assert_close(' y', ref_y, tri_y, 1e-3)
+    assert_close('dx', ref_dx, tri_dx, 1e-3)
+    if ref.weight is not None:
+        assert_close('dw', ref_dw, tri_dw, 1e-3)
+    if ref.bias is not None:
+        assert_close('db', ref_db, tri_db, 1e-3)
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("T", [512])
+@pytest.mark.parametrize("D", [64, 128, 512, 1024, 2048])
+@pytest.mark.parametrize("G", [1, 4])
+@pytest.mark.parametrize("is_rms_norm", [True, False])
+def test_groupnorm(B: int, T: int, D: int, G: int, is_rms_norm: bool):
+    torch.manual_seed(42)
+    x = torch.randn(B, T, D).to(device).requires_grad_(True)
+    if is_rms_norm:
+        ref = GroupNormRef(num_groups=G, hidden_size=D, bias=True, is_rms_norm=True).to(device)
+    else:
+        ref = nn.GroupNorm(G, D).to(device)
+    tri = GroupNorm(G, D, bias=True, is_rms_norm=is_rms_norm).to(device)
+    nn.init.normal_(ref.weight)
+    nn.init.normal_(ref.bias)
+    tri.weight.data.copy_(ref.weight.data)
+    tri.bias.data.copy_(ref.bias.data)
+    ref = ref.to(dtype=torch.float32)
+
+    ref_x = rearrange(x, 'b t d -> (b t) d').to(dtype=torch.float32)
+    ref_y = rearrange(ref(ref_x), '(b t) d -> b t d', b=B)
+    tri_y = tri(x)
+    ref_dx = torch.autograd.grad(ref(ref_x).sum(), x)[0]
+    tri_dx = torch.autograd.grad(tri(x).sum(), x)[0]
+    ref_dw = torch.autograd.grad(ref(ref_x).sum(), ref.weight)[0]
+    tri_dw = torch.autograd.grad(tri(x).sum(), tri.weight)[0]
+    ref_db = torch.autograd.grad(ref(ref_x).sum(), ref.bias)[0]
+    tri_db = torch.autograd.grad(tri(x).sum(), tri.bias)[0]
+
+    assert_close(' y', ref_y, tri_y, 1e-3)
+    assert_close('dx', ref_dx, tri_dx, 1e-3)
+    assert_close('dw', ref_dw, tri_dw, 1e-3)
+    assert_close('db', ref_db, tri_db, 1e-3)
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("H", [2])
+@pytest.mark.parametrize("T", [512])
+@pytest.mark.parametrize("D", [50, 64, 128])
+def test_rmsnorm(B: int, H: int, T: int, D: int):
+    x = torch.randn(B, H, T, D).to(device).requires_grad_(True)
+    ref = LlamaRMSNorm(D, eps=0).to(device)
+    tri = RMSNorm(D, eps=0).to(device)
+    nn.init.normal_(ref.weight)
+    tri.weight.data.copy_(ref.weight.data)
+
+    ref_y = ref(x)
+    tri_y = tri(x)
+    ref_dx = torch.autograd.grad(ref(x).sum(), x)[0]
+    tri_dx = torch.autograd.grad(tri(x).sum(), x)[0]
+
+    ref_dw = torch.autograd.grad(ref(x).sum(), ref.weight)[0]
+    tri_dw = torch.autograd.grad(tri(x).sum(), tri.weight)[0]
+
+    assert_close(' y', ref_y, tri_y, 1e-3)
+    assert_close('dx', ref_dx, tri_dx, 1e-3)
+    assert_close('dw', ref_dw, tri_dw, 1e-3)
+
+
+@pytest.mark.parametrize("N", [1, 16, 128])
+@pytest.mark.parametrize("D", [50, 64, 128])
+def test_layernorm_linear(N: int, D: int):
+    torch.manual_seed(1)
+    x = torch.randn(N, D).to(device).requires_grad_(True)
+    ref = nn.Sequential(nn.LayerNorm(D, elementwise_affine=True, bias=True), nn.Linear(D, D)).to(device)
+    tri = LayerNormLinear(D, elementwise_affine=True, bias=True).to(device)
+    nn.init.normal_(ref[0].weight)
+    nn.init.normal_(ref[0].bias)
+    nn.init.normal_(ref[1].weight, mean=0.0, std=0.01)
+    nn.init.normal_(ref[1].bias, mean=0.0, std=0.01)
+    tri.weight.data.copy_(ref[0].weight.data)
+    tri.bias.data.copy_(ref[0].bias.data)
+    weight, bias = ref[1].weight.clone(), ref[1].bias.clone()
+
+    ref_y = ref(x)
+    tri_y = tri(x, weight, bias)
+    ref_dx = torch.autograd.grad(ref(x).sum(), x)[0]
+    tri_dx = torch.autograd.grad(tri(x, weight, bias).sum(), x)[0]
+    ref_dw = torch.autograd.grad(ref(x).sum(), ref[0].weight)[0]
+    tri_dw = torch.autograd.grad(tri(x, weight, bias).sum(), tri.weight)[0]
+    ref_db = torch.autograd.grad(ref(x).sum(), ref[0].bias)[0]
+    tri_db = torch.autograd.grad(tri(x, weight, bias).sum(), tri.bias)[0]
+    ref_dlw = torch.autograd.grad(ref(x).sum(), ref[1].weight)[0]
+    tri_dlw = torch.autograd.grad(tri(x, weight, bias).sum(), weight)[0]
+    ref_dlb = torch.autograd.grad(ref(x).sum(), ref[1].bias)[0]
+    tri_dlb = torch.autograd.grad(tri(x, weight, bias).sum(), bias)[0]
+
+    assert_close('  y', ref_y, tri_y, 1e-3)
+    assert_close(' dx', ref_dx, tri_dx, 1e-3)
+    assert_close(' dw', ref_dw, tri_dw, 1e-3)
+    assert_close(' db', ref_db, tri_db, 1e-3)
+    assert_close('dlw', ref_dlw, tri_dlw, 1e-3)
+    assert_close('dlb', ref_dlb, tri_dlb, 1e-3)
+
+
+@pytest.mark.parametrize("N", [1, 16, 128])
+@pytest.mark.parametrize("D", [64, 128, 512])
+@pytest.mark.parametrize("G", [1, 4])
+@pytest.mark.parametrize("is_rms_norm", [True, False])
+def test_groupnorm_linear(N: int, D: int, G: int, is_rms_norm: bool):
+    torch.manual_seed(1)
+    x = torch.randn(N, D).to(device).requires_grad_(True)
+    if is_rms_norm:
+        ref = nn.Sequential(
+            GroupNormRef(num_groups=G, hidden_size=D, bias=True, is_rms_norm=True),
+            nn.Linear(D, D),
+        ).to(device)
+    else:
+        ref = nn.Sequential(nn.GroupNorm(G, D), nn.Linear(D, D)).to(device)
+    tri = GroupNormLinear(G, D, bias=True, is_rms_norm=is_rms_norm).to(device)
+    nn.init.normal_(ref[0].weight)
+    nn.init.normal_(ref[0].bias)
+    nn.init.normal_(ref[1].weight, mean=0.0, std=0.01)
+    nn.init.normal_(ref[1].bias, mean=0.0, std=0.01)
+    tri.weight.data.copy_(ref[0].weight.data)
+    tri.bias.data.copy_(ref[0].bias.data)
+    weight, bias = ref[1].weight.clone(), ref[1].bias.clone()
+
+    ref_y = ref(x)
+    tri_y = tri(x, weight, bias)
+    ref_dx = torch.autograd.grad(ref(x).sum(), x)[0]
+    tri_dx = torch.autograd.grad(tri(x, weight, bias).sum(), x)[0]
+    ref_dw = torch.autograd.grad(ref(x).sum(), ref[0].weight)[0]
+    tri_dw = torch.autograd.grad(tri(x, weight, bias).sum(), tri.weight)[0]
+    ref_db = torch.autograd.grad(ref(x).sum(), ref[0].bias)[0]
+    tri_db = torch.autograd.grad(tri(x, weight, bias).sum(), tri.bias)[0]
+    ref_dlw = torch.autograd.grad(ref(x).sum(), ref[1].weight)[0]
+    tri_dlw = torch.autograd.grad(tri(x, weight, bias).sum(), weight)[0]
+    ref_dlb = torch.autograd.grad(ref(x).sum(), ref[1].bias)[0]
+    tri_dlb = torch.autograd.grad(tri(x, weight, bias).sum(), bias)[0]
+
+    assert_close('  y', ref_y, tri_y, 1e-3)
+    assert_close(' dx', ref_dx, tri_dx, 1e-3)
+    assert_close(' dw', ref_dw, tri_dw, 1e-3)
+    assert_close(' db', ref_db, tri_db, 1e-3)
+    assert_close('dlw', ref_dlw, tri_dlw, 1e-3)
+    assert_close('dlb', ref_dlb, tri_dlb, 1e-3)
+
+
+@pytest.mark.parametrize("N", [1, 16, 128])
+@pytest.mark.parametrize("D", [50, 64, 128])
+def test_rmsnorm_linear(N: int, D: int):
+    torch.manual_seed(1)
+    x = torch.randn(N, D).to(device).requires_grad_(True)
+    ref = nn.Sequential(LlamaRMSNorm(D, eps=0), nn.Linear(D, D)).to(device)
+    tri = RMSNormLinear(D, eps=0).to(device)
+    nn.init.normal_(ref[0].weight)
+    nn.init.normal_(ref[1].weight, mean=0.0, std=0.01)
+    nn.init.normal_(ref[1].bias, mean=0.0, std=0.01)
+    tri.weight.data.copy_(ref[0].weight.data)
+    weight, bias = ref[1].weight.clone(), ref[1].bias.clone()
+
+    ref_y = ref(x)
+    tri_y = tri(x, weight, bias)
+    ref_dx = torch.autograd.grad(ref(x).sum(), x)[0]
+    tri_dx = torch.autograd.grad(tri(x, weight, bias).sum(), x)[0]
+    ref_dw = torch.autograd.grad(ref(x).sum(), ref[0].weight)[0]
+    tri_dw = torch.autograd.grad(tri(x, weight, bias).sum(), tri.weight)[0]
+    ref_dlw = torch.autograd.grad(ref(x).sum(), ref[1].weight)[0]
+    tri_dlw = torch.autograd.grad(tri(x, weight, bias).sum(), weight)[0]
+    ref_dlb = torch.autograd.grad(ref(x).sum(), ref[1].bias)[0]
+    tri_dlb = torch.autograd.grad(tri(x, weight, bias).sum(), bias)[0]
+
+    assert_close('  y', ref_y, tri_y, 1e-3)
+    assert_close(' dx', ref_dx, tri_dx, 1e-3)
+    assert_close(' dw', ref_dw, tri_dw, 1e-3)
+    assert_close('dlw', ref_dlw, tri_dlw, 1e-3)
+    assert_close('dlb', ref_dlb, tri_dlb, 1e-3)

code/flash-linear-attention/tests/modules/test_layernorm_gated.py

@@ -0,0 +1,92 @@
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fla.modules import FusedLayerNormGated, FusedRMSNormGated
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'H', 'T', 'D', 'elementwise_affine', 'activation', 'bias'),
+    [
+        pytest.param(*test, id=f"B{test[0]}_H{test[1]}_T{test[2]}_D{test[3]}_affine{test[4]}_{test[5]}_bias{test[6]}")
+        for test in [
+            (2, 2, 1,    64,  False, "silu",   False),
+            (2, 2, 512,  128, True,  "silu",   True),
+            (2, 2, 2048, 1200, True,  "sigmoid", False),
+            (2, 2, 50,   50,  False, "sigmoid", False),
+        ]
+    ],
+)
+def test_layernorm_gated(B: int, H: int, T: int, D: int, elementwise_affine: bool, activation: str, bias: bool):
+    torch.manual_seed(42)
+    x = torch.randn(B, H, T, D).to(device).requires_grad_(True)
+    g = torch.randn(B, H, T, D).to(device).requires_grad_(True)
+
+    ref = nn.LayerNorm(D, elementwise_affine=elementwise_affine, bias=bias).to(device)
+    tri = FusedLayerNormGated(D, elementwise_affine=elementwise_affine, bias=bias, activation=activation).to(device)
+    if ref.weight is not None:
+        nn.init.normal_(ref.weight)
+        tri.weight.data.copy_(ref.weight.data)
+    if ref.bias is not None:
+        nn.init.normal_(ref.bias)
+        tri.bias.data.copy_(ref.bias.data)
+
+    act_fn = F.silu if activation == "silu" else F.sigmoid
+    ref_y = ref(x) * act_fn(g)
+    tri_y = tri(x, g)
+    ref_dx, ref_dg = torch.autograd.grad((ref(x) * act_fn(g)).sum(), (x, g))
+    tri_dx, tri_dg = torch.autograd.grad(tri_y.sum(), (x, g))
+
+    if ref.weight is not None:
+        ref_dw = torch.autograd.grad((ref(x) * act_fn(g)).sum(), ref.weight)[0]
+        tri_dw = torch.autograd.grad(tri(x, g).sum(), tri.weight)[0]
+    if ref.bias is not None:
+        ref_db = torch.autograd.grad((ref(x) * act_fn(g)).sum(), ref.bias)[0]
+        tri_db = torch.autograd.grad(tri(x, g).sum(), tri.bias)[0]
+
+    assert_close(' y', ref_y, tri_y, 1e-3)
+    assert_close('dx', ref_dx, tri_dx, 1e-3)
+    assert_close('dg', ref_dg, tri_dg, 1e-3)
+    if ref.weight is not None:
+        assert_close('dw', ref_dw, tri_dw, 1e-3)
+    if ref.bias is not None:
+        assert_close('db', ref_db, tri_db, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'H', 'T', 'D', 'activation'),
+    [
+        pytest.param(*test, id=f"B{test[0]}_H{test[1]}_T{test[2]}_D{test[3]}_{test[4]}")
+        for test in [
+            (2, 2, 1,    64,  "silu"),
+            (2, 2, 512,  128, "sigmoid"),
+            (2, 2, 2048, 1200, "silu"),
+            (2, 2, 50,   50,  "sigmoid"),
+        ]
+    ],
+)
+def test_rmsnorm_gated(B: int, H: int, T: int, D: int, activation: str):
+    torch.manual_seed(42)
+    x = torch.randn(B, H, T, D).to(device).requires_grad_(True)
+    g = torch.randn(B, H, T, D).to(device).requires_grad_(True)
+    ref = nn.RMSNorm(D, eps=0).to(device)
+    tri = FusedRMSNormGated(D, eps=0, activation=activation).to(device)
+    nn.init.normal_(ref.weight)
+    tri.weight.data.copy_(ref.weight.data)
+
+    act_fn = F.silu if activation == "silu" else F.sigmoid
+    ref_y = ref(x) * act_fn(g)
+    tri_y = tri(x, g)
+    ref_dx, ref_dg = torch.autograd.grad((ref(x) * act_fn(g)).sum(), (x, g))
+    tri_dx, tri_dg = torch.autograd.grad(tri_y.sum(), (x, g))
+
+    ref_dw = torch.autograd.grad((ref(x) * act_fn(g)).sum(), ref.weight)[0]
+    tri_dw = torch.autograd.grad(tri(x, g).sum(), tri.weight)[0]
+
+    assert_close(' y', ref_y, tri_y, 1e-3)
+    assert_close('dx', ref_dx, tri_dx, 1e-3)
+    assert_close('dg', ref_dg, tri_dg, 1e-3)
+    assert_close('dw', ref_dw, tri_dw, 1e-3)

code/flash-linear-attention/tests/modules/test_rotary.py

@@ -0,0 +1,122 @@
+
+import pytest
+import torch
+
+from fla.modules.rotary import RotaryEmbedding, rotary_embedding_ref
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("T", [2048, 4096])
+@pytest.mark.parametrize("H", [4])
+@pytest.mark.parametrize("G", [1, 4])
+@pytest.mark.parametrize("D", [128, 256])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+def test_rotary(B: int, T: int, H: int, G: int, D: int, dtype: torch.dtype):
+    torch.manual_seed(42)
+    q = torch.randn(B, T, H, D).to(device).to(dtype=dtype).requires_grad_()
+    k = torch.randn(B, T, H//G, D).to(device).to(dtype=dtype).requires_grad_()
+    rotary = RotaryEmbedding(D).to(device)
+
+    tri_q, tri_k = rotary(q, k)
+    tri_dq = torch.autograd.grad(tri_q.sum(), q, retain_graph=True)[0]
+    tri_dk = torch.autograd.grad(tri_k.sum(), k, retain_graph=True)[0]
+
+    ref_q = rotary_embedding_ref(q.float(), rotary._cos_cached, rotary._sin_cached).to(dtype=dtype)
+    ref_k = rotary_embedding_ref(k.float(), rotary._cos_cached, rotary._sin_cached).to(dtype=dtype)
+    ref_dq = torch.autograd.grad(ref_q.sum(), q, retain_graph=True)[0]
+    ref_dk = torch.autograd.grad(ref_k.sum(), k, retain_graph=True)[0]
+
+    assert_close(" q", ref_q, tri_q, ratio=1e-5)
+    assert_close(" k", ref_k, tri_k, ratio=1e-5)
+    assert_close("dq", ref_dq, tri_dq, ratio=1e-5)
+    assert_close("dk", ref_dk, tri_dk, ratio=1e-5)
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("T", [2048, 4096])
+@pytest.mark.parametrize("H", [4])
+@pytest.mark.parametrize("G", [1, 4])
+@pytest.mark.parametrize("D", [128, 256])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_rotary_with_offsets(B: int, T: int, H: int, G: int, D: int, dtype: torch.dtype):
+    torch.manual_seed(42)
+    q = torch.randn(B, T, H, D).to(device).to(dtype=dtype).requires_grad_()
+    k = torch.randn(B, T, H//G, D).to(device).to(dtype=dtype).requires_grad_()
+    seqlen_offset = torch.randint(0, T//2, (B,)).to(device)
+    max_seqlen = T + seqlen_offset.max().item()
+    rotary = RotaryEmbedding(D).to(device)
+
+    tri_q, tri_k = rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen)
+    tri_dq = torch.autograd.grad(tri_q.sum(), q, retain_graph=True)[0]
+    tri_dk = torch.autograd.grad(tri_k.sum(), k, retain_graph=True)[0]
+
+    ref_q = torch.cat([
+        rotary_embedding_ref(
+            q[i:i+1].float(),
+            rotary._cos_cached[offset:offset+T],
+            rotary._sin_cached[offset:offset+T],
+        )
+        for i, offset in enumerate(seqlen_offset.tolist())
+    ]).to(dtype=dtype)
+    ref_k = torch.cat([
+        rotary_embedding_ref(
+            k[i:i+1].float(),
+            rotary._cos_cached[offset:offset+T],
+            rotary._sin_cached[offset:offset+T],
+        )
+        for i, offset in enumerate(seqlen_offset.tolist())
+    ]).to(dtype=dtype)
+    ref_dq = torch.autograd.grad(ref_q.sum(), q, retain_graph=True)[0]
+    ref_dk = torch.autograd.grad(ref_k.sum(), k, retain_graph=True)[0]
+
+    assert_close(" q", ref_q, tri_q, ratio=1e-5)
+    assert_close(" k", ref_k, tri_k, ratio=1e-5)
+    assert_close("dq", ref_dq, tri_dq, ratio=1e-5)
+    assert_close("dk", ref_dk, tri_dk, ratio=1e-5)
+
+
+@pytest.mark.parametrize("N", [4])
+@pytest.mark.parametrize("T", [2048, 4096])
+@pytest.mark.parametrize("H", [4])
+@pytest.mark.parametrize("G", [1, 4])
+@pytest.mark.parametrize("D", [128, 256])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_rotary_varlen(N: int, T: int, H: int, G: int, D: int, dtype: torch.dtype):
+    torch.manual_seed(42)
+    q = torch.randn(1, T, H, D).to(device).to(dtype=dtype).requires_grad_()
+    k = torch.randn(1, T, H//G, D).to(device).to(dtype=dtype).requires_grad_()
+    cu_seqlens = torch.cat([
+        torch.tensor([0], dtype=torch.long),
+        torch.arange(1, T)[torch.randperm(T - 1)[:N-1]],
+        torch.tensor([T], dtype=torch.long),
+    ], 0).to(device).sort()[0]
+    rotary = RotaryEmbedding(D).to(device)
+
+    tri_q, tri_k = rotary(q, k, cu_seqlens=cu_seqlens)
+    tri_dq = torch.autograd.grad(tri_q.sum(), q, retain_graph=True)[0]
+    tri_dk = torch.autograd.grad(tri_k.sum(), k, retain_graph=True)[0]
+
+    ref_q = torch.cat([
+        rotary_embedding_ref(
+            q[0, start:end].float(),
+            rotary._cos_cached[:end-start],
+            rotary._sin_cached[:end-start],
+        )
+        for start, end in zip(cu_seqlens.tolist(), cu_seqlens[1:].tolist(), strict=False)
+    ]).to(dtype=dtype).unsqueeze(0)
+    ref_k = torch.cat([
+        rotary_embedding_ref(
+            k[0, start:end].float(),
+            rotary._cos_cached[:end-start],
+            rotary._sin_cached[:end-start],
+        )
+        for start, end in zip(cu_seqlens.tolist(), cu_seqlens[1:].tolist(), strict=False)
+    ]).to(dtype=dtype).unsqueeze(0)
+    ref_dq = torch.autograd.grad(ref_q.sum(), q, retain_graph=True)[0]
+    ref_dk = torch.autograd.grad(ref_k.sum(), k, retain_graph=True)[0]
+
+    assert_close(" q", ref_q, tri_q, ratio=1e-5)
+    assert_close(" k", ref_k, tri_k, ratio=1e-5)
+    assert_close("dq", ref_dq, tri_dq, ratio=1e-5)
+    assert_close("dk", ref_dk, tri_dk, ratio=1e-5)

code/flash-linear-attention/tests/modules/test_token_shift.py

@@ -0,0 +1,133 @@
+
+import pytest
+import torch
+
+from fla.modules.token_shift import token_shift, token_shift_ref
+from fla.utils import assert_close, device
+
+test_b_list = [4]
+test_t_list = [512, 4100, 8192]
+test_h_list = [2560, 4096]
+test_cu_seqlens_list = [
+    None,
+    [0, 4, 7, 40, 128],
+    [0, 10, 20, 64],
+    [0, 32],
+    [0, 1, 3, 4],
+]
+test_dtype_list = [torch.float]
+
+
+@pytest.mark.parametrize('B', test_b_list)
+@pytest.mark.parametrize('T', test_t_list)
+@pytest.mark.parametrize('H', test_h_list)
+@pytest.mark.parametrize('cu_seqlens_val', test_cu_seqlens_list)
+@pytest.mark.parametrize('dtype', test_dtype_list)
+def test_token_shift(B, T, H, cu_seqlens_val, dtype):
+    if cu_seqlens_val is not None:
+        B = 1
+        T = cu_seqlens_val[-1]
+        cu_seqlens_tensor = torch.tensor(cu_seqlens_val, dtype=torch.int32, device=device)
+    else:
+        cu_seqlens_tensor = None
+
+    torch.manual_seed(42)
+
+    x = torch.randn(B, T, H, device=device).to(dtype).requires_grad_(True)
+    dy = torch.randn_like(x)
+
+    ref = token_shift_ref(x, cu_seqlens_tensor)
+    tri = token_shift(x, cu_seqlens_tensor)
+
+    ref.backward(dy)
+    ref_dx, x.grad = x.grad, None
+
+    tri.backward(dy)
+    tri_dx, x.grad = x.grad, None
+
+    assert_close(' x', ref, tri, 1e-3)
+    assert_close('dx', ref_dx, tri_dx, 1e-3)
+
+
+def _split_for_passing(
+    x: torch.Tensor,
+    cu_seqlens,
+    split_at: int = 1,
+):
+    assert x.size(0) == 1
+    assert 0 < split_at < len(cu_seqlens) - 1
+
+    cu0 = [t - cu_seqlens[0] for t in cu_seqlens[: split_at + 1]]
+    cu1 = [t - cu_seqlens[split_at] for t in cu_seqlens[split_at:]]
+    T0, T1 = cu0[-1], cu1[-1]
+
+    x0 = x[:, :T0].contiguous()
+    x1 = x[:, T0: T0 + T1].contiguous()
+    cache1 = x[:, T0 - 1: T0].contiguous()
+    return x0, x1, \
+        torch.tensor(cu0, dtype=torch.int32, device=x.device), \
+        torch.tensor(cu1, dtype=torch.int32, device=x.device), \
+        cache1
+
+
+def _check_passing_vs_whole(
+    B: int,
+    T: int,
+    H: int,
+    cu_seqlens: list[int] | None,
+    dtype: torch.dtype,
+    split_at: int = 1,
+):
+    torch.manual_seed(42)
+
+    if cu_seqlens is None:
+        x = torch.randn(B, T, H, device=device, dtype=dtype, requires_grad=True)
+        cu_seqlens_tensor = None
+    else:
+        B = 1
+        T = cu_seqlens[-1]
+        x = torch.randn(1, T, H, device=device, dtype=dtype, requires_grad=True)
+        cu_seqlens_tensor = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    dy = torch.randn_like(x)
+    ref_out = token_shift(x, cu_seqlens_tensor)
+    ref_out.backward(dy)
+    ref_dx = x.grad.clone()
+    x.grad.zero_()
+
+    if cu_seqlens is None:
+        T0 = T // 2
+        x0 = x[:, :T0].contiguous()
+        x1 = x[:, T0:].contiguous()
+        cu0, cu1 = None, None
+    else:
+        if split_at >= len(cu_seqlens) - 1:
+            pytest.skip("invalid split_at")
+        x0, x1, cu0, cu1, cache1 = _split_for_passing(x, cu_seqlens, split_at)
+
+    out0, cache_out0 = token_shift(x0, cu0, output_cache=True)
+    out1, cache_out1 = token_shift(x1, cu1, cache=cache_out0, output_cache=True)
+
+    cat_out = torch.cat([out0, out1], dim=1)
+    cat_out.backward(dy)
+
+    cat_dx = x.grad.clone()
+
+    assert_close("do", ref_out, cat_out, 1e-3)
+    assert_close("dx", ref_dx, cat_dx, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ("B", "T", "H", "cu_seqlens", "split_at"),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-cu{}-split{}".format(*test))
+        for test in [
+            (2, 512, 1024, None, 1),
+            (1, 8192, 1024, None, 2),
+        ]
+    ],
+)
+def test_all_with_and_without_varlen(B, T, H, cu_seqlens, split_at):
+    dtype = torch.float
+    assert cu_seqlens is None, "This test is for cu_seqlens=None case"
+    _check_passing_vs_whole(B, T, H, cu_seqlens, dtype, split_at)

code/flash-linear-attention/tests/ops/test_attn.py

@@ -0,0 +1,125 @@
+
+import os
+
+import pytest
+import torch
+
+from fla.ops.attn.parallel import parallel_attn
+from fla.ops.utils import prepare_lens
+from fla.utils import assert_close, check_shared_mem, device
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    HAS_FLASH = True
+except Exception:
+    HAS_FLASH = False
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'HQ', 'D', 'scale'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-HQ{}-D{}-scale{}".format(*test))
+        for test in [
+            (1, 63, 1, 1, 64, 1.0),
+            (3, 111, 2, 2, 100, 1.0),
+            (3, 1024, 2, 8, 60, 0.1),
+            (3, 1024, 2, 8, 128, 0.1),
+            (4, 2048, 2, 8, 64, 0.1),
+        ]
+    ],
+)
+def test_parallel(
+    B: int,
+    T: int,
+    H: int,
+    HQ: int,
+    D: int,
+    scale: float,
+):
+    if not check_shared_mem('hopper') and D > 128:
+        pytest.skip(reason="Skip test, do not have enough shard mem")
+    if not HAS_FLASH:
+        pytest.skip(reason="Skipping test because flash-attn is not installed")
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    q = torch.randn((B, T, HQ, D), dtype=torch.float16, device=device).requires_grad_(True)
+    k = torch.randn((B, T, H, D), dtype=torch.float16, device=device).requires_grad_(True)
+    v = torch.randn((B, T, H, D), dtype=torch.float16, device=device).requires_grad_(True)
+    do = torch.randn((B, T, HQ, D), dtype=torch.float16, device=device)
+
+    ref = flash_attn_func(q=q, k=k, v=v, softmax_scale=scale, causal=True)
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+
+    tri = parallel_attn(q=q, k=k, v=v, scale=scale)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    assert_close(" o", ref, tri, 0.005)
+    assert_close("dq", ref_dq, tri_dq, 0.005)
+    assert_close("dk", ref_dk, tri_dk, 0.005)
+    assert_close("dv", ref_dv, tri_dv, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'HQ', 'D', 'cu_seqlens'),
+    [
+        pytest.param(*test, id="H{}-HQ{}-D{}-cu_seqlens{}".format(*test))
+        for test in [
+            (2, 2, 64, [0, 15]),
+            (2, 8, 64, [0, 256, 500, 1000]),
+            (2, 2, 100, [0, 15, 100, 300, 1200, 2000]),
+        ]
+    ],
+)
+def test_parallel_varlen(
+    H: int,
+    HQ: int,
+    D: int,
+    cu_seqlens: list[int],
+):
+    if not HAS_FLASH:
+        pytest.skip(reason="Skipping test because flash-attn is not installed")
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+    dtype = torch.float16
+
+    q = torch.randn((1, T, HQ, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn((1, T, HQ, D), dtype=dtype, device=device)
+
+    ref = flash_attn_varlen_func(
+        q=q.squeeze(0),
+        k=k.squeeze(0),
+        v=v.squeeze(0),
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        max_seqlen_q=prepare_lens(cu_seqlens).max(),
+        max_seqlen_k=prepare_lens(cu_seqlens).max(),
+        causal=True,
+    )
+    ref.backward(do.squeeze(0))
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+
+    tri = parallel_attn(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens=cu_seqlens,
+    )
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    assert_close(" o", ref, tri, 0.004)
+    assert_close("dq", ref_dq.squeeze(), tri_dq.squeeze(), 0.005)
+    assert_close("dk", ref_dk.squeeze(), tri_dk.squeeze(), 0.005)
+    assert_close("dv", ref_dv.squeeze(), tri_dv.squeeze(), 0.005)

code/flash-linear-attention/tests/ops/test_based.py

@@ -0,0 +1,63 @@
+
+import pytest
+import torch
+
+from fla.ops.based import fused_chunk_based, parallel_based
+from fla.ops.based.naive import naive_parallel_based
+from fla.utils import device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 60, torch.float16),
+            (3, 111, 2, 64, torch.float16),
+            (3, 1024, 4, 100, torch.float16),
+            (3, 1024, 8, 128, torch.float16),
+            (4, 2048, 8, 256, torch.float16),
+        ]
+    ],
+)
+def test_based(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn((B, H, T, 16), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, H, T, 16), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, H, T, D), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    ref = naive_parallel_based(q, k, v, use_norm=True)
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+
+    tri = parallel_based(q, k, v, use_norm=True)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    if dtype == torch.float32:
+        assert ref.allclose(tri, 0, 1e-4)
+        assert ref_dq.allclose(tri_dq, 0, 1e-4)
+        assert ref_dk.allclose(tri_dk, 0, 1e-4)
+        assert ref_dv.allclose(tri_dv, 0, 1e-4)
+
+    tri = fused_chunk_based(q, k, v, use_norm=True)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    if dtype == torch.float32:
+        assert ref.allclose(tri, 0, 1e-4)
+        assert ref_dq.allclose(tri_dq, 0, 1e-4)
+        assert ref_dk.allclose(tri_dk, 0, 1e-4)
+        assert ref_dv.allclose(tri_dv, 0, 1e-4)

code/flash-linear-attention/tests/ops/test_comba.py

@@ -0,0 +1,368 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.ops.comba import chunk_comba, fused_recurrent_comba
+from fla.ops.comba.utils import chunk_comba_cumsum_scalar_fwd
+from fla.utils import assert_close, device, is_intel_alchemist
+
+
+def cumsum_comba_local_fwd_reference(s, reverse=False, chunk_size=128):
+    o_0 = torch.zeros_like(s)
+    o_1 = torch.zeros_like(s)
+    T = s.size(1)
+    fn = torch.cumsum
+    for i in range(0, T, chunk_size):
+        s_chunk = s[:, i:i+chunk_size]
+        o_1[:, i:i+chunk_size] = fn(s_chunk.float(), dim=1).to(o_1)
+        o_0[:, i:i+chunk_size] = o_1[:, i:i+chunk_size] - s_chunk
+
+    return o_0, o_1
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'chunk_size', 'dtype'),
+    [
+        pytest.param(*test, id='B{}-T{}-H{}-chunk_size{}-{}'.format(*test))
+        for test in [
+            (32, 200, 4, 64, torch.float),
+            (32, 1000, 4, 64, torch.float),
+            (32, 2048, 8, 128, torch.float),
+        ]
+    ],
+)
+def test_cumsum_local_scalar_fwd(
+    B: int,
+    T: int,
+    H: int,
+    chunk_size: int,
+    dtype: torch.dtype,
+):
+    s = torch.randn((B, T, H), dtype=dtype, device=device).requires_grad_()
+    ref_0, ref_1 = cumsum_comba_local_fwd_reference(s, chunk_size=chunk_size)
+    tri_0, tri_1 = chunk_comba_cumsum_scalar_fwd(s, chunk_size=chunk_size)
+    assert_close("local cumsum scalar", ref_0, tri_0, 0.001 if dtype == torch.float else 0.003)
+    assert_close("local cumsum scalar", ref_1, tri_1, 0.001 if dtype == torch.float else 0.003)
+
+
+def chunk_comba_ref(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    p: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    chunk_size: int = 64,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+):
+    BT = chunk_size
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+    # Calculate padding needed to make T a multiple of BT
+    q, k, v, p, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, p, beta, g])
+
+    T = q.shape[-2]
+    pad_len = (BT - (T % BT)) % BT
+    if pad_len > 0:
+        # Pad all tensors
+        q = F.pad(q, (0, 0, 0, pad_len))
+        k = F.pad(k, (0, 0, 0, pad_len))
+        v = F.pad(v, (0, 0, 0, pad_len))
+        p = F.pad(p, (0, 0, 0, pad_len))
+        beta = F.pad(beta, (0, pad_len))
+        g = F.pad(g, (0, pad_len))
+    decay = g
+    chunk_size = BT
+    b, h, l, d_k = q.shape
+    d_v = v.shape[-1]
+    q = q * scale
+    v = v * beta[..., None]
+    p_beta = p * beta[..., None]
+    assert l % chunk_size == 0
+    # note that diagonal is masked.
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)
+    q, k, v, p_beta, decay, g = map(
+        lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size),
+        [q, k, v, p_beta, decay.unsqueeze(-1), g.unsqueeze(-1)],
+    )
+    decay = decay.squeeze(-1).cumsum(-1)  # [B, H, n, c]
+    decay_0 = decay - g.squeeze(-1)  # [B, H, n, c]
+    L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
+    L_mask_0 = ((decay_0.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
+    # [B, H, n, c, d] @ [B, H, n, d, c] -> [B, H, n, c, c]
+    attn = -((p_beta @ k.transpose(-1, -2)) * L_mask_0).masked_fill(mask, 0)
+    for i in range(1, chunk_size):
+        attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
+    # for U
+    k_cumsum = attn @ v
+    # for W
+    k_cumdecay = attn @ (p_beta * decay_0[..., None].exp())
+    v = k_cumsum
+    S = k.new_zeros(b, h, d_k, d_v)
+    if initial_state is not None:
+        S += initial_state
+    o = torch.zeros_like(v)
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)
+    for i in range(0, l // chunk_size):
+        q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i]
+        attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0)
+        v_prime = k_cumdecay[:, :, i] @ S
+        v_new = v_i - v_prime
+        o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S
+        o[:, :, i] = o_inter + attn @ v_new
+        S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()
+                                                        [..., None]).transpose(-1, -2) @ v_new
+    if not output_final_state:
+        S = None
+    # unpad
+    o = rearrange(o, 'b h n c d -> b h (n c) d')
+    o = o[:, :, :T]
+    o = o.transpose(1, 2)
+    return o, S
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, 1, torch.float),
+            (2, 1024, 4, 60, 1, 1, torch.float),
+            (2, 1024, 8, 128, 1, 0.1, torch.float),
+            (2, 1024, 8, 128, 0.1, 1, torch.float),
+            (2, 1024, 8, 128, 1, 10, torch.float),
+            (4, 2048, 8, 64, 0.1, 1, torch.float),
+            (2, 1024, 8, 128, 1, 0.1, torch.float16),
+            (2, 1024, 8, 128, 1, 10, torch.float16),
+        ]
+    ],
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    p = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    beta = torch.rand(B, T, H, dtype=dtype).sigmoid()
+    g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32))
+    g = g / gate_logit_normalizer
+    h0 = torch.randn(B, H, D, D, dtype=torch.float32)
+    q, k, v, p, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, p, beta, g, h0))
+    ref, ref_ht = chunk_comba_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        p=p.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    tri, tri_ht = fused_recurrent_comba(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        p=p.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    assert_close('o', ref, tri, 0.002)
+    assert_close('ht', ref_ht, tri_ht, 0.002)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype'),
+    [
+        pytest.param(
+            *test,
+            id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}".format(*test),
+        )
+        for test in [
+            (1, 63, 1, 64, 1, 1, 0, False, torch.float16),
+            (2, 1000, 3, 60, 1, 1, 0, False, torch.float16),
+            (2, 1024, 3, 64, 0.1, 1, 0.5, False, torch.float16),
+            (2, 1024, 4, 100, 1, 0.1, 0, False, torch.float16),
+            (2, 1024, 4, 128, 0.1, 1, 0, True, torch.float16),
+            (2, 1024, 4, 128, 0.1, 1, 0.5, False, torch.float16),
+            (2, 1024, 4, 128, 0.1, 10, 0, False, torch.float16),
+            (4, 2048, 8, 64, 0.1, 1, 0, True, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    mask_p: float,
+    use_qk_l2norm_in_kernel: bool,
+    dtype: torch.dtype,
+):
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T, H, D, dtype=dtype)
+    p = torch.randn(B, T, H, D, dtype=dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    beta = torch.rand(B, T, H, dtype=dtype).sigmoid()
+    g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32))
+    g = g / gate_logit_normalizer
+    g = g * (torch.rand_like(g) > mask_p)
+    h0 = torch.zeros(B, H, D, D, dtype=torch.float32)
+    q, k, v, p, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, p, beta, g, h0))
+
+    tri, tri_ht = chunk_comba(
+        q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(),
+        k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(),
+        p=F.normalize(p.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else p.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dp, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, p.grad, beta.grad, g.grad, h0.grad
+    q.grad = k.grad = v.grad = p.grad = beta.grad = g.grad = h0.grad = None
+
+    ref, ref_ht = chunk_comba_ref(
+        q=F.normalize(q.clone(), p=2, dim=-1),
+        k=F.normalize(k.clone(), p=2, dim=-1),
+        p=F.normalize(p.clone(), p=2, dim=-1),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, ref_dk, ref_dv, ref_dp, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, p.grad, beta.grad, g.grad, h0.grad
+
+    assert_close("  o", ref, tri, 0.005)
+    assert_close(" ht", ref_ht, tri_ht, 0.005)
+    assert_close(" dq", ref_dq, tri_dq, 0.005)
+    assert_close(" dk", ref_dk, tri_dk, 0.008)
+    assert_close(" dv", ref_dv, tri_dv, 0.005)
+    assert_close(" dp", ref_dp, tri_dp, 0.008)
+    assert_close(" dg", ref_dg, tri_dg, 0.02)
+    assert_close(" db", ref_dbeta, tri_dbeta, 0.005)
+    assert_close("dh0", ref_dh0, tri_dh0, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, 0, [0, 15], torch.float16),
+            (4, 64, 0, [0, 256, 500, 1000], torch.float16),
+            (4, 64, 0.5, [0, 256, 500, 1000], torch.float16),
+            (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    mask_p: float,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.randn((1, T, H, D), dtype=dtype)
+    k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn((1, T, H, D), dtype=dtype)
+    p = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype))
+    g = g * (torch.rand_like(g) > mask_p)
+    beta = torch.rand(1, T, H, dtype=dtype).sigmoid()
+    h0 = torch.randn((N, H, D, D), dtype=dtype)
+
+    q, k, v, p, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, p, beta, g, h0))
+    do = torch.randn_like(v)
+    dht = torch.rand_like(h0)
+
+    tri, tri_ht = chunk_comba(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        p=p.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        output_final_state=True,
+        initial_state=h0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None
+
+    ref = []
+    ref_ht = []
+    for i in range(N):
+        ref_i, ref_ht_i = chunk_comba_ref(
+            q=q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k=k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v=v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            p=p[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            g=g[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            initial_state=h0[i],
+            output_final_state=True,
+        )
+        ref.append(ref_i)
+        ref_ht.append(ref_ht_i)
+    ref = torch.cat(ref, 1)
+    ref_ht = torch.cat(ref_ht, 0)
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.007)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.007)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.015)
+    assert_close('dg', ref_dg, tri_dg, 0.015)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.007)

code/flash-linear-attention/tests/ops/test_delta.py

@@ -0,0 +1,152 @@
+
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.delta_rule import chunk_delta_rule, fused_recurrent_delta_rule
+from fla.utils import assert_close, device, device_platform
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'use_qk_l2norm_in_kernel', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, False, torch.float16),
+            (2, 100, 4, 60, 0.1, False, torch.float16),
+            (2, 1000, 3, 128, 0.1, False, torch.float16),
+            (2, 1024, 4, 128, 1, True, torch.float16),
+            (3, 2000, 4, 128, 0.1, False, torch.float16),
+            (4, 2048, 8, 64, 0.1, False, torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    use_qk_l2norm_in_kernel: bool,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T, H, D, dtype=dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    beta = torch.randn(B, T, H, dtype=dtype).sigmoid()
+    h0 = torch.randn(B, H, D, D, dtype=torch.float32)
+    q, k, v, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, h0))
+    do = torch.rand_like(v)
+    dht = torch.rand_like(h0)
+
+    tri, tri_ht = chunk_delta_rule(
+        q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(),
+        k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = h0.grad = None
+
+    ref, ref_ht = fused_recurrent_delta_rule(
+        q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(),
+        k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad
+
+    assert_close('o', ref, tri, 0.006)
+    assert_close('ht', ref_ht, tri_ht, 0.006)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.008)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 64, [0, 15], torch.float16),
+            (3, 60, [0, 111, 500], torch.float16),
+            (3, 64, [0, 256, 500, 900, 1000], torch.float16),
+            (4, 100, [0, 15, 100, 300, 1200, 1599, 1800, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, D), dtype=dtype)
+    k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn((1, T, H, D), dtype=dtype)
+    beta = torch.randn(1, T, H, dtype=dtype).sigmoid()
+    h0 = torch.randn(N, H, D, D, dtype=dtype)
+    q, k, v, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, h0))
+    do = torch.randn_like(v)
+    dht = torch.rand_like(h0)
+
+    ref, ref_ht = fused_recurrent_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        output_final_state=True,
+        initial_state=h0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad
+
+    tri, tri_ht = chunk_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        output_final_state=True,
+        initial_state=h0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = h0.grad = None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.008)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.008)

code/flash-linear-attention/tests/ops/test_delta_product.py

@@ -0,0 +1,186 @@
+
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.gated_delta_product import chunk_gated_delta_product
+from fla.ops.gated_delta_product.chunk_ref import chunk_gated_delta_product_ref
+from fla.ops.gated_delta_product.naive import naive_recurrent_gated_delta_product
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'num_householder', 'use_qk_l2norm_in_kernel', 'dtype'),
+    [
+        pytest.param(
+            *test,
+            id="B{}-T{}-H{}-D{}-scale{}-num_householder{}-l2norm{}-{}".format(*test),
+        )
+        for test in [
+            (1, 63, 1, 64, 0.1, 1, False, torch.float16),
+            (2, 200, 3, 60, 0.1, 1, False, torch.float16),
+            (2, 1000, 4, 64, 0.1, 2, False, torch.float16),
+            (2, 1024, 4, 64, 1, 2, True, torch.float16),
+            (2, 1024, 6, 100, 1, 2, False, torch.float16),
+            (4, 1500, 8, 128, 0.1, 3, False, torch.float16),
+            (2, 2048, 8, 128, 1, 3, False, torch.float16),
+            (2, 2048, 8, 128, 1, 3, True, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    num_householder: int,
+    use_qk_l2norm_in_kernel: bool,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T * num_householder, H, D, dtype=dtype)
+    v = torch.randn(B, T * num_householder, H, D, dtype=dtype)
+    beta = torch.rand(B, T * num_householder, H, dtype=dtype).sigmoid()
+    h0 = torch.zeros(B, H, D, D, dtype=torch.float32)
+    q, k, v, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, h0))
+
+    tri, tri_ht = chunk_gated_delta_product(
+        q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(),
+        k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(),
+        v=v.clone(),
+        g=None,
+        beta=beta.clone(),
+        num_householder=num_householder,
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+    do = torch.randn_like(q)
+    dht = torch.randn_like(h0)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = h0.grad = None
+
+    ref, ref_ht = chunk_gated_delta_product_ref(
+        q=F.normalize(q.clone(), p=2, dim=-1),
+        k=F.normalize(k.clone(), p=2, dim=-1),
+        v=v.clone(),
+        g=None,
+        beta=beta.clone(),
+        num_householder=num_householder,
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.02)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'num_householder', 'cu_seqlens', 'dtype'),
+    [
+        (2, 64, 3, [0, 63 ], torch.float16),
+        (2, 100, 2, [0, 63, 100, 500, 1000], torch.float16),
+        (2, 128, 2, [0, 100, 300, 800, 1500, 2000], torch.float16),
+        (2, 256, 3, [0, 100, 123, 300, 500, 800, 1000, 1500, 2048], torch.float16),
+    ],
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    num_householder: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+    cu_seqlens = torch.LongTensor(cu_seqlens).to(device)
+    scale = 1.0
+
+    q = torch.nn.functional.normalize(torch.randn((1, T, H, D), dtype=dtype), dim=-1, p=2)
+    k = torch.nn.functional.normalize(torch.randn(1, T*num_householder, H, D, dtype=dtype), dim=-1, p=2)
+    v = torch.randn((1, T*num_householder, H, D), dtype=dtype)
+    beta = torch.rand(1, T*num_householder, H, dtype=dtype).sigmoid()
+    h0 = torch.randn((N, H, D, D), dtype=dtype)
+
+    q, k, v, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, h0))
+    do = torch.randn_like(q)
+    dht = torch.rand_like(h0)
+
+    tri, tri_ht = chunk_gated_delta_product(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=None,
+        scale=scale,
+        output_final_state=True,
+        num_householder=num_householder,
+        initial_state=h0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = h0.grad = None
+
+    ref, ref_ht = chunk_gated_delta_product_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=None,
+        scale=scale,
+        output_final_state=True,
+        num_householder=num_householder,
+        initial_state=h0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.007)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.007)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.015)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.007)
+    q.grad = k.grad = v.grad = beta.grad = h0.grad = None
+
+    torch_ref = torch.zeros_like(ref)
+    torch_ref_ht = torch.zeros_like(ref_ht)
+    for i in range(len(cu_seqlens) - 1):
+        start, end = cu_seqlens[i], cu_seqlens[i+1]
+        q_i = q[:, start:end, :, :]
+        k_i = k[:, start*num_householder:end*num_householder, :, :]
+        v_i = v[:, start*num_householder:end*num_householder, :, :]
+        beta_i = beta[:, start*num_householder:end*num_householder, :]
+        o3_i, h3_i = naive_recurrent_gated_delta_product(
+            q_i, k_i, v_i, None, beta_i, scale=scale, cu_seqlens=None, output_final_state=True, num_householder=num_householder,
+        )
+        torch_ref[:, start:end, :, :] = o3_i
+        torch_ref_ht[i, :, :, :] = h3_i.squeeze(0)
+
+    ((torch_ref * do).sum() + (torch_ref_ht * dht).sum()).backward(retain_graph=True)
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.007)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.007)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.015)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.007)

code/flash-linear-attention/tests/ops/test_deltaformer.py

@@ -0,0 +1,130 @@
+
+
+import pytest
+import torch
+
+from fla.ops.deltaformer import deltaformer_attn
+from fla.ops.deltaformer.naive import naive_deltaformer_attn
+from fla.utils import assert_close, device, is_intel_alchemist
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (2, 128, 2, 64, torch.float16),
+            (1, 256, 4, 64, torch.float16),
+            (2, 512, 4, 64, torch.float16),
+            (4, 1024, 4, 128, torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    is_intel_alchemist,
+    reason="Skipping test on Intel Alchemist due to known issues with SRAM.",
+)
+def test_deltaformer_attn(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    """
+    Test DeltaFormer pre-attention by comparing fused implementation with naive reference.
+    """
+    torch.manual_seed(42)
+
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    beta = torch.randn((B, T, H), dtype=dtype, device=device).sigmoid().requires_grad_(True)
+
+    do = torch.randn((B, T, H, D), dtype=dtype, device=device)
+
+    ref = naive_deltaformer_attn(q, k, v, beta)
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dbeta, beta.grad = beta.grad.clone(), None
+
+    tri = deltaformer_attn(q, k, v, beta)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dbeta, beta.grad = beta.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.006)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('dbeta', ref_dbeta, tri_dbeta, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 64, [0, 63], torch.float16),
+            (4, 64, [0, 256, 500, 1000], torch.float16),
+            (4, 128, [0, 15, 100, 300, 1200, 2000], torch.float16),
+            (2, 128, [0, 100, 123, 300, 500, 800, 1000, 1500, 2048], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    is_intel_alchemist,
+    reason="Skipping test on Intel Alchemist due to known issues with SRAM.",
+)
+def test_deltaformer_attn_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    beta = torch.randn((1, T, H), dtype=dtype, device=device).sigmoid().requires_grad_()
+
+    do = torch.randn_like(q)
+
+    refs = []
+    for i in range(N):
+        ref = naive_deltaformer_attn(
+            q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            beta[:, cu_seqlens[i]:cu_seqlens[i+1]],
+        )
+        refs.append(ref)
+    ref = torch.cat(refs, dim=1)
+
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dbeta, beta.grad = beta.grad.clone(), None
+
+    tri = deltaformer_attn(q, k, v, beta, cu_seqlens=cu_seqlens)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dbeta, beta.grad = beta.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.006)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('dbeta', ref_dbeta, tri_dbeta, 0.008)

code/flash-linear-attention/tests/ops/test_dplr_delta.py

@@ -0,0 +1,432 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.ops.generalized_delta_rule.dplr import chunk_dplr_delta_rule, fused_recurrent_dplr_delta_rule
+from fla.utils import assert_close, device, device_platform
+
+
+def recurrent_dplr_delta_rule_ref(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    gk: torch.Tensor,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+):
+    q, k, v, a, b, gk = map(lambda x: x.transpose(1, 2).to(torch.float), (q, k, v, a, b, gk))
+
+    B, H, T, K, V = *q.shape, v.shape[-1]
+    o = torch.zeros_like(v)
+    S = torch.zeros(B, H, K, V).to(v)
+    if initial_state is not None:
+        S = initial_state
+    if scale is None:
+        scale = K ** -0.5
+    q = q * scale
+
+    for i in range(T):
+        _q = q[:, :, i]
+        _k = k[:, :, i]
+        _v = v[:, :, i].clone()
+        a_i = a[:, :, i]
+        b_i = b[:, :, i]
+        # first matmul then decay in DPLR.
+        _v2 = (S.clone() * a_i[..., None]).sum(-2)
+        S = S.clone() * gk[:, :, i].exp()[..., None]
+        S = S.clone() + _k.unsqueeze(-1) * _v.unsqueeze(-2) + b_i.unsqueeze(-1) * _v2.unsqueeze(-2)
+        o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S)
+    if not output_final_state:
+        S = None
+    o = o.transpose(1, 2)
+    return o, S
+
+
+def chunk_dplr_delta_rule_ref(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    gk: torch.Tensor,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = True,
+    scale: float = None,
+    chunk_size: int = 64,
+):
+    q, k, v, a, b, gk = map(lambda x: x.transpose(1, 2).to(torch.float), (q, k, v, a, b, gk))
+    BT = chunk_size
+    T = q.shape[-2]
+    pad_len = (BT - (T % BT)) % BT
+
+    q, k, v, a, b, gk = map(lambda x: F.pad(x, (0, 0, 0, pad_len)).to(torch.float), [q, k, v, a, b, gk])
+    B, H, _, K, V = *q.shape, v.shape[-1]
+    NT = q.shape[-2] // BT
+    if scale is None:
+        scale = K ** -0.5
+    q = q * scale
+
+    S = k.new_zeros(B, H, K, V)
+    if initial_state is not None:
+        S += initial_state
+
+    # note that diagonal is masked.
+    mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
+    q, k, v, a, b, gk = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BT), [q, k, v, a, b, gk])
+    gk_cumsum = gk.cumsum(-2)
+    A_ab = torch.zeros(B, H, NT, BT, BT).to(q.device)
+    A_qk = torch.zeros(B, H, NT, BT, BT).to(q.device)
+    A_ak = torch.zeros(B, H, NT, BT, BT).to(q.device)
+    A_qb = torch.zeros(B, H, NT, BT, BT).to(q.device)
+
+    for i in range(BT):
+        a_i = a[:, :, :, i, None]
+        q_i = q[:, :, :, i, None]
+        gk_i = gk_cumsum[:, :, :, i, None]
+        mask = (torch.arange(BT) <= i).to(q.device)
+        attn_i = (gk_i - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp()
+        A_qk[:, :, :, i, :] = (q_i * k * attn_i).sum(-1).clone()
+        A_qb[:, :, :, i, :] = (q_i * b * attn_i).sum(-1).clone()
+        mask = (torch.arange(BT) < i).to(q.device)
+        # shift by one.
+        attn_i = (gk_i - gk[:, :, :, i, None] - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp()
+        A_ab[:, :, :, i, :] = (a_i * b * attn_i).sum(-1).clone()
+        A_ak[:, :, :, i, :] = (a_i * k * attn_i).sum(-1).clone()
+
+    A_ab = A_ab
+    for i in range(1, BT):
+        A_ab[..., i, :i] = A_ab[..., i, :i].clone() + (A_ab[..., i, :, None].clone() * A_ab[..., :, :i].clone()).sum(-2)
+
+    A_ab = A_ab + torch.eye(BT, dtype=torch.float, device=q.device)
+    u = A_ab @ (A_ak @ v)
+    w = A_ab @ ((gk_cumsum-gk).exp() * a)
+
+    o = torch.zeros_like(v)
+    mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
+    for i in range(0, NT):
+        q_i, k_i, v_i, u_i, w_i, b_i = q[:, :, i], k[:, :, i], v[:, :, i], u[:, :, i], w[:, :, i], b[:, :, i]
+        v2_i = u_i + w_i @ S
+        o_1 = A_qk[:, :, i] @ v_i
+        o_2 = A_qb[:, :, i] @ v2_i
+        o_3 = (q_i * gk_cumsum[:, :, i].exp()) @ S
+        o[:, :, i] = o_1 + o_2 + o_3
+        decay = (gk_cumsum[:, :, i, -1, None] - gk_cumsum[:, :, i]).exp()
+        S = S*gk_cumsum[:, :, i, -1, :, None].exp() + (k_i * decay).transpose(-1, -2) @ v_i + \
+            (b_i * decay).transpose(-1, -2) @ v2_i
+
+    S = None if output_final_state is False else S
+    o = rearrange(o, 'b h n c d -> b h (n c) d')
+    o = o[:, :, :T].transpose(1, 2)
+    return o, S
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float),
+            (2, 1024, 4, 60, 1, torch.float),
+            (2, 1024, 8, 128, 1, torch.float),
+            (2, 1024, 8, 128, 0.1, torch.float),
+            (4, 2048, 8, 64, 0.1, torch.float),
+            (2, 1024, 8, 128, 1, torch.float16),
+        ]
+    ],
+)
+def test_recurrent_fwd(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    os.environ['TORCH_CUDA_MATMUL_PRECISION'] = 'highest'
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T, H, D, dtype=dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    a = torch.rand(B, T, H, D, dtype=dtype)
+    gk = torch.randn(B, T, H, D, dtype=torch.float)
+
+    a = F.normalize(a, p=2, dim=-1)
+    b = -a
+    gk = F.logsigmoid(gk) / 16
+
+    h0 = torch.randn(B, H, D, D, dtype=torch.float)
+    q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(False), (q, k, v, a, b, gk, h0))
+    ref, ref_ht = chunk_dplr_delta_rule_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        gk=gk.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    tri, tri_ht = recurrent_dplr_delta_rule_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        gk=gk.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    assert_close('o', ref, tri, 0.001)
+    assert_close('ht', ref_ht, tri_ht, 0.001)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float),
+            (2, 1024, 4, 60, 1, torch.float),
+            (2, 1024, 8, 100, 1, torch.float),
+            (2, 1024, 8, 128, 0.1, torch.float),
+            (4, 2048, 8, 64, 0.1, torch.float),
+        ]
+    ],
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T, H, D, dtype=dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    a = torch.rand(B, T, H, D, dtype=dtype)
+    gk = torch.randn(B, T, H, D, dtype=torch.float)
+
+    a = F.normalize(a, p=2, dim=-1)
+    b = -a
+    gk = F.logsigmoid(gk) / 4
+
+    h0 = torch.randn(B, H, D, D, dtype=torch.float)
+    q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(False), (q, k, v, a, b, gk, h0))
+    ref, ref_ht = recurrent_dplr_delta_rule_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        gk=gk.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+
+    tri, tri_ht = fused_recurrent_dplr_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        gk=gk.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    assert_close('o', ref, tri, 0.002)
+    assert_close('ht', ref_ht, tri_ht, 0.002)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, 1, 0, torch.float16),
+            (2, 1000, 3, 60, 1, 1, 0, torch.float16),
+            (2, 1024, 3, 64, 0.1, 1, 0.5, torch.float16),
+            (2, 1024, 4, 100, 1, 0.1, 0, torch.float16),
+            (2, 1024, 4, 128, 0.1, 1, 0, torch.float16),
+            (2, 1024, 4, 128, 0.1, 1, 0.5, torch.float16),
+            (2, 1024, 4, 128, 0.1, 10, 0, torch.float16),
+            (4, 2048, 8, 64, 0.1, 1, 0, torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    mask_p: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T, H, D, dtype=dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    a = torch.rand(B, T, H, D, dtype=dtype)
+    gk = torch.randn(B, T, H, D, dtype=torch.float)
+
+    a = F.normalize(a, p=2, dim=-1)
+    b = -a
+    gk = F.logsigmoid(gk)
+    gk = gk / gate_logit_normalizer
+    gk = gk * (torch.rand_like(gk) > mask_p)
+
+    h0 = torch.randn(B, H, D, D, dtype=torch.float)
+    q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, a, b, gk, h0))
+    ref, ref_ht = chunk_dplr_delta_rule_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        gk=gk.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_da, ref_db, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad
+    q.grad = k.grad = v.grad = a.grad = b.grad = gk.grad = h0.grad = None
+
+    tri, tri_ht = chunk_dplr_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        gk=gk.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_da, tri_db, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad
+    q.grad = k.grad = v.grad = a.grad = b.grad = gk.grad = h0.grad = None
+
+    assert_close('o', ref, tri, 0.007)
+    assert_close('ht', ref_ht, tri_ht, 0.008)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('da', ref_da, tri_da, 0.008)
+    assert_close('db', ref_db, tri_db, 0.008)
+    if gate_logit_normalizer >= 1 and ref_dg.norm() > 0.01:  # otherwise it is meaningless
+        assert_close('dg', ref_dg, tri_dg, 0.008)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, 0, [0, 15], torch.float16),
+            (4, 64, 0, [0, 256, 500, 1000], torch.float16),
+            (4, 64, 0.5, [0, 256, 500, 1000], torch.float16),
+            (4, 100, 0, [0, 15, 100, 300, 1111, 1599, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    mask_p: float,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn(1, T, H, D, dtype=dtype)
+    k = torch.randn(1, T, H, D, dtype=dtype)
+    v = torch.randn(1, T, H, D, dtype=dtype)
+    a = torch.rand(1, T, H, D, dtype=dtype)
+    gk = torch.randn(1, T, H, D, dtype=torch.float)
+    a = F.normalize(a, p=2, dim=-1)
+    b = -a
+    gk = F.logsigmoid(gk)
+    gk = gk * (torch.rand_like(gk) > mask_p)
+    h0 = torch.randn(N, H, D, D, dtype=torch.float)
+    q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, a, b, gk, h0))
+
+    tri, tri_ht = chunk_dplr_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        gk=gk.clone(),
+        output_final_state=True,
+        initial_state=h0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_da, tri_db, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad
+    q.grad = k.grad = v.grad = a.grad = b.grad = gk.grad = h0.grad = None
+
+    ref = []
+    ref_ht = []
+    for i in range(N):
+        ref_i, ref_ht_i = chunk_dplr_delta_rule_ref(
+            q=q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k=k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v=v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            a=a[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            b=b[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            gk=gk[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            initial_state=h0[i, None],
+            output_final_state=True,
+        )
+        ref.append(ref_i)
+        ref_ht.append(ref_ht_i)
+
+    ref = torch.cat(ref, 1)
+    ref_ht = torch.cat(ref_ht, 0)
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_da, ref_db, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad
+
+    assert_close('o', ref, tri, 0.007)
+    assert_close('ht', ref_ht, tri_ht, 0.008)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('da', ref_da, tri_da, 0.008)
+    assert_close('db', ref_db, tri_db, 0.008)
+    assert_close('dg', ref_dg, tri_dg, 0.008)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.008)

code/flash-linear-attention/tests/ops/test_forgetting_attn.py

@@ -0,0 +1,151 @@
+
+
+import pytest
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from fla.ops.forgetting_attn.parallel import parallel_forgetting_attn
+from fla.utils import assert_close, check_shared_mem, device, is_intel_alchemist
+
+
+def naive_forgetting_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    scale: float | None = None,
+):
+    _, T, HQ, D = q.shape
+    H = k.shape[2]
+    G = HQ // H
+    if scale is None:
+        scale = D ** -0.5
+    gc = g.float().cumsum(1)
+    mask = torch.tril(torch.ones((T, T), dtype=torch.bool, device=device))
+    ref = torch.einsum("bqhd,bkhd->bhqk", q.float() * scale, repeat(k, "b t h d -> b t (h g) d", g=G).float())
+    ref = ref + rearrange(gc, "b t h -> b h t 1") - rearrange(gc, "b t h -> b h 1 t")
+    ref = ref.masked_fill(~mask.unsqueeze(0).unsqueeze(0), -float('inf'))
+    ref = torch.einsum("bhqk,bkhd->bqhd", F.softmax(ref, dim=-1), repeat(v, "b t h d -> b t (h g) d", g=G).float())
+    return ref
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'HQ', 'D', 'scale'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-HQ{}-D{}-scale{}".format(*test))
+        for test in [
+            (1, 63, 1, 1, 64, 1.0),
+            (3, 111, 2, 2, 100, 1.0),
+            (3, 1024, 2, 8, 60, 0.1),
+            (3, 1024, 2, 8, 128, 0.1),
+            (4, 2048, 2, 8, 64, 0.1),
+        ]
+    ],
+)
+def test_parallel(
+    B: int,
+    T: int,
+    H: int,
+    HQ: int,
+    D: int,
+    scale: float,
+):
+    torch.manual_seed(42)
+    dtype = torch.float16
+    if not check_shared_mem('hopper') and D > 128:
+        # maybe we can enable this test on Triton 3.3.0
+        pytest.skip("Skipping test because global shared memory is not available")
+
+    q = torch.randn((B, T, HQ, D), dtype=dtype, device=device).requires_grad_(True)
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+
+    g = torch.randn((B, T, HQ), dtype=dtype, device=device).uniform_(-0.1, -0.01).requires_grad_(True)
+
+    do = torch.randn((B, T, HQ, D), dtype=dtype, device=device)
+    ref = naive_forgetting_attn(q, k, v, g, scale)
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+
+    tri = parallel_forgetting_attn(q=q, k=k, v=v, g=g, scale=scale)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+
+    assert_close(" o", ref, tri, 0.005)
+    assert_close("dq", ref_dq, tri_dq, 0.005)
+    assert_close("dk", ref_dk, tri_dk, 0.005)
+    assert_close("dv", ref_dv, tri_dv, 0.005)
+    assert_close("dg", ref_dg, tri_dg, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'HQ', 'D', 'cu_seqlens'),
+    [
+        pytest.param(*test, id="H{}-HQ{}-D{}-cu_seqlens{}".format(*test))
+        for test in [
+            (2, 2, 64, [0, 15]),
+            (2, 8, 64, [0, 256, 500, 1000]),
+            (2, 2, 100, [0, 15, 100, 300, 1200, 2000]),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    is_intel_alchemist,
+    reason="Intel Triton Failure",
+)
+def test_parallel_varlen(
+    H: int,
+    HQ: int,
+    D: int,
+    cu_seqlens: list[int],
+):
+    torch.manual_seed(42)
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+    dtype = torch.float16
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, HQ, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = torch.rand((1, T, HQ), dtype=dtype, device=device).uniform_(-0.1, -0.01).requires_grad_(True)
+    do = torch.randn((1, T, HQ, D), dtype=dtype, device=device)
+
+    ref = q.new_empty(1, T, HQ, D)
+    for bos, eos in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False):
+        ref[:, bos:eos] = naive_forgetting_attn(
+            q=q[:, bos:eos],
+            k=k[:, bos:eos],
+            v=v[:, bos:eos],
+            g=g[:, bos:eos],
+        )
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+
+    tri = parallel_forgetting_attn(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        cu_seqlens=cu_seqlens,
+    )
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+
+    assert_close("  o", ref, tri, 0.004)
+    assert_close(" dq", ref_dq.squeeze(), tri_dq.squeeze(), 0.005)
+    assert_close(" dk", ref_dk.squeeze(), tri_dk.squeeze(), 0.005)
+    assert_close(" dv", ref_dv.squeeze(), tri_dv.squeeze(), 0.005)
+    assert_close(" dg", ref_dg.squeeze(), tri_dg.squeeze(), 0.005)

code/flash-linear-attention/tests/ops/test_gated_delta.py

@@ -0,0 +1,353 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
+from fla.utils import assert_close, device, is_intel_alchemist
+
+
+def recurrent_gated_delta_rule_ref(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    g: torch.Tensor,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+):
+    q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    o = torch.zeros(B, H, T, V).to(v)
+    h = torch.zeros(B, H, K, V).to(v)
+    if initial_state is not None:
+        h = initial_state
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+    q = q * scale
+    for i in range(T):
+        b_q = q[:, :, i]
+        b_k = k[:, :, i]
+        b_v = v[:, :, i].clone()
+        h = h.clone() * g[:, :, i].exp()[..., None, None]
+        b_beta = beta[:, :, i]
+        b_v = b_v - (h.clone() * b_k[..., None]).sum(-2)
+        b_v = b_v * b_beta[..., None]
+        h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2)
+        o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h)
+    if not output_final_state:
+        h = None
+    o = o.transpose(1, 2).contiguous()
+    return o, h
+
+
+def chunk_gated_delta_rule_ref(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    chunk_size: int = 64,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+):
+    BT = chunk_size
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+    # Calculate padding needed to make T a multiple of BT
+    q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
+
+    T = q.shape[-2]
+    pad_len = (BT - (T % BT)) % BT
+    if pad_len > 0:
+        # Pad all tensors
+        q = F.pad(q, (0, 0, 0, pad_len))
+        k = F.pad(k, (0, 0, 0, pad_len))
+        v = F.pad(v, (0, 0, 0, pad_len))
+        beta = F.pad(beta, (0, pad_len))
+        g = F.pad(g, (0, pad_len))
+    q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g])
+    decay = g
+    chunk_size = BT
+    b, h, l, d_k = q.shape
+    d_v = v.shape[-1]
+    q = q * scale
+    v = v * beta[..., None]
+    k_beta = k * beta[..., None]
+    assert l % chunk_size == 0
+    # note that diagonal is masked.
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)
+    q, k, v, k_beta, decay = map(
+        lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size),
+        [q, k, v, k_beta, decay.unsqueeze(-1)],
+    )
+    decay = decay.squeeze(-1).cumsum(-1)
+    decay_exp = decay.exp()[..., None]
+    L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
+    attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0)
+    for i in range(1, chunk_size):
+        attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
+    attn = attn
+    k_cumsum = attn @ v
+    k_cumdecay = attn @ (k_beta * decay_exp)
+    v = k_cumsum
+    S = k.new_zeros(b, h, d_k, d_v)
+    if initial_state is not None:
+        S = initial_state
+    o = torch.zeros_like(v)
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)
+    for i in range(0, l // chunk_size):
+        q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i]
+        attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0)
+        v_prime = (k_cumdecay[:, :, i]) @ S
+        v_new = v_i - v_prime
+        o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S
+        o[:, :, i] = o_inter + attn @ v_new
+        S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()
+                                                        [..., None]).transpose(-1, -2) @ v_new
+    if not output_final_state:
+        S = None
+    # unpad
+    o = rearrange(o, 'b h n c d -> b h (n c) d')
+    o = o[:, :, :T]
+    o = o.transpose(1, 2)
+    return o, S
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'HV', 'D', 'scale', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-HV{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 1, 64, 1, 1, torch.float),
+            (2, 500, 4, 4, 60, 1, 1, torch.float),
+            (2, 1000, 2, 8, 128, 1, 0.1, torch.float),
+            (3, 1024, 2, 2, 128, 0.1, 1, torch.float),
+            (4, 1024, 3, 3, 128, 1, 10, torch.float),
+            (4, 2048, 4, 4, 64, 0.1, 1, torch.float),
+            (2, 1024, 4, 4, 128, 1, 0.1, torch.float16),
+            (2, 1024, 4, 8, 128, 1, 10, torch.float16),
+        ]
+    ],
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    HV: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn(B, T, H, D, dtype=torch.float32)
+    k = torch.randn(B, T, H, D, dtype=torch.float32)
+    v = torch.randn(B, T, HV, D, dtype=dtype)
+    beta = torch.rand(B, T, HV, dtype=dtype).sigmoid()
+    g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.float32))
+    g = g / gate_logit_normalizer
+    h0 = torch.randn(B, HV, D, D, dtype=torch.float32)
+    q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, g, h0))
+    ref, ref_ht = recurrent_gated_delta_rule_ref(
+        q=F.normalize(repeat(q.clone(), 'b t h d -> b t (h g) d', g=HV // H), p=2, dim=-1).to(dtype),
+        k=F.normalize(repeat(k.clone(), 'b t h d -> b t (h g) d', g=HV // H), p=2, dim=-1).to(dtype),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    tri, tri_ht = fused_recurrent_gated_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        use_qk_l2norm_in_kernel=True,
+        output_final_state=True,
+    )
+    assert_close('o', ref, tri, 0.002)
+    assert_close('ht', ref_ht, tri_ht, 0.002)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype'),
+    [
+        pytest.param(
+            *test,
+            id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}".format(*test),
+        )
+        for test in [
+            (1, 63, 1, 64, 1, 1, 0, False, torch.float16),
+            (2, 500, 3, 60, 1, 1, 0, False, torch.float16),
+            (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16),
+            (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16),
+            (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16),
+            (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16),
+            (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16),
+            (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    mask_p: float,
+    use_qk_l2norm_in_kernel: bool,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+
+    q = torch.rand(B, T, H, D, dtype=dtype)
+    k = torch.rand(B, T, H, D, dtype=dtype)
+    v = torch.rand(B, T, H, D, dtype=dtype)
+    beta = torch.rand(B, T, H, dtype=dtype).sigmoid()
+    g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32))
+    g = g / gate_logit_normalizer
+    g = g * (torch.rand_like(g) > mask_p)
+    h0 = torch.zeros(B, H, D, D, dtype=torch.float32)
+    q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, h0))
+
+    tri, tri_ht = chunk_gated_delta_rule(
+        q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(),
+        k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None
+
+    ref, ref_ht = recurrent_gated_delta_rule_ref(
+        q=F.normalize(q.clone(), p=2, dim=-1),
+        k=F.normalize(k.clone(), p=2, dim=-1),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+    )
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.02)
+    assert_close('dg', ref_dg, tri_dg, 0.02)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 60, 0, [0, 15], torch.float16),
+            (4, 64, 0, [0, 256, 500, 1000], torch.float16),
+            (4, 64, 0.5, [0, 256, 500, 1000], torch.float16),
+            (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    mask_p: float,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    # randomly split the sequence into N segments
+    cu_seqlens = torch.LongTensor(cu_seqlens).to(device)
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, D), dtype=dtype)
+    k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn((1, T, H, D), dtype=dtype)
+    g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype))
+    g = g * (torch.rand_like(g) > mask_p)
+    beta = torch.rand(1, T, H, dtype=dtype).sigmoid()
+    h0 = torch.randn((N, H, D, D), dtype=dtype)
+
+    q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, g, h0))
+    do = torch.randn_like(v)
+    dht = torch.rand_like(h0)
+
+    tri, tri_ht = chunk_gated_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        initial_state=h0.clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None
+
+    ref = []
+    ref_ht = []
+    for i in range(N):
+        ref_i, ref_ht_i = recurrent_gated_delta_rule_ref(
+            q=q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k=k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v=v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            g=g[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            initial_state=h0[i],
+            output_final_state=True,
+        )
+        ref.append(ref_i)
+        ref_ht.append(ref_ht_i)
+    ref = torch.cat(ref, 1)
+    ref_ht = torch.cat(ref_ht, 0)
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.007)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.007)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.015)
+    assert_close('dg', ref_dg, tri_dg, 0.015)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.007)

code/flash-linear-attention/tests/ops/test_gated_delta_product.py

@@ -0,0 +1,210 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.gated_delta_product import chunk_gated_delta_product
+from fla.ops.gated_delta_product.chunk_ref import chunk_gated_delta_product_ref
+from fla.ops.gated_delta_product.naive import naive_recurrent_gated_delta_product
+from fla.utils import assert_close, device, is_intel_alchemist
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'num_householder', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype'),
+    [
+        pytest.param(
+            *test,
+            id="B{}-T{}-H{}-D{}-scale{}-num_householder{}-gate_logit_normalizer{}-mask_p{}-l2norm{}-{}".format(*test),
+        )
+        for test in [
+            (1, 63, 1, 64, 0.1, 1, 1, 0, False, torch.float16),
+            (2, 200, 3, 60, 0.1, 1, 1, 0, False, torch.float16),
+            (2, 1000, 4, 64, 0.1, 2, 0.1, 0.5, False, torch.float16),
+            (2, 1024, 4, 64, 1, 2, 1, 0, True, torch.float16),
+            (2, 1024, 6, 100, 1, 2, 10, 0, False, torch.float16),
+            (4, 1500, 8, 128, 0.1, 3, 1, 0.5, False, torch.float16),
+            (2, 2048, 8, 128, 1, 3, 1, 0, False, torch.float16),
+            (2, 2048, 8, 128, 1, 3, 1, 0, True, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    num_householder: int,
+    gate_logit_normalizer: float,
+    mask_p: float,
+    use_qk_l2norm_in_kernel: bool,
+    dtype: torch.dtype,
+):
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T * num_householder, H, D, dtype=dtype)
+    v = torch.randn(B, T * num_householder, H, D, dtype=dtype)
+    beta = torch.rand(B, T * num_householder, H, dtype=dtype).sigmoid()
+    g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32))
+    h0 = torch.zeros(B, H, D, D, dtype=torch.float32)
+    g = g / gate_logit_normalizer
+    g = g * (torch.rand_like(g) > mask_p)
+    q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, h0))
+
+    tri, tri_ht = chunk_gated_delta_product(
+        q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(),
+        k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        num_householder=num_householder,
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+    do = torch.randn_like(q)
+    dht = torch.randn_like(h0)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None
+
+    ref, ref_ht = chunk_gated_delta_product_ref(
+        q=F.normalize(q.clone(), p=2, dim=-1),
+        k=F.normalize(k.clone(), p=2, dim=-1),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        num_householder=num_householder,
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.02)
+    assert_close('dg', ref_dg, tri_dg, 0.02)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'num_householder', 'mask_p', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-num_householder{}-mask_p{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 64, 3, 0, [0, 63], torch.float16),
+            (2, 100, 2, 0, [0, 63, 100, 500, 1000], torch.float16),
+            (2, 100, 2, 0, [0, 100, 256, 512, 1500, 1500], torch.float16),
+            (2, 128, 2, 0, [0, 100, 300, 800, 1500, 2000], torch.float16),
+            (2, 128, 2, 0.5, [0, 31, 111, 799, 1000, 1500, 1800, 2000], torch.float16),
+            (2, 128, 2, 0.5, [0, 63, 300, 800, 1000, 1399, 2048], torch.float16),
+            (2, 256, 3, 0, [0, 100, 123, 300, 500, 800, 1000, 1500, 2048], torch.float16),
+        ]
+    ],
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    num_householder: int,
+    mask_p: float,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    cu_seqlens = torch.LongTensor(cu_seqlens).to(device)
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+
+    q = torch.nn.functional.normalize(torch.randn((1, T, H, D), dtype=dtype), dim=-1, p=2)
+    k = torch.nn.functional.normalize(torch.randn(1, T*num_householder, H, D, dtype=dtype), dim=-1, p=2)
+    v = torch.randn((1, T*num_householder, H, D), dtype=dtype)
+    g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype))
+    g = g * (torch.rand_like(g) > mask_p)
+    beta = torch.rand(1, T*num_householder, H, dtype=dtype).sigmoid()
+    h0 = torch.randn((N, H, D, D), dtype=dtype)
+
+    q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, g, h0))
+    do = torch.randn_like(q)
+    dht = torch.rand_like(h0)
+    scale = D ** -0.5
+
+    tri, tri_ht = chunk_gated_delta_product(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        scale=scale,
+        output_final_state=True,
+        num_householder=num_householder,
+        initial_state=h0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None
+
+    ref, ref_ht = chunk_gated_delta_product_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        scale=scale,
+        output_final_state=True,
+        num_householder=num_householder,
+        initial_state=h0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.007)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.007)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.015)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.007)
+    assert_close('dg', ref_dg, tri_dg, 0.015)
+    q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None
+
+    torch_ref = torch.zeros_like(ref)
+    torch_ref_ht = torch.zeros_like(ref_ht)
+    for i in range(len(cu_seqlens) - 1):
+        start, end = cu_seqlens[i], cu_seqlens[i+1]
+        q_i = q[:, start:end, :, :]
+        k_i = k[:, start*num_householder:end*num_householder, :, :]
+        v_i = v[:, start*num_householder:end*num_householder, :, :]
+        g_i = g[:, start:end, :]
+        beta_i = beta[:, start*num_householder:end*num_householder, :]
+        o3_i, h3_i = naive_recurrent_gated_delta_product(
+            q_i, k_i, v_i, g_i, beta_i, scale=scale, cu_seqlens=None, output_final_state=True, num_householder=num_householder,
+        )
+        torch_ref[:, start:end, :, :] = o3_i
+        torch_ref_ht[i, :, :, :] = h3_i.squeeze(0)
+
+    ((torch_ref * do).sum() + (torch_ref_ht * dht).sum()).backward(retain_graph=True)
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.007)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.007)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.015)
+    assert_close('dg', ref_dg, tri_dg, 0.015)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.007)

code/flash-linear-attention/tests/ops/test_gla.py

@@ -0,0 +1,325 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.gla import chunk_gla, fused_recurrent_gla
+from fla.ops.gla.naive import naive_recurrent_gla
+from fla.utils import assert_close, device, device_platform
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float),
+            (2, 1024, 4, 60, 1, torch.float),
+            (2, 1024, 8, 128, 0.1, torch.float),
+            (2, 1024, 8, 128, 1, torch.float),
+            (2, 1024, 8, 128, 10, torch.float),
+            (4, 2048, 8, 64, 1, torch.float),
+            (2, 1024, 8, 128, 0.1, torch.float16),
+            (2, 1024, 8, 128, 10, torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    q = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = (F.logsigmoid(torch.rand((B, T, H, D), dtype=dtype, device=device)) / gate_logit_normalizer).requires_grad_()
+    h0 = torch.rand(B, H, D, D, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.randn((B, H, D, D), dtype=dtype, device=device)
+
+    ref, ref_ht = naive_recurrent_gla(
+        q=q,
+        k=k,
+        v=v,
+        gk=g,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_recurrent_gla(
+        q=q,
+        k=k,
+        v=v,
+        gk=g,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, [0, 15], torch.float),
+            (4, 64, [0, 256, 500, 1000], torch.float),
+            (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 64, [0, 1, 100, 300, 1200, 2048], torch.float16),
+            (4, 128, [0, 200, 512, 1200, 2048], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_fused_recurrent_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = F.logsigmoid(torch.rand((1, T, H, D), dtype=dtype, device=device)).requires_grad_()
+    h0 = torch.rand(N, H, D, D, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.randn((N, H, D, D), dtype=dtype, device=device)
+
+    refs, ref_hts = [], []
+    for i in range(N):
+        ref, ref_ht = naive_recurrent_gla(
+            q=q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k=k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v=v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            gk=g[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            initial_state=h0[i],
+            output_final_state=True,
+        )
+        refs.append(ref)
+        ref_hts.append(ref_ht)
+    ref = torch.cat(refs, dim=1)
+    ref_ht = torch.cat(ref_hts, dim=0)
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_recurrent_gla(
+        q=q,
+        k=k,
+        v=v,
+        gk=g,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float16),
+            (2, 1024, 4, 60, 1, torch.float16),
+            (2, 1024, 8, 128, 0.1, torch.float16),
+            (2, 1024, 8, 128, 1, torch.float16),
+            (2, 1024, 8, 128, 10, torch.float16),
+            (4, 2048, 8, 64, 1, torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+    gate_logit_normalizer: float,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    # [B, T, H, D]
+    q = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = (F.logsigmoid(torch.rand((B, T, H, D), dtype=dtype, device=device)) / gate_logit_normalizer).requires_grad_()
+    h0 = torch.rand((B, H, D, D), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.randn((B, H, D, D), dtype=dtype, device=device)
+
+    tri, tri_ht = chunk_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    ref, ref_ht = fused_recurrent_gla(
+        q=q,
+        k=k,
+        v=v,
+        gk=g,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, [0, 15], torch.float16),
+            (4, 64, [0, 256, 500, 1000], torch.float16),
+            (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = F.logsigmoid(torch.rand((1, T, H, D), dtype=dtype, device=device)).requires_grad_()
+    h0 = torch.rand((N, H, D, D), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.rand((N, H, D, D), dtype=dtype, device=device)
+
+    ref, ref_ht = fused_recurrent_gla(
+        q=q,
+        k=k,
+        v=v,
+        gk=g,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = chunk_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)

code/flash-linear-attention/tests/ops/test_gsa.py

@@ -0,0 +1,432 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.gsa import chunk_gsa, fused_recurrent_gsa
+from fla.ops.gsa.naive import naive_recurrent_gsa
+from fla.utils import assert_close, check_shared_mem, device, device_platform
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'M', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-M{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 32, 1, torch.float),
+            (2, 1024, 4, 60, 64, 1, torch.float),
+            (2, 1024, 8, 128, 64, 0.1, torch.float),
+            (2, 1024, 8, 128, 32, 1, torch.float),
+            (2, 1024, 8, 128, 64, 1, torch.float),
+            (2, 1024, 8, 128, 64, 10, torch.float),
+            (4, 2048, 8, 64, 64, 1, torch.float),
+            (2, 1024, 8, 128, 64, 0.1, torch.float16),
+            (2, 1024, 8, 128, 64, 10, torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    M: int,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    s = torch.randn((B, T, H, M), dtype=dtype, device=device).requires_grad_()
+    g = (F.logsigmoid(torch.randn((B, T, H, M), dtype=dtype, device=device)) / gate_logit_normalizer).requires_grad_()
+    hk0 = torch.randn(B, H, D, M, device=device).requires_grad_()
+    hv0 = torch.randn(B, H, M, D, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dhkt = torch.randn_like(hk0)
+    dhvt = torch.randn_like(hv0)
+
+    ref, (ref_hkt, ref_hvt) = naive_recurrent_gsa(q, k, v, s, g, initial_state=(hk0, hv0), output_final_state=True)
+    ((ref * do).sum() + (ref_hkt * dhkt).sum() + (ref_hvt * dhvt).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_ds, s.grad = s.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dhk0, hk0.grad = hk0.grad.clone(), None
+    ref_dhv0, hv0.grad = hv0.grad.clone(), None
+
+    tri, (tri_hkt, tri_hvt) = fused_recurrent_gsa(
+        q=q,
+        k=k,
+        v=v,
+        s=s,
+        g=g,
+        initial_state=(hk0, hv0),
+        output_final_state=True,
+    )
+    ((tri * do).sum() + (tri_hkt * dhkt).sum() + (tri_hvt * dhvt).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_ds, s.grad = s.grad.clone(), None
+    tri_dg, s.grad = g.grad.clone(), None
+    tri_dhk0, hk0.grad = hk0.grad.clone(), None
+    tri_dhv0, hv0.grad = hv0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('hkt', ref_hkt, tri_hkt, 0.005)
+    assert_close('hvt', ref_hvt, tri_hvt, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('ds', ref_ds, tri_ds, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005)
+    assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'M', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-M{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, 64, [0, 15], torch.float),
+            (4, 64, 64, [0, 256, 500, 1000], torch.float),
+            (4, 100, 64, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 64, 64, [0, 1, 100, 300, 1200, 2048], torch.float16),
+            (4, 128, 64, [0, 200, 512, 1200, 2048], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_fused_recurrent_varlen(
+    H: int,
+    D: int,
+    M: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    s = torch.randn((1, T, H, M), dtype=dtype, device=device).requires_grad_()
+    g = F.logsigmoid(torch.randn((1, T, H, M), dtype=dtype, device=device)).requires_grad_()
+    hk0 = torch.randn(N, H, D, M, device=device).requires_grad_()
+    hv0 = torch.randn(N, H, M, D, device=device).requires_grad_()
+    dhkt = torch.randn(N, H, D, M, device=device).requires_grad_()
+    dhvt = torch.randn(N, H, M, D, device=device).requires_grad_()
+
+    do = torch.randn_like(v)
+    refs, ref_hkts, ref_hfts = [], [], []
+    for i in range(N):
+        ref, (ref_hkt, ref_hvt) = naive_recurrent_gsa(
+            q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            s[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            g[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            initial_state=(hk0[i:i+1], hv0[i:i+1]),
+            output_final_state=True,
+        )
+        refs.append(ref)
+        ref_hkts.append(ref_hkt)
+        ref_hfts.append(ref_hvt)
+    ref = torch.cat(refs, 1)
+    ref_hkt = torch.cat(ref_hkts, 0)
+    ref_hvt = torch.cat(ref_hfts, 0)
+    ((ref * do).sum() + (ref_hkt * dhkt).sum() + (ref_hvt * dhvt).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_ds, s.grad = s.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dhk0, hk0.grad = hk0.grad.clone(), None
+    ref_dhv0, hv0.grad = hv0.grad.clone(), None
+
+    tri, (tri_hkt, tri_hvt) = fused_recurrent_gsa(
+        q=q,
+        k=k,
+        v=v,
+        s=s,
+        g=g,
+        initial_state=(hk0, hv0),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_hkt * dhkt).sum() + (tri_hvt * dhvt).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_ds, s.grad = s.grad.clone(), None
+    tri_dg, s.grad = g.grad.clone(), None
+    tri_dhk0, hk0.grad = hk0.grad.clone(), None
+    tri_dhv0, hv0.grad = hv0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('hkt', ref_hkt, tri_hkt, 0.005)
+    assert_close('hvt', ref_hvt, tri_hvt, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('ds', ref_ds, tri_ds, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005)
+    assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'M', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-M{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 32, 1, torch.float16),
+            (2, 1024, 4, 60, 64, 1, torch.float16),
+            (2, 1024, 4, 256, 64, 1, torch.float16),
+            (2, 1024, 4, 128, 64, 0.1, torch.float),
+            (2, 1024, 4, 128, 128, 1, torch.float16),
+            (2, 1024, 4, 128, 64, 10, torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    M: int,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    if (D > 64 or M > 64) and check_shared_mem('hopper') is False:
+        pytest.skip(reason='Current CI do not support this config')
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    s = torch.randn((B, T, H, M), dtype=dtype, device=device).requires_grad_()
+    g = (F.logsigmoid(torch.randn((B, T, H, M), dtype=dtype, device=device)) / gate_logit_normalizer).requires_grad_()
+    hk0 = torch.randn(B, H, D, M, device=device).requires_grad_()
+    hv0 = torch.randn(B, H, M, D, device=device).requires_grad_()
+    dhkt = torch.randn(B, H, D, M, device=device).requires_grad_()
+    dhvt = torch.randn(B, H, M, D, device=device).requires_grad_()
+
+    do = torch.randn_like(v)
+    ref, (ref_hkt, ref_hvt) = fused_recurrent_gsa(
+        q=q,
+        k=k,
+        v=v,
+        s=s,
+        g=g,
+        scale=D**-0.5,
+        initial_state=(hk0, hv0),
+        output_final_state=True)
+    ((ref * do).sum() + (ref_hkt * dhkt).sum() + (ref_hvt * dhvt).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_ds, s.grad = s.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dhk0, hk0.grad = hk0.grad.clone(), None
+    ref_dhv0, hv0.grad = hv0.grad.clone(), None
+
+    tri, (tri_hkt, tri_hvt) = chunk_gsa(
+        q=q,
+        k=k,
+        v=v,
+        s=s,
+        g=g,
+        scale=D**-0.5,
+        initial_state=(hk0, hv0),
+        output_final_state=True,
+    )
+    ((tri * do).sum() + (tri_hkt * dhkt).sum() + (tri_hvt * dhvt).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_ds, s.grad = s.grad.clone(), None
+    tri_dg, s.grad = g.grad.clone(), None
+    tri_dhk0, hk0.grad = hk0.grad.clone(), None
+    tri_dhv0, hv0.grad = hv0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('hkt', ref_hkt, tri_hkt, 0.005)
+    assert_close('hvt', ref_hvt, tri_hvt, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('ds', ref_ds, tri_ds, 0.008)
+    assert_close('dg', ref_dg, tri_dg, 0.008)
+    assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005)
+    assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'M', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-M{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, 64, [0, 15], torch.float16),
+            (4, 64, 64, [0, 256, 500, 1000], torch.float16),
+            (4, 100, 64, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    M: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    if (D > 64 or M > 64) and check_shared_mem('hopper') is False:
+        pytest.skip(reason='Current CI do not support this config')
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    s = torch.randn((1, T, H, M), dtype=dtype, device=device).requires_grad_()
+    g = F.logsigmoid(torch.randn((1, T, H, M), dtype=dtype, device=device)).requires_grad_()
+    hk0 = torch.randn(N, H, D, M, device=device).requires_grad_()
+    hv0 = torch.randn(N, H, M, D, device=device).requires_grad_()
+    dhkt = torch.randn(N, H, D, M, device=device).requires_grad_()
+    dhvt = torch.randn(N, H, M, D, device=device).requires_grad_()
+
+    do = torch.randn_like(v)
+
+    ref, (ref_hkt, ref_hvt) = fused_recurrent_gsa(
+        q=q,
+        k=k,
+        v=v,
+        s=s,
+        g=g,
+        scale=D**-0.5,
+        initial_state=(hk0, hv0),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((ref * do).sum() + (ref_hkt * dhkt).sum() + (ref_hvt * dhvt).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_ds, s.grad = s.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dhk0, hk0.grad = hk0.grad.clone(), None
+    ref_dhv0, hv0.grad = hv0.grad.clone(), None
+
+    tri, (tri_hkt, tri_hvt) = chunk_gsa(
+        q=q,
+        k=k,
+        v=v,
+        s=s,
+        g=g,
+        scale=D**-0.5,
+        initial_state=(hk0, hv0),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_hkt * dhkt).sum() + (tri_hvt * dhvt).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_ds, s.grad = s.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dhk0, hk0.grad = hk0.grad.clone(), None
+    tri_dhv0, hv0.grad = hv0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('hkt', ref_hkt, tri_hkt, 0.005)
+    assert_close('hvt', ref_hvt, tri_hvt, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('ds', ref_ds, tri_ds, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005)
+    assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'HQ', 'H', 'D', 'M', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-HQ{}-H{}-D{}-M{}-{}".format(*test))
+        for test in [
+            (2, 63, 2, 1, 64, 32, torch.float),
+            (2, 200, 8, 2, 64, 64, torch.float),
+            (2, 256, 16, 4, 128, 64, torch.float),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_inference(
+    B: int,
+    T: int,
+    HQ: int,
+    H: int,
+    D: int,
+    M: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    q = torch.randn((B, T, HQ, D), dtype=dtype, device=device)
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device)
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device)
+    s = torch.randn((B, T, H, M), dtype=dtype, device=device)
+    g = F.logsigmoid(torch.randn((B, T, H, M), dtype=dtype, device=device))
+    h0 = (torch.randn(B, H, D, M, dtype=dtype, device=device),
+          torch.randn(B, H, M, D, dtype=dtype, device=device))
+
+    ref, _ = naive_recurrent_gsa(q, k, v, s, g, initial_state=h0)
+    tri = torch.empty_like(ref)
+    for i in range(T):
+        o, ht = fused_recurrent_gsa(
+            q[:, i:i+1],
+            k[:, i:i+1],
+            v[:, i:i+1],
+            s[:, i:i+1],
+            g[:, i:i+1],
+            initial_state=h0,
+            output_final_state=True,
+        )
+        tri[:, i] = o.squeeze(1)
+        assert_close(f'o{i}', ref[:, i], tri[:, i], 0.005)
+        h0 = ht

code/flash-linear-attention/tests/ops/test_hgrn.py

@@ -0,0 +1,164 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.hgrn import chunk_hgrn, fused_recurrent_hgrn
+from fla.ops.hgrn.naive import naive_recurrent_hgrn
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 500, torch.float),
+            (2, 1024, 500, torch.float),
+            (2, 1024, 512, torch.float),
+            (2, 1024, 1000, torch.float),
+            (4, 2048, 2048, torch.float),
+        ]
+    ],
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    x = torch.randn((B, T, D), dtype=dtype, device=device)
+    g = torch.randn((B, T, D), dtype=dtype, device=device)
+    h0 = torch.randn_like(x[:, 0])
+    x, g = (1 - g.sigmoid()) * x, F.logsigmoid(g)
+    x, g, h0 = (i.detach().clone().to(dtype).requires_grad_() for i in (x, g, h0))
+
+    do = torch.randn_like(x)
+    dht = torch.randn_like(h0)
+    ref, ref_ht = naive_recurrent_hgrn(x, g, h0, output_final_state=True)
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dx, x.grad = x.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_recurrent_hgrn(x, g, h0, output_final_state=True)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dx, x.grad = x.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dx', ref_dx, tri_dx, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (500, [0, 15], torch.float),
+            (512, [0, 256, 500, 1000], torch.float),
+            (1000, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (2048, [0, 200, 512, 1200, 2048], torch.float16),
+        ]
+    ],
+)
+def test_fused_recurrent_varlen(
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    x = torch.randn((1, T, D), dtype=dtype, device=device)
+    g = torch.randn((1, T, D), dtype=dtype, device=device)
+    h0 = torch.randn(N, D, dtype=dtype, device=device)
+    x, g = (1 - g.sigmoid()) * x, F.logsigmoid(g)
+    x, g, h0 = (i.detach().clone().to(dtype).requires_grad_() for i in (x, g, h0))
+
+    do = torch.randn_like(x)
+    dht = torch.randn_like(h0)
+    refs, ref_hts = [], []
+    for i in range(N):
+        ref, ref_ht = naive_recurrent_hgrn(
+            x[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            g[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            h0[i:i+1],
+            output_final_state=True,
+        )
+        refs.append(ref)
+        ref_hts.append(ref_ht)
+    ref = torch.cat(refs, 1)
+    ref_ht = torch.cat(ref_hts, 0)
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dx, x.grad = x.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_recurrent_hgrn(x, g, h0, output_final_state=True, cu_seqlens=cu_seqlens)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dx, x.grad = x.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dx', ref_dx, tri_dx, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 500, torch.float16),
+            (2, 500, 1000, torch.float16),
+            (2, 1000, 1024, torch.float16),
+            (4, 2048, 2048, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    x = torch.randn((B, T, D), dtype=dtype, device=device)
+    g = torch.randn((B, T, D), dtype=dtype, device=device)
+    x, g = (1 - g.sigmoid()) * x, F.logsigmoid(g)
+    x, g = (i.detach().clone().to(dtype).requires_grad_() for i in (x, g))
+
+    do = torch.randn_like(x)
+    h0 = torch.randn_like(x[:, 0])
+    ref, _ = fused_recurrent_hgrn(x, g, h0, output_final_state=True)
+    ref.backward(do)
+    ref_dx, x.grad = x.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+
+    tri, _ = chunk_hgrn(x, g, h0, output_final_state=True)
+    tri.backward(do)
+    tri_dx, x.grad = x.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('dx', ref_dx, tri_dx, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)

code/flash-linear-attention/tests/ops/test_iplr_delta.py

@@ -0,0 +1,235 @@
+
+
+import pytest
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.ops.generalized_delta_rule.iplr.chunk import chunk_iplr_delta_rule
+from fla.ops.generalized_delta_rule.iplr.fused_recurrent import fused_recurrent_iplr_delta_rule
+from fla.utils import assert_close, device
+
+
+def chunk_iplr_delta_rule_ref(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = True,
+    scale: float = None,
+    chunk_size: int = 64,
+):
+    BT = chunk_size
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+
+    q, k, v, a, b = map(lambda x: x.transpose(1, 2), (q, k, v, a, b))
+    T = q.shape[-2]
+    pad_len = (BT - (T % BT)) % BT
+    if pad_len > 0:
+        q = F.pad(q, (0, 0, 0, pad_len))
+        k = F.pad(k, (0, 0, 0, pad_len))
+        v = F.pad(v, (0, 0, 0, pad_len))
+        a = F.pad(a, (0, 0, 0, pad_len))
+        b = F.pad(b, (0, 0, 0, pad_len))
+    q, k, v, a, b = map(lambda x: x.to(torch.float32), [q, k, v, a, b])
+
+    B, H, L, DK = q.shape
+    DV = v.shape[-1]
+    q = q * scale
+
+    S = k.new_zeros(B, H, DK, DV)
+    if initial_state is not None:
+        S += initial_state
+
+    # note that diagonal is masked.
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)
+    q, k, v, a, b = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, a, b])
+
+    v2 = (a @ k.transpose(-1, -2)).masked_fill_(mask, 0) @ v
+    attn = (a @ b.transpose(-1, -2)).masked_fill(mask, 0)
+    for i in range(1, chunk_size):
+        attn[..., i, :i] = attn[..., i, :i] + (attn[..., i, :, None].clone() * attn[..., :, :i].clone()).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
+    u = attn @ v2
+    w = attn @ a
+    o = torch.zeros_like(v)
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)
+    for i in range(0, L // chunk_size):
+        current_chunk_size = min(chunk_size, L - i * chunk_size)  # to handle the last chunk with possibly padding
+        q_i = q[:, :, i, :current_chunk_size]
+        k_i = k[:, :, i, :current_chunk_size]
+        v_i = v[:, :, i, :current_chunk_size]
+        u_i = u[:, :, i, :current_chunk_size]
+        w_i = w[:, :, i, :current_chunk_size]
+        b_i = b[:, :, i, :current_chunk_size]
+        o_1 = (q_i @ k_i.transpose(-1, -2)).masked_fill_(mask, 0) @ v_i
+        v2_i = u_i + w_i @ S
+        o_2 = (q_i @ b_i.transpose(-1, -2)).masked_fill_(mask, 0) @ v2_i
+        o_3 = q_i @ S
+        o[:, :, i, :current_chunk_size] = o_1 + o_2 + o_3
+        S = S + k_i.transpose(-1, -2) @ v_i + b_i.transpose(-1, -2) @ v2_i
+    S = None if output_final_state is False else S
+    o = rearrange(o, 'b h n c d -> b h (n c) d')
+    o = o[:, :, :T]
+    o = o.transpose(1, 2)
+    return o, S
+
+
+def recurrence_iplr_delta_rule_ref(
+    q,
+    k,
+    v,
+    a,
+    b,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = True,
+    scale: float | None = None,
+):
+    orig_dtype = q.dtype
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+    q, k, v, a, b = map(lambda x: x.transpose(1, 2).to(torch.float32), [q, k, v, a, b])
+    q = q * scale
+    B, H, L, DK = q.shape
+    DV = v.shape[-1]
+    o = torch.zeros_like(v)
+    S = torch.zeros(B, H, DK, DV).to(v)
+    if initial_state is not None:
+        S += initial_state
+
+    for i in range(q.shape[-2]):
+        _k = k[:, :, i]
+        _q = q[:, :, i]
+        _v = v[:, :, i]
+        _a = a[:, :, i]
+        _b = b[:, :, i]
+        _kv = _k[..., None] * _v[..., None, :] + (S.clone() * _a[..., None]).sum(-2, keepdim=True) * _b[..., None]
+        S = S + _kv
+        o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S)
+    S = None if output_final_state is False else S
+    o = o.transpose(1, 2)
+    return o.to(orig_dtype), S
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float),
+            (2, 1024, 4, 60, 1, torch.float),
+            (2, 1024, 8, 100, 1, torch.float),
+            (2, 1024, 8, 128, 0.1, torch.float),
+            (4, 2048, 8, 64, 0.1, torch.float),
+        ]
+    ],
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    dtype: torch.dtype,
+):
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T, H, D, dtype=dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    a = torch.rand(B, T, H, D, dtype=dtype)
+
+    a = F.normalize(a, p=2, dim=-1)
+    b = -a
+    h0 = torch.zeros(B, H, D, D, dtype=torch.float32)
+    q, k, v, a, b, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, a, b, h0))
+    ref, ref_ht = recurrence_iplr_delta_rule_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    dht = torch.rand_like(h0)
+    do = torch.rand_like(ref)
+    ((dht * ref_ht).sum() + (do * ref).sum()).backward()
+    dq, dk, dv, da, db, dh0 = map(lambda x: x.grad, (q, k, v, a, b, h0))
+    q.grad, k.grad, v.grad, a.grad, b.grad, h0.grad = None, None, None, None, None, None
+    tri, tri_ht = fused_recurrent_iplr_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    ((dht * tri_ht).sum() + (do * tri).sum()).backward()
+    assert_close('o', ref, tri, 0.003)
+    assert_close('ht', ref_ht, tri_ht, 0.003)
+    assert_close('dq', dq, q.grad, 0.003)
+    assert_close('dk', dk, k.grad, 0.003)
+    assert_close('dv', dv, v.grad, 0.003)
+    assert_close('da', da, a.grad, 0.003)
+    assert_close('db', db, b.grad, 0.003)
+    assert_close('dh0', dh0, h0.grad, 0.003)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float16),
+            (2, 500, 3, 60, 1, torch.float16),
+            (2, 1000, 3, 64, 0.1, torch.float16),
+            (2, 1024, 4, 100, 1, torch.float16),
+            (3, 1024, 4, 128, 0.1, torch.float16),
+            (4, 2048, 8, 64, 0.1, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    dtype: torch.dtype,
+):
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = torch.randn(B, T, H, D, dtype=dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    a = torch.rand(B, T, H, D, dtype=dtype)
+
+    a = F.normalize(a, p=2, dim=-1)
+    b = -a
+    h0 = torch.zeros(B, H, D, D, dtype=torch.float32)
+    q, k, v, a, b, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, a, b, h0))
+    ref, ref_ht = recurrence_iplr_delta_rule_ref(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    tri, tri_ht = chunk_iplr_delta_rule(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    assert_close('o', ref, tri, 0.007)
+    assert_close('ht', ref_ht, tri_ht, 0.008)

code/flash-linear-attention/tests/ops/test_kda.py

@@ -0,0 +1,379 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.kda import chunk_kda, fused_recurrent_kda
+from fla.ops.kda.gate import fused_kda_gate, kda_gate_ref
+from fla.ops.kda.naive import naive_chunk_kda, naive_recurrent_kda
+from fla.utils import assert_close, device, is_intel_alchemist
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(
+            *test,
+            id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test),
+        )
+        for test in [
+            (1, 64, 1, 64, 1, 1, torch.float),
+            (2, 512, 3, 60, 1, 1, torch.float),
+            (4, 1024, 4, 128, 0.1, 1, torch.float),
+            (4, 1024, 4, 128, 1, 10, torch.float),
+        ]
+    ],
+)
+def test_naive_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+
+    q = torch.rand(B, T, H, D, dtype=dtype)
+    k = torch.rand(B, T, H, D, dtype=dtype)
+    v = torch.rand(B, T, H, D, dtype=dtype)
+    g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer
+    beta = torch.randn(B, T, H, dtype=dtype).sigmoid()
+    h0 = torch.randn(B, H, D, D, dtype=torch.float32)
+    q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0))
+
+    ref, ref_ht = naive_recurrent_kda(
+        q=F.normalize(q.clone(), p=2, dim=-1),
+        k=F.normalize(k.clone(), p=2, dim=-1),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+
+    tri, tri_ht = naive_chunk_kda(
+        q=F.normalize(q.clone(), p=2, dim=-1),
+        k=F.normalize(k.clone(), p=2, dim=-1),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'use_qk_l2norm_in_kernel', 'dtype'),
+    [
+        pytest.param(
+            *test,
+            id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-use_qk_l2norm_in_kernel{}-{}".format(*test),
+        )
+        for test in [
+            (1, 64, 1, 64, 1, 1, False, torch.float),
+            (2, 512, 3, 60, 1, 1, False, torch.float),
+            (3, 1000, 4, 100, 0.1, 1, True, torch.float),
+            (4, 1024, 4, 128, 0.1, 1, False, torch.float),
+        ]
+    ],
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    use_qk_l2norm_in_kernel: bool,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+
+    q = torch.rand(B, T, H, D, dtype=dtype)
+    k = torch.rand(B, T, H, D, dtype=dtype)
+    v = torch.rand(B, T, H, D, dtype=dtype)
+    g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer
+    beta = torch.randn(B, T, H, dtype=dtype).sigmoid()
+    h0 = torch.randn(B, H, D, D, dtype=torch.float32)
+    q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0))
+
+    ref, ref_ht = naive_recurrent_kda(
+        q=F.normalize(q.clone(), p=2, dim=-1),
+        k=F.normalize(k.clone(), p=2, dim=-1),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+
+    tri, tri_ht = fused_recurrent_kda(
+        q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(),
+        k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype', 'tma'),
+    [
+        pytest.param(
+            *test,
+            id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}-tma{}".format(*test),
+        )
+        for test in [
+            (1, 63, 1, 64, 1, 1, 0, False, torch.float16, True),
+            (2, 500, 3, 60, 1, 1, 0, False, torch.float16, True),
+            (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16, False),
+            (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16, False),
+            (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16, True),
+            (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16, True),
+            (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16, False),
+            (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16, True),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    mask_p: float,
+    use_qk_l2norm_in_kernel: bool,
+    dtype: torch.dtype,
+    tma: bool,
+):
+    torch.manual_seed(42)
+    if not tma:
+        os.environ['FLA_USE_TMA'] = '0'
+    else:
+        os.environ['FLA_USE_TMA'] = '1'
+    q = torch.rand(B, T, H, D, dtype=dtype)
+    k = torch.rand(B, T, H, D, dtype=dtype)
+    v = torch.rand(B, T, H, D, dtype=dtype)
+    g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer
+    g = g * (torch.rand_like(g) > mask_p)
+    beta = torch.randn(B, T, H, dtype=dtype).sigmoid()
+    h0 = torch.randn(B, H, D, D, dtype=torch.float32)
+    q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0))
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+
+    ref, ref_ht = naive_recurrent_kda(
+        q=F.normalize(q.clone(), p=2, dim=-1),
+        k=F.normalize(k.clone(), p=2, dim=-1),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad
+    q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None
+
+    tri, tri_ht = chunk_kda(
+        q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(),
+        k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad
+    q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('dg', ref_dg, tri_dg, 0.02)
+    assert_close('db', ref_db, tri_db, 0.02)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 60, 0, [0, 15], torch.float16),
+            (4, 64, 0, [0, 256, 500, 1000], torch.float16),
+            (4, 128, 0.5, [0, 256, 500, 1000], torch.float16),
+            (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16),
+            (4, 256, 0, [0, 15, 100, 300, 1200, 4096], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    mask_p: float,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    # randomly split the sequence into N segments
+    cu_seqlens = torch.LongTensor(cu_seqlens).to(device)
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, D), dtype=dtype)
+    k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn((1, T, H, D), dtype=dtype)
+    g = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float))
+    g = g * (torch.rand_like(g) > mask_p)
+    beta = torch.rand(1, T, H, dtype=dtype).sigmoid()
+    h0 = torch.randn((N, H, D, D), dtype=dtype)
+
+    q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, g, beta, h0))
+    do = torch.randn_like(v)
+    dht = torch.rand_like(h0)
+
+    tri, tri_ht = chunk_kda(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        initial_state=h0.clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad
+    q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None
+
+    ref = []
+    ref_ht = []
+    for i in range(N):
+        ref_i, ref_ht_i = naive_recurrent_kda(
+            q=q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k=k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v=v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            g=g[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            initial_state=h0[i],
+            output_final_state=True,
+        )
+        ref.append(ref_i)
+        ref_ht.append(ref_ht_i)
+    ref = torch.cat(ref, 1)
+    ref_ht = torch.cat(ref_ht, 0)
+
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.007)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.007)
+    assert_close('dg', ref_dg, tri_dg, 0.015)
+    assert_close('db', ref_db, tri_db, 0.015)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.007)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'use_bias'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-bias{}".format(*test))
+        for test in [
+            (1, 2, 2, 12, False),
+            (1, 32, 2, 16, False),
+            (2, 64, 4, 32, False),
+            (4, 128, 8, 64, False),
+            (4, 128, 8, 128, False),
+            # Add bias tests
+            (1, 2, 2, 12, True),
+            (1, 32, 2, 16, True),
+            (2, 64, 4, 32, True),
+            (4, 128, 8, 64, True),
+            (4, 128, 8, 128, True),
+        ]
+    ],
+)
+def test_kda_gate(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    use_bias: bool,
+):
+    """Test kda gate forward and backward pass - reference vs Triton implementation"""
+    torch.manual_seed(42)
+
+    g = torch.randn(B, T, H * D, dtype=torch.float32)
+    # Ensure some values are > 20 to test the threshold logic in softplus
+    g = g * 30  # Scale up to get values > 20
+    A = torch.log(torch.randn(1, 1, H, 1, dtype=torch.float32).uniform_(1, 16))
+    g_bias = torch.randn(H * D, dtype=torch.float32) if use_bias else None
+
+    # Move to device and set requires_grad
+    g, A = map(lambda x: x.to(device).requires_grad_(True), (g, A))
+    if g_bias is not None:
+        g_bias = g_bias.to(device).requires_grad_(True)
+
+    # Create gradient output
+    do = torch.randn_like(g).view(B, T, H, D)
+
+    # Reference implementation
+    ref = kda_gate_ref(g.clone(), A.clone(), D, g_bias.clone() if g_bias is not None else None)
+    # Triton implementation
+    tri = fused_kda_gate(g.clone(), A.clone(), D, g_bias.clone() if g_bias is not None else None)
+
+    # Backward pass
+    ((ref * do).sum()).backward(retain_graph=True)
+    ref_dg, ref_dA = g.grad, A.grad
+    ref_dgbias = g_bias.grad if g_bias is not None else None
+    g.grad = A.grad = None
+    if g_bias is not None:
+        g_bias.grad = None
+
+    ((tri * do).sum()).backward(retain_graph=True)
+    tri_dg, tri_dA = g.grad, A.grad
+    tri_dgbias = g_bias.grad if g_bias is not None else None
+    g.grad = A.grad = None
+    if g_bias is not None:
+        g_bias.grad = None
+
+    assert_close('o', ref, tri, 1e-4)
+    assert_close('dg', ref_dg, tri_dg, 1e-4)
+    assert_close('dA', ref_dA, tri_dA, 1e-4)
+    if use_bias:
+        assert_close('dgbias', ref_dgbias, tri_dgbias, 1e-4)

code/flash-linear-attention/tests/ops/test_linear_attn.py

@@ -0,0 +1,189 @@
+
+
+import pytest
+import torch
+
+from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn
+from fla.ops.linear_attn.naive import naive_recurrent_linear_attn
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 64, 1, 64, None, torch.float),
+            (2, 512, 4, 60, None, torch.float),
+            (3, 1024, 8, 128, 1., torch.float),
+            (3, 1024, 8, 128, 0.1, torch.float),
+            (3, 1024, 8, 128, None, torch.float),
+            (2, 2048, 8, 256, None, torch.float16),
+            (2, 2048, 4, 256, None, torch.float16),
+        ]
+    ],
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float | None,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    h0 = torch.randn((B, H, D, D), dtype=torch.float, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+
+    ref, ref_ht = naive_recurrent_linear_attn(q, k, v, scale=scale, initial_state=h0, output_final_state=True, normalize=False)
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_recurrent_linear_attn(q, k, v, scale=scale, initial_state=h0, output_final_state=True, normalize=False)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.001)
+    assert_close('ht', ref_ht, tri_ht, 0.001)
+    assert_close('dq', ref_dq, tri_dq, 0.001)
+    assert_close('dk', ref_dk, tri_dk, 0.001)
+    assert_close('dv', ref_dv, tri_dv, 0.001)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.001)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, torch.float16),
+            (2, 500, 3, 60, torch.float16),
+            (2, 1000, 3, 128, torch.float16),
+            (3, 1000, 4, 64, torch.float16),
+            (2, 2048, 4, 256, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    h0 = torch.randn((B, H, D, D), dtype=torch.float, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+
+    ref, ref_ht = fused_recurrent_linear_attn(
+        q.to(torch.float32),
+        k.to(torch.float32),
+        v.to(torch.float32),
+        initial_state=h0,
+        output_final_state=True,
+        normalize=False,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = chunk_linear_attn(
+        q=q,
+        k=k,
+        v=v,
+        initial_state=h0,
+        output_final_state=True,
+        normalize=False,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.001)
+    assert_close('ht', ref_ht, tri_ht, 0.001)
+    assert_close('dq', ref_dq, tri_dq, 0.001)
+    assert_close('dk', ref_dk, tri_dk, 0.001)
+    assert_close('dv', ref_dv, tri_dv, 0.001)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.001)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, torch.float16),
+            (2, 500, 3, 60, torch.float16),
+            (2, 1000, 3, 128, torch.float16),
+            (3, 1000, 4, 64, torch.float16),
+            (2, 2048, 4, 256, torch.float16),
+        ]
+    ],
+)
+def test_fused_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    h0 = torch.randn((B, H, D, D), dtype=torch.float, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+
+    ref, ref_ht = fused_recurrent_linear_attn(
+        q.to(torch.float32),
+        k.to(torch.float32),
+        v.to(torch.float32),
+        initial_state=h0,
+        output_final_state=True,
+        normalize=False,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_chunk_linear_attn(
+        q=q,
+        k=k,
+        v=v,
+        initial_state=h0,
+        output_final_state=True,
+        normalize=False,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.001)
+    assert_close('ht', ref_ht, tri_ht, 0.001)
+    assert_close('dq', ref_dq, tri_dq, 0.001)
+    assert_close('dk', ref_dk, tri_dk, 0.001)
+    assert_close('dv', ref_dv, tri_dv, 0.001)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.001)

code/flash-linear-attention/tests/ops/test_log_linear_attn.py

@@ -0,0 +1,148 @@
+import os
+
+import numpy as np
+import pytest
+import torch
+
+from fla.ops.log_linear_attn import chunk_log_linear_attn
+from fla.ops.log_linear_attn.naive import naive_log_linear_attn
+from fla.utils import assert_close, device, device_platform
+
+
+@pytest.mark.parametrize(
+    ("B", "T", "H", "D", "dtype"),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [(2, 1024, 8, 128, torch.float32), (4, 2048, 8, 64, torch.float32)]
+    ],
+)
+@pytest.mark.skipif(device_platform == "intel", reason="Intel Triton Failure")
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ["TRITON_F32_DEFAULT"] = "ieee"
+
+    L = int(np.log2(T) + 1)
+    x = torch.randn(B, T, H, D, dtype=dtype, device=device)
+    dt = torch.nn.functional.softplus(
+        torch.randn(B, T, H, dtype=torch.float32, device=device) - 4,
+    )
+    a = -torch.exp(torch.rand(H, dtype=torch.float32, device=device))
+    q = torch.randn(B, T, 1, D, dtype=dtype, device=device)
+    k = torch.randn(B, T, 1, D, dtype=dtype, device=device)
+    level_scales = torch.randn(B, T, H, L, dtype=dtype, device=device)
+    v = (x * dt.unsqueeze(-1)).to(dtype=dtype)
+    g = a * dt
+
+    out, _ = chunk_log_linear_attn(q, k, v, g, level_scales)
+
+    ref = naive_log_linear_attn(q, k, v, g, level_scales)
+
+    assert_close("o", ref, out, 0.004)
+
+
+@pytest.mark.parametrize(
+    ("B", "T", "H", "D", "dtype"),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [(2, 512, 8, 64, torch.float32), (2, 1024, 8, 128, torch.float32)]
+    ],
+)
+@pytest.mark.skipif(device_platform == "intel", reason="Intel Triton Failure")
+def test_chunk_bwd(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ["TRITON_F32_DEFAULT"] = "ieee"
+
+    L = int(np.log2(T) + 1)
+    x = torch.randn(B, T, H, D, dtype=dtype, device=device)
+    dt = torch.nn.functional.softplus(
+        torch.randn(B, T, H, dtype=torch.float32, device=device) - 4,
+    )
+    a = -torch.exp(torch.rand(H, dtype=torch.float32, device=device))
+    q = torch.randn(B, T, 1, D, dtype=dtype, device=device)
+    k = torch.randn(B, T, 1, D, dtype=dtype, device=device)
+    level_scales = torch.randn(B, T, H, L, dtype=dtype, device=device)
+    v = (x * dt.unsqueeze(-1)).to(dtype=dtype)
+    g = a * dt
+    do = torch.randn_like(v)
+    q, k, v, g, level_scales = map(lambda x: x.to(device).requires_grad_(), (q, k, v, g, level_scales))
+
+    out, _ = chunk_log_linear_attn(q, k, v, g, level_scales)
+    (out * do).sum().backward()
+    tri_dq, tri_dk, tri_dv, tri_dg, tri_dl = q.grad, k.grad, v.grad, g.grad, level_scales.grad
+    q.grad = k.grad = v.grad = g.grad = level_scales.grad = None
+
+    ref = naive_log_linear_attn(q, k, v, g, level_scales)
+    (ref * do).sum().backward()
+    ref_dq, ref_dk, ref_dv, ref_dg, ref_dl = q.grad, k.grad, v.grad, g.grad, level_scales.grad
+
+    assert_close("o", ref, out, 0.004)
+    assert_close("dq", ref_dq, tri_dq, 0.007)
+    assert_close("dk", ref_dk, tri_dk, 0.008)
+    assert_close("dv", ref_dv, tri_dv, 0.007)
+    assert_close("dg", ref_dg, tri_dg, 0.015)
+    assert_close("dl", ref_dl, tri_dl, 0.015)
+
+
+@pytest.mark.parametrize(
+    ("H", "D", "cu_seqlens", "dtype"),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, [0, 15], torch.float32),
+            (4, 64, [0, 256, 500, 1000], torch.float32),
+            (4, 128, [0, 15, 100, 300, 1200, 2000], torch.float32),
+        ]
+    ],
+)
+@pytest.mark.skipif(device_platform == "intel", reason="Intel Triton Failure")
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ["TRITON_F32_DEFAULT"] = "ieee"
+
+    cu_seqlens = torch.LongTensor(cu_seqlens).to(device)
+    T = cu_seqlens[-1].item()
+
+    L = int(np.ceil(np.log2(T)) + 1)
+    x = torch.randn(1, T, H, D, dtype=dtype, device=device)
+    dt = torch.nn.functional.softplus(
+        torch.randn(1, T, H, dtype=torch.float32, device=device) - 4,
+    )
+    a = -torch.exp(torch.rand(H, dtype=torch.float32, device=device))
+    q = torch.randn(1, T, 1, D, dtype=dtype, device=device)
+    k = torch.randn(1, T, 1, D, dtype=dtype, device=device)
+    level_scales = torch.randn(1, T, H, L, dtype=dtype, device=device)
+    v = (x * dt.unsqueeze(-1)).to(dtype=dtype)
+    g = a * dt
+
+    out, _ = chunk_log_linear_attn(q, k, v, g, level_scales, cu_seqlens=cu_seqlens)
+
+    o = []
+    for i in range(cu_seqlens.shape[0] - 1):
+        bos, eos = cu_seqlens[i], cu_seqlens[i + 1]
+        v_s = v[:, bos:eos]
+        g_s = g[:, bos:eos]
+        k_s = k[:, bos:eos]
+        q_s = q[:, bos:eos]
+        level_scales_s = level_scales[:, bos:eos]
+
+        o.append(naive_log_linear_attn(q_s, k_s, v_s, g_s, level_scales_s))
+    ref = torch.cat(o, dim=1)
+
+    assert_close("o", ref, out, 0.004)

code/flash-linear-attention/tests/ops/test_mesa.py

@@ -0,0 +1,276 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.mesa_net import chunk_mesa_net, mesa_net_decoding_one_step, naive_mesa_net_decoding_one_step, naive_mesa_net_exact
+from fla.utils import assert_close, device, device_platform, is_intel_alchemist
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'gate_range', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-gate_range{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, [0.8, 0.99], torch.float16),
+            (2, 500, 4, 60, [0.8, 0.99], torch.float16),
+            (2, 1024, 8, 128, [0.8, 0.99], torch.float16),
+            (2, 1024, 8, 128, [0.01, 0.1], torch.float16),
+            (2, 1024, 8, 128, [1, 1], torch.float16),
+            (4, 2048, 8, 64, [0.8, 0.99], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Triton Failure',
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    gate_range: tuple[float, float],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    q = torch.rand(B, T, H, D, dtype=dtype) / 10
+    k = F.normalize(torch.rand(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.rand(B, T, H, D, dtype=dtype) / 10
+    beta = torch.rand(B, T, H, dtype=dtype).sigmoid()
+    lower_gate, upper_gate = gate_range
+    g = torch.rand(B, T, H, dtype=dtype).float().uniform_(lower_gate, upper_gate).log()
+    lamb = torch.rand(H, D, dtype=dtype).sigmoid() * 0.75 + 0.25
+    q, k, v, beta, g, lamb = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, lamb))
+    do = torch.rand_like(v)
+
+    k_init_rand = torch.nn.functional.normalize(torch.rand(B, H, D, device=device, dtype=dtype), dim=-1, p=2)
+    h_kk_init = (k_init_rand.unsqueeze(-1) * k_init_rand.unsqueeze(-2)).detach().clone().float().requires_grad_(True)
+    h_kv_init = torch.rand(B, H, D, D, dtype=torch.float32, device=device).requires_grad_(True)
+    d_h_kk_final = torch.rand_like(h_kk_init)
+    d_h_kv_final = torch.rand_like(h_kv_init)
+
+    tri, tri_kk_final, tri_kv_final = chunk_mesa_net(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        lamb=lamb.clone(),
+        max_CG_iteration=D,
+        h_kk_init=h_kk_init.clone(),
+        h_kv_init=h_kv_init.clone(),
+        output_final_state=True,
+    )
+
+    ((tri * do).sum() + (tri_kk_final * d_h_kk_final).sum() + (tri_kv_final * d_h_kv_final).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dlamb = q.grad, k.grad, v.grad, beta.grad, g.grad, lamb.grad
+    tri_dh_kk_init, tri_dh_kv_init = h_kk_init.grad, h_kv_init.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = lamb.grad = h_kk_init.grad = h_kv_init.grad = None
+
+    ref, ref_hkk_final, ref_hkv_final = naive_mesa_net_exact(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        lamb=lamb.clone(),
+        h_kk_init=h_kk_init.clone(),
+        h_kv_init=h_kv_init.clone(),
+    )
+
+    ((ref * do).sum() +
+     (ref_hkk_final * d_h_kk_final).sum() + (ref_hkv_final * d_h_kv_final).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dlamb = q.grad, k.grad, v.grad, beta.grad, g.grad, lamb.grad
+    ref_dh_kk_init, ref_dh_kv_init = h_kk_init.grad, h_kv_init.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = lamb.grad = h_kk_init.grad = h_kv_init.grad = None
+
+    assert_close('o', ref, tri, 0.006)
+    assert_close('h_kk_final', ref_hkk_final, tri_kk_final, 0.008)
+    assert_close('h_kv_final', ref_hkv_final, tri_kv_final, 0.008)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.008)
+    assert_close('dg', ref_dg, tri_dg, 0.008)
+    assert_close('dlamb', ref_dlamb, tri_dlamb, 0.015)
+    assert_close('dh_kk_init', ref_dh_kk_init, tri_dh_kk_init, 0.008)
+    assert_close('dh_kv_init', ref_dh_kv_init, tri_dh_kv_init, 0.008)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'gate_range', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-gate_range{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (3, 50, [0.8, 0.99], [0, 15], torch.float16),
+            (4, 64, [0.8, 0.99], [0, 14, 121, 421, 500], torch.float16),
+            (4, 64, [0.01, 0.1], [0, 256, 500, 1000], torch.float16),
+            (4, 100, [1, 1], [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    gate_range: tuple[float, float],
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    # randomly split the sequence into N segments
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.long, device=device)
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, D), dtype=dtype) / 10
+    k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn((1, T, H, D), dtype=dtype) / 10
+    lower_gate, upper_gate = gate_range
+    g = torch.rand(1, T, H, dtype=dtype).float().uniform_(lower_gate, upper_gate).log()
+    beta = torch.rand(1, T, H, dtype=dtype).sigmoid()
+    lamb = torch.rand(H, D, dtype=dtype).sigmoid() * 0.75 + 0.25
+
+    k_init_rand = torch.nn.functional.normalize(torch.rand(N, H, D, device=device, dtype=dtype), dim=-1, p=2)
+    h_kk_init = (k_init_rand.unsqueeze(-1) * k_init_rand.unsqueeze(-2)).detach().clone().float().requires_grad_(True)
+    h_kv_init = torch.rand(N, H, D, D, dtype=torch.float32, device=device).requires_grad_(True)
+
+    q, k, v, beta, g, lamb, h_kk_init, h_kv_init = map(lambda x: x.to(
+        device).requires_grad_(), (q, k, v, beta, g, lamb, h_kk_init, h_kv_init))
+    do = torch.rand_like(v) / 10
+    d_h_kk_final = torch.rand_like(h_kk_init)
+    d_h_kv_final = torch.rand_like(h_kv_init)
+
+    tri, tri_h_kk_final, tri_h_kv_final = chunk_mesa_net(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        beta=beta.clone(),
+        g=g.clone(),
+        lamb=lamb.clone(),
+        h_kk_init=h_kk_init.clone(),
+        h_kv_init=h_kv_init.clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+
+    ((tri * do).sum() +
+     (tri_h_kk_final * d_h_kk_final).sum() + (tri_h_kv_final * d_h_kv_final).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dlamb, tri_dh_kk_init, tri_dh_kv_init = \
+        q.grad, k.grad, v.grad, beta.grad, g.grad, lamb.grad, h_kk_init.grad, h_kv_init.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = lamb.grad = h_kk_init.grad = h_kv_init.grad = None
+
+    ref = []
+    ref_h_kk_t = []
+    ref_h_kv_t = []
+    for i in range(N):
+        ref_i, ref_h_kk_i, ref_h_kv_i = naive_mesa_net_exact(
+            q=q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k=k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v=v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            g=g[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            lamb=lamb,
+            h_kk_init=h_kk_init[i],
+            h_kv_init=h_kv_init[i],
+        )
+        ref.append(ref_i)
+        ref_h_kk_t.append(ref_h_kk_i)
+        ref_h_kv_t.append(ref_h_kv_i)
+    ref = torch.cat(ref, 1)
+    ref_h_kk_t = torch.cat(ref_h_kk_t, 0)
+    ref_h_kv_t = torch.cat(ref_h_kv_t, 0)
+
+    ((ref * do).sum() + (ref_h_kk_t * d_h_kk_final).sum() + (ref_h_kv_t * d_h_kv_final).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dlamb, ref_dh_kk_init, ref_dh_kv_init = \
+        q.grad, k.grad, v.grad, beta.grad, g.grad, lamb.grad, h_kk_init.grad, h_kv_init.grad
+    q.grad = k.grad = v.grad = beta.grad = g.grad = lamb.grad = h_kk_init.grad = h_kv_init.grad = None
+
+    assert_close('o', ref, tri, 0.006)
+    assert_close('h_kk_final', ref_h_kk_t, tri_h_kk_final, 0.008)
+    assert_close('h_kv_final', ref_h_kv_t, tri_h_kv_final, 0.008)
+    assert_close('dq', ref_dq, tri_dq, 0.008)
+    assert_close('dk', ref_dk, tri_dk, 0.008)
+    assert_close('dv', ref_dv, tri_dv, 0.008)
+    assert_close('db', ref_dbeta, tri_dbeta, 0.015)
+    assert_close('dlamb', ref_dlamb, tri_dlamb, 0.015)
+    assert_close('dg', ref_dg, tri_dg, 0.015)
+    assert_close('dh_kk_0', ref_dh_kk_init, tri_dh_kk_init, 0.007)
+    assert_close('dh_kv_0', ref_dh_kv_init, tri_dh_kv_init, 0.007)
+
+
+@pytest.mark.parametrize(
+    ('B', 'H', 'D', 'gate_range', 'max_CG_step', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-H{}-D{}-gate_range{}-max_CG_step{}-{}".format(*test))
+        for test in [
+            (1, 3, 50, [0.95, 0.99], 1, torch.float16),
+            (2, 4, 60, [0.95, 0.99], 5, torch.float16),
+            (2, 8, 128, [0.95, 0.99], 1, torch.float16),
+            (2, 8, 128, [0.95, 0.99], 5, torch.float16),
+            (2, 8, 128, [0.95, 0.99], 30, torch.float16),
+        ]
+    ],
+)
+def test_decoding_one_step(
+    B: int,
+    H: int,
+    D: int,
+    gate_range: tuple[float, float],
+    max_CG_step: int,
+    dtype: torch.dtype,
+):
+    if is_intel_alchemist and D > 128:
+        pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128')
+    torch.manual_seed(42)
+    torch.set_default_device(device)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    # randomly split the sequence into N segments
+    q = torch.rand((B, H, D), dtype=dtype)
+    k = F.normalize(torch.randn(B, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.rand((B, H, D), dtype=dtype)
+    lower_gate, upper_gate = gate_range
+    g = torch.rand(B, H, dtype=dtype).float().uniform_(lower_gate, upper_gate).log()
+    beta = torch.rand(B, H, dtype=dtype).sigmoid()
+    lamb = torch.rand(H, D, dtype=dtype).sigmoid() * 0.75 + 0.25
+
+    k_init_rand = torch.nn.functional.normalize(torch.rand(B, H, D, device=device, dtype=dtype), dim=-1, p=2)
+    prev_h_kk = (k_init_rand.unsqueeze(-1) * k_init_rand.unsqueeze(-2)).detach().clone().float().requires_grad_(True)
+    prev_h_kv = torch.rand(B, H, D, D, dtype=torch.float32, device=device).requires_grad_(True)
+
+    o, curr_h_kk, curr_h_kv = mesa_net_decoding_one_step(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        lamb=lamb.clone(),
+        beta=beta.clone(),
+        prev_h_kk=prev_h_kk.clone(),
+        prev_h_kv=prev_h_kv.clone(),
+        max_CG_iteration=max_CG_step,
+    )
+
+    o_ref, curr_h_kk_re, curr_h_kv_re = naive_mesa_net_decoding_one_step(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        lamb=lamb.clone(),
+        beta=beta.clone(),
+        prev_h_kk=prev_h_kk.clone(),
+        prev_h_kv=prev_h_kv.clone(),
+        max_CG_iteration=max_CG_step,
+    )
+
+    assert_close('o', o, o_ref, 0.005)
+    assert_close('curr_h_kk', curr_h_kk, curr_h_kk_re, 0.005)
+    assert_close('curr_h_kv', curr_h_kv, curr_h_kv_re, 0.005)

code/flash-linear-attention/tests/ops/test_nsa.py

@@ -0,0 +1,148 @@
+
+import os
+
+import pytest
+import torch
+import triton
+
+from fla.ops.nsa.naive import naive_nsa
+from fla.ops.nsa.parallel import parallel_nsa
+from fla.ops.utils import prepare_token_indices
+from fla.utils import assert_close, device
+
+
+# FIXME
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'HQ', 'D', 'S', 'block_size', 'scale', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-HQ{}-D{}-S{}-block_size{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 16, 64, 16, 32, 1.0, torch.float16),
+            (3, 111, 1, 32, 100, 16, 32, 1.0, torch.float16),
+            (3, 1024, 2, 32, 60, 16, 32, 0.1, torch.float16),
+            (3, 1024, 2, 32, 128, 16, 32, 0.1, torch.float16),
+            (4, 2048, 2, 32, 64, 16, 32, 0.1, torch.float16),
+        ]
+    ],
+)
+def test_parallel(
+    B: int,
+    T: int,
+    H: int,
+    HQ: int,
+    D: int,
+    S: int,
+    block_size: int,
+    scale: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    q = torch.randn((B, T, HQ, D), dtype=dtype, device=device).requires_grad_(True)
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    do = torch.randn((B, T, HQ, D), dtype=dtype, device=device)
+
+    block_indices = torch.full((B, T, H, S), T, dtype=torch.long, device=device)
+    for b in range(B):
+        for t in range(T):
+            for h in range(H):
+                i_i = torch.randperm(max(1, triton.cdiv(t, block_size)))[:S]
+                block_indices[b, t, h, :len(i_i)] = i_i
+    block_indices = block_indices.sort(-1)[0]
+
+    ref = naive_nsa(q=q, k=k, v=v, block_indices=block_indices, block_size=block_size, scale=scale)
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+
+    tri = parallel_nsa(q=q, k=k, v=v, block_indices=block_indices, block_size=block_size, scale=scale)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    assert_close(" o", ref, tri, 0.005)
+    assert_close("dq", ref_dq, tri_dq, 0.005)
+    assert_close("dk", ref_dk, tri_dk, 0.005)
+    assert_close("dv", ref_dv, tri_dv, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'HQ', 'D', 'S', 'block_size', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-HQ{}-D{}-S{}-block_size{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (1, 16, 64, 16, 32, [0, 15], torch.float16),
+            (2, 32, 64, 16, 32, [0, 256, 500, 1000], torch.float16),
+            (2, 32, 100, 16, 32, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_parallel_varlen(
+    H: int,
+    HQ: int,
+    D: int,
+    S: int,
+    block_size: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, HQ, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn((1, T, HQ, D), dtype=dtype, device=device)
+
+    block_indices = torch.full((1, T, H, S), T, dtype=torch.long, device=device)
+    seq_indices = prepare_token_indices(cu_seqlens).tolist()
+
+    for i in range(T):
+        _, t = seq_indices[i]
+        for h in range(H):
+            i_i = torch.randperm(max(1, triton.cdiv(t, block_size)))[:S]
+            block_indices[0, i, h, :len(i_i)] = i_i
+    block_indices = block_indices.sort(-1)[0]
+
+    ref = naive_nsa(
+        q=q,
+        k=k,
+        v=v,
+        block_indices=block_indices,
+        block_size=block_size,
+        cu_seqlens=cu_seqlens,
+    )
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+
+    tri = parallel_nsa(
+        q=q,
+        k=k,
+        v=v,
+        block_indices=block_indices,
+        block_size=block_size,
+        cu_seqlens=cu_seqlens,
+    )
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)

code/flash-linear-attention/tests/ops/test_path_attn.py

@@ -0,0 +1,211 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.ops.path_attn.parallel import parallel_path_attention
+from fla.utils import assert_close, device, is_intel_alchemist
+
+
+def naive_path_attn(q, k, v, w, beta, g, scale, BT=64):
+    original_dtype = q.dtype
+    HQ = q.shape[2]
+    H = k.shape[2]
+    q, k, v, w, beta, g = map(lambda x: x.to(torch.float).transpose(1, 2), [q, k, v, w, beta, g])
+    g_cumsum = g.cumsum(-1)
+    q = q.unsqueeze(2).expand(-1, -1, HQ//HQ, -1, -1).flatten(1, 2)
+    k = k.unsqueeze(2).expand(-1, -1, HQ//H, -1, -1).flatten(1, 2)
+    v = v.unsqueeze(2).expand(-1, -1, HQ//H, -1, -1).flatten(1, 2)
+    w = w.unsqueeze(2).expand(-1, -1, HQ//H, -1, -1).flatten(1, 2)
+    beta = beta.unsqueeze(2).expand(-1, -1, HQ//H, -1).flatten(1, 2)
+
+    b, h, l, _ = q.shape
+    if l % BT != 0:
+        padding_size = BT - l % BT
+        q, k, w = map(lambda x: F.pad(x, (0, 0, 0, padding_size)), [q, k, w])
+        beta = F.pad(beta, (0, padding_size))
+    seq_len = q.shape[2]
+    w_beta = w * beta[..., None]
+    q, k, w, w_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BT), [q, k, w, w_beta])
+    mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
+    T = -(w_beta @ w.transpose(-1, -2)).masked_fill(mask, 0)
+    for i in range(1, BT):
+        T[..., i, :i] = T[..., i, :i].clone() + (T[..., i, :, None].clone() * T[..., :, :i].clone()).sum(-2)
+    T = T + torch.eye(BT, dtype=q.dtype, device=q.device)
+    Twbk = T @ (w_beta @ k.transpose(-1, -2)).masked_fill(mask, 0)
+    qw = (q @ w.transpose(-1, -2)).tril()
+    Twb = T @ w_beta
+    A_local = (q @ k.transpose(-1, -2)).tril() - qw @ Twbk
+    q = q - qw @ Twb
+    k = k - Twbk.transpose(-1, -2) @ w
+    H = w.transpose(-1, -2) @ Twb
+    A = torch.zeros(b, h, seq_len, seq_len, device=q.device)
+    q, k, w, w_beta = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, w, w_beta])
+    for i in range(0, seq_len, BT):
+        q_i = q[:, :, i:i+BT].clone()
+        for j in range(i - BT, -BT, -BT):
+            k_j = k[:, :, j:j+BT]
+            A_ij = q_i @ k_j.transpose(-1, -2)
+            A[:, :, i:i+BT, j:j+BT] = A_ij
+            q_i = q_i - q_i @ H[:, :, j // BT]
+    for i in range(0, seq_len//BT):
+        A[:, :, i*BT:i*BT+BT, i*BT:i*BT+BT] = A_local[:, :, i]
+    A = A.masked_fill_(~torch.tril(torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool)), float("-inf"))
+    A = A[:, :, :l, :l]
+    A = A + g_cumsum[..., None] - g_cumsum[..., None, :]
+    ref_o = (A * scale).softmax(-1).to(v) @ v
+    return ref_o.to(original_dtype).transpose(1, 2)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'HQ', 'D', 'use_forget_gate', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-HQ{}-D{}-use_forget_gate{}-{}".format(*test))
+        for test in [
+            # SY (2025/07/08): It somehow failed on Hopper with error msg: Aborted (core dumped)
+            # (10, 62, 2, 8, 128, True, torch.bfloat16),
+            (5, 512, 2, 8, 128, True, torch.bfloat16),
+            (3, 1024, 2, 8, 64, True, torch.bfloat16),
+            (2, 2000, 1, 4, 64, False, torch.bfloat16),
+            (1, 4000, 1, 2, 128, False, torch.bfloat16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    is_intel_alchemist,
+    reason="Intel Triton Failure",
+)
+def test_parallel(
+    B: int,
+    H: int,
+    HQ: int,
+    T: int,
+    D: int,
+    use_forget_gate: bool,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    q = torch.randn((B, T, HQ, D), dtype=dtype, device=device).requires_grad_(True)
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    w = F.normalize(torch.randn((B, T, H, D), dtype=torch.float, device=device), dim=-1, p=2).requires_grad_(True)
+    beta = torch.empty((B, T, H), dtype=torch.float, device=device).uniform_(1.5, 2.0).requires_grad_(True)
+    if use_forget_gate:
+        g = torch.empty((B, T, HQ), dtype=torch.float, device=device).uniform_(
+            0.95, 1).log().requires_grad_(True)
+    else:
+        g = None
+    do = torch.rand((B, T, HQ, D), dtype=dtype, device=device)
+    scale = D ** -0.5
+    ref = naive_path_attn(q, k, v, w, beta, torch.zeros(B, T, HQ, device=device, dtype=torch.float) if g is None else g, scale)
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    if use_forget_gate:
+        ref_dg, g.grad = g.grad.clone(), None
+    ref_dw, w.grad = w.grad.clone(), None
+    ref_db, beta.grad = beta.grad.clone(), None
+
+    tri, _ = parallel_path_attention(q=q, k=k, v=v, w=w, beta=beta, g=g, scale=scale)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    if use_forget_gate:
+        tri_dg, g.grad = g.grad.clone(), None
+    tri_dw, w.grad = w.grad.clone(), None
+    tri_db, beta.grad = beta.grad.clone(), None
+
+    assert_close(" o", ref, tri, 0.005)
+    assert_close("dq", ref_dq, tri_dq, 0.008)
+    assert_close("dk", ref_dk, tri_dk, 0.008)
+    assert_close("dv", ref_dv, tri_dv, 0.008)
+    if use_forget_gate:
+        assert_close("dg", ref_dg, tri_dg, 0.02)
+    assert_close("dw", ref_dw, tri_dw, 0.015)
+    assert_close("db", ref_db, tri_db, 0.02)
+
+
+@pytest.mark.parametrize(
+    ('H', 'HQ', 'D', 'use_forget_gate', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-HQ{}-D{}-use_forget_gate{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 4, 128, False, [0, 15, 333, 2048], torch.float16),
+            (2, 4, 128, True, [0, 15, 333, 2048], torch.float16),
+            (2, 4, 64, True, [0, 841, 889, 4096], torch.float16),
+            (2, 4, 64, False, [0, 841, 889, 2000, 3000, 4096], torch.float16),
+            (2, 16, 128, True, [0, 500, 1023, 2000, 3000, 4096], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0",
+    reason="Skipping test because TEST_CHUNK_VARLEN is enabled",
+)
+@pytest.mark.skipif(
+    is_intel_alchemist,
+    reason="Intel Triton Failure",
+)
+def test_parallel_varlen(
+    H: int,
+    HQ: int,
+    D: int,
+    use_forget_gate: bool,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.randn((1, T, HQ, D), dtype=dtype, device=device).requires_grad_(True)
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    w = F.normalize(torch.randn((1, T, H, D), dtype=torch.float, device=device), dim=-1, p=2).requires_grad_(True)
+    beta = torch.rand((1, T, H), dtype=torch.float, device=device).sigmoid().requires_grad_(True)
+    if use_forget_gate:
+        g = torch.empty((1, T, HQ), dtype=torch.float, device=device).uniform_(0.95, 1).log().requires_grad_(True)
+    else:
+        g = None
+    do = torch.randn((1, T, HQ, D), dtype=dtype, device=device)
+    scale = D ** -0.5
+    ref = torch.zeros(1, T, HQ, D, device=device, dtype=dtype)
+    for bos, eos in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False):
+        g_segment = torch.zeros(1, eos - bos, HQ, device=device, dtype=torch.float) if g is None else g[:, bos:eos]
+        ref[:, bos:eos] = naive_path_attn(
+            q[:, bos:eos], k[:, bos:eos], v[:, bos:eos],
+            w[:, bos:eos], beta[:, bos:eos], g_segment, scale,
+        )
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    if use_forget_gate:
+        ref_dg, g.grad = g.grad.clone(), None
+    ref_dw, w.grad = w.grad.clone(), None
+    ref_db, beta.grad = beta.grad.clone(), None
+    tri, _ = parallel_path_attention(q=q, k=k, v=v, w=w, beta=beta, g=g, scale=scale, cu_seqlens=cu_seqlens)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    if use_forget_gate:
+        tri_dg, g.grad = g.grad.clone(), None
+    tri_dw, w.grad = w.grad.clone(), None
+    tri_db, beta.grad = beta.grad.clone(), None
+    assert_close(" o", ref, tri, 0.005)
+    assert_close("dq", ref_dq, tri_dq, 0.005)
+    assert_close("dk", ref_dk, tri_dk, 0.005)
+    assert_close("dv", ref_dv, tri_dv, 0.005)
+    if use_forget_gate:
+        assert_close("dg", ref_dg, tri_dg, 0.005)
+    assert_close("dw", ref_dw, tri_dw, 0.005)
+    assert_close("db", ref_db, tri_db, 0.005)

code/flash-linear-attention/tests/ops/test_retention.py

@@ -0,0 +1,308 @@
+
+import os
+
+import pytest
+import torch
+
+from fla.ops.retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'K', 'expand_ratio', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-K{}-expand_ratio{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float16),
+            (2, 500, 3, 60, 1, torch.float16),
+            (2, 1000, 3, 100, 1, torch.float16),
+            (2, 1000, 3, 128, 2, torch.float16),
+            (3, 1024, 4, 256, 2, torch.float16),
+            (4, 2048, 4, 64, 2, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    K: int,
+    expand_ratio: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    V = K * expand_ratio
+
+    q = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, V), dtype=dtype, device=device).requires_grad_()
+    h0 = torch.randn((B, H, K, V), dtype=dtype, device=device).requires_grad_()
+
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+    ref, ref_ht = fused_recurrent_retention(q, k, v, initial_state=h0, output_final_state=True)
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+
+    tri, tri_ht = chunk_retention(q, k, v, initial_state=h0, output_final_state=True)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'K', 'expand_ratio', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-K{}-expand_ratio{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, 1, [0, 15], torch.float16),
+            (4, 64, 2, [0, 256, 500, 1000], torch.float16),
+            (4, 100, 2, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_chunk_varlen(
+    H: int,
+    K: int,
+    expand_ratio: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    V = K * expand_ratio
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.long, device=device)
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, K), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, K), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, V), dtype=dtype, device=device).requires_grad_()
+    h0 = torch.randn((N, H, K, V), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+
+    ref, ref_ht = fused_recurrent_retention(
+        q=q,
+        k=k,
+        v=v,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = chunk_retention(
+        q=q,
+        k=k,
+        v=v,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'K', 'expand_ratio', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-K{}-expand_ratio{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float16),
+            (2, 500, 3, 60, 1, torch.float16),
+            (2, 1000, 3, 100, 1, torch.float16),
+            (2, 1000, 3, 128, 2, torch.float16),
+            (3, 1024, 4, 256, 2, torch.float16),
+            (4, 2048, 4, 64, 2, torch.float16),
+        ]
+    ],
+)
+def test_fused_chunk(
+    B: int,
+    T: int,
+    H: int,
+    K: int,
+    expand_ratio: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    V = K * expand_ratio
+
+    q = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, V), dtype=dtype, device=device).requires_grad_()
+    h0 = torch.randn((B, H, K, V), dtype=dtype, device=device).requires_grad_()
+
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+    ref, ref_ht = fused_recurrent_retention(q, k, v, initial_state=h0, output_final_state=True)
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+
+    tri, tri_ht = fused_chunk_retention(q, k, v, initial_state=h0, output_final_state=True)
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'K', 'expand_ratio', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-K{}-expand_ratio{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, 1, [0, 15], torch.float16),
+            (4, 64, 2, [0, 256, 500, 1000], torch.float16),
+            (4, 100, 2, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_fused_chunk_varlen(
+    H: int,
+    K: int,
+    expand_ratio: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    V = K * expand_ratio
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.long, device=device)
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, K), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, K), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, V), dtype=dtype, device=device).requires_grad_()
+    h0 = torch.randn((N, H, K, V), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn_like(v)
+    dht = torch.randn_like(h0)
+
+    ref, ref_ht = fused_recurrent_retention(
+        q=q,
+        k=k,
+        v=v,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_chunk_retention(
+        q=q,
+        k=k,
+        v=v,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'K', 'expand_ratio', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-K{}-expand_ratio{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float16),
+            (2, 500, 4, 60, 1, torch.float16),
+            (2, 1024, 8, 128, 1, torch.float16),
+            (3, 1024, 8, 128, 2, torch.float16),
+            (3, 1024, 8, 256, 2, torch.float16),
+            (4, 2048, 8, 64, 2, torch.float16),
+        ]
+    ],
+)
+def test_parallel(
+    B: int,
+    T: int,
+    H: int,
+    K: int,
+    expand_ratio: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    V = K * expand_ratio
+
+    q = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, V), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn_like(v)
+
+    ref, _ = fused_recurrent_retention(q, k, v)
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+
+    tri, _ = parallel_retention(q, k, v)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)

code/flash-linear-attention/tests/ops/test_rwkv6.py

@@ -0,0 +1,192 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.rwkv6 import chunk_rwkv6
+from fla.ops.rwkv6.fused_recurrent import fused_recurrent_rwkv6
+from fla.utils import assert_close, device, device_platform
+
+
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason="Intel Triton Failure",
+)
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 15, 2, 60, 1.0, torch.float16),
+            (3, 60, 3, 64, 0.1, torch.float16),
+            (3, 64, 2, 64, 1, torch.float16),
+            (4, 500, 3, 256, 1, torch.float16),
+            (4, 1000, 4, 64, 10, torch.float16),
+            (4, 2048, 4, 64, 1, torch.float16),
+            (4, 2048, 4, 256, 1, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    w = F.logsigmoid(torch.randn((B, T, H, D), dtype=dtype, device=device)) / gate_logit_normalizer
+
+    u = torch.randn(H, D, dtype=dtype, device=device).requires_grad_(True)
+    h0 = torch.randn(B, H, D, D, dtype=dtype, device=device).requires_grad_()
+    w = w.requires_grad_()
+    do = torch.randn_like(v)
+
+    ref, ref_ht = fused_recurrent_rwkv6(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        u.clone(),
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    ref, _ = fused_recurrent_rwkv6(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        u.clone(),
+        initial_state=h0.clone(),
+        output_final_state=False,
+    )
+
+    ((ref * do).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dw, w.grad = w.grad.clone(), None
+    ref_du, u.grad = u.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    # triton implementation
+    tri, tri_ht = chunk_rwkv6(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        u.clone(),
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    ((tri * do).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dw, w.grad = w.grad.clone(), None
+    tri_du, u.grad = u.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dw', ref_dw, tri_dw, 0.005)
+    assert_close('du', ref_du, tri_du, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, [0, 15], torch.float16),
+            (4, 64, [0, 256, 500, 1000], torch.float16),
+            (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    w = F.logsigmoid(torch.randn((1, T, H, D), dtype=dtype, device=device)).requires_grad_(True)
+    u = torch.randn(H, D, dtype=dtype, device=device).requires_grad_(True)
+    h0 = torch.randn((N, H, D, D), dtype=dtype, device=device).requires_grad_()
+    do = torch.randn_like(v)
+
+    ref, ref_ht = fused_recurrent_rwkv6(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        u.clone(),
+        initial_state=h0.clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ref, _ = fused_recurrent_rwkv6(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        u.clone(),
+        initial_state=h0.clone(),
+        output_final_state=False,
+        cu_seqlens=cu_seqlens,
+    )
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dw, w.grad = w.grad.clone(), None
+    ref_du, u.grad = u.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = chunk_rwkv6(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        u.clone(),
+        initial_state=h0.clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dw, w.grad = w.grad.clone(), None
+    tri_du, u.grad = u.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dw', ref_dw, tri_dw, 0.005)
+    assert_close('du', ref_du, tri_du, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)

code/flash-linear-attention/tests/ops/test_rwkv7.py

@@ -0,0 +1,300 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.generalized_delta_rule.dplr.fused_recurrent import fused_recurrent_dplr_delta_rule
+from fla.ops.rwkv7.channel_mixing import channel_mixing_rwkv7, channel_mixing_rwkv7_torch
+from fla.ops.rwkv7.fused_addcmul import fused_addcmul_rwkv7, torch_addcmul_rwkv7
+from fla.ops.rwkv7.fused_k_update import fused_k_rwkv7, k_update_ref
+from fla.ops.rwkv7.fused_recurrent import fused_mul_recurrent_rwkv7
+from fla.ops.rwkv7.gate_output_correction import gate_output_correction, gate_output_correction_ref
+from fla.utils import assert_close, device, is_nvidia_hopper
+
+
+@pytest.mark.parametrize("B", [2])
+@pytest.mark.parametrize("T", [1024])
+@pytest.mark.parametrize("n_embd", [1024])
+@pytest.mark.parametrize("dim_ffn", [4096])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("inplace", [True, False])
+@pytest.mark.parametrize("xprevdim", [2, 3])
+@pytest.mark.skipif(
+    os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0",
+    reason="Skipping test because TEST_CHUNK_VARLEN is enabled",
+)
+def test_channel_mixing_gradients(B, T, n_embd, dim_ffn, dtype, inplace, xprevdim):
+    torch.manual_seed(42)
+    torch._dynamo.config.cache_size_limit = 512
+
+    x = torch.randn(
+        B, T, n_embd, device=device, dtype=dtype, requires_grad=True,
+    )
+    if xprevdim == 3:
+        x_prev = torch.randn(
+            B, 1, n_embd, device=device, dtype=dtype, requires_grad=True,
+        )
+    else:
+        x_prev = torch.randn(
+            B, n_embd, device=device, dtype=dtype, requires_grad=True,
+        )
+    x_k = torch.randn(1, 1, n_embd, device=device, dtype=dtype, requires_grad=True)
+    K_ = torch.randn(n_embd, dim_ffn, device=device, dtype=dtype, requires_grad=True)
+    V_ = torch.randn(dim_ffn, n_embd, device=device, dtype=dtype, requires_grad=True)
+
+    x2 = x.clone().detach().requires_grad_(True)
+    x_prev2 = x_prev.clone().detach().requires_grad_(True)
+    x_k2 = x_k.clone().detach().requires_grad_(True)
+    K_2 = K_.clone().detach().requires_grad_(True)
+    V_2 = V_.clone().detach().requires_grad_(True)
+
+    o1, last1 = channel_mixing_rwkv7_torch(
+        x.to(torch.float32),
+        x_prev.to(torch.float32),
+        x_k.to(torch.float32),
+        K_.to(torch.float32),
+        V_.to(torch.float32),
+    )
+    loss1 = o1.mean() + last1.mean()
+    loss1.backward()
+
+    o2, last2 = channel_mixing_rwkv7(x2, x_prev2, x_k2, K_2, V_2, inplace)
+    loss2 = o2.mean() + last2.mean()
+    loss2.backward()
+
+    assert_close(" dx", x.grad, x2.grad, ratio=5e-3)
+    assert_close(" dxprev", x_prev.grad, x_prev2.grad, ratio=5e-3)
+    assert_close(" dx_k", x_k.grad, x_k2.grad, ratio=5e-3)
+    assert_close(" dK_", K_.grad, K_2.grad, ratio=5e-3)
+    assert_close(" dV_", V_.grad, V_2.grad, ratio=5e-3)
+
+
+@pytest.mark.parametrize('B', [2])
+@pytest.mark.parametrize('T', [1, 1024])
+@pytest.mark.parametrize('H', [1])
+@pytest.mark.parametrize('D', [64])
+@pytest.mark.parametrize('scale', [None, 1])
+@pytest.mark.parametrize('dtype', [torch.float32])
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0',
+    reason='Skipping test because TEST_CHUNK_VARLEN is enabled',
+)
+def test_fused_mul_recurrent_fwd(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    r = torch.empty(B, T, H, D, device=device).uniform_(-8, -6).to(dtype=dtype)
+    k = torch.empty(B, T, H, D, device=device).uniform_(-8, -6).to(dtype=dtype)
+    v = torch.empty(B, T, H, D, device=device).uniform_(-8, -6).to(dtype=dtype)
+    w = torch.empty(B, T, H, D, device=device).uniform_(-8, -6).to(dtype=dtype)
+
+    kk = torch.empty(B, T, H, D, device=device).uniform_(-1, 1)
+    kk = F.normalize(kk, dim=-1).to(dtype=dtype)
+
+    a = -kk.clone()
+    a_scale = torch.empty(B, T, H, D, device=device).uniform_(0, 0.1).to(dtype=dtype)
+    b = (kk * a_scale).requires_grad_(False)  # kk*a
+    h0 = torch.randn(B, H, D, D, dtype=torch.float)
+    r, k, v, a, a_scale, b, w, h0 = map(lambda x: x.to(device).requires_grad_(False),
+                                        (r, k, v, a, a_scale, b, w, h0))
+    ref, ref_ht = fused_recurrent_dplr_delta_rule(
+        q=r.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        a=a.clone(),
+        b=b.clone(),
+        gk=w.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+
+    tri, tri_ht = fused_mul_recurrent_rwkv7(
+        r=r.clone(),
+        w=w.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        kk=kk.clone(),
+        a=a_scale.clone(),
+        scale=scale,
+        initial_state=h0.clone(),
+        output_final_state=True,
+    )
+    assert_close('o', ref, tri, 0.002)
+    assert_close('ht', ref_ht, tri_ht, 0.002)
+
+
+@pytest.mark.parametrize("B", [1])
+@pytest.mark.parametrize("T", [20, 1024, 4100, 131072])
+@pytest.mark.parametrize("H", [2])
+@pytest.mark.parametrize("D", [64])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("use_g", [True, False])
+@pytest.mark.skipif(
+    os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0",
+    reason="Skipping test because TEST_CHUNK_VARLEN is enabled",
+)
+def test_fused_rwkv7_addcmul(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+    use_g: bool,
+):
+    if T == 128 * 1024 and not is_nvidia_hopper:
+        pytest.skip("Skipping test for T=131072 on non-Hopper GPUs")
+    hidden_size = H*D
+    hidden_states = torch.randn(B, T, hidden_size).to(device).to(dtype).requires_grad_()
+    xx = torch.randn(B, T, hidden_size).to(device).to(dtype).requires_grad_()
+    x_r = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_()
+    x_w = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_()
+    x_k = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_()
+    x_v = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_()
+    x_a = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_()
+    if use_g:
+        x_g = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_()
+    else:
+        x_g = None
+    xr0, xw0, xk0, xv0, xa0, xg0 = fused_addcmul_rwkv7(hidden_states, xx, x_r, x_w, x_k, x_v, x_a, x_g)
+    xr1, xw1, xk1, xv1, xa1, xg1 = torch_addcmul_rwkv7(hidden_states.float(),
+                                                       xx.float(), x_r.float(),
+                                                       x_w.float(), x_k.float(),
+                                                       x_v.float(), x_a.float(),
+                                                       x_g.float() if use_g else None)
+    ratio = 1e-5 if dtype == torch.float32 else 0.002
+    assert_close("xr0", xr0, xr1, ratio=ratio)
+    assert_close("xw0", xw0, xw1, ratio=ratio)
+    assert_close("xk0", xk0, xk1, ratio=ratio)
+    assert_close("xv0", xv0, xv1, ratio=ratio)
+    assert_close("xa0", xa0, xa1, ratio=ratio)
+    if use_g:
+        assert_close("xg0", xg0, xg1, ratio=ratio)
+        (xr0 + xw0 + xk0 + xv0 + xa0 + xg0).sum().backward()
+    else:
+        (xr0 + xw0 + xk0 + xv0 + xa0).sum().backward()
+    d_ixr = x_r.grad.clone()
+    d_ixw = x_w.grad.clone()
+    d_ixk = x_k.grad.clone()
+    d_ixv = x_v.grad.clone()
+    d_ixa = x_a.grad.clone()
+    d_hidden = hidden_states.grad.clone()
+    d_xx = xx.grad.clone()
+
+    x_r.grad.zero_()
+    x_w.grad.zero_()
+    x_k.grad.zero_()
+    x_v.grad.zero_()
+    x_a.grad.zero_()
+    if use_g:
+        d_ixg = x_g.grad.clone()
+        x_g.grad.zero_()
+    hidden_states.grad.zero_()
+    xx.grad.zero_()
+
+    if use_g:
+        (xr1 + xw1 + xk1 + xv1 + xa1 + xg1).sum().backward()
+    else:
+        (xr1 + xw1 + xk1 + xv1 + xa1).sum().backward()
+    d_ixr1 = x_r.grad.clone()
+    d_ixw1 = x_w.grad.clone()
+    d_ixk1 = x_k.grad.clone()
+    d_ixv1 = x_v.grad.clone()
+    d_ixa1 = x_a.grad.clone()
+    if use_g:
+        d_ixg1 = x_g.grad.clone()
+    d_hidden1 = hidden_states.grad.clone()
+    d_xx1 = xx.grad.clone()
+
+    assert_close("d_ixr", d_ixr, d_ixr1, ratio=ratio)
+    assert_close("d_ixw", d_ixw, d_ixw1, ratio=ratio)
+    assert_close("d_ixk", d_ixk, d_ixk1, ratio=ratio)
+    assert_close("d_ixv", d_ixv, d_ixv1, ratio=ratio)
+    assert_close("d_ixa", d_ixa, d_ixa1, ratio=ratio)
+    if use_g:
+        assert_close("d_ixg", d_ixg, d_ixg1, ratio=ratio)
+    assert_close("d_hidden", d_hidden, d_hidden1, ratio=ratio)
+    assert_close("d_xx", d_xx, d_xx1, ratio=ratio)
+
+
+@pytest.mark.parametrize("B", [4])
+@pytest.mark.parametrize("T", [13, 4096, 8000])
+@pytest.mark.parametrize("H", [64])
+@pytest.mark.parametrize("D", [64])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("ka_shape", [1, 3])
+def test_fused_k_update(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+    ka_shape: int,
+):
+    k = torch.randn(B, T, H*D).uniform_(-8, 8).to(device).to(dtype).requires_grad_()
+    a = torch.randn(B, T, H*D).uniform_(-8, 8).to(device).to(dtype).requires_grad_()
+    if ka_shape == 1:
+        ka = torch.randn(H*D).uniform_(-8, 8).to(device).to(dtype).requires_grad_()
+    else:
+        ka = torch.randn(1, 1, H*D).uniform_(-8, 8).to(device).to(dtype).requires_grad_()
+
+    ref = k_update_ref(k.float(), a.float(), ka.float())
+    ref.sum().backward()
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_da, a.grad = a.grad.clone(), None
+    ref_dka, ka.grad = ka.grad.clone(), None
+    tri = fused_k_rwkv7(k, a, ka)
+    tri.sum().backward()
+    ratio = 5e-5 if dtype == torch.float32 else 0.002
+    assert_close("  o", tri, ref, ratio=ratio)
+    assert_close(" dk", ref_dk, k.grad, ratio=ratio)
+    assert_close(" da", ref_da, a.grad, ratio=ratio)
+    assert_close("dka", ref_dka, ka.grad, ratio=ratio)
+
+
+@pytest.mark.parametrize("B", [4])
+@pytest.mark.parametrize("T", [4096])
+@pytest.mark.parametrize("H", [64])
+@pytest.mark.parametrize("D", [64])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+def test_gate_output_correction(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    value_dim = H * D
+    torch.manual_seed(0)
+
+    o_ref = torch.randn(B, T, value_dim, device=device, dtype=dtype, requires_grad=True)
+    r_ref = torch.randn(B, T, H, D, device=device, dtype=dtype, requires_grad=True)
+    k_ref = torch.randn(B, T, H, D, device=device, dtype=dtype, requires_grad=True)
+    r_k_ref = torch.randn(H, D, device=device, dtype=dtype, requires_grad=True)
+    v_ref = torch.randn(B, T, H, D, device=device, dtype=dtype, requires_grad=True)
+    g_ref = torch.randn(B, T, value_dim, device=device, dtype=dtype, requires_grad=True)
+
+    tensors_cus = [t.clone().detach().requires_grad_(True) for t in [o_ref, r_ref, k_ref, r_k_ref, v_ref, g_ref]]
+    o_cus, r_cus, k_cus, r_k_cus, v_cus, g_cus = tensors_cus
+
+    output_ref = gate_output_correction_ref(o_ref.float(), r_ref.float(), k_ref.float(),
+                                            r_k_ref.float(), v_ref.float(), g_ref.float())
+    output_ref.sum().backward()
+
+    output_cus = gate_output_correction(o_cus, r_cus, k_cus, r_k_cus, v_cus, g_cus)
+    output_cus.sum().backward()
+
+    assert_close(" o", output_ref, output_cus, 0.002)
+    assert_close("do", o_ref.grad, o_cus.grad, 0.002)
+    assert_close("dr", r_ref.grad, r_cus.grad, 0.002)
+    assert_close("dk", k_ref.grad, k_cus.grad, 0.002)
+    assert_close("drk", r_k_ref.grad, r_k_cus.grad, 0.002)
+    assert_close("dv", v_ref.grad, v_cus.grad, 0.002)
+    assert_close("dg", g_ref.grad, g_cus.grad, 0.002)

code/flash-linear-attention/tests/ops/test_simple_gla.py

@@ -0,0 +1,670 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.simple_gla.chunk import chunk_simple_gla
+from fla.ops.simple_gla.fused_chunk import fused_chunk_simple_gla
+from fla.ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla
+from fla.ops.simple_gla.naive import naive_parallel_simple_gla, naive_recurrent_simple_gla
+from fla.ops.simple_gla.parallel import parallel_simple_gla
+from fla.utils import assert_close, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, 1, torch.float),
+            (2, 500, 4, 60, 1, 1, torch.float),
+            (2, 1024, 8, 128, 1, 0.1, torch.float),
+            (2, 1024, 8, 128, 0.1, 1, torch.float),
+            (2, 1024, 8, 128, 1, 10, torch.float),
+            (4, 2048, 8, 64, 0.1, 1, torch.float),
+            (2, 1024, 8, 128, 1, 0.1, torch.float16),
+            (2, 1024, 8, 128, 1, 10, torch.float16),
+        ]
+    ],
+)
+def test_fused_recurrent(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = torch.randn((B, T, H), dtype=dtype, device=device)
+    g = (F.logsigmoid(g) / gate_logit_normalizer).requires_grad_()
+    h0 = torch.randn(B, H, D, D, device=device).requires_grad_()
+    dht = torch.randn_like(h0)
+    do = torch.randn_like(v)
+    ref, ref_ht = naive_recurrent_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        scale=scale,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_recurrent_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        scale=scale,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005, err_atol=2e-4)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'scale', 'gate_logit_normalizer', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-scale{}-gate_logit_normalizer{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, 1, 1, [0, 15], torch.float),
+            (4, 64, 1, 1, [0, 256, 500, 1000], torch.float),
+            (4, 100, 0.1, 1, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 100, 1, 1, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 100, 1, 10, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 64, 1, 1, [0, 1, 100, 300, 1200, 2048], torch.float16),
+            (4, 128, 1, 1, [0, 200, 512, 1200, 2048], torch.float16),
+        ]
+    ],
+)
+def test_fused_recurrent_varlen(
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = torch.randn((1, T, H), dtype=dtype, device=device)
+    g = (F.logsigmoid(g) / gate_logit_normalizer).requires_grad_()
+    h0 = torch.randn(N, H, D, D, device=device).requires_grad_()
+    dht = torch.randn_like(h0)
+    do = torch.randn_like(v)
+
+    refs, ref_hts = [], []
+    for i, (bos, eos) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)):
+        ref, ref_ht = naive_recurrent_simple_gla(
+            q=q[:, bos:eos],
+            k=k[:, bos:eos],
+            v=v[:, bos:eos],
+            g=g[:, bos:eos],
+            scale=scale,
+            initial_state=h0[i],
+            output_final_state=True,
+        )
+        refs.append(ref)
+        ref_hts.append(ref_ht)
+    ref = torch.cat(refs, 1)
+    ref_ht = torch.cat(ref_hts, 0)
+    ((ref * do).sum() + (ref_ht * dht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_recurrent_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        scale=scale,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.005)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005, err_atol=2e-4)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, 1, torch.float16),
+            (2, 500, 3, 60, 1, 1, torch.float16),
+            (1, 1000, 4, 128, 1, 0.1, torch.float16),
+            (2, 1000, 4, 128, 0.1, 1, torch.float16),
+            (3, 1000, 4, 128, 0.1, 10, torch.float16),
+            (4, 2048, 8, 64, 0.1, 1, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    g = torch.randn((B, T, H), dtype=torch.float32, device=device)
+    h0 = torch.rand((B, H, D, D), dtype=torch.float32, device=device).requires_grad_(True)
+    dht = torch.randn_like(h0)
+    g = (F.logsigmoid(g) / gate_logit_normalizer).requires_grad_(True)
+    do = torch.randn_like(v)
+
+    ref, ref_ht = fused_recurrent_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        scale=scale,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((ref * do).sum() + (dht * ref_ht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = chunk_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        scale=scale,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((tri * do).sum() + (dht * tri_ht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, [0, 15], torch.float16),
+            (4, 64, [0, 256, 500, 1000], torch.float16),
+            (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = F.logsigmoid(torch.randn((1, T, H), dtype=dtype, device=device)).requires_grad_()
+    h0 = torch.randn((N, H, D, D), dtype=torch.float32, device=device).requires_grad_()
+    dht = torch.randn_like(h0)
+    do = torch.randn_like(v)
+
+    ref, ref_ht = fused_recurrent_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((ref * do).sum() + (dht * ref_ht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = chunk_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (dht * tri_ht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, 1, torch.float16),
+            (2, 500, 3, 60, 1, 1, torch.float16),
+            (1, 1000, 4, 128, 1, 0.1, torch.float16),
+            (2, 1000, 4, 128, 0.1, 1, torch.float16),
+            (3, 1000, 4, 128, 0.1, 10, torch.float16),
+            (4, 2048, 8, 64, 0.1, 1, torch.float16),
+        ]
+    ],
+)
+def test_fused_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+    scale: float,
+    gate_logit_normalizer: float,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    g = torch.randn((B, T, H), dtype=torch.float32, device=device)
+    h0 = torch.rand((B, H, D, D), dtype=torch.float32, device=device).requires_grad_(True)
+    dht = torch.randn_like(h0)
+    g = (F.logsigmoid(g) / gate_logit_normalizer).requires_grad_(True)
+    do = torch.randn_like(v)
+
+    ref, ref_ht = fused_recurrent_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        scale=scale,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((ref * do).sum() + (dht * ref_ht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_chunk_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        scale=scale,
+        initial_state=h0,
+        output_final_state=True,
+    )
+    ((tri * do).sum() + (dht * tri_ht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, [0, 15], torch.float16),
+            (4, 64, [0, 256, 500, 1000], torch.float16),
+            (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_fused_chunk_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    N = len(cu_seqlens) - 1
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = F.logsigmoid(torch.randn((1, T, H), dtype=dtype, device=device)).requires_grad_()
+    h0 = torch.randn((N, H, D, D), dtype=torch.float32, device=device).requires_grad_()
+    dht = torch.randn_like(h0)
+    do = torch.randn_like(v)
+
+    ref, ref_ht = fused_recurrent_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((ref * do).sum() + (dht * ref_ht).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+    ref_dh0, h0.grad = h0.grad.clone(), None
+
+    tri, tri_ht = fused_chunk_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum() + (dht * tri_ht).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+    tri_dh0, h0.grad = h0.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('ht', ref_ht, tri_ht, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+    assert_close('dh0', ref_dh0, tri_dh0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, 1, torch.float16),
+            (2, 500, 3, 60, 1, 1, torch.float16),
+            (2, 1024, 4, 128, 0.1, 1, torch.float16),
+            (3, 1024, 4, 128, 0.1, 10, torch.float16),
+            (3, 1024, 4, 256, 0.1, 0.1, torch.float16),
+            (4, 2048, 4, 64, 0.1, 0.1, torch.float16),
+        ]
+    ],
+)
+def test_parallel(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    gate_logit_normalizer: float,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    USE_G = gate_logit_normalizer > 0
+    q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True)
+    g = F.logsigmoid(torch.randn((B, T, H), dtype=dtype, device=device)) if USE_G else None
+    g = (g / gate_logit_normalizer).requires_grad_(True) if USE_G else None
+    do = torch.randn_like(v)
+
+    ref, _ = fused_recurrent_simple_gla(q=q, k=k, v=v, g=g, scale=scale, output_final_state=True)
+    _, ref_A = naive_parallel_simple_gla(q=q, k=k, v=v, g=g, scale=scale)
+    ref.backward(do)
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    if USE_G:
+        ref_dg, g.grad = g.grad.clone(), None
+
+    tri, tri_A = parallel_simple_gla(q=q, k=k, v=v, g=g, scale=scale, output_attentions=True)
+    tri.backward(do)
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    if USE_G:
+        tri_dg, g.grad = g.grad.clone(), None
+    assert_close('o', ref, tri, 0.005)
+    assert_close('A', ref_A, tri_A, 0.005)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    if USE_G:
+        assert_close('dg', ref_dg, tri_dg, 0.015)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (4, 64, [0, 15], torch.float16),
+            (4, 64, [0, 256, 500, 1000], torch.float16),
+            (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_parallel_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_()
+    g = F.logsigmoid(torch.randn((1, T, H), dtype=dtype, device=device)).requires_grad_()
+    do = torch.randn_like(v)
+
+    ref, _ = fused_recurrent_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        output_final_state=False,
+        cu_seqlens=cu_seqlens,
+    )
+    ((ref * do).sum()).backward()
+    ref_dq, q.grad = q.grad.clone(), None
+    ref_dk, k.grad = k.grad.clone(), None
+    ref_dv, v.grad = v.grad.clone(), None
+    ref_dg, g.grad = g.grad.clone(), None
+
+    tri, _ = parallel_simple_gla(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        cu_seqlens=cu_seqlens,
+    )
+    ((tri * do).sum()).backward()
+    tri_dq, q.grad = q.grad.clone(), None
+    tri_dk, k.grad = k.grad.clone(), None
+    tri_dv, v.grad = v.grad.clone(), None
+    tri_dg, g.grad = g.grad.clone(), None
+
+    assert_close('o', ref, tri, 0.004)
+    assert_close('dq', ref_dq, tri_dq, 0.005)
+    assert_close('dk', ref_dk, tri_dk, 0.005)
+    assert_close('dv', ref_dv, tri_dv, 0.005)
+    assert_close('dg', ref_dg, tri_dg, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('vary_A', 'dtype'),
+    [
+        pytest.param(True, torch.float, id=f'vary_A{True}-dtype{torch.float}'),
+        pytest.param(False, torch.float, id=f'vary_A{False}-dtype{torch.float}'),
+        pytest.param(True, torch.float16, id=f'vary_A{True}-dtype{torch.float16}'),
+        pytest.param(False, torch.float16, id=f'vary_A{False}-dtype{torch.float16}'),
+    ],
+)
+def test_simple_gla_to_mamba2(vary_A, dtype):
+    try:
+        from mamba_ssm.modules.ssd_minimal import ssd_minimal_discrete
+        from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
+    except ImportError:
+        pytest.skip('mamba_ssm is not installed.')
+    torch.manual_seed(42)
+
+    # Dimensions, Denoted (B, T, Q, D, P) in Mamba2 paper
+    batch, seq_len, chunk_size, dim, headdim = 2, 512, 8, 64, 16
+    n_heads = dim // headdim  # (H) in the paper
+    ngroups = n_heads  # (G) in the paper; NOTE: do not use group-query here
+    dstate = 64  # (N) in the paper
+    atol = 5e-4 if dtype == torch.float else 1e-2
+
+    x = 0.1 * torch.randn(batch, seq_len, n_heads, headdim, dtype=dtype, device=device)
+    dt = torch.ones(batch, seq_len, n_heads, dtype=dtype, device=device)  # dt=1 can be ignored
+
+    if vary_A:
+        A = -0.1 * torch.rand(1, seq_len, n_heads, dtype=dtype, device=device)
+    else:  # constant A for all position
+        A = -0.1 * torch.rand(n_heads, dtype=dtype, device=device)
+
+    B = 0.1 * torch.randn(batch, seq_len, ngroups, dstate, dtype=dtype, device=device)
+    C = 0.1 * torch.randn(batch, seq_len, ngroups, dstate, dtype=dtype, device=device)
+
+    y_ssd, final_ssd = ssd_minimal_discrete(x * dt.unsqueeze(-1), A * dt, B, C, chunk_size)
+
+    if not vary_A:
+        # NOTE: fused kernel does not support varying A with time
+        y_fuse, final_fuse = mamba_chunk_scan_combined(x, dt, A, B, C, chunk_size, D=None, return_final_states=True)
+        assert y_ssd.allclose(y_fuse, 0, atol), f'y diff: {torch.abs(y_ssd - y_fuse).max()}'
+        # fused kernel upcasts state to float32
+        # https://github.com/state-spaces/mamba/blob/v2.2.2/mamba_ssm/ops/triton/ssd_combined.py#L650
+        final_fuse = final_fuse.to(dtype)
+        assert final_ssd.allclose(final_fuse, 0, atol), f'final diff: {torch.abs(final_ssd - final_fuse).max()}'
+
+    # mapping inputs Mamba2 -> FLA
+    # FLA Now use head_first = False, therefore there is no need to transpose inputs
+    q = C
+    k = B
+    v = x
+    g = (A * dt)
+
+    # mapping outputs Mamba2 -> FLA
+    y_rearrange = y_ssd
+    final_rearrange = final_ssd.transpose(2, 3)
+
+    # comparing output results between FLA kernel and Mamba2 kernel
+    # final_gla_fuse :[N, H, K, V]
+    outputs_gla_fuse, final_gla_fuse = chunk_simple_gla(q, k, v, g, scale=1.0, output_final_state=True)
+    assert y_rearrange.allclose(outputs_gla_fuse, 0, atol), f'y diff: {torch.abs(y_rearrange - outputs_gla_fuse).max()}'
+    final_gla_fuse = final_gla_fuse.to(dtype)  # states hard-coded to float32 in FLA kernel
+    assert final_rearrange.allclose(final_gla_fuse, 0, atol), f'final diff: {torch.abs(final_ssd - final_gla_fuse).max()}'

code/flash-linear-attention/tests/ops/test_solve_tril.py

@@ -0,0 +1,91 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
+from fla.ops.utils.solve_tril import solve_tril
+from fla.utils import assert_close, device, device_platform
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'chunk_size'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-chunk_size{}".format(*test))
+        for test in [
+            (1, 63, 1, 16),
+            (2, 500, 4, 32),
+            (2, 1000, 5, 64),
+            (3, 1024, 6, 64),
+            (4, 2048, 8, 64),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Pytorch Failure',
+)
+def test_solve_tril(B, T, H, chunk_size):
+    # do not randomly intiialize A otherwise the inverse is not stable
+    k = F.normalize(torch.randn((B, H, T, 64), dtype=torch.float32, device=device), dim=-1)
+    # Pad the second-to-last dimension (T) to be a multiple of chunk_size
+    padding_size = (chunk_size - T % chunk_size) % chunk_size
+    k_padded = F.pad(k, (0, 0, 0, padding_size, 0, 0, 0, 0))
+    k_padded = k_padded.reshape(B, H, -1, chunk_size, 64)
+    A = (k_padded @ k_padded.transpose(-1, -2)).tril(-1)
+
+    ref = torch.inverse(A + torch.eye(A.shape[-1], device=A.device)[None, None, None, ...])
+    ref = ref.reshape(B, H, -1, chunk_size)[:, :, :T, :]
+
+    tri = solve_tril(A.reshape(B, H, -1, chunk_size)[:, :, :T, :].transpose(1, 2)).transpose(1, 2)
+
+    assert_close('solve_tril', ref, tri, 0.0001)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'chunk_size', 'cu_seqlens'),
+    [
+        pytest.param(*test, id="H{}-D{}-chunk_size{}-cu_seqlens{}".format(*test))
+        for test in [
+            (4, 64, 16, [0, 15]),
+            (4, 64, 32, [0, 256, 500, 1000]),
+            (4, 100, 64, [0, 15, 100, 300, 1200, 2000]),
+            (4, 64, 16, [0, 1, 100, 300, 1200, 2048]),
+            (4, 128, 32, [0, 200, 512, 1200, 2048]),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+@pytest.mark.skipif(
+    device_platform == 'intel',
+    reason='Intel Pytorch Failure',
+)
+def test_solve_tril_varlen(
+    H: int,
+    D: int,
+    chunk_size: int,
+    cu_seqlens: list[int],
+):
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+    # Construct the input. otherwise inverse's condition number might be too large to measure the error
+    k = F.normalize(torch.randn((1, T, H, D), dtype=torch.bfloat16, device=device), dim=-1)
+    beta = torch.randn((1, T, H), dtype=torch.bfloat16, device=device).sigmoid()
+    A = chunk_scaled_dot_kkt_fwd(k=k, beta=beta, cu_seqlens=cu_seqlens, chunk_size=chunk_size)
+
+    ref = torch.zeros_like(A)
+    for i in range(len(cu_seqlens) - 1):
+        for j in range(cu_seqlens[i], cu_seqlens[i+1], chunk_size):
+            actual_size = min(chunk_size, cu_seqlens[i+1] - j)
+            ref[:, j:j+actual_size, :, :actual_size] = torch.inverse(
+                A[:, j:j+actual_size, :, :actual_size].transpose(1, 2) +
+                torch.eye(actual_size, device=A.device, dtype=A.dtype)[None, None, ...],
+            ).transpose(1, 2)
+
+    tri = solve_tril(A, cu_seqlens=cu_seqlens)
+    assert_close('solve_tril_varlen', ref, tri, 0.0001)

code/flash-linear-attention/tests/ops/test_titans.py

@@ -0,0 +1,122 @@
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.titans.naive import chunk_titans_linear_ref
+from fla.utils import assert_close, device
+
+
+def initialize_chunked_param(B, H, T, BT, dtype=torch.float32):
+    # Calculate number of complete chunks and remaining elements
+    num_complete_chunks = T // BT
+    remainder = T % BT
+
+    # Initialize for complete chunks
+    if num_complete_chunks > 0:
+        theta_chunks = torch.rand(B, H, num_complete_chunks, 1, dtype=dtype)
+        theta_main = theta_chunks.repeat_interleave(
+            BT, dim=2,
+        )  # Shape: (B, H, num_complete_chunks*BT, 1)
+    else:
+        theta_main = torch.empty(B, H, 0, 1, dtype=dtype)
+
+    # Handle remaining elements if any
+    if remainder > 0:
+        theta_remainder = torch.rand(B, H, 1, 1, dtype=dtype)
+        theta_remainder = theta_remainder.repeat_interleave(
+            remainder, dim=2,
+        )  # Shape: (B, H, remainder, 1)
+
+        # Concatenate main chunks with remainder
+        theta = torch.cat([theta_main, theta_remainder], dim=2)
+    else:
+        theta = theta_main
+
+    return theta
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, torch.float16),
+            (2, 100, 4, 60, torch.float16),
+            (2, 1024, 3, 128, torch.float16),
+            (3, 2000, 4, 128, torch.float16),
+            (4, 2048, 8, 64, torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    True, reason='FIXME',
+)
+def test_naive_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    BT = 64
+    # set seed
+    torch.manual_seed(1)
+    # we don't use such initialization in the original code
+    # theta = initialize_chunked_param(B, H, T, BT, dtype)
+    # alpha = initialize_chunked_param(B, H, T, BT, dtype)
+    # eta = initialize_chunked_param(B, H, T, BT, dtype)
+    theta = torch.rand(B, H, T, 1, dtype=dtype)
+    alpha = torch.rand(B, H, T, 1, dtype=dtype)
+    eta = torch.rand(B, H, T, 1, dtype=dtype)
+
+    # titans normalize queries and keys using ℓ2-normalization
+    q = F.normalize(torch.randn(B, H, T, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    k = F.normalize(torch.randn(B, H, T, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn(B, H, T, D, dtype=dtype)
+    w = torch.randn(H, D, dtype=dtype)
+    b = torch.randn(H, D, dtype=dtype)
+    h0 = torch.randn(B, H, D, D, dtype=torch.float32)
+    q = q.permute(0, 2, 1, 3)
+    k = k.permute(0, 2, 1, 3)
+    v = v.permute(0, 2, 1, 3)
+    theta = theta.permute(0, 2, 1, 3)
+    alpha = alpha.permute(0, 2, 1, 3)
+    eta = eta.permute(0, 2, 1, 3)
+    q, k, v, w, b, theta, alpha, eta = map(
+        lambda x: x.to(device).requires_grad_(False), (q, k, v, w, b, theta, alpha, eta),
+    )
+    # in titans paper, h0 is not learnable
+    h0 = h0.to(device)
+
+    ref_naive, ref_ht_naive = chunk_titans_linear_ref(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        b.clone(),
+        theta.clone(),
+        alpha.clone(),
+        eta.clone(),
+        output_final_state=True,
+        chunk_size=BT,
+        initial_state=h0.clone(),
+        use_chunk=False,
+    )
+    ref, ref_ht = chunk_titans_linear_ref(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        b.clone(),
+        theta.clone(),
+        alpha.clone(),
+        eta.clone(),
+        output_final_state=True,
+        chunk_size=BT,
+        initial_state=h0.clone(),
+        use_chunk=True,
+    )
+
+    assert_close(" o", ref, ref_naive, 0.006)
+    assert_close("ht", ref_ht, ref_ht_naive, 0.005)

code/flash-linear-attention/tests/ops/test_ttt.py

@@ -0,0 +1,268 @@
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fla.ops.ttt import chunk_ttt_linear, fused_chunk_ttt_linear
+from fla.ops.ttt.naive import chunk_ttt_linear_ref
+from fla.utils import assert_close, check_shared_mem, device
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float16),
+            (2, 100, 4, 60, 0.1, torch.float16),
+            (2, 1024, 3, 128, 0.1, torch.float16),
+            (2, 1024, 4, 128, 1, torch.float16),
+            (3, 2000, 4, 128, 0.1, torch.float16),
+            (4, 2048, 8, 64, 0.1, torch.float16),
+        ]
+    ],
+)
+def test_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    dtype: torch.dtype,
+):
+    if D > 64 and check_shared_mem('hopper') is False:
+        pytest.skip(reason="Current CI do not support this config")
+    if T > 1000:
+        pytest.skip(reason="Current CI do not support this config")
+    eta_base = 5e-3
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    w = torch.randn(H, D, dtype=dtype)
+    b = torch.randn(H, D, dtype=dtype)
+    eta = torch.randn(B, T, H, 1, dtype=dtype) * eta_base
+    h0 = torch.randn(B, H, D, D, dtype=torch.float32)
+    hb0 = torch.randn(B, H, 1, D, dtype=torch.float32)
+
+    q, k, v, w, b, eta, h0, hb0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, w, b, eta, h0, hb0))
+    do = torch.rand_like(v)
+    dht = torch.rand_like(h0)
+    dhbt = torch.rand_like(hb0)
+
+    tri, tri_ht, tri_hbt = chunk_ttt_linear(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        b.clone(),
+        eta.clone(),
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+        initial_state_bias=hb0.clone(),
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum() + (tri_hbt * dhbt).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dw, tri_db, tri_deta, \
+        tri_dh0, tri_dhb0 = q.grad, k.grad, v.grad, w.grad, b.grad, eta.grad, h0.grad, hb0.grad
+    q.grad = k.grad = v.grad = w.grad = b.grad = eta.grad = h0.grad = hb0.grad = None
+
+    ref, ref_ht, ref_hbt = chunk_ttt_linear_ref(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        b.clone(),
+        eta.clone(),
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+        initial_state_bias=hb0.clone(),
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum() + (ref_hbt * dhbt).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dw, ref_db, ref_deta, \
+        ref_dh0, ref_dhb0 = q.grad, k.grad, v.grad, w.grad, b.grad, eta.grad, h0.grad, hb0.grad
+
+    assert_close("   o", ref, tri, 0.005)
+    assert_close("  ht", ref_ht, tri_ht, 0.005)
+    assert_close(" hbt", ref_hbt, tri_hbt, 0.005)
+    assert_close("  dq", ref_dq, tri_dq, 0.005)
+    assert_close("  dk", ref_dk, tri_dk, 0.010)
+    assert_close("  dv", ref_dv, tri_dv, 0.007)
+    assert_close("  dw", ref_dw, tri_dw, 0.006)
+    assert_close("  db", ref_db, tri_db, 0.006)
+    assert_close("  de", ref_deta, tri_deta, 0.030)  # because the last element of the chunk
+    assert_close(" de0", ref_deta[:, :14, :, :], tri_deta[:, :14, :, :], 0.010)
+    assert_close(" dh0", ref_dh0, tri_dh0, 0.007)
+    assert_close("dhb0", ref_dhb0, tri_dhb0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'scale', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 64, 1, torch.float16),
+            (2, 100, 4, 60, 0.1, torch.float16),
+            (2, 1024, 3, 128, 0.1, torch.float16),
+            (2, 1024, 4, 128, 1, torch.float16),
+            (3, 2000, 4, 128, 0.1, torch.float16),
+            (4, 2048, 8, 64, 0.1, torch.float16),
+        ]
+    ],
+)
+def test_fused_chunk(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    scale: float,
+    dtype: torch.dtype,
+):
+    if D > 64 and check_shared_mem('hopper') is False:
+        pytest.skip(reason="Current CI do not support this config")
+    if T > 1000:
+        pytest.skip(reason="Current CI do not support this config")
+    eta_base = 5e-3
+    q = torch.randn(B, T, H, D, dtype=dtype)
+    k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn(B, T, H, D, dtype=dtype)
+    w = torch.randn(H, D, dtype=dtype)
+    b = torch.randn(H, D, dtype=dtype)
+    eta = torch.randn(B, T, H, 1, dtype=dtype) * eta_base
+    h0 = torch.randn(B, H, D, D, dtype=torch.float32)
+    hb0 = torch.randn(B, H, 1, D, dtype=torch.float32)
+
+    q, k, v, w, b, eta, h0, hb0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, w, b, eta, h0, hb0))
+    do = torch.rand_like(v)
+    dht = torch.rand_like(h0)
+    dhbt = torch.rand_like(hb0)
+
+    tri, tri_ht, tri_hbt = fused_chunk_ttt_linear(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        b.clone(),
+        eta.clone(),
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+        initial_state_bias=hb0.clone(),
+    )
+    ((tri * do).sum() + (tri_ht * dht).sum() + (tri_hbt * dhbt).sum()).backward(retain_graph=True)
+    tri_dq, tri_dk, tri_dv, tri_dw, tri_db, tri_deta, \
+        tri_dh0, tri_dhb0 = q.grad, k.grad, v.grad, w.grad, b.grad, eta.grad, h0.grad, hb0.grad
+    q.grad = k.grad = v.grad = w.grad = b.grad = eta.grad = h0.grad = hb0.grad = None
+
+    ref, ref_ht, ref_hbt = chunk_ttt_linear_ref(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        b.clone(),
+        eta.clone(),
+        scale=scale,
+        output_final_state=True,
+        initial_state=h0.clone(),
+        initial_state_bias=hb0.clone(),
+    )
+    ((ref * do).sum() + (ref_ht * dht).sum() + (ref_hbt * dhbt).sum()).backward(retain_graph=True)
+    ref_dq, ref_dk, ref_dv, ref_dw, ref_db, ref_deta, \
+        ref_dh0, ref_dhb0 = q.grad, k.grad, v.grad, w.grad, b.grad, eta.grad, h0.grad, hb0.grad
+
+    assert_close("   o", ref, tri, 0.005)
+    assert_close("  ht", ref_ht, tri_ht, 0.005)
+    assert_close(" hbt", ref_hbt, tri_hbt, 0.005)
+    assert_close("  dq", ref_dq, tri_dq, 0.005)
+    assert_close("  dk", ref_dk, tri_dk, 0.010)
+    assert_close("  dv", ref_dv, tri_dv, 0.007)
+    assert_close("  dw", ref_dw, tri_dw, 0.005)
+    assert_close("  db", ref_db, tri_db, 0.005)
+    assert_close("  de", ref_deta, tri_deta, 0.03)  # because the last element of the chunk
+    assert_close(" de0", ref_deta[:, :14, :, :], tri_deta[:, :14, :, :], 0.008)
+    assert_close(" dh0", ref_dh0, tri_dh0, 0.006)
+    assert_close("dhb0", ref_dhb0, tri_dhb0, 0.005)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 64, [0, 15], torch.float16),
+            (3, 60, [0, 111, 500], torch.float16),
+            (3, 64, [0, 256, 500, 900, 1000], torch.float16),
+            (4, 100, [0, 15, 100, 300, 1200, 1599, 1800, 2000], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1",
+    reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set",
+)
+def test_chunk_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    if D > 64 and check_shared_mem('hopper') is False:
+        pytest.skip(reason="Current CI do not support this config")
+    torch.manual_seed(42)
+    os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    eta_base = 5e-3
+    # seq-first required for inputs with variable lengths
+    q = torch.randn((1, T, H, D), dtype=dtype)
+    k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
+    v = torch.randn((1, T, H, D), dtype=dtype)
+    eta = torch.randn(1, T, H, 1, dtype=dtype) * eta_base
+    w = torch.randn(H, D, dtype=dtype)
+    b = torch.randn(H, D, dtype=dtype)
+    h0 = torch.randn((N, H, D, D), dtype=torch.float32)
+    hb0 = torch.randn((N, H, 1, D), dtype=torch.float32)
+    q, k, v, w, b, eta, h0, hb0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, w, b, eta, h0, hb0))
+
+    tri, tri_ht, tri_hbt = chunk_ttt_linear(
+        q.clone(),
+        k.clone(),
+        v.clone(),
+        w.clone(),
+        b.clone(),
+        eta.clone(),
+        output_final_state=True,
+        initial_state=h0.clone(),
+        initial_state_bias=hb0.clone(),
+        cu_seqlens=cu_seqlens,
+    )
+
+    ref = []
+    ref_ht = []
+    ref_hbt = []
+    for i in range(N):
+        ref_i, ref_ht_i, ref_hbt_i = chunk_ttt_linear_ref(
+            q=q[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            k=k[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            v=v[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            w=w,
+            b=b,
+            eta=eta[:, cu_seqlens[i]:cu_seqlens[i+1]],
+            initial_state=h0[i],
+            initial_state_bias=hb0[i],
+            output_final_state=True,
+        )
+        ref.append(ref_i)
+        ref_ht.append(ref_ht_i)
+        ref_hbt.append(ref_hbt_i)
+    ref = torch.cat(ref, 1)
+    ref_ht = torch.cat(ref_ht, 0)
+    ref_hbt = torch.cat(ref_hbt, 0)
+
+    assert_close("  o", ref, tri, 0.005)
+    assert_close(" ht", ref_ht, tri_ht, 0.005)
+    assert_close("hbt", ref_hbt, tri_hbt, 0.005)

code/flash-linear-attention/tests/ops/test_utils.py

@@ -0,0 +1,410 @@
+
+import os
+
+import pytest
+import torch
+
+from fla.ops.utils import chunk_global_cumsum, chunk_local_cumsum, mean_pooling
+from fla.ops.utils.index import prepare_lens
+from fla.ops.utils.pack import pack_sequence, unpack_sequence
+from fla.utils import assert_close, device
+
+
+def reversed_cumsum(x, dim=-1):
+    dtype = x.dtype
+    x = x.float()
+    c = x.cumsum(dim)
+    y = x + c.index_select(dim, x.new_tensor([c.shape[dim]-1], dtype=torch.long)) - c
+    return y.to(dtype)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 30, torch.float),
+            (2, 500, 4, 60, torch.float),
+            (2, 1000, 5, 128, torch.float),
+            (3, 1024, 6, 500, torch.float),
+            (4, 2048, 8, 1024, torch.float),
+        ]
+    ],
+)
+def test_global_cumsum(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    s = torch.randn(B, T, H, dtype=dtype).to(device)
+    ref = s.float().cumsum(1).to(dtype)
+    tri = chunk_global_cumsum(s)
+    assert_close('global_cumsum', ref, tri, 1e-3)
+
+    s = torch.randn(B, T, H, D, dtype=dtype).to(device)
+    ref = s.float().cumsum(1).to(dtype)
+    tri = chunk_global_cumsum(s)
+    assert_close('global_cumsum', ref, tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 60, [0, 15], torch.float),
+            (3, 100, [0, 256, 500, 1000], torch.float),
+            (4, 256, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 500, [0, 1, 100, 300, 1200, 2048], torch.float16),
+            (2, 1024, [0, 200, 512, 1200, 2048], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_global_cumsum_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    s = torch.randn(1, T, H, dtype=dtype).to(device)
+    ref = torch.cat([s[:, start:end].float().cumsum(1) for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)], 1).to(dtype)
+    tri = chunk_global_cumsum(s, cu_seqlens=cu_seqlens)
+    assert_close('global_cumsum', ref, tri, 1e-3)
+
+    s = torch.randn(1, T, H, D, dtype=dtype).to(device)
+    ref = torch.cat([s[:, start:end].float().cumsum(1) for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)], 1).to(dtype)
+    tri = chunk_global_cumsum(s, cu_seqlens=cu_seqlens)
+    assert_close('global_cumsum', ref, tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 30, torch.float),
+            (2, 500, 4, 60, torch.float),
+            (2, 1000, 5, 128, torch.float),
+            (3, 1024, 6, 500, torch.float),
+            (4, 2048, 8, 1024, torch.float),
+        ]
+    ],
+)
+def test_global_reversed_cumsum(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    s = torch.randn(B, T, H, dtype=dtype).to(device)
+    ref = reversed_cumsum(s, dim=(1)).to(dtype)
+    tri = chunk_global_cumsum(s, reverse=True)
+    assert_close('global_cumsum', ref, tri, 1e-3)
+
+    s = torch.randn(B, T, H, D, dtype=dtype).to(device)
+    ref = reversed_cumsum(s, dim=(1)).to(dtype)
+    tri = chunk_global_cumsum(s, reverse=True)
+    assert_close('global_cumsum', ref, tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('H', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 60, [0, 15], torch.float),
+            (3, 100, [0, 256, 500, 1000], torch.float),
+            (4, 256, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 500, [0, 1, 100, 300, 1200, 2048], torch.float16),
+            (2, 1024, [0, 200, 512, 1200, 2048], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_global_reversed_cumsum_varlen(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    s = torch.randn(1, T, H, dtype=dtype).to(device)
+    ref = torch.cat([reversed_cumsum(s[:, start:end], 1) for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)], 1).to(dtype)
+    tri = chunk_global_cumsum(s, reverse=True, cu_seqlens=cu_seqlens)
+    assert_close('global_reversed_cumsum', ref, tri, 1e-3)
+
+    s = torch.randn(1, T, H, D, dtype=dtype).to(device)
+    ref = torch.cat([reversed_cumsum(s[:, start:end], 1) for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)], 1).to(dtype)
+    tri = chunk_global_cumsum(s, reverse=True, cu_seqlens=cu_seqlens)
+    assert_close('global_reversed_cumsum', ref, tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'C', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-C{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 16, 30, torch.float),
+            (2, 500, 4, 32, 60, torch.float),
+            (2, 1000, 5, 64, 128, torch.float),
+            (3, 1024, 6, 64, 500, torch.float),
+            (4, 2048, 8, 128, 1024, torch.float),
+        ]
+    ],
+)
+def test_local_cumsum(
+    B: int,
+    T: int,
+    H: int,
+    C: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    s = torch.randn(B, T, H, dtype=dtype).to(device)
+    ref = torch.cat([s[:, i:i+C, :].float().cumsum(1) for i in range(0, T, C)], 1)
+    tri = chunk_local_cumsum(s, chunk_size=C)
+    assert_close('local_cumsum', ref, tri, 1e-3)
+
+    s = torch.randn(B, T, H, D, dtype=dtype).to(device)
+    ref = torch.cat([s[:, i:i+C, :].float().cumsum(1) for i in range(0, T, C)], 1)
+    tri = chunk_local_cumsum(s, chunk_size=C)
+    assert_close('local_cumsum', ref, tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('H', 'C', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-C{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 32, 60, [0, 15], torch.float),
+            (3, 64, 100, [0, 256, 500, 1000], torch.float),
+            (4, 64, 256, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 128, 500, [0, 1, 100, 300, 1200, 2048], torch.float16),
+            (2, 128, 1024, [0, 200, 512, 1200, 2048], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_local_cumsum_varlen(
+    H: int,
+    C: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    s = torch.randn(1, T, H, dtype=dtype).to(device)
+    ref = torch.cat([
+        torch.cat([s[:, i:min(end, i+C), :].float().cumsum(1) for i in range(start, end, C)], 1)
+        for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)
+    ], 1)
+    tri = chunk_local_cumsum(s, chunk_size=C, cu_seqlens=cu_seqlens)
+    assert_close('local_cumsum', ref, tri, 1e-3)
+
+    s = torch.randn(1, T, H, D, dtype=dtype).to(device)
+    ref = torch.cat([
+        torch.cat([s[:, i:min(end, i+C), :].float().cumsum(1) for i in range(start, end, C)], 1)
+        for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)
+    ], 1)
+    tri = chunk_local_cumsum(s, chunk_size=C, cu_seqlens=cu_seqlens)
+    assert_close('local_cumsum', ref, tri, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'C', 'D', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-C{}-D{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 16, 30, torch.float),
+            (2, 500, 4, 32, 60, torch.float),
+            (2, 1000, 5, 64, 128, torch.float),
+            (3, 1024, 6, 64, 500, torch.float),
+            (4, 2048, 8, 128, 1024, torch.float),
+        ]
+    ],
+)
+def test_mean_pooling(
+    B: int,
+    T: int,
+    H: int,
+    C: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    x = torch.randn(B, T, H, D, dtype=dtype).to(device)
+    x.requires_grad = True
+    ref = torch.cat([x[:, i:i+C, :].float().mean(1, True) for i in range(0, T, C)], 1).to(dtype)
+    do = torch.randn_like(ref)
+    ref.backward(do)
+    ref_dx, x.grad = x.grad.clone(), None
+
+    tri = mean_pooling(x, chunk_size=C)
+    tri.backward(do)
+    tri_dx, x.grad = x.grad.clone(), None
+
+    assert_close('mean_pooling', ref, tri, 1e-3)
+    assert_close('mean_pooling', ref_dx, tri_dx, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('H', 'C', 'D', 'cu_seqlens', 'dtype'),
+    [
+        pytest.param(*test, id="H{}-C{}-D{}-cu_seqlens{}-{}".format(*test))
+        for test in [
+            (2, 32, 60, [0, 15], torch.float),
+            (3, 64, 100, [0, 256, 500, 1000], torch.float),
+            (4, 64, 256, [0, 15, 100, 300, 1200, 2000], torch.float),
+            (4, 128, 500, [0, 1, 100, 300, 1200, 2048], torch.float16),
+            (2, 128, 1024, [0, 200, 512, 1200, 2048], torch.float16),
+        ]
+    ],
+)
+@pytest.mark.skipif(
+    os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
+    reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
+)
+def test_mean_pooling_varlen(
+    H: int,
+    C: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    T = cu_seqlens[-1]
+    cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
+
+    x = torch.randn(1, T, H, D, dtype=dtype).to(device).requires_grad_(True)
+    ref = torch.cat([
+        torch.cat([x[:, i:min(end, i+C), :].float().mean(1, True) for i in range(start, end, C)], 1)
+        for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)
+    ], 1).to(dtype)
+    do = torch.randn_like(ref)
+    ref.backward(do)
+    ref_dx, x.grad = x.grad.clone(), None
+
+    tri = mean_pooling(x, chunk_size=C, cu_seqlens=cu_seqlens)
+    tri.backward(do)
+    tri_dx, x.grad = x.grad.clone(), None
+
+    torch.testing.assert_close(ref, tri.to(ref.dtype), rtol=1.6e-2, atol=3e-5)
+    torch.testing.assert_close(ref_dx, tri_dx.to(ref_dx.dtype), rtol=1.6e-2, atol=3e-5)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'padding_side', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-padding_side{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 30, 'left', torch.float),
+            (2, 500, 4, 60, 'right', torch.float),
+            (2, 1000, 5, 128, 'left', torch.float),
+            (3, 1024, 6, 500, 'right', torch.float),
+            (4, 2048, 8, 1024, 'left', torch.float),
+        ]
+    ],
+)
+def test_pack_sequence(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    padding_side: str,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    x = torch.randn(B, T, H, D, dtype=dtype).to(device).requires_grad_(True)
+    cu_seqlens = torch.cat(
+        [torch.tensor([0])]+[torch.randint(0, T, (1,)).clamp(min=1) for _ in range(B)],
+    ).cumsum(-1).to(device)
+    lens = prepare_lens(cu_seqlens)
+
+    if padding_side == 'left':
+        ref = torch.cat([x[i, -length:] for i, length in enumerate(lens.tolist())], 0)
+    else:
+        ref = torch.cat([x[i, :length] for i, length in enumerate(lens.tolist())], 0)
+    dy = torch.randn_like(ref)
+    ref.backward(dy)
+    ref_dx, x.grad = x.grad.clone(), None
+
+    tri = pack_sequence(x, cu_seqlens, padding_side=padding_side)
+    tri.backward(dy)
+    tri_dx, x.grad = x.grad.clone(), None
+
+    assert_close('y', ref, tri, 1e-3)
+    assert_close('dx', ref_dx, tri_dx, 1e-3)
+
+
+@pytest.mark.parametrize(
+    ('B', 'T', 'H', 'D', 'padding_side', 'dtype'),
+    [
+        pytest.param(*test, id="B{}-T{}-H{}-D{}-padding_side{}-{}".format(*test))
+        for test in [
+            (1, 63, 1, 30, 'left', torch.float),
+            (2, 500, 4, 60, 'right', torch.float),
+            (2, 1000, 5, 128, 'left', torch.float),
+            (3, 1024, 6, 500, 'right', torch.float),
+            (4, 2048, 8, 1024, 'left', torch.float),
+        ]
+    ],
+)
+def test_unpack_sequence(
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    padding_side: str,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(42)
+    cu_seqlens = torch.cat(
+        [torch.tensor([0])]+[torch.randint(0, T, (1,)).clamp(min=1) for _ in range(B)],
+    ).cumsum(-1).to(device)
+    lens = prepare_lens(cu_seqlens)
+    desired_shape = (B, lens.max().item() + torch.randint(0, 10, (1,)).item(), H, D)
+
+    x = torch.randn(cu_seqlens[-1].item(), H, D, dtype=dtype).to(device).requires_grad_(True)
+    ref = torch.zeros(desired_shape, device=device, dtype=dtype)
+    dy = torch.randn_like(ref)
+    for i, (bos, eos) in enumerate(zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist(), strict=False)):
+        length = eos - bos
+        if padding_side == 'left':
+            ref[i, -length:] = x[bos:eos]
+        else:
+            ref[i, :length] = x[bos:eos]
+    ref.backward(dy)
+    ref_dx, x.grad = x.grad.clone(), None
+
+    tri = unpack_sequence(x, cu_seqlens, padding_side=padding_side, desired_shape=desired_shape)
+    tri.backward(dy)
+    tri_dx, x.grad = x.grad.clone(), None
+
+    assert_close('y', ref, tri, 1e-3)
+    assert_close('dx', ref_dx, tri_dx, 1e-3)

code/flash-linear-attention/utils/convert_from_llama.py

@@ -0,0 +1,189 @@
+
+# scripts for converting pretrained hf model weights to fla style
+# calling the code to make conversions for mistralai/Mistral-7B-v0.1 would achieve the following results:
+# |    Tasks     |Version|Filter|n-shot|  Metric  |Value |   |Stderr|
+# |--------------|------:|------|-----:|----------|-----:|---|-----:|
+# |arc_challenge |      1|none  |     0|acc       |0.5043|±  |0.0146|
+# |              |       |none  |     0|acc_norm  |0.5392|±  |0.0146|
+# |arc_easy      |      1|none  |     0|acc       |0.8081|±  |0.0081|
+# |              |       |none  |     0|acc_norm  |0.7946|±  |0.0083|
+# |boolq         |      2|none  |     0|acc       |0.8373|±  |0.0065|
+# |copa          |      1|none  |     0|acc       |0.9300|±  |0.0256|
+# |hellaswag     |      1|none  |     0|acc       |0.6127|±  |0.0049|
+# |              |       |none  |     0|acc_norm  |0.8100|±  |0.0039|
+# |lambada_openai|      1|none  |     0|perplexity|3.1810|±  |0.0583|
+# |              |       |none  |     0|acc       |0.7563|±  |0.0060|
+# |openbookqa    |      1|none  |     0|acc       |0.3260|±  |0.0210|
+# |              |       |none  |     0|acc_norm  |0.4380|±  |0.0222|
+# |piqa          |      1|none  |     0|acc       |0.8069|±  |0.0092|
+# |              |       |none  |     0|acc_norm  |0.8215|±  |0.0089|
+# |sciq          |      1|none  |     0|acc       |0.9580|±  |0.0063|
+# |              |       |none  |     0|acc_norm  |0.9390|±  |0.0076|
+# |winogrande    |      1|none  |     0|acc       |0.7395|±  |0.0123|
+
+
+import argparse
+import warnings
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+import fla  # noqa
+
+
+def sizeof_fmt(num, suffix='B'):
+    for unit in ('', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'):
+        if abs(num) < 1024.0:
+            return f'{num:.2f}{unit}{suffix}'
+        num /= 1024.0
+    return f'{num:.2f}Yi{suffix}'
+
+
+def convert(
+    llama: str,
+    config: str,
+    output: str,
+    precision: str = 'float32',
+):
+    AutoTokenizer.from_pretrained(llama).save_pretrained(output)
+    llama = AutoModelForCausalLM.from_pretrained(llama, torch_dtype=precision)
+    print(f"Loading Llama ...\n{llama}")
+
+    config = AutoConfig.from_pretrained(config)
+    config.torch_dtype = precision
+    model = AutoModelForCausalLM.from_config(config)
+    if precision in ['float16', 'fp16']:
+        model = model.to(torch.float16)
+    elif precision in ['bfloat16', 'bf16']:
+        model = model.to(torch.bfloat16)
+    num_parameters = model.num_parameters()
+    print(f"Initializing the model from the config:\n{config}\n{model}")
+    print(f"Number of parameters in total: {num_parameters} ({sizeof_fmt(num_parameters)})")
+
+    print("Copying the weights from Llama to the model ...")
+    vocab_size = llama.model.embed_tokens.weight.shape[0]
+    if model.model.embeddings.weight.shape[0] != vocab_size:
+        warnings.warn(f"Llama and the model have different embedding sizes "
+                      f"({vocab_size} vs {model.model.embeddings.weight.shape[0]}), "
+                      f"the model embeddings will be extended with randomly initialized values or truncated")
+        vocab_size = min(model.model.embeddings.weight.shape[0], vocab_size)
+    print("llama.model.embed_tokens                        -> model.model.embeddings")
+    model.model.embeddings.weight.data[:vocab_size].copy_(llama.model.embed_tokens.weight[:vocab_size])
+    torch.testing.assert_close(model.model.embeddings.weight[:vocab_size], llama.model.embed_tokens.weight[:vocab_size])
+    for i in range(config.num_hidden_layers):
+        if hasattr(model.model.layers[i], 'attn_norm'):
+            if model.model.layers[i].attn_norm.weight is not None:
+                print(f"llama.model.layers{i}.input_layernorm.weight -> model.model.layers{i}.attn_norm.weight")
+                model.model.layers[i].attn_norm.weight.data.copy_(llama.model.layers[i].input_layernorm.weight)
+                torch.testing.assert_close(model.model.layers[i].attn_norm.weight,
+                                           llama.model.layers[i].input_layernorm.weight)
+            if model.model.layers[i].attn_norm.bias is not None:
+                print(f"llama.model.layers{i}.input_layernorm.bias -> model.model.layers{i}.attn_norm.bias")
+                model.model.layers[i].attn_norm.bias.data.copy_(llama.model.layers[i].input_layernorm.bias)
+                torch.testing.assert_close(model.model.layers[i].attn_norm.bias,
+                                           llama.model.layers[i].input_layernorm.bias)
+            model.model.layers[i].attn_norm.eps = llama.model.layers[i].input_layernorm.variance_epsilon
+        if hasattr(model.model.layers[i].attn, 'norm'):
+            if model.model.layers[i].attn.norm.weight is not None:
+                print(f"llama.model.layers{i}.input_layernorm.weight -> model.model.layers{i}.attn.norm.weight")
+                model.model.layers[i].attn.norm.weight.data.copy_(llama.model.layers[i].input_layernorm.weight)
+                torch.testing.assert_close(model.model.layers[i].attn.norm.weight,
+                                           llama.model.layers[i].input_layernorm.weight)
+            if model.model.layers[i].attn.norm.bias is not None:
+                print(f"llama.model.layers{i}.input_layernorm.bias -> model.model.layers{i}.attn.norm.bias")
+                model.model.layers[i].attn.norm.bias.data.copy_(llama.model.layers[i].input_layernorm.bias)
+                torch.testing.assert_close(model.model.layers[i].attn.norm.bias,
+                                           llama.model.layers[i].input_layernorm.bias)
+            model.model.layers[i].attn.norm.eps = llama.model.layers[i].input_layernorm.variance_epsilon
+
+        print(f"llama.model.layers{i}.attn.q_proj.weight  -> model.model.layers{i}.attn.q_proj.weight")
+        model.model.layers[i].attn.q_proj.weight.data.copy_(llama.model.layers[i].self_attn.q_proj.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.q_proj.weight, llama.model.layers[i].self_attn.q_proj.weight)
+        if hasattr(llama.model.layers[i].self_attn.q_proj, 'bias') and hasattr(model.model.layers[i].attn.q_proj, 'bias'):
+            print(f"llama.model.layers{i}.attn.q_proj.bias  -> model.model.layers{i}.attn.q_proj.bias")
+            model.model.layers[i].attn.q_proj.bias.data.copy_(llama.model.layers[i].self_attn.q_proj.bias)
+            torch.testing.assert_close(model.model.layers[i].attn.q_proj.bias, llama.model.layers[i].self_attn.q_proj.bias)
+        print(f"llama.model.layers.{i}.attn.k_proj.weight -> model.model.layers.{i}.attn.k_proj.weight")
+        model.model.layers[i].attn.k_proj.weight.data.copy_(llama.model.layers[i].self_attn.k_proj.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.k_proj.weight, llama.model.layers[i].self_attn.k_proj.weight)
+        if hasattr(llama.model.layers[i].self_attn.k_proj, 'bias') and hasattr(model.model.layers[i].attn.k_proj, 'bias'):
+            print(f"llama.model.layers{i}.attn.k_proj.bias  -> model.model.layers{i}.attn.k_proj.bias")
+            model.model.layers[i].attn.k_proj.bias.data.copy_(llama.model.layers[i].self_attn.k_proj.bias)
+            torch.testing.assert_close(model.model.layers[i].attn.k_proj.bias, llama.model.layers[i].self_attn.k_proj.bias)
+        print(f"llama.model.layers.{i}.attn.v_proj.weight -> model.model.layers.{i}.attn.v_proj.weight")
+        model.model.layers[i].attn.v_proj.weight.data.copy_(llama.model.layers[i].self_attn.v_proj.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.v_proj.weight, llama.model.layers[i].self_attn.v_proj.weight)
+        if hasattr(llama.model.layers[i].self_attn.v_proj, 'bias') and hasattr(model.model.layers[i].attn.v_proj, 'bias'):
+            print(f"llama.model.layers{i}.attn.v_proj.bias  -> model.model.layers{i}.attn.v_proj.bias")
+            model.model.layers[i].attn.v_proj.bias.data.copy_(llama.model.layers[i].self_attn.v_proj.bias)
+            torch.testing.assert_close(model.model.layers[i].attn.v_proj.bias, llama.model.layers[i].self_attn.v_proj.bias)
+
+        print(f"llama.model.layers.{i}.attn.o_proj.weight -> model.model.layers.{i}.attn.o_proj.weight")
+        model.model.layers[i].attn.o_proj.weight.data.copy_(llama.model.layers[i].self_attn.o_proj.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.o_proj.weight, llama.model.layers[i].self_attn.o_proj.weight)
+
+        if hasattr(model.model.layers[i], 'mlp_norm'):
+            if model.model.layers[i].mlp_norm.weight is not None:
+                print(f"llama.model.layers{i}.post_attention_layernorm.weight -> model.model.layers{i}.mlp_norm.weight")
+                model.model.layers[i].mlp_norm.weight.data.copy_(llama.model.layers[i].post_attention_layernorm.weight)
+                torch.testing.assert_close(model.model.layers[i].mlp_norm.weight,
+                                           llama.model.layers[i].post_attention_layernorm.weight)
+            if model.model.layers[i].mlp_norm.bias is not None:
+                print(f"llama.model.layers{i}.post_attention_layernorm.bias -> model.model.layers{i}.mlp_norm.bias")
+                model.model.layers[i].mlp_norm.bias.data.copy_(llama.model.layers[i].post_attention_layernorm.bias)
+                torch.testing.assert_close(model.model.layers[i].mlp_norm.bias,
+                                           llama.model.layers[i].post_attention_layernorm.bias)
+            model.model.layers[i].mlp_norm.eps = llama.model.layers[i].post_attention_layernorm.variance_epsilon
+        if hasattr(model.model.layers[i].mlp, 'norm'):
+            if model.model.layers[i].mlp.norm.weight is not None:
+                print(f"llama.model.layers{i}.post_attention_layernorm.weight -> model.model.layers{i}.mlp.norm.weight")
+                model.model.layers[i].mlp.norm.weight.data.copy_(llama.model.layers[i].post_attention_layernorm.weight)
+                torch.testing.assert_close(model.model.layers[i].mlp.norm.weight,
+                                           llama.model.layers[i].post_attention_layernorm.weight)
+            if model.model.layers[i].mlp.norm.bias is not None:
+                print(f"llama.model.layers{i}.post_attention_layernorm.bias -> model.model.layers{i}.mlp.norm.bias")
+                model.model.layers[i].mlp.norm.bias.data.copy_(llama.model.layers[i].post_attention_layernorm.bias)
+                torch.testing.assert_close(model.model.layers[i].mlp.norm.bias,
+                                           llama.model.layers[i].post_attention_layernorm.bias)
+            model.model.layers[i].mlp.norm.eps = llama.model.layers[i].post_attention_layernorm.variance_epsilon
+
+        print(f"llama.model.layers.{i}.mlp.gate_proj.weight -> model.model.layers.{i}.mlp.gate_proj.weight")
+        model.model.layers[i].mlp.gate_proj.weight.data.copy_(llama.model.layers[i].mlp.gate_proj.weight)
+        torch.testing.assert_close(model.model.layers[i].mlp.gate_proj.weight, llama.model.layers[i].mlp.gate_proj.weight)
+        print(f"llama.model.layers.{i}.mlp.up_proj.weight -> model.model.layers.{i}.mlp.up_proj.weight")
+        model.model.layers[i].mlp.up_proj.weight.data.copy_(llama.model.layers[i].mlp.up_proj.weight)
+        torch.testing.assert_close(model.model.layers[i].mlp.up_proj.weight, llama.model.layers[i].mlp.up_proj.weight)
+
+        print(f"llama.model.layers.{i}.mlp.down_proj.weight -> model.model.layers.{i}.mlp.down_proj.weight")
+        model.model.layers[i].mlp.down_proj.weight.data.copy_(llama.model.layers[i].mlp.down_proj.weight)
+        torch.testing.assert_close(model.model.layers[i].mlp.down_proj.weight,
+                                   llama.model.layers[i].mlp.down_proj.weight)
+
+    if model.model.norm.weight is not None:
+        print("llama.model.norm.weight -> model.model.norm.weight")
+        model.model.norm.weight.data.copy_(llama.model.norm.weight)
+        torch.testing.assert_close(model.model.norm.weight, llama.model.norm.weight)
+    if model.model.norm.bias is not None:
+        print("llama.model.norm.bias -> model.model.norm.bias")
+        model.model.norm.bias.data.copy_(llama.model.norm.bias)
+        torch.testing.assert_close(model.model.norm.bias, llama.model.norm.bias)
+    model.model.norm.eps = llama.model.norm.variance_epsilon
+
+    if not model.config.tie_word_embeddings:
+        print("llama.model.lm_head.weight -> model.lm_head.weight")
+        model.lm_head.weight.data[:vocab_size].copy_(llama.lm_head.weight[:vocab_size])
+        torch.testing.assert_close(model.lm_head.weight[:vocab_size], llama.lm_head.weight[:vocab_size])
+    model.config.rope_theta = llama.config.rope_theta
+
+    print(f"Saving converted model to {output} ...\n{model}")
+    model.save_pretrained(output)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default='mistralai/Mistral-7B-v0.1')
+    parser.add_argument("--config", default='configs/transformer_7B.json')
+    parser.add_argument("--output", default='converted/transformer-7B')
+    parser.add_argument('--precision', type=str, default='float32')
+    args = parser.parse_args()
+    convert(args.model, args.config, args.output, precision=args.precision)

code/flash-linear-attention/utils/convert_from_rwkv6.py

@@ -0,0 +1,213 @@
+
+# scripts for converting pretrained hf model weights to fla style
+# calling the code to make conversions for RWKV/rwkv-6-world-7b would achieve the following results:
+# |    Tasks     |Version|Filter|n-shot|    Metric     | Value |   |Stderr|
+# |--------------|------:|------|-----:|---------------|------:|---|------|
+# |arc_challenge |      1|none  |     0|acc            | 0.4130|±  |0.0144|
+# |              |       |none  |     0|acc_norm       | 0.4403|±  |0.0145|
+# |arc_easy      |      1|none  |     0|acc            | 0.7382|±  |0.0090|
+# |              |       |none  |     0|acc_norm       | 0.7079|±  |0.0093|
+# |boolq         |      2|none  |     0|acc            | 0.6823|±  |0.0081|
+# |copa          |      1|none  |     0|acc            | 0.8700|±  |0.0338|
+# |hellaswag     |      1|none  |     0|acc            | 0.5508|±  |0.0050|
+# |              |       |none  |     0|acc_norm       | 0.7171|±  |0.0045|
+# |lambada_openai|      1|none  |     0|perplexity     | 3.2989|±  |0.0634|
+# |              |       |none  |     0|acc            | 0.7493|±  |0.0060|
+# |openbookqa    |      1|none  |     0|acc            | 0.3200|±  |0.0209|
+# |              |       |none  |     0|acc_norm       | 0.4440|±  |0.0222|
+# |piqa          |      1|none  |     0|acc            | 0.7753|±  |0.0097|
+# |              |       |none  |     0|acc_norm       | 0.7894|±  |0.0095|
+# |sciq          |      1|none  |     0|acc            | 0.9370|±  |0.0077|
+# |              |       |none  |     0|acc_norm       | 0.8860|±  |0.0101|
+# |winogrande    |      1|none  |     0|acc            | 0.6867|±  |0.0130|
+
+import argparse
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+import fla  # noqa
+from fla.utils import device
+
+
+def sizeof_fmt(num, suffix='B'):
+    for unit in ('', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'):
+        if abs(num) < 1024.0:
+            return f'{num:.2f}{unit}{suffix}'
+        num /= 1024.0
+    return f'{num:.2f}Yi{suffix}'
+
+
+def convert(
+    rwkv6: str,
+    config: str,
+    output: str,
+):
+    torch.manual_seed(1)
+    AutoTokenizer.from_pretrained(rwkv6, trust_remote_code=True).save_pretrained(output)
+    rwkv6 = AutoModelForCausalLM.from_pretrained(rwkv6, trust_remote_code=True).to(device)
+    print(f"Loading rwkv6 ...\n{rwkv6}")
+
+    config = AutoConfig.from_pretrained(config)
+    model = AutoModelForCausalLM.from_config(config).to(device)
+    num_parameters = model.num_parameters()
+    print(f"Initializing the model from the config:\n{config}\n{model}")
+    print(f"Number of parameters in total: {num_parameters} ({sizeof_fmt(num_parameters)})")
+
+    print("Copying the weights from rwkv6 to the model ...")
+    print("rwkv6.rwkv.embeddings                        -> model.model.embeddings")
+    model.model.embeddings.weight.data.copy_(rwkv6.rwkv.embeddings.weight)
+    torch.testing.assert_close(model.model.embeddings.weight, rwkv6.rwkv.embeddings.weight)
+    for i in range(config.num_hidden_layers):
+        if hasattr(model.model.layers[i], 'pre_norm'):
+            if model.model.layers[i].pre_norm.weight is not None:
+                print(f"rwkv6.rwkv.blocks{i}.pre_ln.weight -> model.model.layers{i}.pre_norm.weight")
+                model.model.layers[i].pre_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].pre_ln.weight)
+                torch.testing.assert_close(model.model.layers[i].pre_norm.weight, rwkv6.rwkv.blocks[i].pre_ln.weight)
+            if model.model.layers[i].pre_norm.bias is not None:
+                print(f"rwkv6.rwkv.blocks{i}.pre_ln.bias -> model.model.layers{i}.pre_norm.bias")
+                model.model.layers[i].pre_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].pre_ln.bias)
+                torch.testing.assert_close(model.model.layers[i].pre_norm.bias, rwkv6.rwkv.blocks[i].pre_ln.bias)
+            model.model.layers[i].pre_norm.eps = rwkv6.rwkv.blocks[i].pre_ln.eps
+        if model.model.layers[i].attn_norm.weight is not None:
+            print(f"rwkv6.rwkv.blocks{i}.ln1.weight -> model.model.layers{i}.attn_norm.weight")
+            model.model.layers[i].attn_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].ln1.weight)
+            torch.testing.assert_close(model.model.layers[i].attn_norm.weight, rwkv6.rwkv.blocks[i].ln1.weight)
+        if model.model.layers[i].attn_norm.bias is not None:
+            print(f"rwkv6.rwkv.blocks{i}.ln1.bias -> model.model.layers{i}.attn_norm.bias")
+            model.model.layers[i].attn_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].ln1.bias)
+            torch.testing.assert_close(model.model.layers[i].attn_norm.bias, rwkv6.rwkv.blocks[i].ln1.bias)
+        model.model.layers[i].attn_norm.eps = rwkv6.rwkv.blocks[i].ln1.eps
+
+        print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_x -> model.model.layers.{i}.attn.x_proj0.mu")
+        model.model.layers[i].attn.x_proj[0].mu.data.copy_(rwkv6.rwkv.blocks[i].attention.time_maa_x.view(-1))
+        torch.testing.assert_close(model.model.layers[i].attn.x_proj[0].mu,
+                                   rwkv6.rwkv.blocks[i].attention.time_maa_x.view(-1))
+        print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_w1.weight  -> model.model.layers{i}.attn.x_proj0.linear.weight")
+        ww, wk, wv, wr, wg = rwkv6.rwkv.blocks[i].attention.time_maa_w1.view(config.hidden_size, 5, -1).unbind(-2)
+        w = torch.cat((wr, ww, wk, wv, wg), -1).t()
+        model.model.layers[i].attn.x_proj[0].linear.weight.data.copy_(w)
+        torch.testing.assert_close(model.model.layers[i].attn.x_proj[0].linear.weight, w)
+
+        print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_w2.weight  -> model.model.layers{i}.attn.x_proj2.weight")
+        ww, wk, wv, wr, wg = rwkv6.rwkv.blocks[i].attention.time_maa_w2.unbind(0)
+        w = torch.cat((wr, ww, wk, wv, wg), 0).t()
+        model.model.layers[i].attn.x_proj[2].weight.data.copy_(w)
+        torch.testing.assert_close(model.model.layers[i].attn.x_proj[2].weight, w)
+
+        print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_wkvrg  -> model.model.layers{i}.attn.x_bias")
+        bias = torch.stack((rwkv6.rwkv.blocks[i].attention.time_maa_r.view(-1),
+                            rwkv6.rwkv.blocks[i].attention.time_maa_w.view(-1),
+                            rwkv6.rwkv.blocks[i].attention.time_maa_k.view(-1),
+                            rwkv6.rwkv.blocks[i].attention.time_maa_v.view(-1),
+                            rwkv6.rwkv.blocks[i].attention.time_maa_g.view(-1)))
+        model.model.layers[i].attn.x_bias.data.copy_(bias)
+        torch.testing.assert_close(model.model.layers[i].attn.x_bias, bias)
+
+        print(f"rwkv6.rwkv.blocks{i}.attention.receptance.weight  -> model.model.layers{i}.attn.r_proj.linear.weight")
+        model.model.layers[i].attn.r_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.receptance.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.r_proj.linear.weight,
+                                   rwkv6.rwkv.blocks[i].attention.receptance.weight)
+        print(f"rwkv6.rwkv.blocks{i}.attention.time_decay_w1  -> model.model.layers{i}.attn.w_proj.linear.lora0.weight")
+        model.model.layers[i].attn.w_proj.linear.lora[0].weight.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay_w1.t())
+        torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[0].weight,
+                                   rwkv6.rwkv.blocks[i].attention.time_decay_w1.t())
+        print(f"rwkv6.rwkv.blocks{i}.attention.time_decay_w2  -> model.model.layers{i}.attn.w_proj.linear.lora2.weight")
+        model.model.layers[i].attn.w_proj.linear.lora[2].weight.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay_w2.t())
+        torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[2].weight,
+                                   rwkv6.rwkv.blocks[i].attention.time_decay_w2.t())
+        print(f"rwkv6.rwkv.blocks{i}.attention.time_decay  -> model.model.layers{i}.attn.w_proj.linear.lora2.bias")
+        model.model.layers[i].attn.w_proj.linear.lora[2].bias.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay.view(-1))
+        torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[2].bias,
+                                   rwkv6.rwkv.blocks[i].attention.time_decay.view(-1))
+
+        print(f"rwkv6.rwkv.blocks{i}.attention.key.weight -> model.model.layers.{i}.attn.k_proj.linear.weight")
+        model.model.layers[i].attn.k_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.key.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.k_proj.linear.weight,
+                                   rwkv6.rwkv.blocks[i].attention.key.weight)
+        print(f"rwkv6.rwkv.blocks{i}.attention.value.weight -> model.model.layers.{i}.attn.v_proj.linear.weight")
+        model.model.layers[i].attn.v_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.value.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.v_proj.linear.weight,
+                                   rwkv6.rwkv.blocks[i].attention.value.weight)
+        print(f"rwkv6.rwkv.blocks{i}.attention.gate.weight -> model.model.layers.{i}.attn.g_proj.linear.weight")
+        model.model.layers[i].attn.g_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.gate.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.g_proj.linear.weight,
+                                   rwkv6.rwkv.blocks[i].attention.gate.weight)
+        print(f"rwkv6.rwkv.blocks{i}.attention.time_faaaa -> model.model.layers.{i}.attn.bonus")
+        bonus = rwkv6.rwkv.blocks[i].attention.time_faaaa.view(config.num_heads, -1)
+        model.model.layers[i].attn.bonus.data.copy_(bonus)
+        torch.testing.assert_close(model.model.layers[i].attn.bonus, bonus)
+
+        if model.model.layers[i].attn.g_norm.weight is not None:
+            print(f"rwkv6.rwkv.blocks{i}.attention.ln_x.weight -> model.model.layers[i].attn.g_norm.weight")
+            model.model.layers[i].attn.g_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.ln_x.weight)
+            torch.testing.assert_close(model.model.layers[i].attn.g_norm.weight, rwkv6.rwkv.blocks[i].attention.ln_x.weight)
+        if model.model.layers[i].attn.g_norm.bias is not None:
+            print(f"rwkv6.rwkv.blocks{i}.attention.ln_x.bias -> model.model.layers[i].attn.g_norm.bias")
+            model.model.layers[i].attn.g_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].attention.ln_x.bias)
+            torch.testing.assert_close(model.model.layers[i].attn.g_norm.bias, rwkv6.rwkv.blocks[i].attention.ln_x.bias)
+        model.model.layers[i].attn.g_norm.eps = rwkv6.rwkv.blocks[i].attention.ln_x.eps
+
+        print(f"rwkv6.rwkv.blocks{i}.attention.output.weight -> model.model.layers.{i}.attn.o_proj.weight")
+        model.model.layers[i].attn.o_proj.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.output.weight)
+        torch.testing.assert_close(model.model.layers[i].attn.o_proj.weight, rwkv6.rwkv.blocks[i].attention.output.weight)
+
+        if model.model.layers[i].ffn_norm.weight is not None:
+            print(f"rwkv6.rwkv.blocks{i}.ln2.weight -> model.model.layers{i}.ffn_norm.weight")
+            model.model.layers[i].ffn_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].ln2.weight)
+            torch.testing.assert_close(model.model.layers[i].ffn_norm.weight, rwkv6.rwkv.blocks[i].ln2.weight)
+        if model.model.layers[i].ffn_norm.bias is not None:
+            print(f"rwkv6.rwkv.blocks{i}.ln2.bias -> model.model.layers{i}.ffn_norm.bias")
+            model.model.layers[i].ffn_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].ln2.bias)
+            torch.testing.assert_close(model.model.layers[i].ffn_norm.bias, rwkv6.rwkv.blocks[i].ln2.bias)
+        model.model.layers[i].ffn_norm.eps = rwkv6.rwkv.blocks[i].ln2.eps
+
+        print(f"rwkv6.rwkv.blocks{i}.feed_forward.key.weight -> model.model.layers.{i}.ffn.key.linear.weight")
+        model.model.layers[i].ffn.key.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.key.weight)
+        torch.testing.assert_close(model.model.layers[i].ffn.key.linear.weight,
+                                   rwkv6.rwkv.blocks[i].feed_forward.key.weight)
+        print(f"rwkv6.rwkv.blocks{i}.feed_forward.time_maa_k -> model.model.layers.{i}.ffn.key.mu")
+        model.model.layers[i].ffn.key.mu.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.time_maa_k.view(-1))
+        torch.testing.assert_close(model.model.layers[i].ffn.key.mu,
+                                   rwkv6.rwkv.blocks[i].feed_forward.time_maa_k.view(-1))
+
+        print(f"rwkv6.rwkv.blocks{i}.feed_forward.value.weight -> model.model.layers.{i}.ffn.value.weight")
+        model.model.layers[i].ffn.value.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.value.weight)
+        torch.testing.assert_close(model.model.layers[i].ffn.value.weight,
+                                   rwkv6.rwkv.blocks[i].feed_forward.value.weight)
+
+        print(f"rwkv6.rwkv.blocks{i}.feed_forward.receptance.weight -> model.model.layers.{i}.ffn.receptance.linear.weight")
+        model.model.layers[i].ffn.receptance.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.receptance.weight)
+        torch.testing.assert_close(model.model.layers[i].ffn.receptance.linear.weight,
+                                   rwkv6.rwkv.blocks[i].feed_forward.receptance.weight)
+        print(f"rwkv6.rwkv.blocks{i}.feed_forward.time_maa_r -> model.model.layers.{i}.ffn.receptance.mu")
+        model.model.layers[i].ffn.receptance.mu.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.time_maa_r.view(-1))
+        torch.testing.assert_close(model.model.layers[i].ffn.receptance.mu,
+                                   rwkv6.rwkv.blocks[i].feed_forward.time_maa_r.view(-1))
+
+    if model.model.norm.weight is not None:
+        print("rwkv6.rwkv.ln_out.weight -> model.model.norm.weight")
+        model.model.norm.weight.data.copy_(rwkv6.rwkv.ln_out.weight)
+        torch.testing.assert_close(model.model.norm.weight, rwkv6.rwkv.ln_out.weight)
+    if model.model.norm.bias is not None:
+        print("rwkv6.rwkv.ln_out.bias -> model.model.norm.bias")
+        model.model.norm.bias.data.copy_(rwkv6.rwkv.ln_out.bias)
+        torch.testing.assert_close(model.model.norm.bias, rwkv6.rwkv.ln_out.bias)
+    model.model.norm.eps = rwkv6.rwkv.ln_out.eps
+
+    if not model.config.tie_word_embeddings:
+        print("rwkv6.rwkv.head.weight -> model.lm_head.weight")
+        model.lm_head.weight.data.copy_(rwkv6.head.weight)
+        torch.testing.assert_close(model.lm_head.weight, rwkv6.head.weight)
+
+    print(f"Saving converted model \n{model}\n to {output} ...")
+    model.save_pretrained(output)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default='RWKV/rwkv-6-world-7b')
+    parser.add_argument("--config", default='configs/rwkv6_7B.json')
+    parser.add_argument("--output", default='converted/rwkv6-7B')
+    args = parser.parse_args()
+    convert(args.model, args.config, args.output)

code/flash-linear-attention/utils/convert_from_rwkv7.py

@@ -0,0 +1,150 @@
+
+# scripts for converting pretrained hf model weights to fla style
+
+import argparse
+import os
+import re
+
+import torch
+from transformers import AutoModelForCausalLM
+
+import fla  # noqa
+from fla.models.rwkv7 import RWKV7Config
+
+
+def convert(
+    rwkv7: str,
+    output: str,
+    precision: str = 'float32',
+):
+    weights = torch.load(rwkv7, weights_only=True, map_location='cpu')
+    config = RWKV7Config()
+    config.vocab_size = weights['emb.weight'].shape[0]  # 50304
+    config.hidden_size = weights['blocks.0.ffn.key.weight'].shape[1]  # 768
+    config.hidden_ratio = weights['blocks.0.ffn.key.weight'].shape[0] / weights['blocks.0.ffn.key.weight'].shape[1]  # 4.0
+    config.intermediate_size = weights['blocks.0.ffn.key.weight'].shape[0]
+    config.num_hidden_layers = 0
+    while f'blocks.{config.num_hidden_layers}.ffn.key.weight' in weights:
+        config.num_hidden_layers += 1
+    # 12
+    config.value_dim = [config.hidden_size] * config.num_hidden_layers
+    config.decay_low_rank_dim = weights['blocks.0.att.w1'].shape[1]  # 64
+    config.gate_low_rank_dim = weights['blocks.0.att.g1'].shape[1]  # 128
+    config.a_low_rank_dim = weights['blocks.0.att.a1'].shape[1]  # 64
+    try:
+        config.v_low_rank_dim = weights['blocks.1.att.v1'].shape[1]  # 32
+    except KeyError:
+        config.v_low_rank_dim = 32
+
+    if precision in ['bf16', 'bfloat16']:
+        precision = 'bfloat16'
+        dtype = torch.bfloat16
+    if precision in ['fp16', 'float16']:
+        precision = 'float16'
+        dtype = torch.float16
+    if precision in ['fp64', 'double', 'float64']:
+        precision = 'float64'
+        dtype = torch.float64
+
+    config.torch_dtype = precision
+    print(f"Creating model with config:\n{config}")
+    model = AutoModelForCausalLM.from_config(config).to(dtype=dtype)
+
+    print(model)
+    model_dict = model.state_dict()
+    model_names = [n for n in model_dict]
+
+    # these parameters may be present in pth file but are never used:
+    unused_names = ['blocks.0.att.v0', 'blocks.0.att.v1', 'blocks.0.att.v2']
+    # these parameters may or may not be present in pth file:
+    possible_absent_weights = [
+        'model.layers.0.pre_norm.weight', 'model.layers.0.pre_norm.bias',
+    ]
+    # other parameters may raise a KeyError
+
+    def translate_into_fla(name):
+        transposed = False
+        emb_head = {
+            'emb.weight': 'model.embeddings.weight',
+            'ln_out.weight': 'model.norm.weight',
+            'ln_out.bias': 'model.norm.bias',
+            'head.weight': 'lm_head.weight',
+        }
+        proj = {
+            'receptance': 'r_proj',
+            'key': 'k_proj',
+            'value': 'v_proj',
+            'ln_x': 'g_norm',
+            'output': 'o_proj',
+        }
+        if name in unused_names:
+            return '', False
+        if name in emb_head:
+            return emb_head[name], False
+        name_compo = name.split('.')
+        assert name_compo[0] == 'blocks'
+        name_compo[0] = 'model.layers'
+        assert int(name_compo[1]) in range(config.num_hidden_layers)
+        name_compo[2] = {
+            'att': 'attn',
+            'ffn': 'ffn',
+            'ln0': 'pre_norm',
+            'ln1': 'attn_norm',
+            'ln2': 'ffn_norm',
+        }[name_compo[2]]
+        if re.match("[wvag][012]", name_compo[3]):
+            typ, num = name_compo[3]
+            name_compo[3] = f'{typ}_lora.lora.' + {
+                '0': '2.bias',
+                '1': '0.weight',
+                '2': '2.weight',
+            }[num]
+            transposed |= (num in ['1', '2'])
+        elif name_compo[2] == 'attn' and name_compo[3] in proj:
+            name_compo[3] = proj[name_compo[3]]
+        return '.'.join(name_compo), transposed
+
+    for name in weights:
+        fla_name, transposed = translate_into_fla(name)
+        print(f'{name:32} -> {fla_name:42}, {transposed}')
+        if not fla_name:
+            print('redundant parameters in source weight: ', name, '\n')
+            continue
+        weight = weights[name]
+        # print shape information
+        shape1 = list(weight.shape)
+        shape2 = list(model_dict[fla_name].shape)
+        print(f'{str(shape1):32}    {str(shape2)}\n')
+
+        if transposed:
+            weight.t_()
+        if shape1 == [1, 1, config.hidden_size]:
+            weight.squeeze_()
+
+        if "attn.x_" in fla_name:
+            assert model_dict[fla_name].shape[2:] == weight.shape, \
+                f"Shape mismatch for {fla_name}: model_dict={model_dict[fla_name].shape}, weight={weight.shape}"
+        else:
+            assert model_dict[fla_name].shape == weight.shape, \
+                f"Shape mismatch for {fla_name}: model_dict={model_dict[fla_name].shape}, weight={weight.shape}"
+
+        model_dict[fla_name].data.copy_(weight)
+        model_names.remove(fla_name)
+
+    print("uninitialized parameters: ", model_names)
+    for n in model_names:
+        if n not in possible_absent_weights:
+            raise KeyError(n)
+
+    os.makedirs(output, exist_ok=True)
+
+    model.save_pretrained(output, max_shard_size="1000GB")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Convert RWKV7')
+    parser.add_argument('--rwkv7', type=str, help='Path to the input model')
+    parser.add_argument('--output', type=str, help='Directory to save model')
+    parser.add_argument('--precision', type=str, default='float32')
+    args = parser.parse_args()
+    convert(args.rwkv7, args.output, precision=args.precision)

code/inference/README.md

@@ -0,0 +1,73 @@
+# Inference Recipes
+
+Bash-level inference scripts mirroring `train/` — one script per memory row, all calling `inference/unified_inference.py`.
+
+## Usage
+
+```bash
+export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
+
+# Single memory type
+CKPT=./ckpts/context_k1/epoch-0.safetensors \
+  bash inference/memory_baselines_basic/run_infer_context_k1.sh
+
+# With custom prompt and context image
+CKPT=./ckpts/context_k1/epoch-0.safetensors \
+PROMPT="A toy bear on a table" \
+CONTEXT_IMAGE=assets/opendomain_revisit/1774363417.png \
+  bash inference/memory_baselines_basic/run_infer_context_k1.sh
+
+# All memory baselines (needs CKPT_DIR with per-row folders)
+CKPT_DIR=./ckpts bash inference/memory_baselines_basic/run_infer_all.sh
+
+# Dynamic SpatialVID row
+CKPT=/path/to/retrained_dynamic_spatial_mem/epoch-0.safetensors \
+  bash inference/dynamic_spatialvid/run_infer_dyn_spatial_mem.sh
+```
+
+## Environment Variables
+
+| Variable | Default | Description |
+|---|---|---|
+| `CKPT` | (required) | Path to `.safetensors` checkpoint |
+| `WAN_BASE_MODEL` | (required) | Wan 2.1 base model directory |
+| `PROMPT` | Generic game scene prompt | Text prompt |
+| `CONTEXT_IMAGE` | (none) | First-frame context image path |
+| `ACTION_PATH` | `env/action_rotation_left_45.json` | Camera trajectory JSON |
+| `SEED` | `0` | Random seed |
+| `HEIGHT` / `WIDTH` | `352` / `640` | Resolution |
+| `NUM_FRAMES` | `81` | Frames per chunk |
+| `NUM_INFERENCE_STEPS` | `50` | Denoising steps |
+| `SIGMA_SHIFT` | `15.0` (memory baselines) / `5.0` (context learning) | Timestep shift |
+| `INFER_OUTPUT_ROOT` | `inference_outputs/` | Output directory |
+
+## Script Mapping
+
+### Memory Baselines (`inference/memory_baselines_basic/`)
+
+| Inference script | `--memory_type` | Training script |
+|---|---|---|
+| `run_infer_no_memory.sh` | `no_memory` | `run_ablation_no_memory_baseline_two_chunk.sh` |
+| `run_infer_framepack_weight.sh` | `framepack_weight` | `run_ablation_framepack_weight_two_chunk.sh` |
+| `run_infer_framepack_len_r2.sh` | `framepack_len_r2` | `run_ablation_framepack_len_r2_two_chunk.sh` |
+| `run_infer_framepack_len_r4.sh` | `framepack_len_r4` | `run_ablation_framepack_len_r4_two_chunk.sh` |
+| `run_infer_framepack_hybrid_r2.sh` | `framepack_hybrid_r2` | `run_ablation_framepack_hybrid_r2_weight_two_chunk.sh` |
+| `run_infer_framepack_hybrid_r4.sh` | `framepack_hybrid_r4` | `run_ablation_framepack_hybrid_r4_weight_two_chunk.sh` |
+| `run_infer_spatial_mem.sh` | `spatial_mem` | `run_spatial_memory_baseline.sh` |
+| `run_infer_spatial_concat_text.sh` | `spatial_concat_text` | `run_ablation_spatial_concat_text_two_chunk.sh` |
+| `run_infer_spatial_inject_none.sh` | `spatial_inject_none` | `run_ablation_spatial_inject_none_two_chunk.sh` |
+| `run_infer_spatial_cross_attn_readout.sh` | `spatial_cross_attn_readout` | `run_ablation_spatial_cross_attn_readout_two_chunk.sh` |
+| `run_infer_videossm_hybrid.sh` | `videossm_hybrid` | `run_videossm_hybrid_baseline.sh` |
+| `run_infer_block_wise_ssm.sh` | `block_wise_ssm` | `run_ablation_block_wise_ssm_two_chunk.sh` |
+
+### Context Learning (`inference/context_learning/`)
+
+| Inference script | `--memory_type` | Training script |
+|---|---|---|
+| `run_infer_ctx1.sh` | `context_k1` | `run_pre_qkv_ctx1.sh` |
+| `run_infer_ctx5.sh` | `context_k5` | `run_pre_qkv_ctx5.sh` |
+| `run_infer_ctx20.sh` | `context_k20` | `run_pre_qkv_ctx20.sh` |
+
+### Dynamic SpatialVID (`inference/dynamic_spatialvid/`)
+
+Dynamic wrappers mirror the six dynamic training rows in `train/dynamic_spatialvid/`. They are intended for qualitative replay and demo generation; dynamic evaluation scripts are TODO.

code/inference/_shared/common_env_infer.sh

@@ -0,0 +1,59 @@
+#!/bin/bash
+# Shared inference environment — mirrors train/_shared/common_env_memory.sh
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[1]:-${BASH_SOURCE[0]}}")" && pwd)"
+REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/../.." && pwd)}"
+export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}"
+
+# ── Base model ────────────────────────────────────────────────────────
+WAN_BASE_MODEL="${WAN_BASE_MODEL:-}"
+if [ -z "${WAN_BASE_MODEL}" ]; then
+  for _d in \
+    "${REPO_ROOT}/checkpoints/Wan2.1-T2V-1.3B" \
+    "${REPO_ROOT}/models/Wan2.1-T2V-1.3B" \
+    "${REPO_ROOT}/Wan2.1-T2V-1.3B"; do
+    [ -d "${_d}" ] && { WAN_BASE_MODEL="${_d}"; break; }
+  done
+fi
+if [ -z "${WAN_BASE_MODEL}" ]; then
+  echo "[common_env_infer] ERROR: WAN_BASE_MODEL not set." >&2
+  echo "[common_env_infer] HINT: export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B" >&2
+  exit 2
+fi
+for _f in diffusion_pytorch_model.safetensors models_t5_umt5-xxl-enc-bf16.pth Wan2.1_VAE.pth; do
+  [ -f "${WAN_BASE_MODEL}/${_f}" ] || {
+    echo "[common_env_infer] ERROR: missing ${WAN_BASE_MODEL}/${_f}" >&2; exit 2; }
+done
+TOKENIZER_PATH="${TOKENIZER_PATH:-}"
+if [ -z "${TOKENIZER_PATH}" ] && [ -d "${WAN_BASE_MODEL}/google/umt5-xxl" ]; then
+  TOKENIZER_PATH="${WAN_BASE_MODEL}/google/umt5-xxl"
+fi
+export WAN_BASE_MODEL
+export TOKENIZER_PATH
+
+# ── Checkpoint ────────────────────────────────────────────────────────
+# CKPT must be set by the caller or the wrapper script.
+# Inference scripts validate this themselves.
+
+# ── Output ────────────────────────────────────────────────────────────
+INFER_OUTPUT_ROOT="${INFER_OUTPUT_ROOT:-${REPO_ROOT}/inference_outputs}"
+mkdir -p "${INFER_OUTPUT_ROOT}"
+export INFER_OUTPUT_ROOT
+
+# ── Defaults ──────────────────────────────────────────────────────────
+PROMPT="${PROMPT:-A game scene, the camera moves through the environment}"
+CONTEXT_IMAGE="${CONTEXT_IMAGE:-}"
+ACTION_PATH="${ACTION_PATH:-${REPO_ROOT}/env/action_rotation_left_45.json}"
+SEED="${SEED:-0}"
+HEIGHT="${HEIGHT:-352}"
+WIDTH="${WIDTH:-640}"
+NUM_FRAMES="${NUM_FRAMES:-81}"
+NUM_INFERENCE_STEPS="${NUM_INFERENCE_STEPS:-50}"
+SIGMA_SHIFT="${SIGMA_SHIFT:-15.0}"
+CFG_SCALE="${CFG_SCALE:-5.0}"
+FPS="${FPS:-15}"
+
+echo "[common_env_infer] WAN_BASE_MODEL=${WAN_BASE_MODEL}"
+echo "[common_env_infer] INFER_OUTPUT_ROOT=${INFER_OUTPUT_ROOT}"
+cd "${REPO_ROOT}"

code/inference/context_learning/common_env.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+# Common inference wrapper for context_learning.
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+SIGMA_SHIFT="${SIGMA_SHIFT:-5.0}"
+source "${SCRIPT_DIR}/../_shared/common_env_infer.sh"

code/inference/context_learning/run_infer_ctx1.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+# Inference: Context K=1
+# Corresponds to: train/context_learning/run_pre_qkv_ctx1.sh
+set -euo pipefail
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "${SCRIPT_DIR}/common_env.sh"
+
+: "${CKPT:?ERROR: set CKPT to your context_k1 checkpoint path}"
+OUTPUT="${INFER_OUTPUT_ROOT}/context_k1_$(date +%Y%m%d_%H%M%S).mp4"
+
+EXTRA_ARGS=()
+if [ -n "${CONTEXT_IMAGE}" ]; then
+  EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}")
+fi
+if [ -n "${ACTION_PATH}" ]; then
+  EXTRA_ARGS+=(--action_path "${ACTION_PATH}")
+fi
+
+python inference/unified_inference.py \
+  --ckpt "${CKPT}" \
+  --memory_type context_k1 \
+  --base_model "${WAN_BASE_MODEL}" \
+  --prompt "${PROMPT}" \
+  "${EXTRA_ARGS[@]}" \
+  --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \
+  --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \
+  --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \
+  --output_path "${OUTPUT}"