Fix cp_gather_cache block table bounds

csrc/libtorch_stable/cache_kernels.cu

@@ -993,8 +993,9 @@ __global__ void gather_and_maybe_dequant_cache(
     const int32_t* __restrict__ cu_seq_lens,   // [BATCH+1]
     const int32_t* __restrict__ token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNK]
     const int32_t num_tokens, const int32_t block_size,
-    const int64_t block_table_stride, const int64_t cache_block_stride,
-    const int64_t cache_entry_stride, const int64_t dst_entry_stride,
+    const int64_t block_table_stride, const int64_t block_table_width,
+    const int64_t cache_block_stride, const int64_t cache_entry_stride,
+    const int64_t dst_entry_stride,
     const float* __restrict__ scale,
     const int32_t* __restrict__ seq_starts) {  // Optional: starting offsets per
                                                // batch
@@ -1011,22 +1012,31 @@ __global__ void gather_and_maybe_dequant_cache(
     int64_t batch_id = token_to_seq[token_id];
     int64_t batch_start = cu_seq_lens[batch_id];
     int64_t batch_end = cu_seq_lens[batch_id + 1];
-    int32_t batch_offset = token_id - batch_start;
+    int64_t batch_offset = token_id - batch_start;
 
     if (token_id >= batch_end) return;
-    int32_t offset = 0;
+    int64_t offset = 0;
     if (seq_starts != nullptr) {
       offset = seq_starts[batch_id];
     }
     batch_offset += offset;
-    int32_t block_table_id = batch_offset / block_size;
+    int64_t block_table_id = batch_offset / block_size;
     int32_t slot_id = batch_offset % block_size;
-    int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
+    scalar_t* dst_ = dst + token_id * dst_entry_stride;
+
+    if (batch_offset < 0 || block_table_id >= block_table_width) {
+#pragma unroll
+      for (int idx = threadIdx.x; idx < ENTRY_SIZE; idx += CTA_SIZE) {
+        dst_[idx] = scalar_t{};
+      }
+      continue;
+    }
+
+    int64_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 =
         block_id * cache_block_stride + slot_id * cache_entry_stride;
     constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size;
-    scalar_t* dst_ = dst + token_id * dst_entry_stride;
     cache_t* src_ = const_cast<cache_t*>(src_cache) + cache_offset;
 
 #pragma unroll
@@ -1076,9 +1086,9 @@ __global__ void gather_and_maybe_dequant_cache(
           block_table.const_data_ptr<int32_t>(),                              \
           cu_seq_lens.const_data_ptr<int32_t>(),                              \
           token_to_seq.const_data_ptr<int32_t>(), num_tokens, block_size,     \
-          block_table_stride, cache_block_stride, cache_entry_stride,         \
-          dst_entry_stride, reinterpret_cast<const float*>(scale.data_ptr()), \
-          seq_starts_ptr);
+          block_table_stride, block_table_width, cache_block_stride,          \
+          cache_entry_stride, dst_entry_stride,                               \
+          reinterpret_cast<const float*>(scale.data_ptr()), seq_starts_ptr);
 
 #define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \
   CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576)
@@ -1137,6 +1147,7 @@ void gather_and_maybe_dequant_cache(
   }
 
   int64_t block_table_stride = block_table.stride(0);
+  int64_t block_table_width = block_table.size(1);
   int64_t cache_block_stride = src_cache.stride(0);
   int64_t cache_entry_stride = src_cache.stride(1);
   int64_t dst_entry_stride = dst.stride(0);
