minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[bug fix] fix ima with get_mla_kv_buffer_kernel overflow (#14224)

Browse Source 
Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>
This commit is contained in:
Xuchun Shang
2025-12-04 01:20:11 -08:00
committed by GitHub
parent cb8df87fc1
commit af35023e65
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/utils.py
+2 -2
View File
@@ -43,7 +43,7 @@ def set_mla_kv_buffer_kernel(
total_dim = nope_dim + rope_dim total_dim = nope_dim + rope_dim
mask = offs < total_dim mask = offs < total_dim
loc = tl.load(loc_ptr + pid_loc) loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
if base + BLOCK <= nope_dim: if base + BLOCK <= nope_dim:
@@ -168,7 +168,7 @@ def get_mla_kv_buffer_kernel(
rope_dim: tl.constexpr, rope_dim: tl.constexpr,
): ):
pid_loc = tl.program_id(0) pid_loc = tl.program_id(0)
loc = tl.load(loc_ptr + pid_loc) loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
loc_src_ptr = kv_buffer_ptr + loc * buffer_stride loc_src_ptr = kv_buffer_ptr + loc * buffer_stride
nope_offs = tl.arange(0, nope_dim) nope_offs = tl.arange(0, nope_dim)