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

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 11, 12:30 UTC

CPU KV Offloading: Use more CUDA streams (#29013)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

Or Ozeri committed Dec 15, 2025 at 01:50 UTC 174e39ead7cb14a0511b0bbdc15dfd4a01ffb5fb

3 files changed +192 -105

tests/v1/kv_offload/test_cpu_gpu.py

+10 -12

@@ -9,7 +9,7 @@ import torch

9 from vllm.platforms import current_platform

10 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

11 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec

12 -from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandler

12 +from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers

14 BACKENDS_TO_TEST = [FlashAttentionBackend]

@@ -82,7 +82,7 @@ def test_transfer(

83 # create handler

84 cpu_block_size = gpu_blocks_per_cpu_block * gpu_block_size

85 - handler = CpuGpuOffloadingHandler(

85 + handlers = CpuGpuOffloadingHandlers(

86 attn_backends=attn_backends,

87 gpu_block_size=gpu_block_size,

88 cpu_block_size=cpu_block_size,

@@ -112,8 +112,7 @@ def test_transfer(

112

113 # set transfer direction

114 if gpu_to_cpu:

115 - src_kv_caches = handler.gpu_tensors

116 - dst_kv_caches = handler.cpu_tensors

115 + handler = handlers.gpu_to_cpu_handler

116 src_spec_class = GPULoadStoreSpec

117 dst_spec_class = CPULoadStoreSpec

118 src_blocks = gpu_blocks

@@ -122,8 +121,7 @@ def test_transfer(

121 dst_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size

122 dst_size_in_gpu_blocks = num_cpu_blocks * gpu_blocks_per_cpu_block

123 else:

125 - src_kv_caches = handler.cpu_tensors

126 - dst_kv_caches = handler.gpu_tensors

124 + handler = handlers.cpu_to_gpu_handler

125 src_spec_class = CPULoadStoreSpec

126 dst_spec_class = GPULoadStoreSpec

127 src_blocks = cpu_blocks

@@ -144,12 +142,12 @@ def test_transfer(

142 dst_spec = dst_spec_class(dst_blocks)

143

144 # clone src and dst tensors before transfer

147 - orig_src_caches = [x.clone() for x in src_kv_caches]

148 - orig_dst_caches = [x.clone() for x in dst_kv_caches]

145 + orig_src_caches = [x.clone() for x in handler.src_tensors]

146 + orig_dst_caches = [x.clone() for x in handler.dst_tensors]

147

148 # call transfer function

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

152 - assert set(handler.transfer_events.keys()) == {1}

150 + assert set({x[0] for x in handler._transfers}) == {1}

151

152 # wait for transfer to complete

153 end_time = time.time() + 10

@@ -161,15 +159,15 @@ def test_transfer(

159 time.sleep(0.1)

160

161 # verify src tensors did not change

164 - for orig_tensor, tensor in zip(orig_src_caches, src_kv_caches):

162 + for orig_tensor, tensor in zip(orig_src_caches, handler.src_tensors):

163 assert torch.equal(orig_tensor, tensor)

164

165 # verify dst tensors

166 for dst_block in range(dst_size_in_gpu_blocks):

167 src_block_candidate = dst_to_src.get(dst_block)

168 for src_cache, dst_cache, orig_dst_cache, kv_dim in zip(

171 - src_kv_caches,

172 - dst_kv_caches,

169 + handler.src_tensors,

170 + handler.dst_tensors,

171 orig_dst_caches,

172 handler.kv_dim_before_num_blocks,

173 ):

vllm/v1/kv_offload/cpu.py

+7 -7

@@ -13,7 +13,7 @@ from vllm.v1.kv_offload.backends.cpu import CPUBackend

13 from vllm.v1.kv_offload.lru_manager import LRUOffloadingManager

14 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec

15 from vllm.v1.kv_offload.spec import OffloadingSpec

16 -from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandler

16 +from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers

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

@@ -32,7 +32,7 @@ class CPUOffloadingSpec(OffloadingSpec):

32 self._manager: OffloadingManager | None = None

34 # worker-side

35 - self._handler: OffloadingHandler | None = None

35 + self._handlers: CpuGpuOffloadingHandlers | None = None

37 self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru")

@@ -67,13 +67,13 @@ class CPUOffloadingSpec(OffloadingSpec):

67 kv_caches: dict[str, torch.Tensor],

68 attn_backends: dict[str, type[AttentionBackend]],

69 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:

70 - if not self._handler:

70 + if not self._handlers:

71 if not current_platform.is_cuda_alike():

72 raise Exception(

73 "CPU Offloading is currently only supported on CUDA-alike GPUs"

74 )

76 - self._handler = CpuGpuOffloadingHandler(

76 + self._handlers = CpuGpuOffloadingHandlers(

77 attn_backends=attn_backends,

78 gpu_block_size=self.gpu_block_size,

79 cpu_block_size=self.offloaded_block_size,

@@ -81,6 +81,6 @@ class CPUOffloadingSpec(OffloadingSpec):

81 gpu_caches=kv_caches,

82 )

84 - assert self._handler is not None

85 - yield GPULoadStoreSpec, CPULoadStoreSpec, self._handler

86 - yield CPULoadStoreSpec, GPULoadStoreSpec, self._handler

84 + assert self._handlers is not None

85 + yield GPULoadStoreSpec, CPULoadStoreSpec, self._handlers.gpu_to_cpu_handler

86 + yield CPULoadStoreSpec, GPULoadStoreSpec, self._handlers.cpu_to_gpu_handler

vllm/v1/kv_offload/worker/cpu_gpu.py

+175 -86

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections import deque

5 import numpy as np

6 import torch

@@ -8,7 +9,7 @@ from vllm import _custom_ops as ops

9 from vllm.attention.backends.abstract import AttentionBackend

10 from vllm.logger import init_logger

11 from vllm.utils.platform_utils import is_pin_memory_available

11 -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec

12 +from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec

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

14 OffloadingHandler,

15 TransferResult,

@@ -51,7 +52,123 @@ def expand_block_ids(

52 output_idx = output_end_idx

54 -class CpuGpuOffloadingHandler(OffloadingHandler):

55 +class SingleDirectionOffloadingHandler(OffloadingHandler):

56 + """

57 + SingleDirectionOffloadingHandler handles transfers for a single direction,

58 + either CPU->GPU or GPU->CPU.

59 + Transfers are guaranteed to be executed in order of their submission.

60 + Each transfer uses a unique CUDA stream, and its stream will start

61 + executing only after the streams of previous transfers have finished.

62 + """

63 +

64 + def __init__(

65 + self,

66 + src_tensors: list[torch.Tensor],

67 + dst_tensors: list[torch.Tensor],

68 + kv_dim_before_num_blocks: list[bool],

69 + src_block_size_factor: int,

70 + dst_block_size_factor: int,

71 + priority: int,

72 + ):

73 + """

74 + Initialize a SingleDirectionOffloadingHandler.

75 +

76 + Args:

77 + src_tensors: list of KV cache tensors to copy from.

78 + dst_tensors: list of KV cache tensors to copy to.

79 + Order should match src_tensors.

80 + kv_dim_before_num_blocks: list of bools, indicating

81 + whether the respective KV cache tensor has a KV

82 + dimension before its num_blocks dimension.

83 + e.g. (2, num_blocks, ...)

84 + src_block_size_factor: The number of kernel blocks

85 + per KV block in a source tensor.

86 + dst_block_size_factor: The number of kernel blocks

87 + per KV block in a destination tensor.

88 + priority: The priority of the backing CUDA streams.

89 + Lower numbers indicate higher priority.

90 + """

91 + assert len(src_tensors) == len(dst_tensors) == len(kv_dim_before_num_blocks)

92 +

93 + self.src_tensors: list[torch.Tensor] = src_tensors

94 + self.dst_tensors: list[torch.Tensor] = dst_tensors

95 + self.kv_dim_before_num_blocks: list[bool] = kv_dim_before_num_blocks

96 + self.src_block_size_factor: int = src_block_size_factor

97 + self.dst_block_size_factor: int = dst_block_size_factor

98 + self.priority = priority

99 +

100 + # queue of transfers (job_id, stream, event)

101 + self._transfers: deque[tuple[int, torch.cuda.Stream, torch.Event]] = deque()

102 + # list of CUDA streams available for re-use

103 + self._stream_pool: list[torch.cuda.Stream] = []

104 + # list of CUDA events available for re-use

105 + self._event_pool: list[torch.Event] = []

106 +

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

108 + src_spec, dst_spec = transfer_spec

109 + assert isinstance(src_spec, BlockIDsLoadStoreSpec)

110 + assert isinstance(dst_spec, BlockIDsLoadStoreSpec)

111 +

112 + src_blocks = src_spec.block_ids

113 + dst_blocks = dst_spec.block_ids

114 + assert src_blocks.ndim == 1

115 + assert dst_blocks.ndim == 1

116 +

117 + src_sub_block_count = src_blocks.size * self.src_block_size_factor

118 + dst_sub_block_count = dst_blocks.size * self.dst_block_size_factor

119 + src_sub_blocks_to_skip = -dst_blocks.size % self.src_block_size_factor

120 +

121 + assert dst_sub_block_count == src_sub_block_count - src_sub_blocks_to_skip

122 +

123 + src_to_dst = np.empty((dst_sub_block_count, 2), dtype=np.int64)

124 + expand_block_ids(

125 + src_blocks,

126 + self.src_block_size_factor,

127 + src_to_dst[:, 0],

128 + skip_count=src_sub_blocks_to_skip,

129 + )

130 + expand_block_ids(dst_blocks, self.dst_block_size_factor, src_to_dst[:, 1])

131 + src_to_dst_tensor = torch.from_numpy(src_to_dst)

132 +

133 + stream = (

134 + self._stream_pool.pop()

135 + if self._stream_pool

136 + else torch.cuda.Stream(priority=self.priority)

137 + )

138 + event = self._event_pool.pop() if self._event_pool else torch.Event()

139 + if self._transfers:

140 + _, _, last_event = self._transfers[-1]

141 + # assure job will start only after the previous one completes

142 + stream.wait_event(last_event)

143 + with torch.cuda.stream(stream):

144 + for src_tensor, dst_tensor, kv_dim in zip(

145 + self.src_tensors, self.dst_tensors, self.kv_dim_before_num_blocks

146 + ):

147 + if kv_dim:

148 + src_key_cache, src_value_cache = src_tensor

149 + dst_key_cache, dst_value_cache = dst_tensor

150 + ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst_tensor)

151 + ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst_tensor)

152 + else:

153 + ops.swap_blocks(src_tensor, dst_tensor, src_to_dst_tensor)

154 + event.record(stream)

155 +

156 + self._transfers.append((job_id, stream, event))

157 +

158 + # success

159 + return True

160 +

161 + def get_finished(self) -> list[TransferResult]:

162 + results: list[TransferResult] = []

163 + while self._transfers and self._transfers[0][2].query():

164 + job_id, stream, event = self._transfers.popleft()

165 + results.append((job_id, True))

166 + self._stream_pool.append(stream)

167 + self._event_pool.append(event)

168 + return results

169 +

170 +

171 +class CpuGpuOffloadingHandlers:

172 def __init__(

173 self,

174 gpu_block_size: int,

@@ -60,27 +177,20 @@ class CpuGpuOffloadingHandler(OffloadingHandler):

177 gpu_caches: dict[str, torch.Tensor],

178 attn_backends: dict[str, type[AttentionBackend]],

179 ):

180 + assert gpu_caches

181 assert cpu_block_size % gpu_block_size == 0

64 - self.block_size_factor = cpu_block_size // gpu_block_size

65 -

66 - # cuda streams for gpu->cpu and cpu->gpu

67 - self.d2h_stream = torch.cuda.Stream()

68 - self.h2d_stream = torch.cuda.Stream()

69 -

70 - # job_id -> transfer cuda event

71 - self.transfer_events: dict[int, torch.Event] = {}

72 - # list of cuda events available for re-use

73 - self.events_pool: list[torch.Event] = []

182 + block_size_factor = cpu_block_size // gpu_block_size

183

184 pin_memory = is_pin_memory_available()

185

186 # allocate cpu tensors

187 logger.info("Allocating %d CPU tensors...", len(gpu_caches))

79 - self.gpu_tensors: list[torch.Tensor] = []

80 - self.cpu_tensors: list[torch.Tensor] = []

81 - self.kv_dim_before_num_blocks: list[bool] = []

188 + gpu_tensors: list[torch.Tensor] = []

189 + cpu_tensors: list[torch.Tensor] = []

190 + kv_dim_before_num_blocks: list[bool] = []

191 + kernel_block_size: int | None = None

192 for layer_name, gpu_tensor in gpu_caches.items():

83 - self.gpu_tensors.append(gpu_tensor)

193 + gpu_tensors.append(gpu_tensor)

194

195 gpu_shape = gpu_tensor.shape

196 attn_backend = attn_backends[layer_name]

@@ -88,16 +198,21 @@ class CpuGpuOffloadingHandler(OffloadingHandler):

198 num_blocks=1234, block_size=16, num_kv_heads=8, head_size=256

199 )

200

201 + has_layers_dim = False

202 if len(gpu_shape) != len(test_shape):

203 # cross-layers tensor

204 # shape is (num_blocks, ...)

205 assert len(gpu_shape) == len(test_shape) + 1

206 num_blocks_idx = 0

96 - self.kv_dim_before_num_blocks.append(False)

207 + has_layers_dim = True

208 + kv_dim_before_num_blocks.append(False)

209 +

210 + # prepend a dummy num_layers=80 to test_shape

211 + test_shape = (80,) + test_shape

212 elif test_shape[0] == 1234:

213 # shape is (num_blocks, ...)

214 num_blocks_idx = 0

100 - self.kv_dim_before_num_blocks.append(False)

215 + kv_dim_before_num_blocks.append(False)

216 else:

217 # shape should be (2, num_blocks, ...)

218 assert test_shape[0] == 2

@@ -105,13 +220,32 @@ class CpuGpuOffloadingHandler(OffloadingHandler):

220 assert gpu_shape[0] == 2

221

222 num_blocks_idx = 1

108 - self.kv_dim_before_num_blocks.append(True)

223 + kv_dim_before_num_blocks.append(True)

224 +

225 + try:

226 + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(

227 + include_num_layers_dimension=has_layers_dim

228 + )

229 + assert len(kv_cache_stride_order) == len(gpu_shape)

230 + except (AttributeError, NotImplementedError):

231 + kv_cache_stride_order = tuple(range(len(gpu_shape)))

232 +

233 + # permute test_shape according to stride_order

234 + test_shape = tuple(test_shape[i] for i in kv_cache_stride_order)

235 +

236 + # find block_size (16) dimension index

237 + block_size_idx = test_shape.index(16)

238 + if kernel_block_size is not None:

239 + assert kernel_block_size == gpu_shape[block_size_idx]

240 + else:

241 + kernel_block_size = gpu_shape[block_size_idx]

242 + assert gpu_block_size % kernel_block_size == 0

243

244 cpu_shape = list(gpu_shape)

111 - cpu_shape[num_blocks_idx] = num_cpu_blocks * self.block_size_factor

245 + cpu_shape[num_blocks_idx] = num_cpu_blocks * block_size_factor

246

247 logger.debug("Allocating CPU tensor of shape %r", cpu_shape)

114 - self.cpu_tensors.append(

248 + cpu_tensors.append(

249 torch.zeros(

250 cpu_shape,

251 dtype=gpu_tensor.dtype,

@@ -120,72 +254,27 @@ class CpuGpuOffloadingHandler(OffloadingHandler):

254 )

255 )

256

123 - def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:

124 - src_spec, dst_spec = spec

125 - if isinstance(src_spec, CPULoadStoreSpec):

126 - assert isinstance(dst_spec, GPULoadStoreSpec)

127 - stream = self.h2d_stream

128 - src_tensors = self.cpu_tensors

129 - dst_tensors = self.gpu_tensors

130 - src_block_size_factor = self.block_size_factor

131 - dst_block_size_factor = 1

132 - else:

133 - assert isinstance(src_spec, GPULoadStoreSpec)

134 - assert isinstance(dst_spec, CPULoadStoreSpec)

135 - stream = self.d2h_stream

136 - src_tensors = self.gpu_tensors

137 - dst_tensors = self.cpu_tensors

138 - src_block_size_factor = 1

139 - dst_block_size_factor = self.block_size_factor

140 -

141 - src_blocks = src_spec.block_ids

142 - dst_blocks = dst_spec.block_ids

143 - assert src_blocks.ndim == 1

144 - assert dst_blocks.ndim == 1

257 + assert kernel_block_size is not None

258 + gpu_block_size_factor = gpu_block_size // kernel_block_size

259 + cpu_block_size_factor = cpu_block_size // kernel_block_size

260

146 - src_sub_block_count = src_blocks.size * src_block_size_factor

147 - dst_sub_block_count = dst_blocks.size * dst_block_size_factor

148 - src_sub_blocks_to_skip = -dst_blocks.size % src_block_size_factor

261 + # TODO (orozery): adapt swap_blocks to support gpu_block_size_factor

262 + assert gpu_block_size_factor == 1

263

150 - assert dst_sub_block_count == src_sub_block_count - src_sub_blocks_to_skip

151 -

152 - src_to_dst = np.empty((dst_sub_block_count, 2), dtype=np.int64)

153 - expand_block_ids(

154 - src_blocks,

155 - src_block_size_factor,

156 - src_to_dst[:, 0],

157 - skip_count=src_sub_blocks_to_skip,

264 + self.gpu_to_cpu_handler = SingleDirectionOffloadingHandler(

265 + src_tensors=gpu_tensors,

266 + dst_tensors=cpu_tensors,

267 + kv_dim_before_num_blocks=kv_dim_before_num_blocks,

268 + src_block_size_factor=gpu_block_size_factor,

269 + dst_block_size_factor=cpu_block_size_factor,

270 + priority=1,

271 )

159 - expand_block_ids(dst_blocks, dst_block_size_factor, src_to_dst[:, 1])

160 - src_to_dst_tensor = torch.from_numpy(src_to_dst)

272

162 - event = self.events_pool.pop() if self.events_pool else torch.Event()

163 - with torch.cuda.stream(stream):

164 - for src_tensor, dst_tensor, kv_dim in zip(

165 - src_tensors, dst_tensors, self.kv_dim_before_num_blocks

166 - ):

167 - if kv_dim:

168 - src_key_cache = src_tensor[0]

169 - dst_key_cache = dst_tensor[0]

170 - ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst_tensor)

171 - src_value_cache = src_tensor[1]

172 - dst_value_cache = dst_tensor[1]

173 - ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst_tensor)

174 - else:

175 - ops.swap_blocks(src_tensor, dst_tensor, src_to_dst_tensor)

176 - event.record(stream)

177 -

178 - self.transfer_events[job_id] = event

179 -

180 - # success

181 - return True

182 -

183 - def get_finished(self) -> list[TransferResult]:

184 - results: list[TransferResult] = []

185 - for job_id, event in self.transfer_events.items():

186 - if event.query():

187 - results.append((job_id, True))

188 - self.events_pool.append(event)

189 - for job_id, _ in results:

190 - del self.transfer_events[job_id]

191 - return results

273 + self.cpu_to_gpu_handler = SingleDirectionOffloadingHandler(

274 + src_tensors=cpu_tensors,

275 + dst_tensors=gpu_tensors,

276 + kv_dim_before_num_blocks=kv_dim_before_num_blocks,

277 + src_block_size_factor=cpu_block_size_factor,

278 + dst_block_size_factor=gpu_block_size_factor,

279 + priority=-1,

280 + )