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 11, 00:00 UTC
[Misc] Remove unused custom ops `copy_blocks` and `copy_blocks_mla` (#30967)
Signed-off-by: rongfu.leng <rongfu.leng@daocloud.io>
rongfu.leng committed
Dec 24, 2025 at 10:22 UTC
4ed11105d7b81fe9005b119c86ef42f198ced2cf
6 files changed
-286
csrc/cache.h
-10
@@ -9,16 +9,6 @@
9
void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
10
const torch::Tensor& block_mapping);
11
12
-// Note: the key_caches and value_caches vectors are constant but
13
-// not the Tensors they contain. The vectors need to be const refs
14
-// in order to satisfy pytorch's C++ operator registration code.
15
-void copy_blocks(std::vector<torch::Tensor> const& key_caches,
16
- std::vector<torch::Tensor> const& value_caches,
17
- const torch::Tensor& block_mapping);
18
-
19
-void copy_blocks_mla(std::vector<torch::Tensor> const& kv_caches,
20
- const torch::Tensor& block_mapping);
21
-
12
void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
13
torch::Tensor& key_cache, torch::Tensor& value_cache,
14
torch::Tensor& slot_mapping,
csrc/cache_kernels.cu
-88
@@ -119,94 +119,6 @@ __global__ void copy_blocks_mla_kernel(
119
120
} // namespace vllm
121
122
-// Note: the key_caches and value_caches vectors are constant but
123
-// not the Tensors they contain. The vectors need to be const refs
124
-// in order to satisfy pytorch's C++ operator registration code.
125
-void copy_blocks(std::vector<torch::Tensor> const& key_caches,
126
- std::vector<torch::Tensor> const& value_caches,
127
- const torch::Tensor& block_mapping) {
128
- int num_layers = key_caches.size();
129
- TORCH_CHECK(num_layers == value_caches.size());
130
- if (num_layers == 0) {
131
- return;
132
- }
133
- torch::Device cache_device = key_caches[0].device();
134
- TORCH_CHECK(cache_device.is_cuda());
135
-
136
- // Create data structures for the kernel.
137
- // Create an array of pointers to the key and value caches.
138
- int64_t key_cache_ptrs[num_layers];
139
- int64_t value_cache_ptrs[num_layers];
140
- for (int layer_idx = 0; layer_idx < num_layers; ++layer_idx) {
141
- key_cache_ptrs[layer_idx] =
142
- reinterpret_cast<int64_t>(key_caches[layer_idx].data_ptr());
143
- value_cache_ptrs[layer_idx] =
144
- reinterpret_cast<int64_t>(value_caches[layer_idx].data_ptr());
145
- }
146
-
147
- // block_mapping is a 2D tensor with shape (num_pairs, 2).
148
- int num_pairs = block_mapping.size(0);
149
-
150
- // Move the data structures to the GPU.
151
- // NOTE: This synchronizes the CPU and GPU.
152
- torch::Tensor key_cache_ptrs_tensor =
153
- torch::from_blob(key_cache_ptrs, {num_layers}, torch::kInt64)
154
- .to(cache_device);
155
- torch::Tensor value_cache_ptrs_tensor =
156
- torch::from_blob(value_cache_ptrs, {num_layers}, torch::kInt64)
157
- .to(cache_device);
158
-
159
- // Launch the kernel.
160
- const int numel_per_block = key_caches[0][0].numel();
161
- dim3 grid(num_layers, num_pairs);
162
- dim3 block(std::min(1024, numel_per_block));
163
- const at::cuda::OptionalCUDAGuard device_guard(cache_device);
164
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
165
- VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(
166
- key_caches[0].scalar_type(), "copy_blocks_kernel", ([&] {
167
- vllm::copy_blocks_kernel<scalar_t><<<grid, block, 0, stream>>>(
168
- key_cache_ptrs_tensor.data_ptr<int64_t>(),
169
- value_cache_ptrs_tensor.data_ptr<int64_t>(),
170
- block_mapping.data_ptr<int64_t>(), numel_per_block);
171
- }));
172
-}
173
-
174
-// copy blocks kernel for MLA (assumes a joint KV-cache)
175
-void copy_blocks_mla(std::vector<torch::Tensor> const& kv_caches,
176
- const torch::Tensor& block_mapping) {
177
- int num_layers = kv_caches.size();
178
- if (num_layers == 0) {
179
- return;
180
- }
181
- torch::Device cache_device = kv_caches[0].device();
182
- TORCH_CHECK(cache_device.is_cuda(), "kv_cache must be on CUDA");
183
-
184
- std::vector<int64_t> cache_ptrs(num_layers);
185
- for (int layer_idx = 0; layer_idx < num_layers; ++layer_idx) {
186
- cache_ptrs[layer_idx] =
187
- reinterpret_cast<int64_t>(kv_caches[layer_idx].data_ptr());
188
- }
189
- torch::Tensor cache_ptrs_tensor =
190
- torch::from_blob(cache_ptrs.data(), {num_layers}, torch::kInt64)
191
- .to(cache_device);
192
-
193
- int num_pairs = block_mapping.size(0);
194
- // We use the stride instead of numel in case the cache is padded for memory
195
- // alignment reasons, we assume the blocks data (inclusive of any padding)
196
- // is contiguous in memory
197
- int mem_footprint_per_block = kv_caches[0].stride(0);
198
- dim3 grid(num_layers, num_pairs);
199
- dim3 block(std::min(1024, mem_footprint_per_block));
200
- const at::cuda::OptionalCUDAGuard device_guard(cache_device);
201
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
202
- VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(
203
- kv_caches[0].scalar_type(), "copy_blocks_mla_kernel", ([&] {
204
- vllm::copy_blocks_mla_kernel<scalar_t><<<grid, block, 0, stream>>>(
205
- cache_ptrs_tensor.data_ptr<int64_t>(),
206
- block_mapping.data_ptr<int64_t>(), mem_footprint_per_block);
207
- }));
208
-}
209
-
122
namespace vllm {
123
124
// Used to copy/convert one element
csrc/torch_bindings.cpp
-10
@@ -685,16 +685,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
685
"swap_blocks(Tensor src, Tensor! dst, Tensor block_mapping) -> ()");
686
cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);
687
688
- // Copy the cache blocks from src to dst.
689
- cache_ops.def(
690
- "copy_blocks(Tensor(a!)[] key_caches, Tensor[](b!) value_caches, "
691
- "Tensor block_mapping) -> ()");
692
- cache_ops.impl("copy_blocks", torch::kCUDA, ©_blocks);
693
-
694
- cache_ops.def(
695
- "copy_blocks_mla(Tensor(a!)[] kv_caches, Tensor block_mapping) -> ()");
696
- cache_ops.impl("copy_blocks_mla", torch::kCUDA, ©_blocks_mla);
697
-
688
// Reshape the key and value tensors and cache them.
689
cache_ops.def(
690
"reshape_and_cache(Tensor key, Tensor value,"
tests/kernels/attention/test_cache.py
-154
@@ -40,93 +40,6 @@ KV_CACHE_DTYPE = ["auto", "fp8"]
40
RESHAPE_FLASH_IMPLEMENTATIONS = ["cuda", "triton"]
41
42
43
-@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
44
-@pytest.mark.parametrize("num_layers", NUM_LAYERS)
45
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
46
-@pytest.mark.parametrize("head_size", HEAD_SIZES)
47
-@pytest.mark.parametrize("block_size", BLOCK_SIZES)
48
-@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
49
-@pytest.mark.parametrize("dtype", DTYPES)
50
-@pytest.mark.parametrize("seed", SEEDS)
51
-@pytest.mark.parametrize("device", CUDA_DEVICES)
52
-@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
53
-@torch.inference_mode()
54
-def test_copy_blocks(
55
- kv_cache_factory,
56
- num_mappings: int,
57
- num_layers: int,
58
- num_heads: int,
59
- head_size: int,
60
- block_size: int,
61
- num_blocks: int,
62
- dtype: torch.dtype,
63
- seed: int,
64
- kv_cache_dtype: str,
65
- device: str,
66
-) -> None:
67
- if kv_cache_dtype == "fp8" and head_size % 16:
68
- pytest.skip()
69
- current_platform.seed_everything(seed)
70
- torch.set_default_device(device)
71
- torch.cuda.set_device(device)
72
- # Generate random block mappings where each source block is mapped to two
73
- # destination blocks.
74
- assert 2 * num_mappings <= num_blocks
75
- src_blocks = random.sample(range(num_blocks), num_mappings)
76
- remaining_blocks = list(set(range(num_blocks)) - set(src_blocks))
77
- dst_blocks = random.sample(remaining_blocks, 2 * num_mappings)
78
- block_mapping: list[tuple[int, int]] = []
79
- for i in range(num_mappings):
80
- src = src_blocks[i]
81
- dst1 = dst_blocks[2 * i]
82
- dst2 = dst_blocks[2 * i + 1]
83
- block_mapping.append((src, dst1))
84
- block_mapping.append((src, dst2))
85
-
86
- # Create the KV caches.
87
- key_caches, value_caches = kv_cache_factory(
88
- num_blocks,
89
- block_size,
90
- num_layers,
91
- num_heads,
92
- head_size,
93
- kv_cache_dtype,
94
- dtype,
95
- seed,
96
- device,
97
- )
98
-
99
- # Clone the KV caches.
100
- cloned_key_caches = [key_cache.clone() for key_cache in key_caches]
101
- cloned_value_caches = [value_cache.clone() for value_cache in value_caches]
102
-
103
- # Call the copy blocks kernel.
104
- block_mapping_tensor = torch.tensor(
105
- block_mapping, dtype=torch.int64, device=device
106
- ).view(-1, 2)
107
-
108
- opcheck(
109
- torch.ops._C_cache_ops.copy_blocks,
110
- (key_caches, value_caches, block_mapping_tensor),
111
- test_utils=DEFAULT_OPCHECK_TEST_UTILS,
112
- cond=(head_size == HEAD_SIZES[0]),
113
- )
114
- ops.copy_blocks(key_caches, value_caches, block_mapping_tensor)
115
-
116
- # Run the reference implementation.
117
- for src, dst in block_mapping:
118
- for cloned_key_cache in cloned_key_caches:
119
- cloned_key_cache[dst].copy_(cloned_key_cache[src])
120
- for cloned_value_cache in cloned_value_caches:
121
- cloned_value_cache[dst].copy_(cloned_value_cache[src])
122
-
123
- # Compare the results.
124
- for key_cache, cloned_key_cache in zip(key_caches, cloned_key_caches):
125
- torch.testing.assert_close(key_cache, cloned_key_cache)
126
- for value_cache, cloned_value_cache in zip(value_caches, cloned_value_caches):
127
- torch.testing.assert_close(value_cache, cloned_value_cache)
128
-
129
-
43
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
44
@pytest.mark.parametrize("num_heads", NUM_HEADS)
45
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -763,73 +676,6 @@ def test_concat_and_cache_ds_mla(
676
torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1)
677
678
766
-@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
767
-@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
768
-@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
769
-@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
770
-@pytest.mark.parametrize("num_layers", NUM_LAYERS)
771
-@pytest.mark.parametrize("dtype", DTYPES)
772
-@pytest.mark.parametrize("seed", SEEDS)
773
-@pytest.mark.parametrize("device", CUDA_DEVICES)
774
-@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
775
-@torch.inference_mode()
776
-def test_copy_blocks_mla(
777
- kv_lora_rank: int,
778
- qk_rope_head_dim: int,
779
- block_size: int,
780
- num_blocks: int,
781
- num_layers: int,
782
- dtype: torch.dtype,
783
- seed: int,
784
- device: str,
785
- kv_cache_dtype: str,
786
-) -> None:
787
- current_platform.seed_everything(seed)
788
- torch.set_default_device(device)
789
- torch.cuda.set_device(device)
790
-
791
- entry_size = kv_lora_rank + qk_rope_head_dim
792
-
793
- kv_caches = []
794
- for _ in range(num_layers):
795
- kv_cache = _create_mla_cache(
796
- num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
797
- )
798
- _fill_mla_cache(kv_cache, kv_cache_dtype=kv_cache_dtype)
799
- kv_caches.append(kv_cache)
800
-
801
- ref_caches = [kv_cache.clone() for kv_cache in kv_caches]
802
-
803
- num_mappings = min(2, num_blocks // 2)
804
- src_blocks = random.sample(range(num_blocks), num_mappings)
805
- remaining = list(set(range(num_blocks)) - set(src_blocks))
806
- dst_blocks = random.sample(remaining, 2 * num_mappings)
807
- block_mapping = []
808
- for i in range(num_mappings):
809
- src = src_blocks[i]
810
- dst1 = dst_blocks[2 * i]
811
- dst2 = dst_blocks[2 * i + 1]
812
- block_mapping.append((src, dst1))
813
- block_mapping.append((src, dst2))
814
- block_mapping_tensor = torch.tensor(
815
- block_mapping, dtype=torch.int64, device=device
816
- ).view(-1, 2)
817
-
818
- for src, dst in block_mapping:
819
- for ref_cache in ref_caches:
820
- ref_cache[dst].copy_(ref_cache[src])
821
-
822
- opcheck(
823
- torch.ops._C_cache_ops.copy_blocks_mla,
824
- (kv_caches, block_mapping_tensor),
825
- test_utils=DEFAULT_OPCHECK_TEST_UTILS,
826
- )
827
- ops.copy_blocks_mla(kv_caches, block_mapping_tensor)
828
-
829
- for kv_cache, ref_cache in zip(kv_caches, ref_caches):
830
- torch.testing.assert_close(kv_cache, ref_cache)
831
-
832
-
679
@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
680
@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
681
@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
vllm/_custom_ops.py
-12
@@ -2328,18 +2328,6 @@ def concat_and_cache_mla(
2328
)
2329
2330
2331
-def copy_blocks(
2332
- key_caches: list[torch.Tensor],
2333
- value_caches: list[torch.Tensor],
2334
- block_mapping: torch.Tensor,
2335
-) -> None:
2336
- torch.ops._C_cache_ops.copy_blocks(key_caches, value_caches, block_mapping)
2337
-
2338
-
2339
-def copy_blocks_mla(kv_caches: list[torch.Tensor], block_mapping: torch.Tensor) -> None:
2340
- torch.ops._C_cache_ops.copy_blocks_mla(kv_caches, block_mapping)
2341
-
2342
-
2331
def swap_blocks(
2332
src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor
2333
) -> None:
vllm/_ipex_ops.py
-12
@@ -383,18 +383,6 @@ class ipex_ops:
383
)
384
return None
385
386
- @staticmethod
387
- def copy_blocks(
388
- key_caches: list[torch.Tensor],
389
- value_caches: list[torch.Tensor],
390
- block_mapping: torch.Tensor,
391
- ) -> None:
392
- torch.xpu.copy_blocks( # type: ignore
393
- key_caches,
394
- value_caches,
395
- block_mapping,
396
- )
397
-
386
@staticmethod
387
def swap_blocks(
388
src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor