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
[Kernel] cuda kernels for upcoming decode context parallel feature (#23791)
Co-authored-by: hongchao <hongchao@msh.team>
yzds committed
Aug 28, 2025 at 15:29 UTC
186aced5ffb62b62b41eb1beaf2a598ada43351b
5 files changed
+374
-1
csrc/cache.h
+16
-1
@@ -36,6 +36,13 @@ void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,
36
const std::string& kv_cache_dtype,
37
torch::Tensor& scale);
38
39
+void cp_fused_concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,
40
+ torch::Tensor& cp_local_token_select_indices,
41
+ torch::Tensor& kv_cache,
42
+ torch::Tensor& slot_mapping,
43
+ const std::string& kv_cache_dtype,
44
+ torch::Tensor& scale);
45
+
46
// Just for unittest
47
void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
48
const double scale, const std::string& kv_cache_dtype);
@@ -47,4 +54,12 @@ void gather_and_maybe_dequant_cache(
54
torch::Tensor const& cu_seq_lens, // [BATCH+1]
55
int64_t batch_size, const std::string& kv_cache_dtype,
56
torch::Tensor const& scale,
50
- std::optional<torch::Tensor> seq_starts = std::nullopt);
\ No newline at end of file
57
+ std::optional<torch::Tensor> seq_starts = std::nullopt);
58
+
59
+// TODO(hc): cp_gather_cache need support scaled kvcahe in the future.
60
+void cp_gather_cache(
61
+ torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
62
+ torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
63
+ torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
64
+ torch::Tensor const& cu_seq_lens, // [BATCH+1]
65
+ int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
csrc/cache_kernels.cu
+247
@@ -1,6 +1,7 @@
1
#include <torch/all.h>
2
#include <ATen/cuda/CUDAContext.h>
3
#include <c10/cuda/CUDAGuard.h>
4
+#include <c10/cuda/CUDAException.h>
5
6
#include "cuda_utils.h"
7
#include "cuda_compat.h"
@@ -395,6 +396,51 @@ __global__ void concat_and_cache_mla_kernel(
396
copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
397
}
398
399
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
400
+__global__ void cp_fused_concat_and_cache_mla_kernel(
401
+ const scalar_t* __restrict__ kv_c, // [num_full_tokens, kv_lora_rank]
402
+ const scalar_t* __restrict__ k_pe, // [num_full_tokens, pe_dim]
403
+ const int64_t* __restrict__ cp_local_token_select_indices, // [num_tokens]
404
+ cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank
405
+ // + pe_dim)]
406
+ const int64_t* __restrict__ slot_mapping, // [num_tokens]
407
+ const int block_stride, //
408
+ const int entry_stride, //
409
+ const int kv_c_stride, //
410
+ const int k_pe_stride, //
411
+ const int kv_lora_rank, //
412
+ const int pe_dim, //
413
+ const int block_size, //
414
+ const float* scale //
415
+) {
416
+ const int64_t token_idx = cp_local_token_select_indices[blockIdx.x];
417
+ const int64_t slot_idx = slot_mapping[blockIdx.x];
418
+ // NOTE: slot_idx can be -1 if the token is padded
419
+ if (slot_idx < 0) {
420
+ return;
421
+ }
422
+ const int64_t block_idx = slot_idx / block_size;
423
+ const int64_t block_offset = slot_idx % block_size;
424
+
425
+ auto copy = [&](const scalar_t* __restrict__ src, cache_t* __restrict__ dst,
426
+ int src_stride, int dst_stride, int size, int offset) {
427
+ for (int i = threadIdx.x; i < size; i += blockDim.x) {
428
+ const int64_t src_idx = token_idx * src_stride + i;
429
+ const int64_t dst_idx =
430
+ block_idx * block_stride + block_offset * entry_stride + i + offset;
431
+ if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
432
+ dst[dst_idx] = src[src_idx];
433
+ } else {
434
+ dst[dst_idx] =
435
+ fp8::scaled_convert<cache_t, scalar_t, kv_dt>(src[src_idx], *scale);
436
+ }
437
+ }
438
+ };
439
+
440
+ copy(kv_c, kv_cache, kv_c_stride, block_stride, kv_lora_rank, 0);
441
+ copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
442
+}
443
+
444
} // namespace vllm
445
446
// KV_T is the data type of key and value tensors.
@@ -508,6 +554,20 @@ void reshape_and_cache_flash(
554
kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
555
reinterpret_cast<const float*>(scale.data_ptr()));
556
557
+// KV_T is the data type of key and value tensors.
558
+// CACHE_T is the stored data type of kv-cache.
559
+// KV_DTYPE is the real data type of kv-cache.
560
+#define CALL_CP_FUSED_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \
561
+ vllm::cp_fused_concat_and_cache_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
562
+ <<<grid, block, 0, stream>>>( \
563
+ reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
564
+ reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
565
+ cp_local_token_select_indices.data_ptr<int64_t>(), \
566
+ reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
567
+ slot_mapping.data_ptr<int64_t>(), block_stride, entry_stride, \
568
+ kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
569
+ reinterpret_cast<const float*>(scale.data_ptr()));
570
+
571
void concat_and_cache_mla(
572
torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
573
torch::Tensor& k_pe, // [num_tokens, pe_dim]
@@ -546,6 +606,50 @@ void concat_and_cache_mla(
606
CALL_CONCAT_AND_CACHE_MLA);
607
}
608
609
+// Note(hc): cp_fused_concat_and_cache_mla fuses the following three kernel
610
+// calls into one:
611
+// k_c_normed.index_select(0, cp_local_token_select_indices) + \
612
+// k_pe.squeeze(1).index_select(0, cp_local_token_select_indices) + \
613
+// concat_and_cache_mla.
614
+void cp_fused_concat_and_cache_mla(
615
+ torch::Tensor& kv_c, // [num_total_tokens, kv_lora_rank]
616
+ torch::Tensor& k_pe, // [num_total_tokens, pe_dim]
617
+ torch::Tensor& cp_local_token_select_indices, // [num_tokens]
618
+ torch::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank +
619
+ // pe_dim)]
620
+ torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
621
+ const std::string& kv_cache_dtype, torch::Tensor& scale) {
622
+ // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
623
+ // slot_mapping.size(0) because of padding for CUDA graphs.
624
+ // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
625
+ // both include padding.
626
+ // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
627
+ // since key includes padding for CUDA graphs, while slot_mapping does not.
628
+ // In this case, slot_mapping.size(0) represents the actual number of tokens
629
+ // before padding.
630
+ // For compatibility with both cases, we use slot_mapping.size(0) as the
631
+ // number of tokens.
632
+ int num_tokens = slot_mapping.size(0);
633
+ int kv_lora_rank = kv_c.size(1);
634
+ int pe_dim = k_pe.size(1);
635
+ int block_size = kv_cache.size(1);
636
+
637
+ TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
638
+
639
+ int kv_c_stride = kv_c.stride(0);
640
+ int k_pe_stride = k_pe.stride(0);
641
+ int block_stride = kv_cache.stride(0);
642
+ int entry_stride = kv_cache.stride(1);
643
+
644
+ dim3 grid(num_tokens);
645
+ dim3 block(std::min(kv_lora_rank, 512));
646
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_c));
647
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
648
+
649
+ DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
650
+ CALL_CP_FUSED_CONCAT_AND_CACHE_MLA);
651
+}
652
+
653
namespace vllm {
654
655
template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
@@ -779,3 +883,146 @@ void gather_and_maybe_dequant_cache(
883
884
DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype, CALL_GATHER_CACHE);
885
}
886
+
887
+namespace vllm {
888
+template <typename scalar_t>
889
+// Note(hc): The cp_gather_cache allows seq_starts to no longer be divisible by
890
+// block_size.
891
+__global__ void cp_gather_cache(
892
+ const scalar_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
893
+ // ENTRY_SIZE]
894
+ scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRY_SIZE]
895
+ const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]
896
+ const int32_t* __restrict__ cu_seq_lens, // [BATCH+1]
897
+ const int32_t block_size, const int32_t entry_size,
898
+ const int64_t block_table_stride, const int64_t cache_block_stride,
899
+ const int64_t cache_entry_stride, const int64_t dst_entry_stride,
900
+ const int32_t* __restrict__ seq_starts // Optional: starting offsets per
901
+ // batch
902
+) {
903
+ const int64_t bid = blockIdx.x; // Batch ID
904
+ const int32_t num_splits = gridDim.y;
905
+ const int32_t split = blockIdx.y;
906
+ const int32_t seq_start = cu_seq_lens[bid];
907
+ const int32_t seq_end = cu_seq_lens[bid + 1];
908
+ const int32_t seq_len = seq_end - seq_start;
909
+ const int32_t tot_slots = seq_len;
910
+ const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits);
911
+
912
+ const int32_t split_start = split * split_slots;
913
+ const int32_t split_end = min((split + 1) * split_slots, tot_slots);
914
+
915
+ const bool is_active_split = (split_start < tot_slots);
916
+ const bool is_last_split = (split_end == tot_slots);
917
+
918
+ if (!is_active_split) return;
919
+
920
+ // Adjust the pointer for the block_table for this batch.
921
+ // If seq_starts is provided, compute an offset based on it
922
+ const int32_t batch_offset = bid * block_table_stride;
923
+ int32_t offset = split_start;
924
+ if (seq_starts != nullptr) {
925
+ offset += seq_starts[bid];
926
+ }
927
+ int32_t offset_div = offset / block_size;
928
+ offset = offset % block_size;
929
+ const int32_t* batch_block_table = block_table + batch_offset;
930
+
931
+ // Adjust dst pointer based on the cumulative sequence lengths.
932
+ dst += seq_start * dst_entry_stride;
933
+
934
+ auto copy_entry = [&](const scalar_t* __restrict__ _src,
935
+ scalar_t* __restrict__ _dst) {
936
+ for (int i = threadIdx.x; i < entry_size; i += blockDim.x)
937
+ _dst[i] = _src[i];
938
+ };
939
+
940
+ for (int pid = split_start; pid < split_end; ++pid) {
941
+ auto block_id = batch_block_table[offset_div];
942
+ auto block_start_ptr = src_cache + block_id * cache_block_stride;
943
+ auto block_dst_ptr = dst + pid * dst_entry_stride;
944
+ copy_entry(block_start_ptr + offset * cache_entry_stride, block_dst_ptr);
945
+ offset += 1;
946
+ // bump to next block
947
+ if (offset == block_size) {
948
+ offset_div += 1;
949
+ offset = 0;
950
+ }
951
+ }
952
+}
953
+} // namespace vllm
954
+
955
+// Macro to dispatch the kernel based on the data type.
956
+#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \
957
+ vllm::cp_gather_cache<CPY_DTYPE><<<grid, block, 0, stream>>>( \
958
+ reinterpret_cast<CPY_DTYPE*>(src_cache.data_ptr()), \
959
+ reinterpret_cast<CPY_DTYPE*>(dst.data_ptr()), \
960
+ block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
961
+ block_size, entry_size, block_table_stride, cache_block_stride, \
962
+ cache_entry_stride, dst_entry_stride, seq_starts_ptr);
963
+
964
+// Gather sequences from the cache into the destination tensor.
965
+// - cu_seq_lens contains the cumulative sequence lengths for each batch
966
+// - block_table contains the cache block indices for each sequence
967
+// - Optionally, seq_starts (if provided) offsets the starting slot index by
968
+// seq_starts[bid]
969
+void cp_gather_cache(
970
+ torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
971
+ torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
972
+ torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
973
+ torch::Tensor const& cu_seq_lens, // [BATCH+1]
974
+ int64_t batch_size,
975
+ std::optional<torch::Tensor> seq_starts = std::nullopt) {
976
+ at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
977
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
978
+
979
+ int32_t block_size = src_cache.size(1);
980
+ int32_t entry_size = src_cache.flatten(2, -1).size(2);
981
+
982
+ TORCH_CHECK(block_table.dtype() == torch::kInt32,
983
+ "block_table must be int32");
984
+ TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32,
985
+ "cu_seq_lens must be int32");
986
+ if (seq_starts.has_value()) {
987
+ TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
988
+ "seq_starts must be int32");
989
+ }
990
+
991
+ TORCH_CHECK(src_cache.device() == dst.device(),
992
+ "src_cache and dst must be on the same device");
993
+ TORCH_CHECK(src_cache.device() == block_table.device(),
994
+ "src_cache and block_table must be on the same device");
995
+ TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
996
+ "src_cache and cu_seq_lens must be on the same device");
997
+ if (seq_starts.has_value()) {
998
+ TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
999
+ "src_cache and seq_starts must be on the same device");
1000
+ }
1001
+
1002
+ int64_t block_table_stride = block_table.stride(0);
1003
+ int64_t cache_block_stride = src_cache.stride(0);
1004
+ int64_t cache_entry_stride = src_cache.stride(1);
1005
+ int64_t dst_entry_stride = dst.stride(0);
1006
+
1007
+ // Decide on the number of splits based on the batch size.
1008
+ int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
1009
+ dim3 grid(batch_size, num_splits);
1010
+ dim3 block(1024);
1011
+
1012
+ TORCH_CHECK(src_cache.dtype() == dst.dtype(),
1013
+ "src_cache and dst must have the same dtype");
1014
+
1015
+ const int dtype_bits = src_cache.element_size() * 8;
1016
+ const int32_t* seq_starts_ptr =
1017
+ seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
1018
+
1019
+ if (dtype_bits == 32) {
1020
+ CALL_CP_GATHER_CACHE(uint32_t);
1021
+ } else if (dtype_bits == 16) {
1022
+ CALL_CP_GATHER_CACHE(uint16_t);
1023
+ } else if (dtype_bits == 8) {
1024
+ CALL_CP_GATHER_CACHE(uint8_t);
1025
+ } else {
1026
+ TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits);
1027
+ }
1028
+}
csrc/torch_bindings.cpp
+15
@@ -686,6 +686,16 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
686
" Tensor scale) -> ()");
687
cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
688
689
+ cache_ops.def(
690
+ "cp_fused_concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
691
+ " Tensor cp_local_token_select_indices,"
692
+ " Tensor! kv_cache,"
693
+ " Tensor slot_mapping,"
694
+ " str kv_cache_dtype,"
695
+ " Tensor scale) -> ()");
696
+ cache_ops.impl("cp_fused_concat_and_cache_mla", torch::kCUDA,
697
+ &cp_fused_concat_and_cache_mla);
698
+
699
// Convert the key and value cache to fp8 data type.
700
cache_ops.def(
701
"convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
@@ -702,6 +712,11 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
712
" Tensor scale, Tensor? seq_starts) -> ()");
713
cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,
714
&gather_and_maybe_dequant_cache);
715
+
716
+ cache_ops.def(
717
+ "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
718
+ "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
719
+ cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache);
720
}
721
722
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
tests/kernels/attention/test_cache.py
+72
@@ -790,6 +790,78 @@ def test_gather_and_maybe_dequant_cache_mla(kv_lora_rank, qk_rope_head_dim,
790
torch.testing.assert_close(dst, expected)
791
792
793
+@pytest.mark.parametrize("kv_lora_rank", [512])
794
+@pytest.mark.parametrize("qk_rope_head_dim", [64])
795
+@pytest.mark.parametrize("block_size", [16])
796
+@pytest.mark.parametrize("num_blocks", [1024])
797
+@pytest.mark.parametrize("max_seq_len", [512])
798
+@pytest.mark.parametrize("batch_size", [8])
799
+@pytest.mark.parametrize("dtype", [torch.float32])
800
+@pytest.mark.parametrize("kv_cache_dtype",
801
+ ["auto"]) # You can also test "fp8" if needed.
802
+@pytest.mark.parametrize("device", CUDA_DEVICES)
803
+@torch.inference_mode()
804
+def test_cp_gather_cache_mla(kv_lora_rank, qk_rope_head_dim, block_size,
805
+ num_blocks, max_seq_len, batch_size, dtype,
806
+ kv_cache_dtype, device):
807
+ entry_size = kv_lora_rank + qk_rope_head_dim
808
+ src_cache = _create_mla_cache(num_blocks, block_size, entry_size, dtype,
809
+ kv_cache_dtype, device)
810
+ _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)
811
+
812
+ seq_len_tensor = torch.randint(0,
813
+ max_seq_len + 1, (batch_size, ),
814
+ device=device)
815
+
816
+ total_tokens = seq_len_tensor.sum()
817
+ cu_seq_lens = torch.empty((batch_size + 1),
818
+ dtype=torch.int32,
819
+ device=device)
820
+ cu_seq_lens[0] = 0
821
+ cu_seq_lens[1:] = seq_len_tensor.cumsum(dim=0).to(dtype=torch.int32)
822
+ print("seq_len_tensor", seq_len_tensor)
823
+
824
+ tot_blocks_tensor = (seq_len_tensor + block_size - 1) // block_size
825
+ block_table = torch.empty((batch_size, num_blocks),
826
+ dtype=torch.int32,
827
+ device=device)
828
+
829
+ for b in range(batch_size):
830
+ perm = torch.randperm(num_blocks, device=device)
831
+ block_table[b, :] = perm
832
+
833
+ dst = torch.zeros((total_tokens, entry_size),
834
+ dtype=src_cache.dtype,
835
+ device=device)
836
+
837
+ expected_batches = []
838
+ for b in range(batch_size):
839
+ s = seq_len_tensor[b]
840
+ if s == 0:
841
+ continue
842
+ tot = tot_blocks_tensor[b]
843
+ blocks = block_table[b, :tot].tolist()
844
+
845
+ gathered_rows = []
846
+ for i in range(tot - 1):
847
+ gathered_rows.append(src_cache[blocks[i]])
848
+ remaining = s - (tot - 1) * block_size
849
+ gathered_rows.append(src_cache[blocks[-1], :remaining, :])
850
+
851
+ batch_expected = torch.cat(gathered_rows, dim=0)
852
+ expected_batches.append(batch_expected)
853
+ expected = torch.cat(expected_batches, dim=0)
854
+
855
+ opcheck(
856
+ torch.ops._C_cache_ops.cp_gather_cache,
857
+ (src_cache, dst, block_table, cu_seq_lens, batch_size, None),
858
+ test_utils=DEFAULT_OPCHECK_TEST_UTILS,
859
+ )
860
+
861
+ ops.cp_gather_cache(src_cache, dst, block_table, cu_seq_lens, batch_size)
862
+ torch.testing.assert_close(dst, expected)
863
+
864
+
865
@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
866
@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
867
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
vllm/_custom_ops.py
+24
@@ -1625,6 +1625,20 @@ def concat_and_cache_mla(
1625
scale)
1626
1627
1628
+def cp_fused_concat_and_cache_mla(
1629
+ kv_c: torch.Tensor,
1630
+ k_pe: torch.Tensor,
1631
+ cp_local_token_select_indices: torch.Tensor,
1632
+ kv_cache: torch.Tensor,
1633
+ slot_mapping: torch.Tensor,
1634
+ kv_cache_dtype: str,
1635
+ scale: torch.Tensor,
1636
+) -> None:
1637
+ torch.ops._C_cache_ops.cp_fused_concat_and_cache_mla(
1638
+ kv_c, k_pe, cp_local_token_select_indices, kv_cache, slot_mapping,
1639
+ kv_cache_dtype, scale)
1640
+
1641
+
1642
def copy_blocks(key_caches: list[torch.Tensor],
1643
value_caches: list[torch.Tensor],
1644
block_mapping: torch.Tensor) -> None:
@@ -1662,6 +1676,16 @@ def gather_and_maybe_dequant_cache(
1676
scale, seq_starts)
1677
1678
1679
+def cp_gather_cache(src_cache: torch.Tensor,
1680
+ dst: torch.Tensor,
1681
+ block_table: torch.Tensor,
1682
+ cu_seq_lens: torch.Tensor,
1683
+ batch_size: int,
1684
+ seq_starts: Optional[torch.Tensor] = None) -> None:
1685
+ torch.ops._C_cache_ops.cp_gather_cache(src_cache, dst, block_table,
1686
+ cu_seq_lens, batch_size, seq_starts)
1687
+
1688
+
1689
def get_device_attribute(attribute: int, device: int) -> int:
1690
return torch.ops._C_cuda_utils.get_device_attribute(attribute, device)
1691