Fix amd rope definition (#14556)

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)

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"]