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

[Kernel] Enable fast random sample on Kunlun P

Browse Source
This commit is contained in:
xyDong0223
2025-12-10 21:52:48 +08:00
committed by GitHub
parent f109a76a39
commit 670c2397b8
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm_kunlun/v1/sample/ops/topk_topp_sampler.py
View File

@@ -151,7 +151,12 @@ def random_sample(
# not have its own seed. Then, we overwrite the values for the requests # not have its own seed. Then, we overwrite the values for the requests
# that have their own seeds. # that have their own seeds.
if len(generators) != probs.shape[0]: if len(generators) != probs.shape[0]:
q.exponential_() if os.getenv('FAST_RANDOM_SAMPLE') == "1":
q.uniform_()
q = -torch.log(q)
q = q.clamp(min=1e-4)
else:
q.exponential_()
if generators: if generators:
# TODO(woosuk): This can be slow because we handle each request # TODO(woosuk): This can be slow because we handle each request
# one by one. Optimize this. # one by one. Optimize this.