diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index 9981e4670..336c3b5a1 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/sgl-kernel/csrc/common_extension_rocm.cc @@ -186,6 +186,15 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { */ m.def("apply_token_bitmask_inplace_cuda(Tensor logits, Tensor bitmask, Tensor? indices=None) -> ()"); m.impl("apply_token_bitmask_inplace_cuda", &ApplyTokenBitmaskInplace); + + /* + * From csrc/elementwise + */ + m.def( + "rotary_embedding(Tensor positions, Tensor! query," + " Tensor!? key, int head_size," + " Tensor cos_sin_cache, bool is_neox) -> ()"); + m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding); } REGISTER_EXTENSION(common_ops) diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index 8cdb7c695..525b9c0be 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -51,6 +51,7 @@ sources = [ "csrc/moe/moe_topk_sigmoid_kernels.cu", "csrc/speculative/eagle_utils.cu", "csrc/kvcacheio/transfer.cu", + "csrc/elementwise/pos_enc.cu", ] cxx_flags = ["-O3"]