chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix(rope): restore K writeback in fused rope + kv store kernel (#19636)

Browse Source
This commit is contained in:
David Wang
2026-03-07 07:41:35 -05:00
committed by GitHub
parent f016738f4c
commit 19c51fe2fa
2 changed files with 17 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/jit_kernel/rope.py
View File

@@ -78,7 +78,7 @@ def apply_rope_inplace(
module.run_rope(q, k, cos_sin_cache, positions)
@register_custom_op(mutates_args=["q", "k_cache", "v_cache"])
@register_custom_op(mutates_args=["q", "k", "k_cache", "v_cache"])
def apply_rope_inplace_with_kvcache(
q: torch.Tensor,
k: torch.Tensor,
@@ -95,8 +95,8 @@ def apply_rope_inplace_with_kvcache(
"""
Fused inplace RoPE + KV cache store.
Applies rotary position embedding to q inplace. For k, applies RoPE and
stores the result in k_cache. The original v is also stored in v_cache.
Applies rotary position embedding to q and k inplace. The rotated k is also
stored in k_cache. The original v is also stored in v_cache.
Args:
q: Query tensor of shape [num_tokens, num_qo_heads, head_dim].