From bcb34da9f9f7bce82a867b2ea30275e83537a128 Mon Sep 17 00:00:00 2001 From: Jay Thakur Date: Thu, 30 Apr 2026 11:09:06 +0530 Subject: [PATCH] Add deterministic mode for XPU operations (#16793) Co-authored-by: Ma Mingfei --- .../batch_invariant_ops.py | 32 +++++++++++-------- python/sglang/srt/layers/layernorm.py | 11 +++++++ python/sglang/srt/utils/common.py | 14 ++++++++ 3 files changed, 44 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py index 80bb66e7e..30a8cdb46 100644 --- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py +++ b/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( - enable_bmm: bool = True, -): +def enable_batch_invariant_mode(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 diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index b582ac5e0..3ac232dd7 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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 diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 3d3dc7b03..b0244416c 100644 --- a/python/sglang/srt/utils/common.py +++ b/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