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, 21:30 UTC
[Bugfix] Fix prefix caching for Mamba 'all' mode (Nemotron models) (#34874)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
Feb 24, 2026 at 00:31 UTC
a2ba6a52443f9823d55a7f779c48c5de98abbc80
2 files changed
+166
tests/v1/attention/test_mamba_update_block_table.py
new
+145
@@ -0,0 +1,145 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Regression test for https://github.com/vllm-project/vllm/issues/34865
5
+
6
+When multiple KV cache groups share the same MambaSpec (as in Nemotron
7
+hybrid models), the metadata caching optimization reuses metadata from
8
+an earlier group via update_block_table(). In 'all' mode with CUDA graphs,
9
+update_block_table() must copy block_idx_last_scheduled_token and
10
+block_idx_last_computed_token to the *current* builder's persistent
11
+buffers, otherwise CUDA graph replay reads stale values from uninitialized
12
+buffers.
13
+"""
14
+
15
+from types import SimpleNamespace
16
+
17
+import torch
18
+
19
+from vllm.config.compilation import CUDAGraphMode
20
+from vllm.v1.attention.backends.mamba_attn import (
21
+ BaseMambaAttentionMetadata,
22
+ BaseMambaAttentionMetadataBuilder,
23
+)
24
+from vllm.v1.kv_cache_interface import MambaSpec
25
+
26
+
27
+class _ConcreteMambaBuilder(
28
+ BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetadata]
29
+):
30
+ """Minimal concrete subclass for testing (base class is ABC)."""
31
+
32
+ metadata_cls = BaseMambaAttentionMetadata
33
+
34
+
35
+def _make_vllm_config(block_size, max_model_len, max_num_seqs):
36
+ """Create a minimal mock VllmConfig with only the fields the builder
37
+ accesses, avoiding any model download / HF config inspection."""
38
+ return SimpleNamespace(
39
+ cache_config=SimpleNamespace(mamba_cache_mode="all"),
40
+ compilation_config=SimpleNamespace(
41
+ cudagraph_mode=CUDAGraphMode.FULL,
42
+ max_cudagraph_capture_size=None,
43
+ ),
44
+ scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
45
+ model_config=SimpleNamespace(max_model_len=max_model_len),
46
+ )
47
+
48
+
49
+def test_update_block_table_copies_block_idx_to_persistent_buffers():
50
+ """update_block_table() must write block_idx tensors to the current
51
+ builder's persistent buffers, not leave them pointing to a different
52
+ builder's buffers."""
53
+
54
+ block_size = 16
55
+ max_model_len = 256
56
+ num_reqs = 4
57
+ device = torch.device("cpu")
58
+
59
+ vllm_config = _make_vllm_config(block_size, max_model_len, num_reqs)
60
+
61
+ spec = MambaSpec(
62
+ block_size=block_size,
63
+ shapes=((1,), (1,)),
64
+ dtypes=(torch.float32,),
65
+ mamba_cache_mode="all",
66
+ )
67
+
68
+ # Two builders simulating two KV cache groups with the same MambaSpec.
69
+ builder_a = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
70
+ builder_b = _ConcreteMambaBuilder(spec, ["layer1"], vllm_config, device)
71
+
72
+ # Sanity: each builder has its own persistent buffer.
73
+ assert (
74
+ builder_a.block_idx_last_scheduled_token.data_ptr()
75
+ != builder_b.block_idx_last_scheduled_token.data_ptr()
76
+ )
77
+
78
+ # Construct decode-only metadata as if builder_a.build() produced it.
79
+ max_blocks = max_model_len // block_size
80
+ seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
81
+ block_idx_vals = (seq_lens - 1) // block_size # [3, 3, 3, 3]
82
+
83
+ builder_a.block_idx_last_scheduled_token[:num_reqs].copy_(block_idx_vals)
84
+ builder_a.block_idx_last_computed_token[:num_reqs].copy_(block_idx_vals)
85
+
86
+ metadata_a = BaseMambaAttentionMetadata(
87
+ num_prefills=0,
88
+ num_prefill_tokens=0,
89
+ num_decodes=num_reqs,
90
+ num_decode_tokens=num_reqs,
91
+ num_reqs=num_reqs,
92
+ has_initial_states_p=None,
93
+ query_start_loc_p=None,
94
+ num_computed_tokens_p=None,
95
+ state_indices_tensor=builder_a.state_indices_tensor[:num_reqs],
96
+ block_idx_last_scheduled_token=(
97
+ builder_a.block_idx_last_scheduled_token[:num_reqs]
98
+ ),
99
+ block_idx_first_scheduled_token_p=None,
100
+ block_idx_last_computed_token=(
101
+ builder_a.block_idx_last_computed_token[:num_reqs]
102
+ ),
103
+ seq_lens=seq_lens,
104
+ )
105
+
106
+ # Call update_block_table on builder_b (simulates the metadata caching
107
+ # optimization reusing metadata from builder_a's group).
108
+ blk_table = torch.randint(
109
+ 0, 100, (num_reqs, max_blocks), dtype=torch.int32, device=device
110
+ )
111
+ slot_mapping = torch.zeros(num_reqs, dtype=torch.int64, device=device)
112
+
113
+ metadata_b = builder_b.update_block_table(metadata_a, blk_table, slot_mapping)
114
+
115
+ # block_idx tensors must live in builder_b's persistent buffers.
116
+ def shares_storage(tensor, buffer):
117
+ return (
118
+ tensor.untyped_storage().data_ptr() == buffer.untyped_storage().data_ptr()
119
+ )
120
+
121
+ assert shares_storage(
122
+ metadata_b.block_idx_last_scheduled_token,
123
+ builder_b.block_idx_last_scheduled_token,
124
+ ), "block_idx_last_scheduled_token not in builder_b's persistent buffer"
125
+
126
+ assert shares_storage(
127
+ metadata_b.block_idx_last_computed_token,
128
+ builder_b.block_idx_last_computed_token,
129
+ ), "block_idx_last_computed_token not in builder_b's persistent buffer"
130
+
131
+ # Must NOT point to builder_a's buffers.
132
+ assert not shares_storage(
133
+ metadata_b.block_idx_last_scheduled_token,
134
+ builder_a.block_idx_last_scheduled_token,
135
+ ), "block_idx_last_scheduled_token still points to builder_a's buffer"
136
+
137
+ # Values must be correct (copied from metadata_a).
138
+ torch.testing.assert_close(
139
+ metadata_b.block_idx_last_scheduled_token,
140
+ block_idx_vals,
141
+ )
142
+ torch.testing.assert_close(
143
+ metadata_b.block_idx_last_computed_token,
144
+ block_idx_vals,
145
+ )
vllm/v1/attention/backends/mamba_attn.py
+21
@@ -331,5 +331,26 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
331
persistent_state_indices_t.copy_(state_indices_t, non_blocking=True)
332
state_indices_t = persistent_state_indices_t
333
334
+ # For 'all' mode, also update prefix caching block indices
335
+ # to use this builder's persistent buffers (required for CUDA
336
+ # graph replay to read from the correct memory addresses).
337
+ if self.vllm_config.cache_config.mamba_cache_mode == "all":
338
+ assert metadata.block_idx_last_scheduled_token is not None
339
+ assert metadata.block_idx_last_computed_token is not None
340
+ self.block_idx_last_scheduled_token[:num_reqs].copy_(
341
+ metadata.block_idx_last_scheduled_token[:num_reqs],
342
+ non_blocking=True,
343
+ )
344
+ new_metadata.block_idx_last_scheduled_token = (
345
+ self.block_idx_last_scheduled_token[: metadata.num_decode_tokens]
346
+ )
347
+ self.block_idx_last_computed_token[:num_reqs].copy_(
348
+ metadata.block_idx_last_computed_token[:num_reqs],
349
+ non_blocking=True,
350
+ )
351
+ new_metadata.block_idx_last_computed_token = (
352
+ self.block_idx_last_computed_token[: metadata.num_decode_tokens]
353
+ )
354
+
355
new_metadata.state_indices_tensor = state_indices_t
356
return new_metadata