From 6abef3862734f4328ee9fd3f1d574f96d3826d93 Mon Sep 17 00:00:00 2001 From: Estrella-xx <820167017@qq.com> Date: Tue, 9 Jun 2026 09:22:09 +0800 Subject: [PATCH] [NPU]adaptation to support deterministic inference (#21197) --- .../batch_invariant_ops.py | 58 ++++++++++++----- .../npu/attention/ascend_backend.py | 1 - .../npu_batch_invariant_ops.py | 62 +++++++++++++++++++ python/sglang/srt/layers/sampler.py | 46 ++++++++++++-- python/sglang/srt/server_args.py | 13 ++-- python/sglang/srt/utils/common.py | 2 + 6 files changed, 155 insertions(+), 27 deletions(-) create mode 100644 python/sglang/srt/hardware_backend/npu/batch_invariant_ops/npu_batch_invariant_ops.py 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 482d75028..51a911eb1 100644 --- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py +++ b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py @@ -10,6 +10,7 @@ import triton import triton.language as tl from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM +from sglang.srt.utils import is_npu from sglang.srt.utils.common import ( calc_diff, get_bool_env_var, @@ -17,6 +18,10 @@ from sglang.srt.utils.common import ( get_dispatch_device_backend, ) +_is_npu = is_npu() +if _is_npu: + import torch_npu + if ENABLE_JIT_DEEPGEMM: import deep_gemm @@ -982,21 +987,46 @@ def enable_batch_invariant_mode(enable_bmm: bool = True): _batch_invariant_MODE = True _batch_invariant_LIB = torch.library.Library("aten", "IMPL") - # 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, dispatch_key - ) - _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, dispatch_key) - _batch_invariant_LIB.impl("aten::rms_norm", _rms_norm_aten_compat, dispatch_key) - _batch_invariant_LIB.impl("aten::mm.dtype", _mm_dtype_compat, dispatch_key) + if not _is_npu: + # 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, dispatch_key + ) + _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, dispatch_key) + _batch_invariant_LIB.impl("aten::rms_norm", _rms_norm_aten_compat, dispatch_key) + _batch_invariant_LIB.impl("aten::mm.dtype", _mm_dtype_compat, dispatch_key) - if enable_bmm: - _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 + if enable_bmm: + _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 + else: + from sglang.srt.hardware_backend.npu.batch_invariant_ops.npu_batch_invariant_ops import ( + npu_add_rms_norm_batch_invariant, + npu_fused_infer_attention_score_batch_invariant, + npu_log_softmax_batch_invariant, + npu_matmul_batch_invariant, + npu_mean_batch_invariant, + npu_mm_batch_invariant, + ) + + _batch_invariant_LIB.impl("aten::mm", npu_mm_batch_invariant, dispatch_key) + _batch_invariant_LIB.impl( + "aten::matmul", npu_matmul_batch_invariant, dispatch_key + ) + _batch_invariant_LIB.impl( + "aten::mean.dim", npu_mean_batch_invariant, dispatch_key + ) + _batch_invariant_LIB.impl( + "aten::_log_softmax", npu_log_softmax_batch_invariant, dispatch_key + ) + torch.ops.npu.npu_fused_infer_attention_score = ( + npu_fused_infer_attention_score_batch_invariant + ) + torch_npu.npu_add_rms_norm = npu_add_rms_norm_batch_invariant def disable_batch_invariant_mode(): diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index c61dd84a6..d0a4beb2d 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -1164,7 +1164,6 @@ class AscendAttnBackend(AttentionBackend): return attn_output if self.use_fia: - """FIA supports multi-bs in the current version of CANN""" q = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) num_token_padding = q.shape[0] if num_token_padding > forward_batch.num_token_non_padded_cpu: diff --git a/python/sglang/srt/hardware_backend/npu/batch_invariant_ops/npu_batch_invariant_ops.py b/python/sglang/srt/hardware_backend/npu/batch_invariant_ops/npu_batch_invariant_ops.py new file mode 100644 index 000000000..b947c3c89 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/batch_invariant_ops/npu_batch_invariant_ops.py @@ -0,0 +1,62 @@ +# Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/batch_invariant_ops/batch_invariant_ops.py + +import batch_invariant_ops # noqa: F401 +import torch +import torch_npu + + +def npu_mm_batch_invariant(a, b): + return torch.ops.batch_invariant_ops.npu_mm_batch_invariant(a, b) + + +def npu_matmul_batch_invariant(a, b): + return torch.ops.batch_invariant_ops.npu_matmul_batch_invariant(a, b) + + +def npu_mean_batch_invariant( + input, dim, keepdim=False, dtype: torch.dtype | None = None +): + assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}" + if len(dim) == 1: + return torch.ops.batch_invariant_ops.npu_reduce_mean_batch_invariant( + input, dim[0], keepdim=keepdim + ) + else: + assert input.dtype in { + torch.float16, + torch.bfloat16, + torch.float32, + }, "only float types supported for now" + n_elems = 1 + for d in dim: + n_elems *= input.shape[d] + return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems + + +def npu_log_softmax_batch_invariant(input, dim, _half_to_float): + assert not _half_to_float, "not implemented" + return torch.ops.batch_invariant_ops.npu_log_softmax_batch_invariant(input, dim=dim) + + +def npu_fused_infer_attention_score_batch_invariant(*args, **kwargs): + return ( + torch.ops.batch_invariant_ops.npu_fused_infer_attention_score_batch_invariant( + *args, **kwargs + ) + ) + + +def npu_add_rms_norm_batch_invariant( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, +): + """ + AclnnAddRmsNorm can't ensure batch invariant, + so we need to split it into add and rms_norm. + """ + x_ = x + residual + residual_ = x_ + x_, _ = torch_npu.npu_rms_norm(x_, weight, eps) + return x_, None, residual_ diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index bc79cca97..84d57b8e2 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -154,7 +154,11 @@ class Sampler(nn.Module): if self.use_ascend_backend: # Ascend backend: sample from logits directly. batch_next_token_ids, logprobs = self._forward_ascend_backend( - logits, sampling_info, simple_sampling_case, return_logprob + logits, + sampling_info, + simple_sampling_case, + return_logprob, + positions, ) elif ( self.use_log_softmax_logprob @@ -281,6 +285,7 @@ class Sampler(nn.Module): logits: torch.Tensor, sampling_info: SamplingBatchInfo, simple_sampling_case: bool, + positions: torch.Tensor, ) -> torch.Tensor: """Sample from temperature-scaled logits without softmax. @@ -288,7 +293,13 @@ class Sampler(nn.Module): """ if simple_sampling_case: probs = torch.softmax(logits, dim=-1) - batch_next_token_ids = torch.multinomial(probs, num_samples=1).view(-1) + if sampling_info.sampling_seed is not None: + probabilities = probs.to(torch.float64).log_() + batch_next_token_ids = multinomial_with_seed( + probabilities, sampling_info.sampling_seed, positions + ).view(-1) + else: + batch_next_token_ids = torch.multinomial(probs, num_samples=1).view(-1) return batch_next_token_ids.to(torch.int32) else: assert ( @@ -300,6 +311,8 @@ class Sampler(nn.Module): sampling_info.top_ps, sampling_info.min_ps, sampling_info.need_min_p_sampling, + sampling_info.sampling_seed, + positions, ) return batch_next_token_ids.to(torch.int32) @@ -309,6 +322,7 @@ class Sampler(nn.Module): sampling_info: SamplingBatchInfo, simple_sampling_case: bool, return_logprob: bool, + positions: torch.Tensor, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Handle the full Ascend backend sampling path. @@ -321,7 +335,7 @@ class Sampler(nn.Module): """ logits.div_(sampling_info.temperatures) batch_next_token_ids = self._sample_from_logits( - logits, sampling_info, simple_sampling_case + logits, sampling_info, simple_sampling_case, positions ) logprobs = None if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB: @@ -517,6 +531,8 @@ def top_k_top_p_min_p_sampling_from_logits_ascend( top_ps: torch.Tensor, min_ps: torch.Tensor, need_min_p_sampling: bool, + sampling_seed: Optional[torch.Tensor], + positions: torch.Tensor, ): """A top-k, top-p and min-p sampling implementation for ascend npu with torch_npu interface. @@ -534,7 +550,17 @@ def top_k_top_p_min_p_sampling_from_logits_ascend( min_p_mask = probs_top_k_top_p < min_p_thresholds.view(-1, 1) probs_top_k_top_p.masked_fill_(min_p_mask, 0.0) - batch_next_token_ids = torch.multinomial(probs_top_k_top_p, num_samples=1) + if sampling_seed is None: + batch_next_token_ids = torch.multinomial(probs_top_k_top_p, num_samples=1) + else: + logprobs_top_k_top_p = probs_top_k_top_p.to( + torch.float64 + ) # Using float64 for numerical stability + del probs_top_k_top_p + logprobs_top_k_top_p.log_() + batch_next_token_ids = multinomial_with_seed( + logprobs_top_k_top_p, sampling_seed, positions + ) else: probs = torch.softmax(logits, dim=-1) probs_sort, probs_idx = probs.sort(dim=-1, descending=True) @@ -556,14 +582,22 @@ def top_k_top_p_min_p_sampling_from_logits_ascend( min_p_mask = probs_sort < min_p_thresholds.view(-1, 1) probs_sort.masked_fill_(min_p_mask, 0.0) - sampled_index = torch.multinomial(probs_sort, num_samples=1) + if sampling_seed is None: + sampled_index = torch.multinomial(probs_sort, num_samples=1) + else: + logprobs = probs_sort.to( + torch.float64 + ) # Using float64 for numerical stability + del probs_sort + logprobs.log_() + sampled_index = multinomial_with_seed(logprobs, sampling_seed, positions) probs_idx = probs_idx.to(torch.int32) batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index) return batch_next_token_ids.view(-1) -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, disable=is_npu()) def multinomial_with_seed( logprobs: torch.Tensor, seed: torch.Tensor, positions: torch.Tensor ) -> torch.Tensor: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 375354b03..9435f9254 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -196,9 +196,9 @@ ATTENTION_BACKEND_CHOICES = [ "intel_xpu", ] -DETERMINISTIC_ATTENTION_BACKEND_CHOICES = ["flashinfer", "fa3", "triton"] +DETERMINISTIC_ATTENTION_BACKEND_CHOICES = ["flashinfer", "fa3", "triton", "ascend"] -RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["fa3", "triton"] +RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["fa3", "triton", "ascend"] DISAGG_TRANSFER_BACKEND_CHOICES = [ "mooncake", @@ -4204,10 +4204,11 @@ class ServerArgs: self.enable_flashinfer_allreduce_fusion = False # Check sampling backend - self.sampling_backend = "pytorch" - logger.warning( - "Sampling backend is set to pytorch for deterministic inference." - ) + if self.sampling_backend != "ascend": + self.sampling_backend = "pytorch" + logger.warning( + "Sampling backend is set to pytorch for deterministic inference." + ) is_deepseek_model = False if parse_connector_type(self.model_path) != ConnectorType.INSTANCE: try: diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 3606bc36a..4556d06b1 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -759,6 +759,8 @@ def get_dispatch_device_backend(): dispatch_key = "CUDA" elif is_xpu(): dispatch_key = "XPU" + elif is_npu(): + dispatch_key = "NPU" else: raise RuntimeError("No supported accelerator (CUDA/XPU) available") return dispatch_key