mirror of
https://github.com/kvcache-ai/sglang.git
synced 2026-07-07 15:57:17 +00:00
15 lines
267 B
Python
15 lines
267 B
Python
import torch
|
|
|
|
from sglang.jit_kernel.add_constant import add_constant
|
|
|
|
|
|
def main():
|
|
c = 1024
|
|
src = torch.arange(0, 1024 + 1, dtype=torch.int32).cuda()
|
|
dst = add_constant(src, c)
|
|
assert torch.all(dst == src + c)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|