[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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user