[KDA] Fix missing beta sigmoid in PTX prefill (#40685)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>

python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py

@@ -19,7 +19,7 @@ Correctness-sensitive cases stay on Triton:
 Single-sequence token counts that are not a multiple of the kernel's 64-token
 chunk are padded up to a bucket (1k/2k/4k/8k/16k/32k) in a persistent staging
 buffer, which bounds the resident workspace set. Pad rows are state-neutral:
-k/v/beta zero => no rank-1 update; raw gate -1000 => transformed decay of
+k/v zero => no rank-1 update, even with beta sigmoid; raw gate -1000 => decay of
 exactly 1. Multi-sequence batches go through the kernel's own varlen grid
 (real cu_seqlens, no padding), so their shapes are whatever the scheduler
 produces and each distinct shape can retain another workspace.
@@ -327,6 +327,7 @@ class PtxKDAKernel(LinearAttnKernelBase):
                 dt_bias=self._flat_param(dt_bias),
                 return_intermediate_states=return_intermediate_states,
                 use_qk_l2norm_in_kernel=True,
+                use_beta_sigmoid_in_kernel=kwargs.get("beta_is_raw", False),
             )
             out, final_state, h = result[0], result[1], result[10]
             ssm_states.index_copy_(0, slot, final_state.to(ssm_states.dtype))

test/registered/kernels/ops/attention/test_kda_prefill.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import patch
 
 import torch
 import torch.nn.functional as F
@@ -10,6 +11,8 @@ from sglang.kernels.ops.attention.linear.kda_nvidia_prefill import (
 from sglang.kernels.ops.attention.linear.kda_ptx_prefill import (
     chunk_kda_fwd as ptx_chunk_kda_fwd,
 )
+from sglang.srt.layers.attention.linear.kernels.kda_ptx import PtxKDAKernel
+from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -79,6 +82,45 @@ def _reference(q, k, v, gate, beta, a_log, dt_bias, state, fused_qk_norm):
 
 
 class TestKdaPrefill(CustomTestCase):
+    @torch.inference_mode()
+    def test_ptx_padded_raw_beta(self):
+        """Raw beta must match Triton, including final state after neutral padding."""
+        if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (
+            10,
+            3,
+        ):
+            self.skipTest("PTX KDA prefill requires GB300")
+        q, k, v, gate, beta, a_log, dt_bias, state = _inputs(2, seq_len=1025)
+        state.fill_(0.1)
+        actual_state = state.clone()
+        inputs = dict(
+            q=q,
+            k=k,
+            v=v,
+            g=gate,
+            beta=beta,
+            cache_indices=torch.zeros(1, device="cuda", dtype=torch.int32),
+            query_start_loc=torch.tensor([0, 1025], device="cuda", dtype=torch.int32),
+            A_log=a_log,
+            dt_bias=dt_bias,
+            lower_bound=-5.0,
+            beta_is_raw=True,
+            extend_seq_lens_cpu=[1025],
+        )
+        kernel = PtxKDAKernel()
+        with patch.object(
+            kernel._triton,
+            "extend",
+            side_effect=AssertionError("PTX unexpectedly fell back to Triton"),
+        ):
+            actual = kernel.extend(**inputs, ssm_states=actual_state)
+        # Triton may mutate inputs, so run the reference last.
+        expected = TritonKDAKernel().extend(**inputs, ssm_states=state)
+        torch.testing.assert_close(
+            actual.float(), expected.float(), rtol=2e-2, atol=3e-2
+        )
+        torch.testing.assert_close(actual_state, state, rtol=2e-2, atol=3e-2)
+
     @torch.inference_mode()
     def test_nvidia_prefill(self):
         if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:

test/registered/unit/layers/attention/linear/kernels/test_kda_ptx.py

@@ -95,8 +95,10 @@ class TestPtxKDATrackRouting(CustomTestCase):
         kernel = self._make_kernel()
         kernel._triton = _RejectTriton()
         h = torch.zeros(3, 2, 128, 128, dtype=torch.float32)
+        beta_flags = []
 
         def fake_fwd(*args, **kwargs):
+            beta_flags.append(kwargs["use_beta_sigmoid_in_kernel"])
             return [
                 args[2].clone(),  # out == v
                 kwargs["initial_state"].clone(),  # final_state
@@ -107,24 +109,16 @@ class TestPtxKDATrackRouting(CustomTestCase):
         kernel._fwd = fake_fwd
         x = self._inputs()
 
+        for beta_kwargs in ({"beta_is_raw": True}, {}, {"beta_is_raw": False}):
             out, h_out = kernel.extend(
-            x["q"],
-            x["k"],
-            x["v"],
-            x["g"],
-            x["beta"],
-            ssm_states=x["ssm_states"],
-            cache_indices=x["cache_indices"],
-            query_start_loc=x["query_start_loc"],
-            A_log=x["A_log"],
-            dt_bias=x["dt_bias"],
+                **x,
+                **beta_kwargs,
                 return_intermediate_states=True,
                 track_ssm_h_src=torch.empty(0, dtype=torch.long),
-            extend_seq_lens_cpu=x["extend_seq_lens_cpu"],
             )
-
             self.assertEqual(tuple(out.shape), (1, 164, 2, 128))
             self.assertIs(h_out, h)
+        self.assertEqual(beta_flags, [True, False, False])
 
 
 if __name__ == "__main__":