[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
@@ -149,6 +149,14 @@ void apply_rope_pos_ids_cos_sin_cache(
    const std::optional<at::Tensor>& v_buffer,
    const std::optional<at::Tensor>& kv_cache_loc);

+void rotary_embedding(
+    torch::Tensor& positions,
+    torch::Tensor& query,
+    std::optional<torch::Tensor> key,
+    int64_t head_size,
+    torch::Tensor& cos_sin_cache,
+    bool is_neox);
+
 void downcast_fp8(
    at::Tensor& k,
    at::Tensor& v,