From af35023e65faec8be6547e3affd8e4aadc83be0e Mon Sep 17 00:00:00 2001 From: Xuchun Shang Date: Thu, 4 Dec 2025 17:20:11 +0800 Subject: [PATCH] [bug fix] fix ima with get_mla_kv_buffer_kernel overflow (#14224) Signed-off-by: Xuchun Shang --- python/sglang/srt/mem_cache/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index d0abf02fe..361accb87 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/python/sglang/srt/mem_cache/utils.py @@ -43,7 +43,7 @@ def set_mla_kv_buffer_kernel( total_dim = nope_dim + rope_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 if base + BLOCK <= nope_dim: @@ -168,7 +168,7 @@ def get_mla_kv_buffer_kernel( rope_dim: tl.constexpr, ): 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 nope_offs = tl.arange(0, nope_dim)