From 10592e9c08b9e1ceaf111ed22de12b37c4bc0111 Mon Sep 17 00:00:00 2001 From: Zhihao Lyu Date: Sat, 15 Nov 2025 21:59:58 +0800 Subject: [PATCH] [Ascend][Feat] Add Ascend sampling backend (#12692) Co-authored-by: ronnie_zheng --- docs/advanced_features/server_arguments.md | 2 +- python/sglang/srt/layers/sampler.py | 68 ++++++++++++- python/sglang/srt/server_args.py | 2 +- .../ascend/test_ascend_sampling_backend.py | 96 +++++++++++++++++++ test/srt/run_suite.py | 1 + 5 files changed, 165 insertions(+), 4 deletions(-) create mode 100644 test/srt/ascend/test_ascend_sampling_backend.py diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index ae17ada40..d483738b7 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -239,7 +239,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--attention-backend` | Choose the kernels for attention layers. | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | | `--prefill-attention-backend` | Choose the kernels for prefill attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | | `--decode-attention-backend` | Choose the kernels for decode attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | -| `--sampling-backend` | Choose the kernels for sampling layers. | `None` | `flashinfer`, `pytorch` | +| `--sampling-backend` | Choose the kernels for sampling layers. | `None` | `flashinfer`, `pytorch`, `ascend` | | `--grammar-backend` | Choose the backend for grammar-guided decoding. | `None` | `xgrammar`, `outlines`, `llguidance`, `none` | | `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `triton_attn`, `ascend_attn`, `aiter_attn` | | `--nsa-prefill` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter` | diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 59a0f3bb9..7f6f6a010 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -12,8 +12,9 @@ from sglang.srt.layers.dp_attention import ( ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo +from sglang.srt.sampling.sampling_params import TOP_K_ALL from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import crash_on_warnings, get_bool_env_var, is_cuda +from sglang.srt.utils import crash_on_warnings, get_bool_env_var, is_cuda, is_npu if is_cuda(): from sgl_kernel import ( @@ -23,6 +24,8 @@ if is_cuda(): top_p_renorm_prob, ) +if is_npu(): + import torch_npu logger = logging.getLogger(__name__) @@ -112,7 +115,11 @@ class Sampler(nn.Module): # Post process logits logits.div_(sampling_info.temperatures) - logits[:] = torch.softmax(logits, dim=-1) + # For ascend backend, softmax is not needed before sampling + if not get_global_server_args().sampling_backend == "ascend" or ( + return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB + ): + logits[:] = torch.softmax(logits, dim=-1) probs = logits del logits @@ -150,6 +157,14 @@ class Sampler(nn.Module): sampling_info.sampling_seed, positions, ) + elif get_global_server_args().sampling_backend == "ascend": + batch_next_token_ids = top_k_top_p_min_p_sampling_from_probs_ascend( + probs, + sampling_info.top_ks, + sampling_info.top_ps, + sampling_info.min_ps, + sampling_info.need_min_p_sampling, + ) else: raise ValueError( f"Invalid sampling backend: {get_global_server_args().sampling_backend}" @@ -288,6 +303,55 @@ def top_k_top_p_min_p_sampling_from_probs_torch( return batch_next_token_ids +def top_k_top_p_min_p_sampling_from_probs_ascend( + probs: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, + need_min_p_sampling: bool, +): + """A top-k, top-p and min-p sampling implementation for ascend npu with torch_npu interface.""" + # 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) + ): + logits_top_k_top_p = torch_npu.npu_top_k_top_p(probs, top_ps, top_ks) + probs_top_k_top_p = logits_top_k_top_p.softmax(dim=-1) + + if need_min_p_sampling: + min_p_thresholds = probs_top_k_top_p.max(dim=-1) * min_ps + 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) + else: + probs = torch.softmax(probs, dim=-1) + probs_sort, probs_idx = probs.sort(dim=-1, descending=True) + + # when top_k is -1 (in which sglang turns it to TOP_K_ALL), make it explicitly equal to logit's size + topk_all_mask = top_ks == TOP_K_ALL + top_ks.masked_fill_(topk_all_mask, probs.shape[1]) + top_k_mask = torch.arange(0, probs.shape[-1], device=probs.device).view( + 1, -1 + ) >= top_ks.view(-1, 1) + probs_sort.masked_fill_(top_k_mask, 0.0) + + probs_sum = torch.cumsum(probs_sort, dim=-1) + top_p_mask = probs_sum - probs_sort > top_ps.view(-1, 1) + probs_sort.masked_fill_(top_p_mask, 0.0) + + if need_min_p_sampling: + min_p_thresholds = probs_sort[:, 0] * min_ps + 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) + 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) + + def multinomial_with_seed( inputs: 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 1aa9e7c8c..7bc1b23b8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2706,7 +2706,7 @@ class ServerArgs: parser.add_argument( "--sampling-backend", type=str, - choices=["flashinfer", "pytorch"], + choices=["flashinfer", "pytorch", "ascend"], default=ServerArgs.sampling_backend, help="Choose the kernels for sampling layers.", ) diff --git a/test/srt/ascend/test_ascend_sampling_backend.py b/test/srt/ascend/test_ascend_sampling_backend.py new file mode 100644 index 000000000..7b6307912 --- /dev/null +++ b/test/srt/ascend/test_ascend_sampling_backend.py @@ -0,0 +1,96 @@ +import unittest +from types import SimpleNamespace + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class TestAscendSamplingBackend(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = "Qwen/Qwen2.5-7B-Instruct" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--sampling-backend", + "ascend", + "--disable-radix-cache", + "--disable-cuda-graph", + "--mem-fraction-static", + 0.85, + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_mmlu(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mmlu", + num_examples=64, + num_threads=32, + temperature=0.1, + ) + + metrics = run_eval(args) + self.assertGreaterEqual(metrics["score"], 0.65) + + def test_greedy(self): + + first_text = None + + # ensure the answer is identical across single response + for _ in range(5): + response_single = requests.post( + self.base_url + "/generate", + json={ + "text": "The capital of Germany is", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + }, + ).json() + text = response_single["text"] + if first_text is None: + first_text = text + + self.assertEqual(text, first_text) + + first_text = None + + response_batch = requests.post( + self.base_url + "/generate", + json={ + "text": ["The capital of Germany is"] * 10, + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + }, + ).json() + + # ensure the answer is identical among the batch + for i in range(10): + text = response_batch[i]["text"] + if first_text is None: + first_text = text + self.assertEqual(text, first_text) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 6f1cb206f..d516feb43 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -524,6 +524,7 @@ suite_ascend = { TestFile("ascend/test_ascend_graph_tp1_bf16.py", 400), TestFile("ascend/test_ascend_tp1_bf16.py", 400), TestFile("ascend/test_ascend_hicache_mha.py", 400), + TestFile("ascend/test_ascend_sampling_backend.py", 400), ], "per-commit-2-ascend-npu": [ TestFile("ascend/test_ascend_graph_tp2_bf16.py", 400),