Validate the independent FLA 0.5 compatibility path

README.md

@@ -33,9 +33,10 @@ Once installed, FlashKDA is auto-dispatched from `flash-linear-attention`'s `chu
 
 **Requirements**
 
-1. Install `flash-linear-attention >= 0.5.0`:
+1. Install a compatible `flash-linear-attention` 0.5 release. FlashKDA's
+   raw-gate/raw-beta dispatch contract requires FLA 0.5.1 or newer:
    ```bash
-   pip install -U flash-linear-attention
+   pip install -U "flash-linear-attention>=0.5.1,<0.6"
    ```
 2. Call `chunk_kda` under `torch.inference_mode()` 
    ```python
@@ -54,7 +55,7 @@ Once installed, FlashKDA is auto-dispatched from `flash-linear-attention`'s `chu
            safe_gate=True,
            A_log=A_log, dt_bias=dt_bias,
            lower_bound=lower_bound,
-           transpose_state_layout=True,
+           state_v_first=True,
            cu_seqlens=cu_seqlens,
        )
    ```

benchmarks/bench_fwd.py

@@ -1,7 +1,13 @@
+import os
 import torch
 import flash_kda
 import torch.nn.functional as F
 import math
+
+# Keep the FLA baseline independent instead of dispatching chunk_kda back to
+# the FlashKDA implementation being benchmarked.
+os.environ["FLA_FLASH_KDA"] = "0"
+
 from fla.ops.kda import chunk_kda
 from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 
@@ -104,9 +110,10 @@ def run_chunk_kda():
             use_gate_in_kernel=True,
             use_qk_l2norm_in_kernel=True,
             use_beta_sigmoid_in_kernel=True,
+            safe_gate=True,
             A_log=A_log, dt_bias=dt_bias,
             lower_bound=LOWER_BOUND,
-            transpose_state_layout=True,
+            state_v_first=True,
             **extra,
         )
 
@@ -124,7 +131,7 @@ def run_chunk_gated_delta_rule():
             initial_state=h0_gdn,
             output_final_state=True,
             use_qk_l2norm_in_kernel=True,
-            transpose_state_layout=True,
+            state_v_first=True,
             **extra,
         )
 

benchmarks/generate_benchmark_md.py

@@ -26,12 +26,13 @@
 FLA_CHUNK_KDA_OPTIONS_MD = (
     "- `fla_chunk_kda` configuration: `use_gate_in_kernel=True`, "
     "`use_qk_l2norm_in_kernel=True`, `use_beta_sigmoid_in_kernel=True`, "
-    "`lower_bound=-5`, `transpose_state_layout=True`"
+    "`safe_gate=True`, `lower_bound=-5`, `state_v_first=True`, "
+    "`FLA_FLASH_KDA=0`"
 )
 FLA_CHUNK_GDN_OPTIONS_MD = (
     "- `fla_chunk_gated_delta_rule` configuration: scalar per-head gate "
     "`g` of shape `(1, T, H)`, `use_qk_l2norm_in_kernel=True`, "
-    "`transpose_state_layout=True`"
+    "`state_v_first=True`"
 )
 
 RE_HEADER_FIXED = re.compile(

tests/test.sh

@@ -1,4 +1,4 @@
 set -e
 pip install -e .
-pip install "flash-linear-attention>=0.5.0" matplotlib
+pip install "flash-linear-attention>=0.5.1,<0.6" matplotlib
 python tests/test_fwd.py

tests/test_fwd.py

@@ -2,6 +2,8 @@
 import torch.nn.functional as F
 import flash_kda
 import math
+import os
+from unittest.mock import patch
 
 from torch_ref import torch_ref
 
@@ -69,7 +71,10 @@ def make_test_cases(H, D, T_shape, dtype, device):
 
 
 def run_fla_gold_reference(q, k, v, g, beta, h0, A_log, dt_bias, scale, lower_bound, cu_seqlens=None):
-    """Run fused_recurrent_kda in fp64 as gold reference, and chunk_kda in bf16.
+    """Run fused_recurrent_kda in fp64 and FLA's Triton chunk_kda in bf16.
+
+    FLA_FLASH_KDA is disabled for chunk_kda so this tests the independent
+    upstream implementation rather than dispatching back to flash_kda.
     beta: [B, T, H] bf16 logits (pre-sigmoid).
     """
     from fla.ops.kda import chunk_kda, fused_recurrent_kda
@@ -97,31 +102,32 @@ def run_fla_gold_reference(q, k, v, g, beta, h0, A_log, dt_bias, scale, lower_bo
         use_qk_l2norm_in_kernel=True,
         use_gate_in_kernel=False,
         lower_bound=None,
-        transpose_state_layout=True,
+        state_v_first=True,
         **fla_kwargs,
     )
     tri = tri.to(torch.float32)
     tri_ht = tri_ht.to(torch.float32)
 
-    # upstream hasn't implemented use_beta_sigmoid_in_kernel; pass post-sigmoid beta explicitly.
-    beta_activated = torch.sigmoid(beta.clone().float()).to(torch.bfloat16)
-    chunk_o, chunk_ht = chunk_kda(
-        q=q.clone().to(torch.bfloat16),
-        k=k.clone().to(torch.bfloat16),
-        v=v.clone(),
-        g=g.clone(),
-        beta=beta_activated,
-        scale=scale,
-        initial_state=h0.clone(),
-        output_final_state=True,
-        use_gate_in_kernel=True,
-        use_qk_l2norm_in_kernel=True,
-        A_log=A_log.clone(),
-        dt_bias=dt_bias.clone(),
-        lower_bound=lower_bound,
-        transpose_state_layout=True,
-        **fla_kwargs,
-    )
+    with patch.dict(os.environ, {"FLA_FLASH_KDA": "0"}):
+        chunk_o, chunk_ht = chunk_kda(
+            q=q.clone().to(torch.bfloat16),
+            k=k.clone().to(torch.bfloat16),
+            v=v.clone(),
+            g=g.clone(),
+            beta=beta.clone(),
+            scale=scale,
+            initial_state=h0.clone(),
+            output_final_state=True,
+            use_gate_in_kernel=True,
+            use_qk_l2norm_in_kernel=True,
+            use_beta_sigmoid_in_kernel=True,
+            safe_gate=True,
+            A_log=A_log.clone(),
+            dt_bias=dt_bias.clone(),
+            lower_bound=lower_bound,
+            state_v_first=True,
+            **fla_kwargs,
+        )
 
     return tri, tri_ht, chunk_o, chunk_ht
 
@@ -359,7 +365,8 @@ def test_fwd_vs_fla():
     for r in results:
         assert_close(f"{r['name']} o", r["tri"], r["out"], 0.005)
         assert_close(f"{r['name']} ht", r["tri_ht"], r["final_state"], 0.005, warning=True)
-        assert_close(f"{r['name']} chunk_kda ht", r["tri_ht"], r["chunk_ht"], 0.005, warning=True)
+        assert_close(f"{r['name']} chunk_kda o", r["chunk_o"], r["out"], 0.006)
+        assert_close(f"{r['name']} chunk_kda ht", r["chunk_ht"], r["final_state"], 0.006)
     print("Assert results: Success")
 
 
@@ -420,7 +427,8 @@ def test_fwd_varlen_vs_fla():
     for r in results:
         assert_close(f"{r['name']} o", r["tri"], r["out"], 0.006)
         assert_close(f"{r['name']} ht", r["tri_ht"], r["final_state"], 0.006, warning=True)
-        assert_close(f"{r['name']} chunk_kda ht", r["tri_ht"], r["chunk_ht"], 0.005, warning=True)
+        assert_close(f"{r['name']} chunk_kda o", r["chunk_o"], r["out"], 0.006)
+        assert_close(f"{r['name']} chunk_kda ht", r["chunk_ht"], r["final_state"], 0.006)
     print("Assert results: Success")