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 14, 05:30 UTC
Add gather_indexer_k_quant_cache kernel (#25931)
Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Signed-off-by: Simon Mo <simon.mo@hey.com> Signed-off-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Simon Mo <simon.mo@hey.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com>
Barry Kang committed
Oct 8, 2025 at 12:58 UTC
127c8b782a63010e56bc33d45e87baf3a6415568
4 files changed
+146
csrc/cache.h
+8
@@ -64,3 +64,11 @@ void indexer_k_quant_and_cache(
64
torch::Tensor& slot_mapping, // [num_tokens]
65
int64_t quant_block_size, // quantization block size
66
const std::string& scale_fmt);
67
+
68
+// Extract function to gather quantized K cache
69
+void cp_gather_indexer_k_quant_cache(
70
+ const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
71
+ torch::Tensor& dst_k, // [num_tokens, head_dim]
72
+ torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4]
73
+ const torch::Tensor& block_table, // [batch_size, num_blocks]
74
+ const torch::Tensor& cu_seq_lens); // [batch_size + 1]
\ No newline at end of file
csrc/cache_kernels.cu
+120
@@ -572,6 +572,70 @@ __global__ void indexer_k_quant_and_cache_kernel(
572
}
573
}
574
575
+template <int BLOCK_Y_SIZE>
576
+__global__ void cp_gather_indexer_k_quant_cache_kernel(
577
+ const char* __restrict__ kv_cache, // [num_blocks, block_size,
578
+ // cache_stride]
579
+ char* __restrict__ dst_k, // [num_tokens, head_dim]
580
+ char* __restrict__ dst_scale, // [num_tokens, head_dim / quant_block_size *
581
+ // 4]
582
+ const int* __restrict__ block_table, // [batch_size, num_blocks]
583
+ const int* __restrict__ cu_seq_lens, // [batch_size + 1]
584
+ const int batch_size, // batch size
585
+ const int64_t token_stride, // stride for each token in dst_k
586
+ const int64_t head_dim, // dimension of each head
587
+ const int64_t block_stride, // stride for each block in kv_cache
588
+ const int64_t cache_token_stride, // stride for each token in kv_cache
589
+ const int64_t cache_block_size, // num_tokens for each block in kv_cache
590
+ const int num_blocks, // number of blocks
591
+ const int num_tokens, // number of tokens
592
+ const int quant_block_size // quantization block size
593
+) {
594
+ constexpr int VEC_SIZE = sizeof(float4) / sizeof(char);
595
+ const int token_idx = blockIdx.x * blockDim.y + threadIdx.y;
596
+ const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
597
+ // Find batch index within a block
598
+ __shared__ int batch_idx[BLOCK_Y_SIZE];
599
+ for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
600
+ iter++) {
601
+ int tid = iter * blockDim.x + threadIdx.x;
602
+ if (tid < batch_size) {
603
+ const int seq_start = cu_seq_lens[tid];
604
+ const int seq_end = cu_seq_lens[tid + 1];
605
+ if (token_idx >= seq_start && token_idx < seq_end) {
606
+ batch_idx[threadIdx.y] = tid;
607
+ }
608
+ }
609
+ }
610
+
611
+#ifndef USE_ROCM
612
+ __syncwarp();
613
+#endif
614
+
615
+ if (head_idx >= head_dim || token_idx >= num_tokens) {
616
+ return;
617
+ }
618
+ const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
619
+ const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
620
+ inbatch_seq_idx / cache_block_size];
621
+ const int64_t src_block_offset = block_idx * block_stride;
622
+ const int64_t cache_inblock_offset =
623
+ (inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
624
+ const int64_t src_inblock_offset = src_block_offset + cache_inblock_offset;
625
+ const int64_t dst_inblock_offset = token_idx * token_stride + head_idx;
626
+
627
+ reinterpret_cast<float4*>(dst_k)[dst_inblock_offset / VEC_SIZE] =
628
+ reinterpret_cast<const float4*>(kv_cache)[src_inblock_offset / VEC_SIZE];
629
+ ;
630
+ if (threadIdx.x == 0) {
631
+ const int64_t src_scale_offset =
632
+ src_block_offset + cache_block_size * head_dim +
633
+ cache_inblock_offset * 4 / quant_block_size;
634
+ reinterpret_cast<float*>(dst_scale)[dst_inblock_offset / quant_block_size] =
635
+ reinterpret_cast<const float*>(kv_cache)[src_scale_offset / 4];
636
+ }
637
+}
638
+
639
} // namespace vllm
640
641
// KV_T is the data type of key and value tensors.
@@ -1173,3 +1237,59 @@ void indexer_k_quant_and_cache(
1237
DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
1238
CALL_INDEXER_K_QUANT_AND_CACHE);
1239
}
1240
+
1241
+// Macro to dispatch the kernel based on the data amount.
1242
+#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \
1243
+ vllm::cp_gather_indexer_k_quant_cache_kernel<BLOCK_Y_SIZE> \
1244
+ <<<dim3((num_tokens + BLOCK_Y_SIZE - 1) / BLOCK_Y_SIZE, \
1245
+ (head_dim + 8 * vec_size - 1) / (8 * vec_size)), \
1246
+ dim3(8, BLOCK_Y_SIZE), 0, stream>>>( \
1247
+ reinterpret_cast<char*>(kv_cache.data_ptr()), \
1248
+ reinterpret_cast<char*>(dst_k.data_ptr()), \
1249
+ reinterpret_cast<char*>(dst_scale.data_ptr()), \
1250
+ block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
1251
+ batch_size, dst_k.stride(0), dst_k.size(1), kv_cache.stride(0), \
1252
+ kv_cache.stride(1), kv_cache.size(1), block_table.size(1), \
1253
+ num_tokens, quant_block_size);
1254
+
1255
+void cp_gather_indexer_k_quant_cache(
1256
+ const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
1257
+ torch::Tensor& dst_k, // [num_tokens, head_dim]
1258
+ torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4]
1259
+ const torch::Tensor& block_table, // [batch_size, num_blocks]
1260
+ const torch::Tensor& cu_seq_lens // [batch_size + 1]
1261
+) {
1262
+ int batch_size = block_table.size(0);
1263
+ int num_tokens = dst_k.size(0);
1264
+ int head_dim = dst_k.size(1);
1265
+ int quant_block_size = head_dim * 4 / dst_scale.size(1);
1266
+
1267
+ TORCH_CHECK(kv_cache.device() == dst_k.device(),
1268
+ "kv_cache and dst_k must be on the same device");
1269
+ TORCH_CHECK(kv_cache.device() == dst_scale.device(),
1270
+ "kv_cache and dst_scale must be on the same device");
1271
+ TORCH_CHECK(kv_cache.device() == block_table.device(),
1272
+ "kv_cache and block_table must be on the same device");
1273
+ TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(),
1274
+ "kv_cache and cu_seq_lens must be on the same device");
1275
+ TORCH_CHECK(head_dim % quant_block_size == 0,
1276
+ "head_dim must be divisible by quant_block_size");
1277
+
1278
+ constexpr int vec_size = 16;
1279
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_cache));
1280
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1281
+
1282
+ if (num_tokens < 32) {
1283
+ CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(1);
1284
+ } else if (num_tokens < 64) {
1285
+ CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(2);
1286
+ } else if (num_tokens < 128) {
1287
+ CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(4);
1288
+ } else if (num_tokens < 256) {
1289
+ CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(8);
1290
+ } else if (num_tokens < 512) {
1291
+ CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(16);
1292
+ } else {
1293
+ CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(32);
1294
+ }
1295
+}
csrc/torch_bindings.cpp
+6
@@ -727,6 +727,12 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
727
"int quant_block_size, str kv_cache_dtype) -> ()");
728
cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
729
&indexer_k_quant_and_cache);
730
+
731
+ cache_ops.def(
732
+ "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
733
+ "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
734
+ cache_ops.impl("cp_gather_indexer_k_quant_cache", torch::kCUDA,
735
+ &cp_gather_indexer_k_quant_cache);
736
}
737
738
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
vllm/_custom_ops.py
+12
@@ -2108,6 +2108,18 @@ def indexer_k_quant_and_cache(
2108
)
2109
2110
2111
+def cp_gather_indexer_k_quant_cache(
2112
+ kv_cache: torch.Tensor,
2113
+ dst_k: torch.Tensor,
2114
+ dst_scale: torch.Tensor,
2115
+ block_table: torch.Tensor,
2116
+ cu_seq_lens: torch.Tensor,
2117
+) -> None:
2118
+ torch.ops._C_cache_ops.cp_gather_indexer_k_quant_cache(
2119
+ kv_cache, dst_k, dst_scale, block_table, cu_seq_lens
2120
+ )
2121
+
2122
+
2123
def get_device_attribute(attribute: int, device: int) -> int:
2124
return torch.ops._C_cuda_utils.get_device_attribute(attribute, device)
2125