From 5beb2fd5528bd39f9bffe16a5e2ed2cb50ab2953 Mon Sep 17 00:00:00 2001 From: Jensen Date: Wed, 16 Sep 2026 15:15:16 +0800 Subject: [PATCH] [NPU] Avoid device synchronization in Ascend sampling (#39404) --- python/sglang/srt/layers/sampler.py | 6 +- .../srt/sampling/sampling_batch_info.py | 8 +++ .../unit/sampling/test_ascend_sampler.py | 57 +++++++++++++++++++ .../unit/sampling/test_sampling_batch_info.py | 19 +++++++ 4 files changed, 87 insertions(+), 3 deletions(-) create mode 100644 test/registered/unit/sampling/test_ascend_sampler.py diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index d5ee8d826..abf59fd1c 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/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) diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index fad02f235..966e91d4d 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/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) diff --git a/test/registered/unit/sampling/test_ascend_sampler.py b/test/registered/unit/sampling/test_ascend_sampler.py new file mode 100644 index 000000000..70786258a --- /dev/null +++ b/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() diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 91cefe30a..7e52f64ef 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/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."""