A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
[ROCm][V1] Update reshape_and_cache to properly work with CUDA graph padding (#13922)
Sage Moore committed
Feb 26, 2025 at 20:04 UTC
378b3ef6f834a2536fc01aa4cd89d09737232bac
1 file changed
+1
-1
csrc/cache_kernels.cu
+1
-1
@@ -375,7 +375,7 @@ void reshape_and_cache(
375
torch::Tensor& slot_mapping, // [num_tokens]
376
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
377
torch::Tensor& v_scale) {
378
- int num_tokens = key.size(0);
378
+ int num_tokens = slot_mapping.size(0);
379
int num_heads = key.size(1);
380
int head_size = key.size(2);
381
int block_size = key_cache.size(3);