Fix Qwen3.5 deterministic batch-invariant logprobs (#27869)

python/sglang/srt/layers/attention/fla/layernorm_gated.py

@@ -15,6 +15,7 @@ import triton.language as tl
 from einops import rearrange
 
 from sglang.jit_kernel.utils import is_arch_support_pdl
+from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
 from sglang.srt.model_executor.cuda_graph_config import (
     Backend,
     Phase,
@@ -192,9 +193,10 @@ def _get_sm_count(device: torch.device) -> int:
 
 
 def calc_rows_per_block(M: int, device: torch.device) -> int:
-    # When piecewise cuda graph is enabled, use a constant value to avoid
+    # Use a constant value when the row count must not affect kernel numerics.
-    # torch.compile creating guards on the dynamic batch dimension.
+    if is_batch_invariant_mode_enabled() or check_cuda_graph_backend(
-    if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
+        Phase.PREFILL, Backend.TC_PIECEWISE
+    ):
         return MAX_ROWS_PER_BLOCK
     sm_count = _get_sm_count(device)
     rows_per_block = next_power_of_2(cdiv(M, 2 * sm_count))

python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py

@@ -13,6 +13,7 @@ import torch
 import torch.nn.functional as F
 import triton.language as tl
 
+from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
 from sglang.srt.environ import envs
 from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
 from sglang.srt.layers.moe.utils import get_moe_padding_size
@@ -88,6 +89,10 @@ if not _is_cuda and not _is_hip and not _is_xpu:
 padding_size = get_moe_padding_size(_use_aiter)
 
 
+def _use_moe_sum_reduce_torch_compile(num_tokens: int) -> bool:
+    return num_tokens <= 32 and not is_batch_invariant_mode_enabled()
+
+
 @register_custom_op(mutates_args=["hidden_states"])
 def inplace_fused_experts(
     hidden_states: torch.Tensor,
@@ -728,7 +733,7 @@ def _fused_moe_kernel_sequence(
             ).squeeze(dim=1)
         else:
             # According to micro benchmark results, torch.compile can get better performance for small token.
-            if num_tokens <= 32:
+            if _use_moe_sum_reduce_torch_compile(num_tokens):
                 moe_sum_reduce_torch_compile(
                     intermediate_cache3.view(*intermediate_cache3.shape),
                     out_hidden_states,
@@ -748,7 +753,7 @@ def _fused_moe_kernel_sequence(
             )
         else:
             # According to micro benchmark results, torch.compile can get better performance for small token.
-            if num_tokens <= 32:
+            if _use_moe_sum_reduce_torch_compile(num_tokens):
                 moe_sum_reduce_torch_compile(
                     intermediate_cache3.view(*intermediate_cache3.shape),
                     out_hidden_states,

test/registered/attention/test_qwen35_deterministic.py

@@ -0,0 +1,44 @@
+"""
+Usage:
+cd test/srt
+python3 -m unittest test_qwen35_deterministic.TestQwen35Fa3Deterministic
+"""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_deterministic_utils import (
+    COMMON_SERVER_ARGS,
+    TestDeterministicBase,
+)
+
+register_cuda_ci(est_time=360, stage="extra-b", runner_config="4-gpu-h100")
+
+QWEN35 = "Qwen/Qwen3.5-35B-A3B"
+
+
+class TestQwen35Fa3Deterministic(TestDeterministicBase):
+    @classmethod
+    def get_model(cls):
+        return QWEN35
+
+    @classmethod
+    def get_server_args(cls):
+        return list(COMMON_SERVER_ARGS) + [
+            "--tp",
+            "4",
+            "--attention-backend",
+            "fa3",
+            "--skip-server-warmup",
+            "--mamba-scheduler-strategy",
+            "extra_buffer",
+            "--enable-flashinfer-allreduce-fusion",
+            "--tokenizer-worker-num",
+            "6",
+            "--mem-fraction-static",
+            "0.8",
+        ]
+
+
+if __name__ == "__main__":
+    unittest.main()