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

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

[perf] Use pinned memory for async H2D transfer in do_mamba_copy_block (#35480)

Signed-off-by: Huamin Li <3ericli@gmail.com>

Huamin Li committed Feb 27, 2026 at 09:50 UTC 157722da756daa6f967433903680745abc0c4861

4 files changed +85 -44

tests/v1/e2e/test_mamba_prefix_cache.py

+10 -10

@@ -325,6 +325,7 @@ def get_fake_process_mamba_fn(

325 requests: dict[str, CachedRequestState],

326 forward_context: dict[str, Any],

327 mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],

328 + copy_bufs: mamba_utils.MambaCopyBuffers,

329 ):

330 nonlocal copy_info

331 copy_info = None

@@ -337,6 +338,7 @@ def get_fake_process_mamba_fn(

338 requests,

339 forward_context,

340 mamba_state_copy_funcs,

341 + copy_bufs,

342 )

343 if cur_step_action is not None:

344 check_copy_info(

@@ -355,6 +357,7 @@ def get_fake_process_mamba_fn(

357 mamba_state_idx: dict[str, int],

358 forward_context: dict[str, Any],

359 mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],

360 + copy_bufs: mamba_utils.MambaCopyBuffers,

361 ):

362 nonlocal copy_info

363 copy_info = None

@@ -366,6 +369,7 @@ def get_fake_process_mamba_fn(

369 mamba_state_idx,

370 forward_context,

371 mamba_state_copy_funcs,

372 + copy_bufs,

373 )

374 if cur_step_action is not None:

375 check_copy_info(

@@ -376,19 +380,15 @@ def get_fake_process_mamba_fn(

380 )

381 return ret

382

379 - def fake_copy_fn(

380 - src_state_list: list[int],

381 - dest_state_list: list[int],

382 - num_elements_list: list[int],

383 - ):

383 + def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers):

384 nonlocal copy_info

385 assert copy_info is None

386 + n = copy_bufs.offset

387 + src_state_list = copy_bufs.src_ptrs.cpu[:n].tolist()

388 + dest_state_list = copy_bufs.dst_ptrs.cpu[:n].tolist()

389 + num_elements_list = copy_bufs.sizes.cpu[:n].tolist()

390 copy_info = (src_state_list, dest_state_list, num_elements_list)

387 - return original_copy_fn(

388 - src_state_list,

389 - dest_state_list,

390 - num_elements_list,

391 - )

391 + return original_copy_fn(copy_bufs)

392

393 return fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn

394

tests/v1/worker/test_mamba_utils.py

@@ -62,6 +62,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():

62 {},

63 {},

64 (),

65 + MagicMock(),

66 )

68 assert mamba_state_idx == {"keep": 99}

vllm/v1/worker/gpu_model_runner.py

+14

