[Fix] Use unit KDA gate scale when A_log is omitted

fla/ops/kda/backends/triton_ascend/fused_recurrent.py

@@ -200,7 +200,7 @@ def fused_recurrent_kda_fwd_kernel_npu(
                 if USE_LOWER_BOUND:
                     b_gk = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_g)
                 else:
-                    b_gk = -exp(b_A) * softplus(b_g)
+                    b_gk = -(exp(b_A) if HAS_A else b_A) * softplus(b_g)
             else:
                 b_gk = b_g
 

fla/ops/kda/backends/triton_ascend/gate.py

@@ -109,7 +109,7 @@ def kda_gate_fwd_kernel_npu(
         p_b = tl.make_block_ptr(dt_bias, (H * D,), (1,), (i_h * D,), (BD,), (0,))
         b_g = b_g + tl.load(p_b, boundary_check=(0,)).to(tl.float32)
     if not USE_LOWER_BOUND:
-        b_yg = -exp(b_A) * softplus(b_g)
+        b_yg = -(exp(b_A) if HAS_A else b_A) * softplus(b_g)
     else:
         b_yg = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_g)
     tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), boundary_check=(0, 1))
@@ -203,7 +203,7 @@ def kda_gate_bwd_kernel_npu(
         b_g = b_g + tl.load(dt_bias + o_b, mask=o_d < D, other=0.0).to(tl.float32)
 
     if not USE_LOWER_BOUND:
-        b_A = -exp(b_A)
+        b_A = -(exp(b_A) if HAS_A else b_A)
         b_yg = b_A * softplus(b_g)
         b_dg = b_A * (b_dyg * tl.sigmoid(b_g))
         b_dA = tl.sum(tl.sum(b_dyg * b_yg, 1), 0)
@@ -328,7 +328,7 @@ def kda_gate_chunk_cumsum_vector_kernel_npu(
 
     b_A = tl.load(A_log + i_h).to(tl.float32) if HAS_A else 1.0
     if not USE_LOWER_BOUND:
-        b_gate = -exp(b_A) * softplus(b_s)
+        b_gate = -(exp(b_A) if HAS_A else b_A) * softplus(b_s)
     else:
         b_gate = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_s)
 

fla/ops/kda/chunk.py

@@ -269,10 +269,9 @@ def chunk_kda(
             Whether to compute the log-space KDA decay internally.
             - If ``True``:
               The passed ``g`` acts as the raw input for ``-exp(A_log) * softplus(g + dt_bias.view(HV, K))``.
-              Note that as part of the input arguments,
-              ``A_log`` (shape ``[HV]``) and the optional ``dt_bias`` (shape ``[HV * K]``) should be provided.
-              When ``lower_bound`` is set, ``A_log`` may be ``None``,
-              in which case the gate is ``lower_bound * sigmoid(g + dt_bias)``.
+              ``A_log`` (shape ``[HV]``) and ``dt_bias`` (shape ``[HV * K]``) are optional.
+              When ``A_log`` is ``None``, the decay scale is 1, equivalent to ``A_log=0``.
+              When ``lower_bound`` is set, the gate is ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``.
             - If ``False``, ``g`` is expected to be the pre-computed decay value.
             Default: ``False``.
         use_beta_sigmoid_in_kernel (bool, Optional):
@@ -467,8 +466,6 @@ def chunk_kda(
     A_log, dt_bias = None, None
     if use_gate_in_kernel:
         A_log, dt_bias = kwargs.get("A_log"), kwargs.get("dt_bias")
-        if A_log is None and lower_bound is None:
-            raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True` and `lower_bound` is not set.")
 
     chunk_size = kwargs.pop("chunk_size", 64)
     if chunk_size not in (32, 64):

fla/ops/kda/fused_recurrent.py

@@ -170,7 +170,7 @@ def fused_recurrent_kda_fwd_kernel(
             if USE_LOWER_BOUND:
                 b_gk = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_g)
             else:
-                b_gk = -exp(b_A) * softplus(b_g)
+                b_gk = -(exp(b_A) if HAS_A else b_A) * softplus(b_g)
         else:
             b_gk = b_g
 
@@ -371,8 +371,7 @@ def fused_recurrent_kda(
             betas of shape `[B, T, HV]`.
         A_log (Optional[torch.Tensor]):
             Decay parameter of shape `[HV]`.
-            When `use_gate_in_kernel=True` together with `lower_bound`,
-            may be `None` to use `lower_bound * sigmoid(g + dt_bias)`.
+            When `None`, uses a unit decay scale, equivalent to `A_log=0`.
         dt_bias (Optional[torch.Tensor]):
             Bias added to `g` before activation, of shape `[HV]`. Only used when `use_gate_in_kernel=True`.
         scale (Optional[float]):

fla/ops/kda/gate.py

@@ -26,20 +26,21 @@
 
 def naive_kda_gate(
     g: torch.Tensor,
-    A_log: torch.Tensor,
+    A_log: torch.Tensor | None = None,
     dt_bias: torch.Tensor | None = None,
     output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
     """
     Torch reference implementation for KDA gate computation.
 
-    Computes: g = -A_log.exp().unsqueeze(-1) * softplus(g + dt_bias.view(g.shape[-2:]))
+    Computes: g = -A_log.exp().unsqueeze(-1) * softplus(g + dt_bias.view(g.shape[-2:])).
+    When ``A_log`` is ``None``: ``g = -softplus(g + dt_bias)``.
 
     Args:
         g (torch.Tensor):
             Input tensor of shape `[..., H, K]`.
-        A_log (torch.Tensor):
-            Parameter tensor with `H` elements.
+        A_log (torch.Tensor | None):
+            Optional parameter tensor with `H` elements. When ``None``, uses a unit decay scale.
         dt_bias (torch.Tensor | None):
             Optional bias tensor added to `g` before activation, shape `[H * K]`.
 
@@ -51,7 +52,10 @@ def naive_kda_gate(
     if dt_bias is not None:
         g = g + dt_bias.view(H, -1)
 
-    g = (-A_log.view(H, 1).float().exp() * F.softplus(g.float())).to(output_dtype)
+    g = F.softplus(g.float())
+    if A_log is not None:
+        g = A_log.view(H, 1).float().exp() * g
+    g = -g.to(output_dtype)
     return g
 
 
@@ -144,7 +148,7 @@ def kda_gate_fwd_kernel(
         o_b = i_h * D + tl.arange(0, BD)
         b_g = b_g + tl.load(dt_bias + o_b, mask=o_b < H * D, other=0.0).to(tl.float32)
     if not USE_LOWER_BOUND:
-        b_yg = -exp(b_A) * softplus(b_g)
+        b_yg = -(exp(b_A) if HAS_A else b_A) * softplus(b_g)
     else:
         b_yg = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_g)
     tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), mask=m_g)
@@ -216,7 +220,7 @@ def kda_gate_bwd_kernel(
 
     # [BT, BD]
     if not USE_LOWER_BOUND:
-        b_A = -exp(b_A)
+        b_A = -(exp(b_A) if HAS_A else b_A)
         b_yg = b_A * softplus(b_g)
         b_dg = b_A * (b_dyg * tl.sigmoid(b_g))
         b_dA = tl.sum(tl.sum(b_dyg * b_yg, 1), 0)
@@ -374,14 +378,14 @@ def fused_kda_gate(
 
     Computes: g = -A_log.exp().unsqueeze(-1) * softplus(g + dt_bias.view(g.shape[-2:]))
     When ``lower_bound`` is set: g = lower_bound * sigmoid(exp(A_log) * (g + dt_bias)).
-    When ``A_log`` is ``None`` (requires ``lower_bound``): g = lower_bound * sigmoid(g + dt_bias).
+    When ``A_log`` is ``None``, the decay scale is 1, equivalent to ``A_log=0``.
 
     Args:
         g (torch.Tensor):
             Input tensor of shape `[..., H, K]`.
         A_log (torch.Tensor | None):
             Optional parameter tensor with `H` elements.
-            When ``None``, the gate reduces to ``lower_bound * sigmoid(g + dt_bias)`` (requires ``lower_bound``).
+            When ``None``, uses a unit decay scale.
         dt_bias (torch.Tensor | None):
             Optional bias tensor added to `g` before activation, shape `[H * K]`.
 
@@ -459,7 +463,7 @@ def kda_gate_chunk_cumsum_vector_kernel(
     b_A = tl.load(A_log + i_h).to(tl.float32) if HAS_A else 1.0
     if not USE_LOWER_BOUND:
         # Apply gate: -exp(A_log) * softplus(g + bias)
-        b_gate = -exp(b_A) * softplus(b_s)
+        b_gate = -(exp(b_A) if HAS_A else b_A) * softplus(b_s)
     else:
         b_gate = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_s)
 

tests/ops/test_kda.py

@@ -292,6 +292,7 @@ def test_fused_recurrent_state_v_first(
             (4, 1024, 4, 4, 128, 0.1, True, True, True, torch.float16),
             (1, 64, 1, 1, 64, 1, False, False, True, torch.float),
             (2, 256, 2, 4, 64, 1, False, True, True, torch.float),
+            (2, 63, 2, 4, 64, 1, False, True, False, torch.bfloat16),
         ]
     ],
 )
@@ -494,6 +495,7 @@ def test_fused_recurrent_vllm_decode(
         "mask_p",
         "use_qk_l2norm_in_kernel",
         "use_gate_in_kernel",
+        "has_a_log",
         "dtype",
         "safe_gate",
         "disable_recompute",
@@ -504,28 +506,29 @@ def test_fused_recurrent_vllm_decode(
             *test,
             id=(
                 "B{}-T{}-H{}-HV{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}"
-                "-use_qk_l2norm{}-use_gate{}-{}-safe_gate{}-disable_recompute{}-chunk_size{}"
+                "-use_qk_l2norm{}-use_gate{}-has_a_log{}-{}-safe_gate{}-disable_recompute{}-chunk_size{}"
             ).format(*test),
         )
         for test in [
-            (1, 63, 1, 1, 64, 1, 1, 0, False, False, torch.float16, True, False, 64),
-            (2, 500, 3, 3, 60, 1, 1, 0, False, False, torch.float16, True, True, 64),
-            (2, 1000, 3, 3, 64, 0.1, 1, 0.5, False, False, torch.float16, False, True, 64),
-            (3, 1024, 4, 4, 100, 1, 0.1, 0, False, False, torch.float16, False, False, 64),
-            (4, 1024, 4, 4, 128, 0.1, 1, 0, False, False, torch.float16, True, True, 64),
-            (4, 1024, 4, 4, 128, 0.1, 1, 0, True, False, torch.float16, True, False, 64),
-            (2, 1500, 4, 4, 128, 0.1, 10, 0, False, True, torch.float16, False, True, 64),
-            (4, 2048, 8, 8, 64, 0.1, 1, 0, False, True, torch.float16, True, True, 64),
-
-            (2, 1024, 2, 4, 64, 0.1, 1, 0, False, False, torch.float16, False, False, 64),
-            (2, 1024, 2, 8, 64, 0.1, 1, 0, False, True, torch.float16, False, False, 64),
-            (2, 1024, 4, 8, 128, 0.1, 1, 0, True, True, torch.float16, False, False, 64),
-            (2, 160, 2, 4, 64, 0.1, 1, 0, False, True, torch.float16, True, True, 32),
-
-            (2, 1024, 2, 4, 64, 0.1, 1, 0, True, True, torch.bfloat16, True, False, 64),
-            (2, 1024, 2, 4, 64, 0.1, 1, 0, False, True, torch.bfloat16, False, False, 64),
-            (2, 160, 2, 4, 64, 0.1, 1, 0, False, True, torch.bfloat16, True, True, 32),
-            (2, 1024, 2, 8, 128, 0.1, 1, 0, True, True, torch.bfloat16, False, False, 64),
+            (1, 63, 1, 1, 64, 1, 1, 0, False, False, False, torch.float16, True, False, 64),
+            (2, 500, 3, 3, 60, 1, 1, 0, False, False, False, torch.float16, True, True, 64),
+            (2, 1000, 3, 3, 64, 0.1, 1, 0.5, False, False, False, torch.float16, False, True, 64),
+            (3, 1024, 4, 4, 100, 1, 0.1, 0, False, False, False, torch.float16, False, False, 64),
+            (4, 1024, 4, 4, 128, 0.1, 1, 0, False, False, False, torch.float16, True, True, 64),
+            (4, 1024, 4, 4, 128, 0.1, 1, 0, True, False, False, torch.float16, True, False, 64),
+            (2, 1500, 4, 4, 128, 0.1, 10, 0, False, True, True, torch.float16, False, True, 64),
+            (4, 2048, 8, 8, 64, 0.1, 1, 0, False, True, True, torch.float16, True, True, 64),
+
+            (2, 1024, 2, 4, 64, 0.1, 1, 0, False, False, False, torch.float16, False, False, 64),
+            (2, 1024, 2, 8, 64, 0.1, 1, 0, False, True, True, torch.float16, False, False, 64),
+            (2, 1024, 4, 8, 128, 0.1, 1, 0, True, True, True, torch.float16, False, False, 64),
+            (2, 160, 2, 4, 64, 0.1, 1, 0, False, True, True, torch.float16, True, True, 32),
+
+            (2, 1024, 2, 4, 64, 0.1, 1, 0, True, True, True, torch.bfloat16, True, False, 64),
+            (2, 1024, 2, 4, 64, 0.1, 1, 0, False, True, True, torch.bfloat16, False, False, 64),
+            (2, 160, 2, 4, 64, 0.1, 1, 0, False, True, True, torch.bfloat16, True, True, 32),
+            (2, 1024, 2, 8, 128, 0.1, 1, 0, True, True, True, torch.bfloat16, False, False, 64),
+            (1, 65, 1, 1, 64, 1, 1, 0, False, True, False, torch.float16, False, False, 64),
         ]
     ],
 )
@@ -540,6 +543,7 @@ def test_chunk(
     mask_p: float,
     use_qk_l2norm_in_kernel: bool,
     use_gate_in_kernel: bool,
+    has_a_log: bool,
     dtype: torch.dtype,
     safe_gate: bool,
     disable_recompute: bool,
@@ -550,8 +554,9 @@ def test_chunk(
     k = torch.rand(B, T, H, D, dtype=dtype)
     v = torch.rand(B, T, HV, D, dtype=dtype)
     g = torch.randn(B, T, HV, D, dtype=torch.float if not use_gate_in_kernel else dtype)
+    A_log = None
     if use_gate_in_kernel:
-        A_log = torch.randn(HV, dtype=torch.float)
+        A_log = torch.randn(HV, dtype=torch.float) if has_a_log else None
         dt_bias = torch.randn(HV * D, dtype=torch.float)
     else:
         g = F.logsigmoid(g) / gate_logit_normalizer
@@ -567,8 +572,10 @@ def test_chunk(
 
     beta = torch.randn(B, T, HV, dtype=dtype).sigmoid()
     h0 = torch.randn(B, HV, D, D, dtype=torch.float32)
+    if A_log is not None:
+        A_log = A_log.to(device).requires_grad_(True)
     if use_gate_in_kernel:
-        A_log, dt_bias = map(lambda x: x.to(device).requires_grad_(True), (A_log, dt_bias))
+        dt_bias = dt_bias.to(device).requires_grad_(True)
     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)
@@ -585,8 +592,9 @@ def test_chunk(
         output_final_state=True,
     )
     ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True)
-    if use_gate_in_kernel:
+    if A_log is not None:
         ref_dA, A_log.grad = A_log.grad, None
+    if use_gate_in_kernel:
         ref_dbias, dt_bias.grad = dt_bias.grad, None
     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
@@ -597,7 +605,7 @@ def test_chunk(
         v=v.clone(),
         g=g.clone(),
         beta=beta.clone(),
-        A_log=(A_log.clone() if use_gate_in_kernel else None),
+        A_log=(A_log.clone() if A_log is not None else None),
         dt_bias=(dt_bias.clone() if use_gate_in_kernel else None),
         scale=scale,
         initial_state=h0.clone(),
@@ -610,8 +618,9 @@ def test_chunk(
         chunk_size=chunk_size,
     )
     ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True)
-    if use_gate_in_kernel:
+    if A_log is not None:
         tri_dA, A_log.grad = A_log.grad, None
+    if use_gate_in_kernel:
         tri_dbias, dt_bias.grad = dt_bias.grad, None
     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
@@ -623,8 +632,9 @@ def test_chunk(
     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)
-    if use_gate_in_kernel:
+    if A_log is not None:
         assert_close("dA", ref_dA, tri_dA, 0.003, warning=True)
+    if use_gate_in_kernel:
         assert_close("dbias", ref_dbias, tri_dbias, 0.008)
     assert_close("dh0", ref_dh0, tri_dh0, 0.008)
 
@@ -1018,6 +1028,7 @@ def test_chunk_varlen_prefill(
             (2, 64, 4, 32, False, False, -5.0),
             (4, 128, 8, 64, False, True, -5.0),
             (4, 128, 8, 128, False, True, -5.0),
+            (1, 32, 2, 16, False, True, None),
         ]
     ],
 )
@@ -1053,7 +1064,11 @@ def test_gate(
             lower_bound=LOWER_BOUND,
         )
     else:
-        ref = naive_kda_gate(g=g.clone(), A_log=A_log.clone(), dt_bias=dt_bias.clone() if dt_bias is not None else None)
+        ref = naive_kda_gate(
+            g=g.clone(),
+            A_log=A_log.clone() if A_log is not None else None,
+            dt_bias=dt_bias.clone() if dt_bias is not None else None,
+        )
     tri = fused_kda_gate(
         g=g.clone(),
         A_log=A_log.clone() if A_log is not None else None,