[NPU]adaptation to support deterministic inference (#21197)

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,6 +987,7 @@ def enable_batch_invariant_mode(enable_bmm: bool = True):
     _batch_invariant_MODE = True
     _batch_invariant_LIB = torch.library.Library("aten", "IMPL")
 
+    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)
@@ -997,6 +1003,30 @@ def enable_batch_invariant_mode(enable_bmm: bool = True):
             # 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():

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:

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_

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,6 +293,12 @@ class Sampler(nn.Module):
         """
         if simple_sampling_case:
             probs = torch.softmax(logits, dim=-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:
@@ -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)
 
+        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)
 
+        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:

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,6 +4204,7 @@ class ServerArgs:
                 self.enable_flashinfer_allreduce_fusion = False
 
             # Check sampling backend
+            if self.sampling_backend != "ascend":
                 self.sampling_backend = "pytorch"
                 logger.warning(
                     "Sampling backend is set to pytorch for deterministic inference."

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