minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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 20:41:35 +08: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
python/sglang/jit_kernel/csrc/elementwise/rope.cuh
+14 -8
View File
@@ -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_x[j] = cast<DType2, fp32x2_t>({out_x0, out_x1});
input_vec_y[j] = cast<DType2, fp32x2_t>({out_y0, out_y1}); 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); store_as<InputStorage>(input, input_vec_x, lane_id);
const auto output_x = load_q ? input : k_out; const auto input_y_out = pointer::offset(input, (kRopeDim / 2) * sizeof(DType));
store_as<InputStorage>(output_x, input_vec_x, lane_id); store_as<InputStorage>(input_y_out, input_vec_y, lane_id);
const auto output_y = pointer::offset(output_x, (kRopeDim / 2) * sizeof(DType)); if (!load_q) {
store_as<InputStorage>(output_y, input_vec_y, lane_id); 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 { } else {
using CacheStorage = AlignedVector<float, kVecSize>; using CacheStorage = AlignedVector<float, kVecSize>;
auto input_vec = load_as<InputStorage>(input, lane_id); 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; const auto out_y = x * sin + y * cos;
input_vec[j] = cast<DType2, fp32x2_t>({out_x, out_y}); 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); store_as<InputStorage>(input, input_vec, lane_id);
const auto output = load_q ? input : k_out; if (!load_q) {
store_as<InputStorage>(output, input_vec, lane_id); 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);
}
} }
} }
python/sglang/jit_kernel/rope.py
+3 -3
View File
@@ -78,7 +78,7 @@ def apply_rope_inplace(
module.run_rope(q, k, cos_sin_cache, positions) 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( def apply_rope_inplace_with_kvcache(
q: torch.Tensor, q: torch.Tensor,
k: torch.Tensor, k: torch.Tensor,
@@ -95,8 +95,8 @@ def apply_rope_inplace_with_kvcache(
""" """
Fused inplace RoPE + KV cache store. Fused inplace RoPE + KV cache store.
Applies rotary position embedding to q inplace. For k, applies RoPE and Applies rotary position embedding to q and k inplace. The rotated k is also
stores the result in k_cache. The original v is also stored in v_cache. stored in k_cache. The original v is also stored in v_cache.
Args: Args:
q: Query tensor of shape [num_tokens, num_qo_heads, head_dim]. q: Query tensor of shape [num_tokens, num_qo_heads, head_dim].