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 9, 19:30 UTC
OffloadingConnector: Support kernel_block_size != block_size (#30692)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Jan 22, 2026 at 14:30 UTC
421012b63ae2e3f26bba0506e9f3951e9dc53ded
7 files changed
+156
-114
csrc/cache.h
+1
@@ -7,6 +7,7 @@
7
#include <vector>
8
9
void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
10
+ int64_t block_size_in_bytes,
11
const torch::Tensor& block_mapping);
12
13
void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
csrc/cache_kernels.cu
+1
-4
@@ -25,6 +25,7 @@ typedef __hip_bfloat16 __nv_bfloat16;
25
#endif
26
27
void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
28
+ int64_t block_size_in_bytes,
29
const torch::Tensor& block_mapping) {
30
torch::Device src_device = src.device();
31
torch::Device dst_device = dst.device();
@@ -49,10 +50,6 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
50
char* src_ptr = static_cast<char*>(src.data_ptr());
51
char* dst_ptr = static_cast<char*>(dst.data_ptr());
52
52
- // We use the stride instead of numel in case the cache is padded for memory
53
- // alignment reasons, we assume the blocks data (inclusive of any padding)
54
- // is contiguous in memory
55
- const int64_t block_size_in_bytes = src.element_size() * src.stride(0);
53
const at::cuda::OptionalCUDAGuard device_guard(
54
src_device.is_cuda() ? src_device : dst_device);
55
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
csrc/torch_bindings.cpp
+2
-1
@@ -692,7 +692,8 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
692
// Cache ops
693
// Swap in (out) the cache blocks from src to dst.
694
cache_ops.def(
695
- "swap_blocks(Tensor src, Tensor! dst, Tensor block_mapping) -> ()");
695
+ "swap_blocks(Tensor src, Tensor! dst,"
696
+ " int block_size_in_bytes, Tensor block_mapping) -> ()");
697
cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);
698
699
// Reshape the key and value tensors and cache them.
tests/kernels/attention/test_cache.py
+29
-6
@@ -405,19 +405,41 @@ def test_swap_blocks(
405
406
# Call the swap_blocks kernel.
407
do_opcheck = head_size == HEAD_SIZES[0]
408
+ src_cache = src_key_caches[0]
409
+ block_size_in_bytes = src_cache.element_size() * src_cache.stride(0)
410
opcheck(
411
torch.ops._C_cache_ops.swap_blocks,
410
- (src_key_caches[0], dist_key_caches[0], block_mapping_tensor),
412
+ (
413
+ src_key_caches[0],
414
+ dist_key_caches[0],
415
+ block_size_in_bytes,
416
+ block_mapping_tensor,
417
+ ),
418
cond=do_opcheck,
419
)
420
opcheck(
421
torch.ops._C_cache_ops.swap_blocks,
415
- (src_value_caches[0], dist_value_caches[0], block_mapping_tensor),
422
+ (
423
+ src_value_caches[0],
424
+ dist_value_caches[0],
425
+ block_size_in_bytes,
426
+ block_mapping_tensor,
427
+ ),
428
cond=do_opcheck,
429
)
430
419
- ops.swap_blocks(src_key_caches[0], dist_key_caches[0], block_mapping_tensor)
420
- ops.swap_blocks(src_value_caches[0], dist_value_caches[0], block_mapping_tensor)
431
+ ops.swap_blocks(
432
+ src_key_caches[0],
433
+ dist_key_caches[0],
434
+ block_size_in_bytes,
435
+ block_mapping_tensor,
436
+ )
437
+ ops.swap_blocks(
438
+ src_value_caches[0],
439
+ dist_value_caches[0],
440
+ block_size_in_bytes,
441
+ block_mapping_tensor,
442
+ )
443
444
for src, dst in block_mapping:
445
torch.testing.assert_close(
@@ -723,13 +745,14 @@ def test_swap_blocks_mla(
745
block_mapping, dtype=torch.int64, device="cpu"
746
).view(-1, 2)
747
748
+ block_size_in_bytes = src_cache.element_size() * src_cache.stride(0)
749
opcheck(
750
torch.ops._C_cache_ops.swap_blocks,
728
- (src_cache, dst_cache, block_mapping_tensor),
751
+ (src_cache, dst_cache, block_size_in_bytes, block_mapping_tensor),
752
test_utils=DEFAULT_OPCHECK_TEST_UTILS,
753
)
754
732
- ops.swap_blocks(src_cache, dst_cache, block_mapping_tensor)
755
+ ops.swap_blocks(src_cache, dst_cache, block_size_in_bytes, block_mapping_tensor)
756
757
for src, dst in block_mapping:
758
torch.testing.assert_close(
tests/v1/kv_offload/test_cpu_gpu.py
+51
-45
@@ -25,8 +25,9 @@ if not current_platform.is_rocm():
25
26
NUM_GPU_BLOCKS = [64]
27
NUM_CPU_BLOCKS = [256]
28
-GPU_BLOCK_SIZES = [16]
29
-GPU_BLOCKS_PER_CPU_BLOCK = [1, 3]
28
+KERNEL_BLOCK_SIZES = [16]
29
+LOGICAL_BLOCK_SIZES = [16, 32]
30
+LOGICAL_BLOCKS_PER_CPU_BLOCK = [1, 3]
31
HEAD_SIZES = [64]
32
NUM_HEADS = [8]
33
NUM_LAYERS = [4]
@@ -40,8 +41,9 @@ NUM_MAPPINGS = [3]
41
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
42
@pytest.mark.parametrize("head_size", HEAD_SIZES)
43
@pytest.mark.parametrize("num_heads", NUM_HEADS)
43
-@pytest.mark.parametrize("gpu_block_size", GPU_BLOCK_SIZES)
44
-@pytest.mark.parametrize("gpu_blocks_per_cpu_block", GPU_BLOCKS_PER_CPU_BLOCK)
44
+@pytest.mark.parametrize("kernel_block_size", KERNEL_BLOCK_SIZES)
45
+@pytest.mark.parametrize("logical_block_size", LOGICAL_BLOCK_SIZES)
46
+@pytest.mark.parametrize("logical_blocks_per_cpu_block", LOGICAL_BLOCKS_PER_CPU_BLOCK)
47
@pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
48
@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
49
@pytest.mark.parametrize("num_layers", NUM_LAYERS)
@@ -55,8 +57,9 @@ def test_transfer(
57
num_mappings: int,
58
head_size: int,
59
num_heads: int,
58
- gpu_block_size: int,
59
- gpu_blocks_per_cpu_block: int,
60
+ kernel_block_size: int,
61
+ logical_block_size: int,
62
+ logical_blocks_per_cpu_block: int,
63
num_gpu_blocks: int,
64
num_cpu_blocks: int,
65
num_layers: int,
@@ -69,6 +72,10 @@ def test_transfer(
72
# create per-layer GPU KV caches based on available attn_backends
73
attn_backends_list = BACKENDS_TO_TEST
74
75
+ assert logical_block_size % kernel_block_size == 0
76
+ kernel_blocks_per_gpu_block = logical_block_size // kernel_block_size
77
+ num_gpu_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
78
+
79
gpu_caches = {}
80
attn_backends = {}
81
for i in range(num_layers):
@@ -78,15 +85,16 @@ def test_transfer(
85
attn_backends[layer_name] = attn_backend
86
87
gpu_cache_shape = attn_backend.get_kv_cache_shape(
81
- num_gpu_blocks, gpu_block_size, num_heads, head_size
88
+ num_gpu_kernel_blocks, kernel_block_size, num_heads, head_size
89
)
90
gpu_caches[layer_name] = torch.rand(gpu_cache_shape, dtype=dtype, device=device)
91
92
# create handler
86
- cpu_block_size = gpu_blocks_per_cpu_block * gpu_block_size
93
+ cpu_block_size = logical_blocks_per_cpu_block * logical_block_size
94
+ kernel_blocks_per_cpu_block = cpu_block_size // kernel_block_size
95
handlers = CpuGpuOffloadingHandlers(
96
attn_backends=attn_backends,
89
- gpu_block_size=gpu_block_size,
97
+ gpu_block_size=logical_block_size,
98
cpu_block_size=cpu_block_size,
99
num_cpu_blocks=num_cpu_blocks,
100
gpu_caches=gpu_caches,
@@ -94,22 +102,34 @@ def test_transfer(
102
103
# select block mappings
104
gpu_blocks = random.sample(
97
- range(num_gpu_blocks), num_mappings * gpu_blocks_per_cpu_block
105
+ range(num_gpu_blocks), num_mappings * logical_blocks_per_cpu_block
106
)
107
cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings)
108
109
+ # convert gpu blocks to kernel block size
110
+ gpu_blocks_in_kernel_block_size = []
111
+ for gpu_block in gpu_blocks:
112
+ base_block_id = gpu_block * kernel_blocks_per_gpu_block
113
+ for i in range(kernel_blocks_per_gpu_block):
114
+ gpu_blocks_in_kernel_block_size.append(i + base_block_id)
115
+
116
# convert cpu blocks to gpu block size
102
- cpu_blocks_in_gpu_block_size = []
117
+ cpu_blocks_in_kernel_block_size = []
118
for cpu_block in cpu_blocks:
104
- base_block_id = cpu_block * gpu_blocks_per_cpu_block
105
- for i in range(gpu_blocks_per_cpu_block):
106
- cpu_blocks_in_gpu_block_size.append(i + base_block_id)
119
+ base_block_id = cpu_block * kernel_blocks_per_cpu_block
120
+ for i in range(kernel_blocks_per_cpu_block):
121
+ cpu_blocks_in_kernel_block_size.append(i + base_block_id)
122
108
- # maybe skip a GPU block to test reading from the middle of a CPU block
123
+ # maybe skip some GPU block to test reading from the middle of a CPU block
124
if not gpu_to_cpu:
110
- gpu_blocks = gpu_blocks[gpu_blocks_per_cpu_block - 1 :]
111
- cpu_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size[
112
- gpu_blocks_per_cpu_block - 1 :
125
+ gpu_blocks_to_skip = logical_blocks_per_cpu_block - 1
126
+ gpu_blocks = gpu_blocks[gpu_blocks_to_skip:]
127
+ kernel_blocks_to_skip = gpu_blocks_to_skip * kernel_blocks_per_gpu_block
128
+ gpu_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size[
129
+ kernel_blocks_to_skip:
130
+ ]
131
+ cpu_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size[
132
+ kernel_blocks_to_skip:
133
]
134
135
# set transfer direction
@@ -119,23 +139,23 @@ def test_transfer(
139
dst_spec_class = CPULoadStoreSpec
140
src_blocks = gpu_blocks
141
dst_blocks = cpu_blocks
122
- src_blocks_in_gpu_block_size = gpu_blocks
123
- dst_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size
124
- dst_size_in_gpu_blocks = num_cpu_blocks * gpu_blocks_per_cpu_block
142
+ src_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
143
+ dst_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
144
+ dst_size_in_kernel_blocks = num_cpu_blocks * kernel_blocks_per_cpu_block
145
else:
146
handler = handlers.cpu_to_gpu_handler
147
src_spec_class = CPULoadStoreSpec
148
dst_spec_class = GPULoadStoreSpec
149
src_blocks = cpu_blocks
150
dst_blocks = gpu_blocks
131
- src_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size
132
- dst_blocks_in_gpu_block_size = gpu_blocks
133
- dst_size_in_gpu_blocks = num_gpu_blocks
151
+ src_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
152
+ dst_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
153
+ dst_size_in_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
154
155
# build dst -> src mapping
156
dst_to_src = {}
157
for src_block, dst_block in zip(
138
- src_blocks_in_gpu_block_size, dst_blocks_in_gpu_block_size
158
+ src_blocks_in_kernel_block_size, dst_blocks_in_kernel_block_size
159
):
160
dst_to_src[dst_block] = src_block
161
@@ -165,29 +185,15 @@ def test_transfer(
185
assert torch.equal(orig_tensor, tensor)
186
187
# verify dst tensors
168
- for dst_block in range(dst_size_in_gpu_blocks):
188
+ for dst_block in range(dst_size_in_kernel_blocks):
189
src_block_candidate = dst_to_src.get(dst_block)
170
- for src_cache, dst_cache, orig_dst_cache, kv_dim in zip(
190
+ for src_cache, dst_cache, orig_dst_cache in zip(
191
handler.src_tensors,
192
handler.dst_tensors,
193
orig_dst_caches,
174
- handler.kv_dim_before_num_blocks,
194
):
176
- if kv_dim:
177
- # iterate over key, value
178
- for i in range(2):
179
- if src_block_candidate is not None:
180
- expected_value = src_cache[i][src_block_candidate]
181
- else:
182
- expected_value = orig_dst_cache[i][dst_block]
183
- torch.testing.assert_close(
184
- dst_cache[i][dst_block].cpu(), expected_value.cpu()
185
- )
195
+ if src_block_candidate is not None:
196
+ expected_value = src_cache[src_block_candidate]
197
else:
187
- if src_block_candidate is not None:
188
- expected_value = src_cache[src_block_candidate]
189
- else:
190
- expected_value = orig_dst_cache[dst_block]
191
- torch.testing.assert_close(
192
- dst_cache[dst_block].cpu(), expected_value.cpu()
193
- )
198
+ expected_value = orig_dst_cache[dst_block]
199
+ torch.testing.assert_close(dst_cache[dst_block].cpu(), expected_value.cpu())
vllm/_custom_ops.py
+25
-2
@@ -2455,9 +2455,32 @@ def concat_and_cache_mla_rope_fused(
2455
2456
2457
def swap_blocks(
2458
- src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor
2458
+ src: torch.Tensor,
2459
+ dst: torch.Tensor,
2460
+ block_size_in_bytes: int,
2461
+ block_mapping: torch.Tensor,
2462
) -> None:
2460
- torch.ops._C_cache_ops.swap_blocks(src, dst, block_mapping)
2463
+ """
2464
+ Copy specific blocks from one tensor to another.
2465
+
2466
+ This method assumes each of the two input tensors is composed of
2467
+ consecutive contiguous blocks, of size block_size_in_bytes.
2468
+ i.e. the memory layout for each tensor is:
2469
+ [block0] [block1] ... [block N]
2470
+
2471
+ block_mapping determines the subset of blocks to copy of the source tensor,
2472
+ and their matching destination block number on the destination tensor.
2473
+ block_mapping is expected to be a tensor of shape (num_blocks_to_copy, 2)
2474
+ where each block_mapping[i] represents a single copy operation, copying
2475
+ block #block_mapping[i][0] from the source tensor
2476
+ to block #block_mapping[i][1] on the destination tensor.
2477
+ block_mapping should have dtype int64.
2478
+
2479
+ The source and the destination tensors can be either on cpu or gpu,
2480
+ but not both on cpu.
2481
+ the block mapping tensor must on cpu.
2482
+ """
2483
+ torch.ops._C_cache_ops.swap_blocks(src, dst, block_size_in_bytes, block_mapping)
2484
2485
2486
def convert_fp8(
vllm/v1/kv_offload/worker/cpu_gpu.py
+47
-56
@@ -65,7 +65,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
65
self,
66
src_tensors: list[torch.Tensor],
67
dst_tensors: list[torch.Tensor],
68
- kv_dim_before_num_blocks: list[bool],
68
src_block_size_factor: int,
69
dst_block_size_factor: int,
70
):
@@ -76,22 +75,23 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
75
src_tensors: list of KV cache tensors to copy from.
76
dst_tensors: list of KV cache tensors to copy to.
77
Order should match src_tensors.
79
- kv_dim_before_num_blocks: list of bools, indicating
80
- whether the respective KV cache tensor has a KV
81
- dimension before its num_blocks dimension.
82
- e.g. (2, num_blocks, ...)
78
src_block_size_factor: The number of kernel blocks
79
per KV block in a source tensor.
80
dst_block_size_factor: The number of kernel blocks
81
per KV block in a destination tensor.
82
"""
88
- assert len(src_tensors) == len(dst_tensors) == len(kv_dim_before_num_blocks)
83
+ assert len(src_tensors) == len(dst_tensors)
84
85
self.src_tensors: list[torch.Tensor] = src_tensors
86
self.dst_tensors: list[torch.Tensor] = dst_tensors
92
- self.kv_dim_before_num_blocks: list[bool] = kv_dim_before_num_blocks
93
- self.src_block_size_factor: int = src_block_size_factor
94
- self.dst_block_size_factor: int = dst_block_size_factor
87
+ min_block_size_factor = min(src_block_size_factor, dst_block_size_factor)
88
+ self.src_block_size_factor: int = src_block_size_factor // min_block_size_factor
89
+ self.dst_block_size_factor: int = dst_block_size_factor // min_block_size_factor
90
+
91
+ self.block_size_in_bytes = [
92
+ tensor.element_size() * tensor.stride(0) * min_block_size_factor
93
+ for tensor in src_tensors
94
+ ]
95
96
assert len(src_tensors) > 0
97
self.gpu_to_cpu: bool = self.src_tensors[0].is_cuda
@@ -142,16 +142,17 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
142
# assure job will start only after the previous one completes
143
stream.wait_event(last_event)
144
with torch.cuda.stream(stream):
145
- for src_tensor, dst_tensor, kv_dim in zip(
146
- self.src_tensors, self.dst_tensors, self.kv_dim_before_num_blocks
145
+ for src_tensor, dst_tensor, block_size_in_bytes in zip(
146
+ self.src_tensors,
147
+ self.dst_tensors,
148
+ self.block_size_in_bytes,
149
):
148
- if kv_dim:
149
- src_key_cache, src_value_cache = src_tensor
150
- dst_key_cache, dst_value_cache = dst_tensor
151
- ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst_tensor)
152
- ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst_tensor)
153
- else:
154
- ops.swap_blocks(src_tensor, dst_tensor, src_to_dst_tensor)
150
+ ops.swap_blocks(
151
+ src_tensor,
152
+ dst_tensor,
153
+ block_size_in_bytes,
154
+ src_to_dst_tensor,
155
+ )
156
event.record(stream)
157
158
self._transfer_events[job_id] = event
@@ -188,19 +189,12 @@ class CpuGpuOffloadingHandlers:
189
):
190
assert gpu_caches
191
assert cpu_block_size % gpu_block_size == 0
191
- block_size_factor = cpu_block_size // gpu_block_size
192
-
193
- pin_memory = is_pin_memory_available()
192
195
- # allocate cpu tensors
196
- logger.info("Allocating %d CPU tensors...", len(gpu_caches))
197
- gpu_tensors: list[torch.Tensor] = []
198
- cpu_tensors: list[torch.Tensor] = []
199
- kv_dim_before_num_blocks: list[bool] = []
193
+ # find kernel block size and determine layout per each gpu tensor
194
kernel_block_size: int | None = None
195
+ # list of (gpu_tensor, split_k_and_v)
196
+ parsed_gpu_tensors: list[tuple[torch.Tensor, bool]] = []
197
for layer_name, gpu_tensor in gpu_caches.items():
202
- gpu_tensors.append(gpu_tensor)
203
-
198
gpu_shape = gpu_tensor.shape
199
attn_backend = attn_backends[layer_name]
200
test_shape = attn_backend.get_kv_cache_shape(
@@ -208,28 +202,20 @@ class CpuGpuOffloadingHandlers:
202
)
203
204
has_layers_dim = False
205
+ split_k_and_v = False
206
if len(gpu_shape) != len(test_shape):
207
# cross-layers tensor
208
# shape is (num_blocks, ...)
209
assert len(gpu_shape) == len(test_shape) + 1
215
- num_blocks_idx = 0
210
has_layers_dim = True
217
- kv_dim_before_num_blocks.append(False)
218
-
211
# prepend a dummy num_layers=80 to test_shape
212
test_shape = (80,) + test_shape
221
- elif test_shape[0] == 1234:
222
- # shape is (num_blocks, ...)
223
- num_blocks_idx = 0
224
- kv_dim_before_num_blocks.append(False)
225
- else:
213
+ elif test_shape[0] != 1234:
214
# shape should be (2, num_blocks, ...)
215
assert test_shape[0] == 2
216
assert test_shape[1] == 1234
217
assert gpu_shape[0] == 2
230
-
231
- num_blocks_idx = 1
232
- kv_dim_before_num_blocks.append(True)
218
+ split_k_and_v = True
219
220
try:
221
kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
@@ -250,30 +236,36 @@ class CpuGpuOffloadingHandlers:
236
kernel_block_size = gpu_shape[block_size_idx]
237
assert gpu_block_size % kernel_block_size == 0
238
253
- cpu_shape = list(gpu_shape)
254
- cpu_shape[num_blocks_idx] = num_cpu_blocks * block_size_factor
255
-
256
- logger.debug("Allocating CPU tensor of shape %r", cpu_shape)
257
- cpu_tensors.append(
258
- torch.zeros(
259
- cpu_shape,
260
- dtype=gpu_tensor.dtype,
261
- device="cpu",
262
- pin_memory=pin_memory,
263
- )
264
- )
239
+ parsed_gpu_tensors.append((gpu_tensor, split_k_and_v))
240
241
assert kernel_block_size is not None
267
- gpu_block_size_factor = gpu_block_size // kernel_block_size
242
cpu_block_size_factor = cpu_block_size // kernel_block_size
243
+ gpu_block_size_factor = gpu_block_size // kernel_block_size
244
+ num_cpu_kernel_blocks = num_cpu_blocks * cpu_block_size_factor
245
+
246
+ # allocate cpu tensors
247
+ pin_memory = is_pin_memory_available()
248
+ logger.info("Allocating %d CPU tensors...", len(parsed_gpu_tensors))
249
+ gpu_tensors: list[torch.Tensor] = []
250
+ cpu_tensors: list[torch.Tensor] = []
251
+ for gpu_tensor, split_k_and_v in parsed_gpu_tensors:
252
+ cpu_shape = list(gpu_tensor.shape)
253
+ cpu_shape[1 if split_k_and_v else 0] = num_cpu_kernel_blocks
254
+
255
+ logger.debug("Allocating CPU tensor of shape %r", cpu_shape)
256
+ cpu_tensor = torch.zeros(
257
+ cpu_shape,
258
+ dtype=gpu_tensor.dtype,
259
+ device="cpu",
260
+ pin_memory=pin_memory,
261
+ )
262
270
- # TODO (orozery): adapt swap_blocks to support gpu_block_size_factor
271
- assert gpu_block_size_factor == 1
263
+ gpu_tensors.extend(gpu_tensor.unbind(0) if split_k_and_v else [gpu_tensor])
264
+ cpu_tensors.extend(cpu_tensor.unbind(0) if split_k_and_v else [cpu_tensor])
265
266
self.gpu_to_cpu_handler = SingleDirectionOffloadingHandler(
267
src_tensors=gpu_tensors,
268
dst_tensors=cpu_tensors,
276
- kv_dim_before_num_blocks=kv_dim_before_num_blocks,
269
src_block_size_factor=gpu_block_size_factor,
270
dst_block_size_factor=cpu_block_size_factor,
271
)
@@ -281,7 +273,6 @@ class CpuGpuOffloadingHandlers:
273
self.cpu_to_gpu_handler = SingleDirectionOffloadingHandler(
274
src_tensors=cpu_tensors,
275
dst_tensors=gpu_tensors,
284
- kv_dim_before_num_blocks=kv_dim_before_num_blocks,
276
src_block_size_factor=cpu_block_size_factor,
277
dst_block_size_factor=gpu_block_size_factor,
278
)