Add deterministic mode for XPU operations (#16793)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py

@@ -10,7 +10,12 @@ import triton
 import triton.language as tl
 
 from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
-from sglang.srt.utils.common import calc_diff, get_bool_env_var
+from sglang.srt.utils.common import (
+    calc_diff,
+    get_bool_env_var,
+    get_device_core_count,
+    get_dispatch_device_backend,
+)
 
 if ENABLE_JIT_DEEPGEMM:
     import deep_gemm
@@ -169,7 +174,7 @@ def _matmul_persistent_triton(
     assert (
         bias is None or bias.dim() == 1
     ), "Currently assuming bias is 1D, let Horace know if you run into this"
-    NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count
+    NUM_SMS = get_device_core_count()
     M, K = a.shape
     K, N = b.shape
     dtype = a.dtype
@@ -490,7 +495,7 @@ def mean_dim(
         Tensor with mean values along specified dimension
     """
     # Validate inputs
-    assert input.is_cuda, "Input must be a CUDA tensor"
+    assert input.is_cuda or input.is_xpu, "Input must be a CUDA or XPU tensor"
     assert (
         -input.ndim <= dim < input.ndim
     ), f"Invalid dimension {dim} for tensor with {input.ndim} dimensions"
@@ -733,7 +738,7 @@ def bmm_batch_invariant(a, b, *, out=None):
         else:
             c = out
 
-        NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count
+        NUM_SMS = get_device_core_count()
 
         # Use fixed kernel configuration for determinism
         configs = {
@@ -938,25 +943,26 @@ def is_batch_invariant_mode_enabled():
     return _batch_invariant_MODE
 
 
-def enable_batch_invariant_mode(
+def enable_batch_invariant_mode(enable_bmm: bool = True):
-    enable_bmm: bool = True,
-):
     global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
     if _batch_invariant_MODE:
         return
 
+    dispatch_key = get_dispatch_device_backend()
+
     _batch_invariant_MODE = True
     _batch_invariant_LIB = torch.library.Library("aten", "IMPL")
-    _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")
+
-    _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
+    # Register for detected device
+    _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, dispatch_key)
+    _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, dispatch_key)
     _batch_invariant_LIB.impl(
-        "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"
+        "aten::_log_softmax", _log_softmax_batch_invariant, dispatch_key
     )
-    _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
+    _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, dispatch_key)
 
     if enable_bmm:
-        _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA")
+        _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, dispatch_key)
-
         # Also monkeypatch torch.bmm directly as a fallback
         _original_torch_bmm = torch.bmm
         torch.bmm = bmm_batch_invariant

python/sglang/srt/layers/layernorm.py

@@ -439,6 +439,17 @@ class RMSNorm(MultiPlatformOp):
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if self.variance_size_override is not None:
             return self.forward_native(x, residual, post_residual_addition)
+        if is_batch_invariant_mode_enabled():
+            if (
+                residual is not None
+                or get_global_server_args().rl_on_policy_target == "fsdp"
+            ):
+                return self.forward_native(x, residual, post_residual_addition)
+            return rms_norm_batch_invariant(
+                x,
+                self.weight.data,
+                self.variance_epsilon,
+            )
         if residual is not None:
             if post_residual_addition is not None:
                 residual = residual + post_residual_addition

python/sglang/srt/utils/common.py

@@ -688,6 +688,16 @@ def make_layers_non_pp(
     return layers
 
 
+def get_dispatch_device_backend():
+    if is_cuda_alike():
+        dispatch_key = "CUDA"
+    elif is_xpu():
+        dispatch_key = "XPU"
+    else:
+        raise RuntimeError("No supported accelerator (CUDA/XPU) available")
+    return dispatch_key
+
+
 @lru_cache(maxsize=1)
 def get_device_module():
     return torch.get_device_module()
@@ -700,6 +710,8 @@ def set_random_seed(seed: int) -> None:
     torch.manual_seed(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
+    if torch.xpu.is_available():
+        torch.xpu.manual_seed_all(seed)
 
 
 def load_audio(
@@ -1948,6 +1960,8 @@ def get_device_count() -> int:
 def get_device_core_count(device_id: int = 0) -> int:
     if (hasattr(torch, "cuda") and torch.cuda.is_available()) or is_musa():
         return torch.cuda.get_device_properties(device_id).multi_processor_count
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        return torch.xpu.get_device_properties(device_id).gpu_eu_count
 
     return 0