EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[0.11.0][ops] npu_top_k_top_p supports k and p only (#4153)

Browse Source 
### What this PR does / why we need it?
With CANN 8.3 and corresponding PTA 2.7.1, `npu_top_k_top_p` supports
passing only k (1<=k<=1024) and p separately.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
E2E performance test with only `top_k` and `p` seperately. This pr gains
0.2ms improvements in TPOT with `batch_size=16`.

Signed-off-by: linfeng-yuan <1102311262@qq.com>
This commit is contained in:
linfeng-yuan
2025-12-09 15:45:40 +08:00
committed by GitHub
parent 4588d1f215
commit 695e5c9ebc
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
vllm_ascend/sample/sampler.py
View File

@@ -24,14 +24,14 @@ class AscendTopKTopPSampler(TopKTopPSampler):
k: torch.Tensor,
p: torch.Tensor,
) -> torch.Tensor:
# npu_top_k_top_p uses the operator aclnnApplyTopKTopP, but aclnnApplyTopKTopP currently does not support 310P
if not is_310p() and p is not None and k is not None and 1 <= int(
k.max()) <= 1024:
# npu_top_k_top_p's parameter order is (logits, p, k), not (logits, k, p)
return torch_npu.npu_top_k_top_p(logits, p, k)
if p is None and k is None:
return logits
# npu_top_k_top_p uses the operator aclnnApplyTopKTopP, but aclnnApplyTopKTopP currently does not support 310P
if not is_310p():
# npu_top_k_top_p requires parameter k ranged from 1 to 1024
if k is None or 1 <= int(k.max()) <= 1024:
# npu_top_k_top_p's parameter order is (logits, p, k), not (logits, k, p)
return torch_npu.npu_top_k_top_p(logits, p, k)
probs = logits.softmax(dim=-1)
probs_sort, _ = probs.sort(dim=-1, descending=False)