From 4c81cd1b0904f8f7778e87c0974b464aab21b74c Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Tue, 22 Sep 2026 06:24:38 +0000 Subject: [PATCH] [KDA] Fix missing beta sigmoid in PTX prefill (#40685) Co-authored-by: Mohammad Angkad --- .../attention/linear/kernels/kda_ptx.py | 3 +- .../kernels/ops/attention/test_kda_prefill.py | 42 +++++++++++++++++++ .../attention/linear/kernels/test_kda_ptx.py | 30 ++++++------- 3 files changed, 56 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py b/python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py index bcd0b707b..1d5e07896 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py +++ b/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)) diff --git a/test/registered/kernels/ops/attention/test_kda_prefill.py b/test/registered/kernels/ops/attention/test_kda_prefill.py index 157d911c3..a15bac22a 100644 --- a/test/registered/kernels/ops/attention/test_kda_prefill.py +++ b/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: diff --git a/test/registered/unit/layers/attention/linear/kernels/test_kda_ptx.py b/test/registered/unit/layers/attention/linear/kernels/test_kda_ptx.py index 231db3cd2..5509c7ef7 100644 --- a/test/registered/unit/layers/attention/linear/kernels/test_kda_ptx.py +++ b/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() - 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"], - 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) + for beta_kwargs in ({"beta_is_raw": True}, {}, {"beta_is_raw": False}): + out, h_out = kernel.extend( + **x, + **beta_kwargs, + return_intermediate_states=True, + track_ssm_h_src=torch.empty(0, dtype=torch.long), + ) + self.assertEqual(tuple(out.shape), (1, 164, 2, 128)) + self.assertIs(h_out, h) + self.assertEqual(beta_flags, [True, False, False]) if __name__ == "__main__":