@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / aad88f8486

Fork 0

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 4, 20:00 UTC

[kv_offload+HMA][8/N]: Support multi-group worker transfer (#38453)

Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Or Ozeri committed Apr 22, 2026 at 08:44 UTC aad88f84862b29917241a4b67774186956459350

4 files changed +345 -77

tests/v1/kv_offload/test_cpu_gpu.py

+233 -13

@@ -27,6 +27,7 @@ SEEDS = [0]

27 DEVICE_TYPE = current_platform.device_type

28 DEVICES = [f"{DEVICE_TYPE}:0"]

29 NUM_MAPPINGS = [3]

30 +NUM_MAPPINGS_PER_GROUP = [2]

33 @pytest.mark.parametrize("gpu_to_cpu", [True, False])

@@ -58,9 +59,7 @@ def test_transfer(

59 # build CanonicalKVCacheTensor list: one per tensor

60 kv_cache_tensors: list[CanonicalKVCacheTensor] = []

61 for i in range(num_tensors):

61 - gpu_tensor = torch.randint(

62 - -128,

63 - 127,

62 + gpu_tensor = torch.zeros(

63 (num_gpu_blocks, gpu_page_size_bytes),

64 dtype=torch.int8,

65 device=device,

@@ -119,26 +118,36 @@ def test_transfer(

118 for j in range(block_size_factor)

119 ]

120

122 - # maybe skip some GPU blocks to test reading from the middle of a CPU block

123 - if not gpu_to_cpu:

124 - blocks_to_skip = block_size_factor - 1

121 + # maybe skip some GPU blocks to test reading/writing from the middle of a CPU block

122 + blocks_to_skip = block_size_factor - 1

123 + if blocks_to_skip > 0:

124 gpu_blocks = gpu_blocks[blocks_to_skip:]

125 cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:]

126

127 # set transfer direction

128 if gpu_to_cpu:

129 handler = handlers.gpu_to_cpu_handler

131 - src_spec = GPULoadStoreSpec(gpu_blocks, group_sizes=(len(gpu_blocks),))

130 + src_spec = GPULoadStoreSpec(

131 + gpu_blocks, group_sizes=(len(gpu_blocks),), block_indices=(blocks_to_skip,)

132 + )

133 dst_spec = CPULoadStoreSpec(cpu_blocks)

134 dst_to_src = dict(zip(cpu_blocks_expanded, gpu_blocks))

134 - num_dst_sub_blocks = num_cpu_blocks * block_size_factor

135 + num_dst_sub_blocks = num_gpu_blocks

136 else:

137 handler = handlers.cpu_to_gpu_handler

138 src_spec = CPULoadStoreSpec(cpu_blocks)

138 - dst_spec = GPULoadStoreSpec(gpu_blocks, group_sizes=(len(gpu_blocks),))

139 + dst_spec = GPULoadStoreSpec(

140 + gpu_blocks, group_sizes=(len(gpu_blocks),), block_indices=(blocks_to_skip,)

141 + )

142 dst_to_src = dict(zip(gpu_blocks, cpu_blocks_expanded))

143 num_dst_sub_blocks = num_gpu_blocks

144

145 + # randomize src and dst tensors before transfer

146 + for tensor in handler.src_tensors:

147 + tensor.random_()

148 + for tensor in handler.dst_tensors:

149 + tensor.random_()

150 +

151 # clone src and dst tensors before transfer

152 orig_src_tensors = [x.clone() for x in handler.src_tensors]

153 orig_dst_tensors = [x.clone() for x in handler.dst_tensors]

@@ -146,7 +155,7 @@ def test_transfer(

155 # call transfer function

156 start_time = time.time()

157 assert handler.transfer_async(1, (src_spec, dst_spec))

149 - assert set({x.job_id for x in handler._transfers}) == {1}

158 + assert {x.job_id for x in handler._transfers} == {1}

159

160 # wait for transfer to complete

161 end_time = time.time() + 10

@@ -155,11 +164,14 @@ def test_transfer(

164 if finished:

165 assert finished[0].job_id == 1

166 assert finished[0].success

158 - assert finished[0].transfer_type == (

159 - ("GPU", "CPU") if gpu_to_cpu else ("CPU", "GPU")

167 + assert (

168 + finished[0].transfer_type == ("GPU", "CPU")

169 + if gpu_to_cpu

170 + else ("CPU", "GPU")

171 )

172 assert finished[0].transfer_size == (

162 - len(gpu_blocks) * handler.group_block_size_in_bytes[0]

173 + len(gpu_blocks)

174 + * sum([x.page_size_bytes for x in handler.kv_cache_groups_data_refs[0]])

175 )

176 assert finished[0].transfer_time > 0

177 assert finished[0].transfer_time < (time.time() - start_time)

@@ -196,3 +208,211 @@ def test_transfer(

208 handlers.gpu_to_cpu_handler.shutdown()

209 if mmap_region:

210 mmap_region.cleanup()

211 +

212 +

213 +@pytest.mark.parametrize("gpu_to_cpu", [True, False])

214 +@pytest.mark.parametrize("num_mappings_per_group", NUM_MAPPINGS_PER_GROUP)

215 +@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)

216 +@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS)

217 +@pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)

218 +@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)

219 +@pytest.mark.parametrize("seed", SEEDS)

220 +@pytest.mark.parametrize("device", DEVICES)

221 +@torch.inference_mode()

222 +def test_transfer_multi_group(

223 + default_vllm_config,

224 + gpu_to_cpu: bool,

225 + num_mappings_per_group: int,

226 + gpu_page_size_bytes: int,

227 + block_size_factor: int,

228 + num_gpu_blocks: int,

229 + num_cpu_blocks: int,

230 + seed: int,

231 + device: str,

232 +) -> None:

233 + """Test transfers with three KV cache groups:

234 + - Group 0: aligned transfer with num_mappings_per_group blocks

235 + - Group 1: zero blocks (empty group)

236 + - Group 2: unaligned CPU->GPU transfer (logical_offset=block_size_factor-1,

237 + causing the implementation to skip source sub-blocks) with

238 + num_mappings_per_group blocks

239 + """

240 + set_random_seed(seed)

241 +

242 + # 3 groups, each with 2 tensors

243 + num_groups = 3

244 + tensors_per_group = 2

245 + num_tensors = num_groups * tensors_per_group

246 + kv_cache_tensors: list[CanonicalKVCacheTensor] = []

247 + for _ in range(num_tensors):

248 + gpu_tensor = torch.zeros(

249 + (num_gpu_blocks, gpu_page_size_bytes),

250 + dtype=torch.int8,

251 + device=device,

252 + )

253 + kv_cache_tensors.append(

254 + CanonicalKVCacheTensor(

255 + tensor=gpu_tensor,

256 + page_size_bytes=gpu_page_size_bytes,

257 + )

258 + )

259 +

260 + kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]] = [

261 + [

262 + CanonicalKVCacheRef(

263 + tensor_idx=g * tensors_per_group + i,

264 + page_size_bytes=gpu_page_size_bytes,

265 + )

266 + for i in range(tensors_per_group)

267 + ]

268 + for g in range(num_groups)

269 + ]

270 +

271 + canonical_kv_caches = CanonicalKVCaches(

272 + tensors=kv_cache_tensors, group_data_refs=kv_cache_groups_data_refs

273 + )

274 +

275 + handlers = CpuGpuOffloadingHandlers(

276 + kv_caches=canonical_kv_caches,

277 + block_size_factor=block_size_factor,

278 + num_cpu_blocks=num_cpu_blocks,

279 + )

280 +

281 + # group 0: aligned, group 1: empty, group 2: unaligned on CPU->GPU

282 + group_sizes_in_cpu_blocks = [num_mappings_per_group, 0, num_mappings_per_group]

283 +

284 + total_cpu_blocks = sum(group_sizes_in_cpu_blocks)

285 + total_gpu_blocks_needed = total_cpu_blocks * block_size_factor

286 + gpu_blocks_all = random.sample(range(num_gpu_blocks), total_gpu_blocks_needed)

287 + cpu_blocks_all = random.sample(range(num_cpu_blocks), total_cpu_blocks)

288 +

289 + # split gpu/cpu blocks per group

290 + gpu_blocks_per_group: list[list[int]] = []

291 + cpu_blocks_per_group: list[list[int]] = []

292 + gpu_offset = 0

293 + cpu_offset = 0

294 + for size in group_sizes_in_cpu_blocks:

295 + gpu_count = size * block_size_factor

296 + gpu_blocks_per_group.append(gpu_blocks_all[gpu_offset : gpu_offset + gpu_count])

297 + cpu_blocks_per_group.append(cpu_blocks_all[cpu_offset : cpu_offset + size])

298 + gpu_offset += gpu_count

299 + cpu_offset += size

300 +

301 + # expand cpu blocks to gpu-page granularity

302 + cpu_blocks_expanded_per_group = [

303 + [

304 + cpu_block * block_size_factor + j

305 + for cpu_block in cpu_blocks

306 + for j in range(block_size_factor)

307 + ]

308 + for cpu_blocks in cpu_blocks_per_group

309 + ]

310 +

311 + # skip sub-blocks from group 2 to test unaligned transfers.

312 + sub_blocks_to_skip = block_size_factor - 1 # e.g. 2 when block_size_factor=3

313 + if sub_blocks_to_skip > 0:

314 + gpu_blocks_per_group[2] = gpu_blocks_per_group[2][

315 + sub_blocks_to_skip:-sub_blocks_to_skip

316 + ]

317 + cpu_blocks_expanded_per_group[2] = cpu_blocks_expanded_per_group[2][

318 + sub_blocks_to_skip:-sub_blocks_to_skip

319 + ]

320 +

321 + # build flat gpu_blocks list and group_sizes in GPU blocks

322 + gpu_blocks: list[int] = []

323 + group_sizes: list[int] = []

324 + for gpu_blks in gpu_blocks_per_group:

325 + gpu_blocks.extend(gpu_blks)

326 + group_sizes.append(len(gpu_blks))

327 +

328 + # build flat cpu_blocks list

329 + cpu_blocks = []

330 + for cpu_blks in cpu_blocks_per_group:

331 + cpu_blocks.extend(cpu_blks)

332 +

333 + # block_indices: only relevant for unaligned transfers

334 + block_indices: list[int] = [0, 0, sub_blocks_to_skip]

335 +

336 + if gpu_to_cpu:

337 + handler = handlers.gpu_to_cpu_handler

338 + src_spec = GPULoadStoreSpec(

339 + gpu_blocks, group_sizes=group_sizes, block_indices=block_indices

340 + )

341 + dst_spec = CPULoadStoreSpec(cpu_blocks)

342 + # per-group mapping: cpu sub-block -> gpu sub-block

343 + dst_to_src_per_group = [

344 + dict(zip(expanded, gpu_blks))

345 + for expanded, gpu_blks in zip(

346 + cpu_blocks_expanded_per_group, gpu_blocks_per_group

347 + )

348 + ]

349 + num_dst_sub_blocks = num_cpu_blocks * block_size_factor

350 + else:

351 + handler = handlers.cpu_to_gpu_handler

352 + src_spec = CPULoadStoreSpec(cpu_blocks)

353 + dst_spec = GPULoadStoreSpec(

354 + gpu_blocks, group_sizes=group_sizes, block_indices=block_indices

355 + )

356 + # per-group mapping: gpu sub-block -> cpu sub-block

357 + dst_to_src_per_group = [

358 + dict(zip(gpu_blks, expanded))

359 + for gpu_blks, expanded in zip(

360 + gpu_blocks_per_group, cpu_blocks_expanded_per_group

361 + )

362 + ]

363 + num_dst_sub_blocks = num_gpu_blocks

364 +

365 + # randomize src and dst tensors before transfer

366 + for tensor in handler.src_tensors:

367 + tensor.random_()

368 + for tensor in handler.dst_tensors:

369 + tensor.random_()

370 +

371 + orig_src_tensors = [x.clone() for x in handler.src_tensors]

372 + orig_dst_tensors = [x.clone() for x in handler.dst_tensors]

373 +

374 + assert handler.transfer_async(1, (src_spec, dst_spec))

375 + assert {x.job_id for x in handler._transfers} == {1}

376 +

377 + end_time = time.time() + 10

378 + while time.time() < end_time:

379 + finished = handler.get_finished()

380 + if finished:

381 + assert finished[0].job_id == 1

382 + assert finished[0].success

383 + expected_bytes = sum(

384 + group_size * sum([x.page_size_bytes for x in data_refs])

385 + for group_size, data_refs in zip(

386 + group_sizes, handler.kv_cache_groups_data_refs

387 + )

388 + )

389 + assert finished[0].transfer_size == expected_bytes

390 + break

391 + time.sleep(0.1)

392 +

393 + # verify src tensors did not change

394 + for orig_tensor, tensor in zip(orig_src_tensors, handler.src_tensors):

395 + assert torch.equal(orig_tensor, tensor)

396 +

397 + # verify dst tensors at gpu-page granularity

398 + for group_idx, dst_to_src in enumerate(dst_to_src_per_group):

399 + group_tensor_offset = group_idx * tensors_per_group

400 + for tensor_idx in range(tensors_per_group):

401 + src_tensor = handler.src_tensors[group_tensor_offset + tensor_idx]

402 + dst_tensor = handler.dst_tensors[group_tensor_offset + tensor_idx]

403 + orig_dst_tensor = orig_dst_tensors[group_tensor_offset + tensor_idx]

404 + src_view = src_tensor.view(-1, gpu_page_size_bytes)

405 + dst_view = dst_tensor.view(-1, gpu_page_size_bytes)

406 + orig_dst_view = orig_dst_tensor.view(-1, gpu_page_size_bytes)

407 + for dst_sub_block in range(num_dst_sub_blocks):

408 + src_sub_block = dst_to_src.get(dst_sub_block)

409 + if src_sub_block is not None:

410 + expected = src_view[src_sub_block]

411 + else:

412 + expected = orig_dst_view[dst_sub_block]

413 + torch.testing.assert_close(

414 + dst_view[dst_sub_block].cpu(), expected.cpu()

415 + )

416 +

417 + handlers.cpu_to_gpu_handler.shutdown()

418 + handlers.gpu_to_cpu_handler.shutdown()

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+3 -1

@@ -381,7 +381,9 @@ class OffloadingConnectorScheduler:

381 for i in range(self.config.block_size_factor):

382 src_block_ids.append(block_ids[gpu_block_idx + i])

383 src_spec = GPULoadStoreSpec(

384 - src_block_ids, group_sizes=(len(src_block_ids),)

384 + src_block_ids,

385 + group_sizes=(len(src_block_ids),),

386 + block_indices=(0,),

387 )

388

389 reqs_to_store[req_id] = (src_spec, dst_spec)

vllm/v1/kv_offload/mediums.py

+4 -6

@@ -34,26 +34,24 @@ class GPULoadStoreSpec(BlockIDsLoadStoreSpec):

34 will correspond to logically contiguous blocks, e.g. blocks 5-10 of a some request.

35 block_indices[i] will represent the block index of the first block in group #i.

36 Thus, len(block_indices) == len(group_sizes) = number of KV cache groups.

37 - This information is required in order to support loading from offloaded blocks

37 + This information is required in order to support off/loading from offloaded blocks

38 which are larger than GPU blocks.

39 In such cases, the first GPU block per each group may be unaligned to the offloaded

40 block size, and so knowing block_indices[i] allows the worker to correctly

41 skip part of the first matching offloaded block.

42 - Offloading from GPU is always aligned to offloaded block size, and so

43 - block_indices will only be set by the offloading connector when loading into GPU.

42 """

44 def __init__(

45 self,

46 block_ids: list[int],

47 group_sizes: Sequence[int],

50 - block_indices: Sequence[int] | None = None,

48 + block_indices: Sequence[int],

49 ):

50 super().__init__(block_ids)

51 assert sum(group_sizes) == len(block_ids)

54 - assert block_indices is None or len(block_indices) == len(group_sizes)

52 + assert len(block_indices) == len(group_sizes)

53 self.group_sizes: Sequence[int] = group_sizes

56 - self.block_indices: Sequence[int] | None = block_indices

54 + self.block_indices: Sequence[int] = block_indices

56 @staticmethod

57 def medium() -> str:

vllm/v1/kv_offload/worker/cpu_gpu.py

+105 -57

@@ -9,9 +9,10 @@ import torch

10 from vllm import _custom_ops as ops

11 from vllm.logger import init_logger

12 +from vllm.utils.math_utils import cdiv

13 from vllm.utils.platform_utils import is_pin_memory_available

14 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion

14 -from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec

15 +from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec, GPULoadStoreSpec

16 from vllm.v1.kv_offload.spec import CanonicalKVCacheRef, CanonicalKVCaches

17 from vllm.v1.kv_offload.worker.worker import (

18 OffloadingHandler,

@@ -135,9 +136,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

136 assert len(gpu_tensors) == len(cpu_tensors)

137 assert len(gpu_tensors) > 0

138

138 - # assert a single KV group until transfer_async supports multiple groups

139 - assert len(kv_cache_groups_data_refs) == 1

140 -

139 # assert input tensors are as expected

140 for gpu_tensor, cpu_tensor in zip(gpu_tensors, cpu_tensors):

141 assert gpu_tensor.dtype == torch.int8

@@ -157,29 +155,13 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

155 cpu_tensors if gpu_to_cpu else gpu_tensors

156 )

157 self.gpu_to_cpu: bool = gpu_to_cpu

158 + self.kv_cache_groups_data_refs = kv_cache_groups_data_refs

159

160 # GPU blocks may be smaller

161 # cpu_page_size = gpu_page_size * block_size_factor.

162 self.src_block_size_factor = 1 if self.gpu_to_cpu else block_size_factor

163 self.dst_block_size_factor = block_size_factor if self.gpu_to_cpu else 1

164

166 - # per-tensor block size in byte

167 - self.tensor_block_size_in_bytes = [

168 - gpu_tensor.shape[1] for gpu_tensor in gpu_tensors

169 - ]

170 -

171 - # per-group block size in bytes

172 - self.group_block_size_in_bytes = []

173 - for kv_cache_group_data_refs in kv_cache_groups_data_refs:

174 - group_block_size_in_bytes = 0

175 - for kv_cache_data_ref in kv_cache_group_data_refs:

176 - # TODO(orozery): use kv_cache_data_ref.page_size_bytes

177 - # once swap_blocks support it

178 - group_block_size_in_bytes += self.tensor_block_size_in_bytes[

179 - kv_cache_data_ref.tensor_idx

180 - ]

181 - self.group_block_size_in_bytes.append(group_block_size_in_bytes)

182 -

165 self.transfer_type = ("GPU", "CPU") if self.gpu_to_cpu else ("CPU", "GPU")

166 # job_id -> event

167 self._transfer_events: dict[int, torch.Event] = {}

@@ -190,11 +172,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

172 # list of CUDA events available for re-use

173 self._event_pool: list[torch.Event] = []

174

193 - # Pre-compute block sizes for batch copies.

194 - self._block_size_in_bytes_arr = np.array(

195 - self.tensor_block_size_in_bytes, dtype=np.int64

196 - )

197 -

175 def transfer_async(self, job_id: int, transfer_spec: TransferSpec) -> bool:

176 src_spec, dst_spec = transfer_spec

177 assert isinstance(src_spec, BlockIDsLoadStoreSpec)

@@ -205,37 +182,108 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

182 assert src_blocks.ndim == 1

183 assert dst_blocks.ndim == 1

184

208 - src_sub_block_count = src_blocks.size * self.src_block_size_factor

209 - dst_sub_block_count = dst_blocks.size * self.dst_block_size_factor

210 - src_sub_blocks_to_skip = -dst_blocks.size % self.src_block_size_factor

211 -

212 - assert dst_sub_block_count == src_sub_block_count - src_sub_blocks_to_skip

213 -

214 - num_pairs = dst_sub_block_count

215 - num_tensors = len(self.src_tensors)

216 - total = num_pairs * num_tensors

217 -

218 - all_src = np.empty(total, dtype=np.int64)

219 - all_dst = np.empty(total, dtype=np.int64)

220 - all_sizes = np.empty(total, dtype=np.int64)

221 -

222 - for t_idx, bsz in enumerate(self._block_size_in_bytes_arr):

223 - start = t_idx * num_pairs

224 - end = start + num_pairs

225 - compute_sub_block_ptrs(

226 - block_ids=src_blocks,

227 - block_size_factor=self.src_block_size_factor,

228 - output=all_src[start:end],

229 - tensor=self.src_tensors[t_idx],

230 - skip_count=src_sub_blocks_to_skip,

185 + num_src_blocks = len(src_blocks)

186 + num_dst_blocks = len(dst_blocks)

187 +

188 + # There are 2 types of transfers:

189 + # 1. GPU -> CPU

190 + # 2. CPU -> GPU

191 + #

192 + # transfers are also to CPU blocks, EXCEPT MAYBE for the first and last block.

193 + # i.e. the first and last CPU blocks in src_blocks can match against

194 + # a smaller (byte-wise) set of GPU blocks in dst_blocks.

195 + # In such cases, we may need to skip some gpu-sized sub-blocks,

196 + # and start reading/writing from the middle of the first CPU block.

197 + # If we have multiple KV cache groups (when using HMA with hybrid models),

198 + # we may have a partial first/last CPU block per each group.

199 + # The group_sizes parameter encodes the size of each group of blocks

200 + # in the GPU dst_blocks.

201 + # If group_sizes is None, we assume all blocks belong to a single group.

202 + # The logical_offset parameter maps each group of blocks to its logical

203 + # offset inside the request, counting in GPU blocks.

204 + # This allows us to find the correct starting position

205 + # in the matching first CPU block.

206 +

207 + # extract group_sizes from the GPU spec

208 + gpu_spec = src_spec if self.gpu_to_cpu else dst_spec

209 + assert isinstance(gpu_spec, GPULoadStoreSpec)

210 + group_sizes = gpu_spec.group_sizes

211 + assert len(group_sizes) == len(self.kv_cache_groups_data_refs)

212 +

213 + # extract block indices from the GPU spec

214 + block_indices = gpu_spec.block_indices

215 + assert len(block_indices) == len(self.kv_cache_groups_data_refs)

216 +

217 + num_copy_ops = 0

218 + for group_size, group_data_refs in zip(

219 + group_sizes, self.kv_cache_groups_data_refs

220 + ):

221 + num_copy_ops += group_size * len(group_data_refs)

222 +

223 + all_src = np.empty(num_copy_ops, dtype=np.int64)

224 + all_dst = np.empty(num_copy_ops, dtype=np.int64)

225 + all_sizes = np.empty(num_copy_ops, dtype=np.int64)

226 +

227 + src_offset = 0

228 + dst_offset = 0

229 + op_idx = 0

230 + # count total number of bytes copied

231 + num_transfer_bytes = 0

232 + for group_size, block_idx, group_data_refs in zip(

233 + group_sizes, block_indices, self.kv_cache_groups_data_refs

234 + ):

235 + if group_size == 0:

236 + continue

237 +

238 + src_logical_blocks_to_skip = block_idx % self.src_block_size_factor

239 + dst_logical_blocks_to_skip = block_idx % self.dst_block_size_factor

240 + src_logical_blocks_count = group_size + src_logical_blocks_to_skip

241 + dst_logical_blocks_count = group_size + dst_logical_blocks_to_skip

242 +

243 + dst_blocks_count = cdiv(

244 + dst_logical_blocks_count, self.dst_block_size_factor

245 )

232 - compute_sub_block_ptrs(

233 - block_ids=dst_blocks,

234 - block_size_factor=self.dst_block_size_factor,

235 - output=all_dst[start:end],

236 - tensor=self.dst_tensors[t_idx],

246 + dst_end_offset = dst_offset + dst_blocks_count

247 + assert dst_end_offset <= num_dst_blocks

248 +

249 + src_blocks_count = cdiv(

250 + src_logical_blocks_count, self.src_block_size_factor

251 )

238 - all_sizes[start:end] = bsz

252 + src_end_offset = src_offset + src_blocks_count

253 + assert src_end_offset <= num_src_blocks

254 +

255 + group_src = src_blocks[src_offset:src_end_offset]

256 + group_dst = dst_blocks[dst_offset:dst_end_offset]

257 +

258 + for data_ref in group_data_refs:

259 + t_idx = data_ref.tensor_idx

260 + end_idx = op_idx + group_size

261 +

262 + compute_sub_block_ptrs(

263 + group_src,

264 + self.src_block_size_factor,

265 + all_src[op_idx:end_idx],

266 + self.src_tensors[t_idx],

267 + skip_count=src_logical_blocks_to_skip,

268 + )

269 + compute_sub_block_ptrs(

270 + group_dst,

271 + self.dst_block_size_factor,

272 + all_dst[op_idx:end_idx],

273 + self.dst_tensors[t_idx],

274 + skip_count=dst_logical_blocks_to_skip,

275 + )

276 +

277 + all_sizes[op_idx:end_idx] = data_ref.page_size_bytes

278 + num_transfer_bytes += group_size * data_ref.page_size_bytes

279 + op_idx = end_idx

280 +

281 + src_offset = src_end_offset

282 + dst_offset = dst_end_offset

283 +

284 + assert src_offset == num_src_blocks

285 + assert dst_offset == num_dst_blocks

286 + assert op_idx == num_copy_ops

287

288 batch_src = torch.from_numpy(all_src)

289 batch_dst = torch.from_numpy(all_dst)

@@ -263,7 +311,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

311 stream.wait_event(last_event)

312 with torch.cuda.stream(stream):

313 start_event.record(stream)

266 - if total > 0:

314 + if num_copy_ops > 0:

315 ops.swap_blocks_batch(batch_src, batch_dst, batch_sizes)

316 end_event.record(stream)

317

@@ -274,7 +322,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

322 stream=stream,

323 start_event=start_event,

324 end_event=end_event,

277 - num_bytes=dst_sub_block_count * self.group_block_size_in_bytes[0],

325 + num_bytes=num_transfer_bytes,

326 )

327 )

328