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

22
python/sglang/jit_kernel/csrc/elementwise/rope.cuh
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_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);
}
}
}

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].