test_sgl_kernel.py

import torch
import sglang
import sgl_kernel

if __name__ == "__main__":
    a = torch.randn((1, 1024), dtype=torch.float32).to("xpu")

    ref_res = a + a
    res3 = torch.ops.sgl_kernel.sgl_test_sycl(a, a)

    print("ref_res is: {}".format(ref_res), flush=True)
    print("res3 is: {}".format(res3), flush=True)

    print(torch.testing.assert_allclose(ref_res, res3), flush=True)
    print(torch.allclose(ref_res, res3), flush=True)

    # print("Res is: {}".format(res), flush=True)
    print("Done ----", flush=True)