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

2026-06-09 09:22:09 +08:00
parent ab70153b62
commit 6abef38627
6 changed files with 155 additions and 27 deletions
@@ -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():
@@ -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:
@@ -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_
@@ -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:
@@ -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:
@@ -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