[Bugfix] Restore gather_and_maybe_dequant_cache OOB guard (#49427)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

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 =

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,