[Minor] Enhance JIT kernel and add dev docs (#14570)
This commit is contained in:
14
python/sglang/jit_kernel/test_add_constant.py
Normal file
14
python/sglang/jit_kernel/test_add_constant.py
Normal file
@@ -0,0 +1,14 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user