[1/2] Add rope kernel in sgl-kernel (#14334)

2025-12-04 16:45:44 +08:00
parent e3ab23c1a6
commit cb8df87fc1
8 changed files with 293 additions and 12 deletions
@@ -90,6 +90,12 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
      "Tensor? v, Tensor!? k_buffer, Tensor!? v_buffer, Tensor? kv_cache_loc) -> ()");
  m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache);

+  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);
+
  m.def(
      "downcast_fp8(Tensor k, Tensor v, Tensor k_out, Tensor v_out, Tensor k_scale, Tensor v_scale, Tensor loc, "
      "int mult, int offset) -> ()");