[Bugfix] Restore `gather_and_maybe_dequant_cache` OOB guard (#49427)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jul 22, 2026 at 21:47 UTC
910cc8543a6907c9cc87c417f8f2420969278bf5
2 files changed
+8
-3
csrc/libtorch_stable/cache_kernels.cu
+3
@@ -1025,6 +1025,9 @@ __global__ void gather_and_maybe_dequant_cache(
1025
batch_offset += offset;
1026
int32_t block_table_id = batch_offset / block_size;
1027
int32_t slot_id = batch_offset % block_size;
1028
+ // seq_starts may push the block index past the end of the batch's block
1029
+ // table row.
1030
+ if (block_table_id >= block_table_stride) continue;
1031
int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
1032
int32_t block_id = block_table[block_table_offset];
1033
int64_t cache_offset =
tests/kernels/test_cache_kernels.py
+5
-3
@@ -21,9 +21,9 @@ def test_gather_cache_oob():
21
seq_starts causes the block_table offset to read out of bounds.
22
"""
23
24
- batch_size = 1
24
block_size = 64
26
- entry_size = 128
25
+ # The kernel only supports the MLA entry sizes.
26
+ entry_size = 576
27
28
block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
29
@@ -34,6 +34,7 @@ def test_gather_cache_oob():
34
35
seq_len = 65
36
cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
37
+ token_to_seq = torch.zeros(seq_len, dtype=torch.int32, device="cuda")
38
39
# src_cache: [num_blocks, block_size, entry_size]
40
num_blocks = 5
@@ -51,7 +52,8 @@ def test_gather_cache_oob():
52
dst,
53
block_table,
54
cu_seq_lens,
54
- batch_size,
55
+ token_to_seq,
56
+ seq_len,
57
"auto", # kv_cache_dtype
58
scale,
59
seq_starts,