@@ -755,6 +755,7 @@ class GPUModelRunner(

755 self.execute_model_state: ExecuteModelState | None = None

756 self.kv_connector_output: KVConnectorOutput | None = None

757 self.mamba_state_idx: dict[str, int] = {}

758 + self._mamba_copy_bufs: mamba_utils.MambaCopyBuffers | None = None

759 self.layerwise_nvtx_hooks_registered = False

760

761 def update_max_model_len(self, max_model_len: int) -> None:

@@ -849,6 +850,16 @@ class GPUModelRunner(

850 with_numpy=numpy,

851 )

852

853 + def _get_mamba_copy_bufs(self) -> mamba_utils.MambaCopyBuffers:

854 + if self._mamba_copy_bufs is None:

855 + self._mamba_copy_bufs = mamba_utils.MambaCopyBuffers.create(

856 + self.max_num_reqs,

857 + self.kv_cache_config,

858 + self.model.get_mamba_state_copy_func(),

859 + self._make_buffer,

860 + )

861 + return self._mamba_copy_bufs

862 +

863 def _init_model_kwargs(self):

864 model_kwargs = dict[str, Any]()

865

@@ -1211,6 +1222,7 @@ class GPUModelRunner(

1222 self.mamba_state_idx,

1223 self.compilation_config.static_forward_context,

1224 self.model.get_mamba_state_copy_func(),

1225 + self._get_mamba_copy_bufs(),

1226 )

1227

1228 def _update_streaming_request(

@@ -3505,6 +3517,7 @@ class GPUModelRunner(

3517 self.requests,

3518 self.compilation_config.static_forward_context,

3519 self.model.get_mamba_state_copy_func(),

3520 + self._get_mamba_copy_bufs(),

3521 )

3522

3523 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0

@@ -5997,6 +6010,7 @@ class GPUModelRunner(

6010 """

6011 kv_cache_config = deepcopy(kv_cache_config)

6012 self.kv_cache_config = kv_cache_config

6013 + self._mamba_copy_bufs = None

6014 self.may_add_encoder_only_layers_to_kv_cache_config()

6015 self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)

6016 self.initialize_attn_backend(kv_cache_config)

vllm/v1/worker/mamba_utils.py

+60 -34

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +import dataclasses

4 import itertools

5 +from collections.abc import Callable

6 from typing import Any

8 import torch

@@ -13,6 +15,7 @@ from vllm.triton_utils import tl, triton

15 from vllm.utils.math_utils import cdiv

16 from vllm.v1.core.sched.output import SchedulerOutput

17 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec

18 +from vllm.v1.utils import CpuGpuBuffer

19 from vllm.v1.worker.gpu_input_batch import CachedRequestState

20 from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch

@@ -59,10 +62,36 @@ def get_mamba_groups(kv_cache_config: KVCacheConfig) -> tuple[list[int], MambaSp

62 return mamba_group_ids, mamba_specs[0]

65 +@dataclasses.dataclass

66 +class MambaCopyBuffers:

67 + src_ptrs: CpuGpuBuffer

68 + dst_ptrs: CpuGpuBuffer

69 + sizes: CpuGpuBuffer

70 + offset: int = 0

71 +

72 + @classmethod

73 + def create(

74 + cls,

75 + max_num_reqs: int,

76 + kv_cache_config: KVCacheConfig,

77 + copy_funcs: tuple[MambaStateCopyFunc, ...],

78 + make_buffer: Callable[..., CpuGpuBuffer],

79 + ) -> "MambaCopyBuffers":

80 + mamba_group_ids, _ = get_mamba_groups(kv_cache_config)

81 + entries_per_req = sum(

82 + len(kv_cache_config.kv_cache_groups[gid].layer_names)

83 + for gid in mamba_group_ids

84 + ) * len(copy_funcs)

85 + n = max_num_reqs * entries_per_req

86 + return cls(

87 + src_ptrs=make_buffer(n, dtype=torch.int64),

88 + dst_ptrs=make_buffer(n, dtype=torch.int64),

89 + sizes=make_buffer(n, dtype=torch.int32),

90 + )

91 +

92 +

93 def collect_mamba_copy_meta(

63 - src_state_list: list[int],

64 - dest_state_list: list[int],

65 - num_elements_list: list[int],

94 + copy_bufs: MambaCopyBuffers,

95 kv_cache_config: KVCacheConfig,

96 mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],

97 mamba_group_ids: list[int],

@@ -71,10 +100,15 @@ def collect_mamba_copy_meta(

100 accept_token_bias: int,

101 req_state: CachedRequestState,

102 forward_context: dict[str, Any],

74 -):

103 +) -> None:

104 if src_block_idx == dest_block_idx and accept_token_bias == 0:

105 return

106

107 + src_ptrs_np = copy_bufs.src_ptrs.np

108 + dst_ptrs_np = copy_bufs.dst_ptrs.np

109 + sizes_np = copy_bufs.sizes.np

110 + offset = copy_bufs.offset

111 +

112 for mamba_group_id in mamba_group_ids:

113 block_ids = req_state.block_ids[mamba_group_id]

114 dest_block_id = block_ids[dest_block_idx]

@@ -87,25 +121,23 @@ def collect_mamba_copy_meta(

121 state, block_ids, src_block_idx, accept_token_bias + 1

122 )

123

90 - src_state_list.append(copy_spec.start_addr)

91 - dest_state_list.append(state[dest_block_id].data_ptr())

92 - num_elements_list.append(copy_spec.num_elements * state.element_size())

124 + src_ptrs_np[offset] = copy_spec.start_addr

125 + dst_ptrs_np[offset] = state[dest_block_id].data_ptr()

126 + sizes_np[offset] = copy_spec.num_elements * state.element_size()

127 + offset += 1

128

129 + copy_bufs.offset = offset

130

95 -def do_mamba_copy_block(

96 - src_state_list: list[int],

97 - dest_state_list: list[int],

98 - num_elements_list: list[int],

99 -):

100 - if len(src_state_list) == 0:

101 - return

102 - assert len(src_state_list) == len(dest_state_list)

103 - assert len(src_state_list) == len(num_elements_list)

104 - src_state_ptrs = torch.tensor(src_state_list, device="cuda", dtype=torch.int64)

105 - dst_state_ptrs = torch.tensor(dest_state_list, device="cuda", dtype=torch.int64)

106 - num_elements = torch.tensor(num_elements_list, device="cuda", dtype=torch.int32)

131

108 - batch_memcpy(src_state_ptrs, dst_state_ptrs, num_elements)

132 +def do_mamba_copy_block(copy_bufs: MambaCopyBuffers):

133 + n = copy_bufs.offset

134 + if n == 0:

135 + return

136 + batch_memcpy(

137 + copy_bufs.src_ptrs.copy_to_gpu(n),

138 + copy_bufs.dst_ptrs.copy_to_gpu(n),

139 + copy_bufs.sizes.copy_to_gpu(n),

140 + )

141

142

143 def preprocess_mamba(

@@ -117,6 +149,7 @@ def preprocess_mamba(

149 requests: dict[str, CachedRequestState],

150 forward_context: dict[str, Any],

151 mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],

152 + copy_bufs: MambaCopyBuffers,

153 ):

154 """

155 Copy the mamba state of previous step to the last

@@ -138,9 +171,7 @@ def preprocess_mamba(

171 for req_id in itertools.chain(finished_req_ids, preempted_req_ids, resumed_req_ids):

172 mamba_state_idx.pop(req_id, None)

173

141 - src_state_list: list[int] = []

142 - dest_state_list: list[int] = []

143 - num_elements_list: list[int] = []

174 + copy_bufs.offset = 0

175 for i, req_id in enumerate(input_batch.req_ids):

176 req_state = requests[req_id]

177 prev_state_idx = mamba_state_idx.get(req_id)

@@ -169,9 +200,7 @@ def preprocess_mamba(

200 mamba_state_idx[req_id] = curr_state_idx

201 if prev_state_idx != -1 and prev_state_idx != curr_state_idx:

202 collect_mamba_copy_meta(

172 - src_state_list,

173 - dest_state_list,

174 - num_elements_list,

203 + copy_bufs,

204 kv_cache_config,

205 mamba_state_copy_funcs,

206 mamba_group_ids,

@@ -182,7 +211,7 @@ def preprocess_mamba(

211 forward_context,

212 )

213 input_batch.num_accepted_tokens_cpu[i] = 1

185 - do_mamba_copy_block(src_state_list, dest_state_list, num_elements_list)

214 + do_mamba_copy_block(copy_bufs)

215

216

217 def postprocess_mamba(

@@ -193,6 +222,7 @@ def postprocess_mamba(

222 mamba_state_idx: dict[str, int],

223 forward_context: dict[str, Any],

224 mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],

225 + copy_bufs: MambaCopyBuffers,

226 ):

227 """

228 If a blocks is converted from partial block to full block in this step, copy the

@@ -203,9 +233,7 @@ def postprocess_mamba(

233 num_accepted_tokens_cpu = input_batch.num_accepted_tokens_cpu

234 # NOTE: can be optimized as this function always returns the same result

235 mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)

206 - src_state_list: list[int] = []

207 - dest_state_list: list[int] = []

208 - num_elements_list: list[int] = []

236 + copy_bufs.offset = 0

237 for i, req_id in enumerate(input_batch.req_ids):

238 req_state = requests[req_id]

239 num_computed_tokens = req_state.num_computed_tokens

@@ -225,9 +253,7 @@ def postprocess_mamba(

253 src_block_idx = mamba_state_idx[req_id]

254 dest_block_idx = aligned_new_computed_tokens // mamba_spec.block_size - 1

255 collect_mamba_copy_meta(

228 - src_state_list,

229 - dest_state_list,

230 - num_elements_list,

256 + copy_bufs,

257 kv_cache_config,

258 mamba_state_copy_funcs,

259 mamba_group_ids,

@@ -239,4 +265,4 @@ def postprocess_mamba(

265 )

266 if src_block_idx == dest_block_idx:

267 num_accepted_tokens_cpu[i] = 1

242 - do_mamba_copy_block(src_state_list, dest_state_list, num_elements_list)

268 + do_mamba_copy_block(copy_bufs)