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
-    # 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))

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,