From a4ffb996db3e823708a53637f851e2677d4455a8 Mon Sep 17 00:00:00 2001
From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Date: Thu, 20 Aug 2026 19:38:53 +0800
Subject: [PATCH] [Fix] Keep deterministic GDN prefill on Triton (#35632)

---
 .../layers/attention/linear/gdn_backend.py    | 15 ++++++
 .../test_gdn_prefill_backend_policy.py        | 47 ++++++++++++++++++-
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py
index 13d026efa..304fe1596 100644
--- a/python/sglang/srt/layers/attention/linear/gdn_backend.py
+++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py
@@ -11,6 +11,7 @@ from sglang.srt.configs.hybrid_arch import hybrid_gdn_config
 from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
 from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
 from sglang.srt.layers.attention.linear.utils import (
+    LinearAttnBackends,
     LinearAttnKernelBackend,
     build_verify_intermediate_state_indices,
 )
@@ -71,6 +72,7 @@ def flashinfer_gdn_prefill_default(model_runner: ModelRunner) -> Optional[str]:
     if (
         get_exec().mamba.linear_attn_prefill_backend is not None
         or get_exec().mamba.linear_attn_backend != "triton"
+        or get_exec().deterministic.enable_deterministic_inference
         or get_memory().enable_page_major_kv_layout
         or sm_major not in (9, 10)
     ):
@@ -113,6 +115,18 @@ def flashinfer_gdn_prefill_default(model_runner: ModelRunner) -> Optional[str]:
     return "flashinfer"
 
 
+def _validate_gdn_linear_attn_backends(backends: LinearAttnBackends) -> None:
+    if (
+        get_exec().deterministic.enable_deterministic_inference
+        and backends.prefill.is_flashinfer()
+    ):
+        raise ValueError(
+            "FlashInfer GDN prefill is not supported with "
+            "--enable-deterministic-inference. Use "
+            "--linear-attn-prefill-backend triton."
+        )
+
+
 class GDNKernelDispatcher:
     """Dispatches GDN kernel calls to the appropriate backend per mode."""
 
@@ -359,6 +373,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
     needs_cpu_seq_lens: bool = False
 
     def __init__(self, model_runner: ModelRunner):
+        _validate_gdn_linear_attn_backends(model_runner.linear_attn_backends)
         super().__init__(model_runner)
         self.conv_states_shape = (
             model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
diff --git a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py
index 96d44a70e..296e834af 100644
--- a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py
+++ b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py
@@ -11,15 +11,20 @@ from sglang.srt.layers.attention.linear import gdn_backend
 from sglang.srt.layers.attention.linear.gdn_backend import (
     GDNAttnBackend,
     GDNKernelDispatcher,
+    _validate_gdn_linear_attn_backends,
     flashinfer_gdn_prefill_default,
 )
 from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
     maybe_build_flashinfer_checkpoint_plan,
 )
 from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
-from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend
+from sglang.srt.layers.attention.linear.utils import (
+    LinearAttnKernelBackend,
+    resolve_linear_attn_backends,
+)
 from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
@@ -73,7 +78,7 @@ def make_runner(
     )
 
 
-class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
+class TestFlashInferGDNPrefillBackendPolicy(CustomTestCase):
     def apply_policy(
         self,
         runner,
@@ -119,6 +124,44 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
                 runner = make_runner(self, linear_attn_prefill_backend=backend)
                 self.assertIsNone(self.apply_policy(runner))
 
+    def test_declines_when_deterministic_inference_is_enabled(self):
+        """Batch-sensitive GDN prefill must not bypass deterministic inference."""
+        runner = make_runner(
+            self,
+            state_dtype=torch.float32,
+            enable_deterministic_inference=True,
+        )
+
+        self.assertIsNone(
+            self.apply_policy(
+                runner,
+                capability=(9, 0),
+                cuda_version="12.9",
+            )
+        )
+
+    def test_rejects_explicit_flashinfer_prefill_in_deterministic_mode(self):
+        """Explicit backend precedence must not bypass deterministic GDN startup."""
+        cases = (
+            {"linear_attn_prefill_backend": "flashinfer"},
+            {"linear_attn_backend": "flashinfer"},
+        )
+        for fields in cases:
+            with self.subTest(fields=fields):
+                make_runner(
+                    self,
+                    enable_deterministic_inference=True,
+                    **fields,
+                )
+                backends = resolve_linear_attn_backends()
+
+                with self.assertRaisesRegex(
+                    ValueError,
+                    "FlashInfer GDN prefill is not supported with "
+                    "--enable-deterministic-inference",
+                ):
+                    _validate_gdn_linear_attn_backends(backends)
+
     def test_rejects_unsupported_capability(self):
         cases = (
             ("non_cuda", {}, {"cuda": False}),