perf: speed up topk_topp_sampler

Author: ZhengWG

vllm_ascend/ops/__init__.py

@@ -23,6 +23,7 @@
 import vllm_ascend.ops.layernorm  # noqa
 import vllm_ascend.ops.rotary_embedding  # noqa
 import vllm_ascend.ops.vocab_parallel_embedding  # noqa
+import vllm_ascend.ops.topk_topp_sampler
 
 
 class dummyFusionOp:

vllm_ascend/ops/topk_topp_sampler.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Optional
+import torch
+from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler, apply_top_k_top_p_tpu, random_sample
+
+
+def forward_npu(
+    self,
+    logits: torch.Tensor,
+    generators: dict[int, torch.Generator],
+    k: Optional[torch.Tensor],
+    p: Optional[torch.Tensor],
+) -> torch.Tensor:
+    logits = apply_top_k_top_p_tpu(logits, k, p)
+    probs = logits.softmax(dim=-1, dtype=torch.float32)
+    return random_sample(probs, generators)
+
+
+TopKTopPSampler.forward_native = forward_npu