[Kernel] Add JIT apply_rope_with_cos_sin_cache_inplace (#18155)

2026-02-05 21:49:37 +08:00
parent 4aa03d91fd
commit 2eb4359ada
4 changed files with 1197 additions and 1 deletions
@@ -37,7 +37,10 @@ _is_xpu = is_xpu()
 _is_musa = is_musa()

 if _is_cuda:
-    from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace
+    from sglang.jit_kernel.rope import (
+        FusedSetKVBufferArg,
+        apply_rope_with_cos_sin_cache_inplace,
+    )
 else:
    FusedSetKVBufferArg = None