@@ -1237,8 +1248,9 @@ __global__ void cp_gather_cache(
     const int32_t* __restrict__ block_table,  // [BATCH, BLOCK_INDICES]
     const int32_t* __restrict__ cu_seq_lens,  // [BATCH+1]
     const int32_t block_size, const int32_t entry_size,
-    const int64_t block_table_stride, const int64_t cache_block_stride,
-    const int64_t cache_entry_stride, const int64_t dst_entry_stride,
+    const int64_t block_table_stride, const int64_t block_table_width,
+    const int64_t cache_block_stride, const int64_t cache_entry_stride,
+    const int64_t dst_entry_stride,
     const int32_t* __restrict__ seq_starts  // Optional: starting offsets per
                                             // batch
 ) {
@@ -1261,12 +1273,13 @@ __global__ void cp_gather_cache(
   // Adjust the pointer for the block_table for this batch.
   // If seq_starts is provided, compute an offset based on it
   const int32_t batch_offset = bid * block_table_stride;
-  int32_t offset = split_start;
+  int64_t offset = split_start;
   if (seq_starts != nullptr) {
     offset += seq_starts[bid];
   }
-  int32_t offset_div = offset / block_size;
-  offset = offset % block_size;
+  const bool starts_before_cache = offset < 0;
+  int64_t offset_div = starts_before_cache ? -1 : offset / block_size;
+  offset = starts_before_cache ? 0 : offset % block_size;
   const int32_t* batch_block_table = block_table + batch_offset;
 
   // Adjust dst pointer based on the cumulative sequence lengths.
@@ -1277,12 +1290,21 @@ __global__ void cp_gather_cache(
     for (int i = threadIdx.x; i < entry_size; i += blockDim.x)
       _dst[i] = _src[i];
   };
+  auto zero_entry = [&](scalar_t* __restrict__ _dst) {
+    for (int i = threadIdx.x; i < entry_size; i += blockDim.x)
+      _dst[i] = scalar_t{};
+  };
 
   for (int pid = split_start; pid < split_end; ++pid) {
-    auto block_id = batch_block_table[offset_div];
-    auto block_start_ptr = src_cache + block_id * cache_block_stride;
     auto block_dst_ptr = dst + pid * dst_entry_stride;
-    copy_entry(block_start_ptr + offset * cache_entry_stride, block_dst_ptr);
+    if (starts_before_cache || offset_div < 0 ||
+        offset_div >= block_table_width) {
+      zero_entry(block_dst_ptr);
+    } else {
+      auto block_id = batch_block_table[offset_div];
+      auto block_start_ptr = src_cache + block_id * cache_block_stride;
+      copy_entry(block_start_ptr + offset * cache_entry_stride, block_dst_ptr);
+    }
     offset += 1;
     // bump to next block
     if (offset == block_size) {
@@ -1300,8 +1322,8 @@ __global__ void cp_gather_cache(
       reinterpret_cast<CPY_DTYPE*>(dst.data_ptr()),                  \
       block_table.const_data_ptr<int32_t>(),                         \
       cu_seq_lens.const_data_ptr<int32_t>(), block_size, entry_size, \
-      block_table_stride, cache_block_stride, cache_entry_stride,    \
-      dst_entry_stride, seq_starts_ptr);
+      block_table_stride, block_table_width, cache_block_stride,     \
+      cache_entry_stride, dst_entry_stride, seq_starts_ptr);
 
 // Gather sequences from the cache into the destination tensor.
 //  - cu_seq_lens contains the cumulative sequence lengths for each batch
@@ -1347,6 +1369,7 @@ void cp_gather_cache(
   }
 
   int64_t block_table_stride = block_table.stride(0);
+  int64_t block_table_width = block_table.size(1);
   int64_t cache_block_stride = src_cache.stride(0);
   int64_t cache_entry_stride = src_cache.stride(1);
   int64_t dst_entry_stride = dst.stride(0);

tests/kernels/test_cache_kernels.py

@@ -42,23 +42,65 @@ def test_gather_cache_oob():
     )
 
     dst = torch.empty((seq_len, entry_size), dtype=torch.float16, device="cuda")
+    dst.fill_(1)
 
     scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
+    token_to_seq = torch.zeros(seq_len, dtype=torch.int32, device="cuda")
 
     # Calling the C++ function gather_and_maybe_dequant_cache
     ops.gather_and_maybe_dequant_cache(
         src_cache,
         dst,
         block_table,
         cu_seq_lens,
-        batch_size,
+        token_to_seq,
+        seq_len,
         "auto",  # kv_cache_dtype
         scale,
         seq_starts,
     )
 
     torch.accelerator.synchronize()
-    assert True
+    torch.testing.assert_close(dst, torch.zeros_like(dst))
+
+
+@pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="Need CUDA device")
+def test_cp_gather_cache_oob():
+    """
+    Tests for OOB read in cp_gather_cache (Issue #45462).
+    This reproduces the case where seq_starts advances beyond the available
+    block_table entries for the row.
+    """
+
+    batch_size = 1
+    block_size = 64
+    entry_size = 128
+
+    block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
+    seq_starts = torch.tensor([128], dtype=torch.int32, device="cuda")
+
+    seq_len = 65
+    cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
+
+    num_blocks = 5
+    src_cache = torch.randn(
+        (num_blocks, block_size, entry_size), dtype=torch.float16, device="cuda"
+    )
+
+    dst = torch.empty((seq_len, entry_size), dtype=torch.float16, device="cuda")
+    dst.fill_(1)
+
+    ops.cp_gather_cache(
+        src_cache,
+        dst,
+        block_table,
+        cu_seq_lens,
+        batch_size,
+        seq_starts,
+    )
+
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(dst, torch.zeros_like(dst))
 
 
 if __name__ == "__main__":