Fix amd rope definition (#14556)

2025-12-07 23:47:03 -08:00
parent cf0478d602
commit aeff0d386b
2 changed files with 10 additions and 0 deletions
@@ -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)