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 6, 05:45 UTC
[Hybrid] calling get_mamba_groups() once at MambaCopyBuffers.create() (#37318)
Signed-off-by: Francesco Fusco <ffu@zurich.ibm.com>
Francesco Fusco committed
Mar 21, 2026 at 10:29 UTC
298e5108482e52fed40de315011c30e08342c979
2 files changed
+11
-5
tests/v1/worker/test_mamba_utils.py
+2
-1
@@ -36,6 +36,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():
36
spec = MagicMock(block_size=64, num_speculative_blocks=0)
37
cache_config = MagicMock(enable_prefix_caching=True)
38
input_batch = MagicMock(req_ids=[])
39
+ copy_bufs = MagicMock(mamba_group_ids=[0], mamba_spec=spec)
40
41
mamba_state_idx = {
42
"finished": 1,
@@ -62,7 +63,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():
63
{},
64
{},
65
(),
65
- MagicMock(),
66
+ copy_bufs,
67
)
68
69
assert mamba_state_idx == {"keep": 99}
vllm/v1/worker/mamba_utils.py
+9
-4
@@ -67,6 +67,8 @@ class MambaCopyBuffers:
67
src_ptrs: CpuGpuBuffer
68
dst_ptrs: CpuGpuBuffer
69
sizes: CpuGpuBuffer
70
+ mamba_group_ids: list[int]
71
+ mamba_spec: MambaSpec
72
offset: int = 0
73
74
@classmethod
@@ -77,7 +79,7 @@ class MambaCopyBuffers:
79
copy_funcs: tuple[MambaStateCopyFunc, ...],
80
make_buffer: Callable[..., CpuGpuBuffer],
81
) -> "MambaCopyBuffers":
80
- mamba_group_ids, _ = get_mamba_groups(kv_cache_config)
82
+ mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
83
entries_per_req = sum(
84
len(kv_cache_config.kv_cache_groups[gid].layer_names)
85
for gid in mamba_group_ids
@@ -87,6 +89,8 @@ class MambaCopyBuffers:
89
src_ptrs=make_buffer(n, dtype=torch.int64),
90
dst_ptrs=make_buffer(n, dtype=torch.int64),
91
sizes=make_buffer(n, dtype=torch.int32),
92
+ mamba_group_ids=mamba_group_ids,
93
+ mamba_spec=mamba_spec,
94
)
95
96
@@ -155,7 +159,8 @@ def preprocess_mamba(
159
Copy the mamba state of previous step to the last
160
(1 + num_speculative_blocks) block.
161
"""
158
- mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
162
+ mamba_group_ids = copy_bufs.mamba_group_ids
163
+ mamba_spec = copy_bufs.mamba_spec
164
num_speculative_blocks = mamba_spec.num_speculative_blocks
165
# TODO(Chen): we need to optimize this function a lot
166
assert cache_config.enable_prefix_caching
@@ -231,8 +236,8 @@ def postprocess_mamba(
236
num_scheduled_tokens_dict = scheduler_output.num_scheduled_tokens
237
scheduled_spec_decode_tokens_dict = scheduler_output.scheduled_spec_decode_tokens
238
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)
239
+ mamba_group_ids = copy_bufs.mamba_group_ids
240
+ mamba_spec = copy_bufs.mamba_spec
241
copy_bufs.offset = 0
242
for i, req_id in enumerate(input_batch.req_ids):
243
req_state = requests[req_id]