From 910cc8543a6907c9cc87c417f8f2420969278bf5 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 22 Jul 2026 21:47:37 +0100 Subject: [PATCH] [Bugfix] Restore `gather_and_maybe_dequant_cache` OOB guard (#49427) Signed-off-by: Nick Hill --- csrc/libtorch_stable/cache_kernels.cu | 3 +++ tests/kernels/test_cache_kernels.py | 8 +++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index 2d4b47b4b4..5c1628537a 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -1025,6 +1025,9 @@ __global__ void gather_and_maybe_dequant_cache( batch_offset += offset; int32_t block_table_id = batch_offset / block_size; int32_t slot_id = batch_offset % block_size; + // seq_starts may push the block index past the end of the batch's block + // table row. + if (block_table_id >= block_table_stride) continue; int32_t block_table_offset = batch_id * block_table_stride + block_table_id; int32_t block_id = block_table[block_table_offset]; int64_t cache_offset = diff --git a/tests/kernels/test_cache_kernels.py b/tests/kernels/test_cache_kernels.py index 25402fe03e..eff4508b0d 100644 --- a/tests/kernels/test_cache_kernels.py +++ b/tests/kernels/test_cache_kernels.py @@ -21,9 +21,9 @@ def test_gather_cache_oob(): seq_starts causes the block_table offset to read out of bounds. """ - batch_size = 1 block_size = 64 - entry_size = 128 + # The kernel only supports the MLA entry sizes. + entry_size = 576 block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda") @@ -34,6 +34,7 @@ def test_gather_cache_oob(): seq_len = 65 cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda") + token_to_seq = torch.zeros(seq_len, dtype=torch.int32, device="cuda") # src_cache: [num_blocks, block_size, entry_size] num_blocks = 5 @@ -51,7 +52,8 @@ def test_gather_cache_oob(): dst, block_table, cu_seq_lens, - batch_size, + token_to_seq, + seq_len, "auto", # kv_cache_dtype scale, seq_starts,