[Kernel] Triton-based Top-k and Top-p sampler kernels (#33538)

Signed-off-by: js_park <cakeng@naver.com>
Signed-off-by: Jongseok Park <37990712+cakeng@users.noreply.github.com>
Signed-off-by: Sunga Kim <sunga.kim@berkeley.edu>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Sunga Kim <sunga.kim@berkeley.edu>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Jongseok Park
2026-02-17 15:14:30 -08:00
committed by GitHub
parent dc5fa77a4e
commit c656ba3b4d
6 changed files with 2002 additions and 32 deletions

View File

@@ -14,16 +14,12 @@ def cdiv(a: int, b: int) -> int:
def next_power_of_2(n: int) -> int:
"""The next power of 2 (inclusive)"""
if n < 1:
return 1
return 1 << (n - 1).bit_length()
return 1 if n < 1 else 1 << (n - 1).bit_length()
def prev_power_of_2(n: int) -> int:
"""The previous power of 2 (inclusive)"""
if n <= 0:
return 0
return 1 << (n.bit_length() - 1)
return 0 if n <= 0 else 1 << (n.bit_length() - 1)
def round_up(x: int, y: int) -> int: