[Feature] Support JIT set kv cache (#16273)

2026-01-11 09:34:09 +08:00
parent a2c2c09d7d
commit d112f6a25b
9 changed files with 516 additions and 33 deletions
--- a/python/sglang/jit_kernel/benchmark/bench_qknorm.py
+++ b/python/sglang/jit_kernel/benchmark/bench_qknorm.py
@@ -53,6 +53,7 @@ def flashinfer_qknorm(
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
 ) -> None:
+    from flashinfer import rmsnorm

    rmsnorm(q, q_weight, out=q)
    rmsnorm(k, k_weight, out=k)