[NPU] Avoid device synchronization in Ascend sampling (#39404)

python/sglang/srt/layers/sampler.py

@@ -615,6 +615,7 @@ class Sampler(nn.Module):
                 sampling_info.need_min_p_sampling,
                 sampling_info.sampling_seed,
                 positions,
+                npu_top_k_top_p_eligible=sampling_info.npu_top_k_top_p_eligible,
             )
             return batch_next_token_ids.to(torch.int32)
 
@@ -782,15 +783,14 @@ def top_k_top_p_min_p_sampling_from_logits_ascend(
     need_min_p_sampling: bool,
     sampling_seed: Optional[torch.Tensor],
     positions: torch.Tensor,
+    npu_top_k_top_p_eligible: bool = False,
 ):
     """A top-k, top-p and min-p sampling implementation for ascend npu with torch_npu interface.
 
     Takes temperature-scaled logits as input (softmax is applied internally).
     """
     # torch_npu.npu_top_k_top_p requires top_k value range in [1, 1024]
-    if hasattr(torch_npu, "npu_top_k_top_p") and torch.all(
-        (top_ks <= 1024) & (top_ks >= 1)
-    ):
+    if hasattr(torch_npu, "npu_top_k_top_p") and npu_top_k_top_p_eligible:
         logits_top_k_top_p = torch_npu.npu_top_k_top_p(logits, top_ps, top_ks)
         probs_top_k_top_p = logits_top_k_top_p.softmax(dim=-1)
 

python/sglang/srt/sampling/sampling_batch_info.py

@@ -85,6 +85,10 @@ class SamplingBatchInfo:
     # Handle logit bias
     logit_bias: Optional[torch.Tensor] = None
 
+    # Host-side eligibility for torch_npu.npu_top_k_top_p. Keeping this off the
+    # device avoids a scalar synchronization in the per-token sampling path.
+    npu_top_k_top_p_eligible: bool = False
+
     @classmethod
     def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
         enable_deterministic = get_exec().deterministic.enable_deterministic_inference
@@ -207,6 +211,9 @@ class SamplingBatchInfo:
             need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs),
             need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs),
             need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs),
+            npu_top_k_top_p_eligible=all(
+                1 <= r.sampling_params.top_k <= 1024 for r in reqs
+            ),
             vocab_size=vocab_size,
             penalizer_orchestrator=penalizer_orchestrator,
             has_custom_logit_processor=has_custom_logit_processor,
@@ -502,6 +509,7 @@ class SamplingBatchInfo:
         self.need_top_p_sampling |= other.need_top_p_sampling
         self.need_top_k_sampling |= other.need_top_k_sampling
         self.need_min_p_sampling |= other.need_min_p_sampling
+        self.npu_top_k_top_p_eligible &= other.npu_top_k_top_p_eligible
 
         self.adjusted_merge_batch(other)
 

test/registered/unit/sampling/test_ascend_sampler.py

@@ -0,0 +1,57 @@
+"""Unit tests for the Ascend sampling dispatch path."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from sglang.srt.layers import sampler as sampler_module
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestAscendSamplerDispatch(unittest.TestCase):
+    def test_top_k_dispatch_does_not_read_device_values_on_cpu(self):
+        logits = torch.tensor([[0.1, 0.2, 0.3, 0.4]])
+        top_ks = torch.tensor([2], dtype=torch.int32)
+        top_ps = torch.ones(1)
+        min_ps = torch.zeros(1)
+        positions = torch.zeros(1, dtype=torch.int64)
+
+        for eligible in (True, False):
+            with self.subTest(eligible=eligible):
+                npu_top_k_top_p = MagicMock(return_value=logits)
+                with (
+                    patch.object(
+                        sampler_module,
+                        "torch_npu",
+                        SimpleNamespace(npu_top_k_top_p=npu_top_k_top_p),
+                        create=True,
+                    ),
+                    patch.object(
+                        sampler_module.torch,
+                        "all",
+                        side_effect=AssertionError("device predicate was inspected"),
+                    ),
+                ):
+                    result = (
+                        sampler_module.top_k_top_p_min_p_sampling_from_logits_ascend(
+                            logits.clone(),
+                            top_ks.clone(),
+                            top_ps,
+                            min_ps,
+                            False,
+                            None,
+                            positions,
+                            npu_top_k_top_p_eligible=eligible,
+                        )
+                    )
+
+                self.assertEqual(tuple(result.shape), (1,))
+                self.assertEqual(npu_top_k_top_p.called, eligible)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/sampling/test_sampling_batch_info.py

@@ -492,18 +492,21 @@ class TestMergeBatch(CustomTestCase):
             need_top_p_sampling=False,
             need_top_k_sampling=False,
             need_min_p_sampling=False,
+            npu_top_k_top_p_eligible=True,
         )
         info2 = _make_info(
             is_all_greedy=False,
             need_top_p_sampling=True,
             need_top_k_sampling=True,
             need_min_p_sampling=True,
+            npu_top_k_top_p_eligible=False,
         )
         info1.merge_batch(info2)
         self.assertFalse(info1.is_all_greedy)  # AND semantics
         self.assertTrue(info1.need_top_p_sampling)  # OR semantics
         self.assertTrue(info1.need_top_k_sampling)  # OR semantics
         self.assertTrue(info1.need_min_p_sampling)  # OR semantics
+        self.assertFalse(info1.npu_top_k_top_p_eligible)  # AND semantics
 
     def test_merge_with_logit_bias(self):
         """Test that merge pads missing logit_bias with zeros before concatenation."""
@@ -675,6 +678,22 @@ class TestFromScheduleBatch(CustomTestCase):
         self.assertTrue(info.need_min_p_sampling)  # 0.1 > 0
         self.assertFalse(info.is_all_greedy)  # top_k=50 > 1
 
+    def test_npu_top_k_top_p_eligibility_uses_request_params(self):
+        cases = (
+            ((1, 1024), True),
+            ((1, 1025), False),
+            ((4, TOP_K_ALL), False),
+        )
+        for top_ks, expected in cases:
+            with self.subTest(top_ks=top_ks):
+                batch = MagicMock()
+                batch.reqs = [self._make_req(top_k=top_k) for top_k in top_ks]
+                batch.device = DEVICE
+
+                info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE)
+
+                self.assertEqual(info.npu_top_k_top_p_eligible, expected)
+
     def test_no_logit_bias_when_all_none(self):
         """Test that logit_bias stays None when no request has logit_bias set."""