From 19c51fe2fafd582ab4f0b29d7696b1fa08f17419 Mon Sep 17 00:00:00 2001
From: David Wang <21328423+dcw02@users.noreply.github.com>
Date: Sat, 7 Mar 2026 07:41:35 -0500
Subject: [PATCH] fix(rope): restore K writeback in fused rope + kv store
 kernel (#19636)

---
 .../jit_kernel/csrc/elementwise/rope.cuh      | 22 ++++++++++++-------
 python/sglang/jit_kernel/rope.py              |  6 ++---
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/python/sglang/jit_kernel/csrc/elementwise/rope.cuh b/python/sglang/jit_kernel/csrc/elementwise/rope.cuh
index 27b4e7ec8..36f583e49 100644
--- a/python/sglang/jit_kernel/csrc/elementwise/rope.cuh
+++ b/python/sglang/jit_kernel/csrc/elementwise/rope.cuh
@@ -201,11 +201,15 @@ __global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStorePa
         input_vec_x[j] = cast<DType2, fp32x2_t>({out_x0, out_x1});
         input_vec_y[j] = cast<DType2, fp32x2_t>({out_y0, out_y1});
       }
-      const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
-      const auto output_x = load_q ? input : k_out;
-      store_as<InputStorage>(output_x, input_vec_x, lane_id);
-      const auto output_y = pointer::offset(output_x, (kRopeDim / 2) * sizeof(DType));
-      store_as<InputStorage>(output_y, input_vec_y, lane_id);
+      store_as<InputStorage>(input, input_vec_x, lane_id);
+      const auto input_y_out = pointer::offset(input, (kRopeDim / 2) * sizeof(DType));
+      store_as<InputStorage>(input_y_out, input_vec_y, lane_id);
+      if (!load_q) {
+        const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
+        store_as<InputStorage>(k_out, input_vec_x, lane_id);
+        const auto k_out_y = pointer::offset(k_out, (kRopeDim / 2) * sizeof(DType));
+        store_as<InputStorage>(k_out_y, input_vec_y, lane_id);
+      }
     } else {
       using CacheStorage = AlignedVector<float, kVecSize>;
       auto input_vec = load_as<InputStorage>(input, lane_id);
@@ -220,9 +224,11 @@ __global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStorePa
         const auto out_y = x * sin + y * cos;
         input_vec[j] = cast<DType2, fp32x2_t>({out_x, out_y});
       }
-      const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
-      const auto output = load_q ? input : k_out;
-      store_as<InputStorage>(output, input_vec, lane_id);
+      store_as<InputStorage>(input, input_vec, lane_id);
+      if (!load_q) {
+        const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
+        store_as<InputStorage>(k_out, input_vec, lane_id);
+      }
     }
   }
 
diff --git a/python/sglang/jit_kernel/rope.py b/python/sglang/jit_kernel/rope.py
index 02c4e5fb6..ebc316094 100644
--- a/python/sglang/jit_kernel/rope.py
+++ b/python/sglang/jit_kernel/rope.py
@@ -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].