From 69090af34f45f42d917c0af09c5e3c37e51ac189 Mon Sep 17 00:00:00 2001 From: RolaoDenthu Date: Thu, 24 Sep 2026 19:11:36 +0000 Subject: [PATCH] Fix int32 offset overflow in Triton DSv4 KV store kernels --- python/sglang/kernels/ops/kvcache/triton_store_cache.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/kvcache/triton_store_cache.py b/python/sglang/kernels/ops/kvcache/triton_store_cache.py index 509171e58d37..1261720610cd 100644 --- a/python/sglang/kernels/ops/kvcache/triton_store_cache.py +++ b/python/sglang/kernels/ops/kvcache/triton_store_cache.py @@ -61,7 +61,8 @@ def _triton_fused_store_flashmla_kernel( if token_id >= N: return - loc = tl.load(indices_ptr + token_id).to(tl.int32) + # Per-layer caches exceed 2 GiB; int32 byte offsets wrap and write outside the buffer. + loc = tl.load(indices_ptr + token_id).to(tl.int64) page = loc // PAGE_SIZE slot = loc % PAGE_SIZE @@ -169,7 +170,7 @@ def _triton_fused_store_indexer_kernel( if token_id >= N: return - loc = tl.load(indices_ptr + token_id).to(tl.int32) + loc = tl.load(indices_ptr + token_id).to(tl.int64) page = loc // PAGE_SIZE slot = loc % PAGE_SIZE