From 171037c3e73d2fce139724fa8cd921afa1060983 Mon Sep 17 00:00:00 2001
From: Yuzhen Zhou <82826991+zyzshishui@users.noreply.github.com>
Date: Sat, 13 Jun 2026 23:23:06 -0700
Subject: [PATCH] Fix Qwen3.5 deterministic batch-invariant logprobs (#27869)

---
 .../layers/attention/fla/layernorm_gated.py   |  8 ++--
 .../moe/moe_runner/triton_utils/fused_moe.py  |  9 +++-
 .../attention/test_qwen35_deterministic.py    | 44 +++++++++++++++++++
 3 files changed, 56 insertions(+), 5 deletions(-)
 create mode 100644 test/registered/attention/test_qwen35_deterministic.py

diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py
index 798033634..a076d161c 100644
--- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py
+++ b/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
-    # torch.compile creating guards on the dynamic batch dimension.
-    if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
+    # Use a constant value when the row count must not affect kernel numerics.
+    if is_batch_invariant_mode_enabled() or check_cuda_graph_backend(
+        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))
diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
index a2f739b2d..b7245e5ec 100644
--- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+++ b/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,
diff --git a/test/registered/attention/test_qwen35_deterministic.py b/test/registered/attention/test_qwen35_deterministic.py
new file mode 100644
index 000000000..bb0d5c37c
--- /dev/null
+++ b/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()