[Mamba] enable prefix cache by default (#50991)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jiangyun Zhu committed
Aug 4, 2026 at 23:20 UTC
f9c74b4b9c25c202cb84e0e9908b82a503a8c7c4
4 files changed
+30
-36
tests/models/language/generation/test_hybrid.py
+24
-17
@@ -42,11 +42,6 @@ HYBRID_MODELS = [
42
"tiny-random/qwen3-next-moe",
43
]
44
45
-HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL = {
46
- "LiquidAI/LFM2-1.2B",
47
- "tiny-random/qwen3-next-moe",
48
-}
49
-
45
FULL_CUDA_GRAPH_MODELS = [
46
"ai21labs/Jamba-tiny-dev",
47
"Zyphra/Zamba2-1.2B-instruct",
@@ -95,15 +90,11 @@ def test_models(
90
example_prompts, max_tokens, num_logprobs
91
)
92
98
- extra_kwargs = {}
99
- if model in HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL:
100
- extra_kwargs["enable_chunked_prefill"] = True
101
-
93
with vllm_runner(
94
model,
95
max_num_seqs=MAX_NUM_SEQS,
96
attention_backend=ATTN_BACKEND,
106
- **extra_kwargs,
97
+ enable_chunked_prefill=True,
98
) as vllm_model:
99
vllm_outputs = vllm_model.generate_greedy_logprobs(
100
example_prompts, max_tokens, num_logprobs
@@ -140,7 +131,9 @@ def test_batching(
131
_set_conv_state_layout(monkeypatch, conv_state_layout)
132
133
for_loop_outputs = []
143
- with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
134
+ with vllm_runner(
135
+ model, max_num_seqs=MAX_NUM_SEQS, enable_chunked_prefill=True
136
+ ) as vllm_model:
137
for prompt in example_prompts:
138
(single_output,) = vllm_model.generate_greedy_logprobs(
139
[prompt], max_tokens, num_logprobs
@@ -224,7 +217,7 @@ def test_mamba_cache_cg_padding(
217
example_prompts.append(example_prompts[0])
218
219
try:
227
- with vllm_runner(model) as vllm_model:
220
+ with vllm_runner(model, enable_chunked_prefill=True) as vllm_model:
221
vllm_model.generate_greedy(example_prompts, max_tokens)
222
except RuntimeError:
223
pytest.fail(
@@ -250,7 +243,9 @@ def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
243
a single step.
244
"""
245
try:
253
- with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
246
+ with vllm_runner(
247
+ model, max_num_seqs=MAX_NUM_SEQS, enable_chunked_prefill=True
248
+ ) as vllm_model:
249
vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
250
except ValueError:
251
pytest.fail(
@@ -272,7 +267,9 @@ def test_state_cleanup(
267
If it's not cleaned, an error would be expected.
268
"""
269
try:
275
- with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
270
+ with vllm_runner(
271
+ model, max_num_seqs=MAX_NUM_SEQS, enable_chunked_prefill=True
272
+ ) as vllm_model:
273
for _ in range(10):
274
vllm_model.generate_greedy([example_prompts[0]] * 100, 1)
275
except ValueError:
@@ -294,14 +291,20 @@ def test_distributed_correctness(
291
num_logprobs: int,
292
) -> None:
293
with vllm_runner(
297
- model, tensor_parallel_size=1, max_num_seqs=MAX_NUM_SEQS
294
+ model,
295
+ tensor_parallel_size=1,
296
+ max_num_seqs=MAX_NUM_SEQS,
297
+ enable_chunked_prefill=True,
298
) as vllm_model:
299
vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
300
example_prompts, max_tokens, num_logprobs
301
)
302
303
with vllm_runner(
304
- model, tensor_parallel_size=2, max_num_seqs=MAX_NUM_SEQS
304
+ model,
305
+ tensor_parallel_size=2,
306
+ max_num_seqs=MAX_NUM_SEQS,
307
+ enable_chunked_prefill=True,
308
) as vllm_model:
309
vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
310
example_prompts, max_tokens, num_logprobs
@@ -340,7 +343,10 @@ def test_full_cuda_graph(
343
)
344
345
with vllm_runner(
343
- model, max_num_seqs=MAX_NUM_SEQS, attention_backend=ATTN_BACKEND
346
+ model,
347
+ max_num_seqs=MAX_NUM_SEQS,
348
+ attention_backend=ATTN_BACKEND,
349
+ enable_chunked_prefill=True,
350
) as vllm_model:
351
vllm_outputs = vllm_model.generate_greedy_logprobs(
352
example_prompts, max_tokens, num_logprobs
@@ -388,6 +394,7 @@ def test_fp32_cache_state(
394
model,
395
max_num_seqs=MAX_NUM_SEQS,
396
gpu_memory_utilization=0.9,
397
+ enable_chunked_prefill=True,
398
**{cache_dtype_param: "float32"},
399
) as vllm_model:
400
vllm_outputs = vllm_model.generate_greedy_logprobs(
vllm/config/cache.py
+3
-3
@@ -135,10 +135,10 @@ class CacheConfig:
135
"""The cache strategy for Mamba layers:
136
137
- "none": set when prefix caching is disabled.
138
- - "all": cache the mamba state of all tokens at position i * block_size. This is
139
- the default behavior (for models that support it) when prefix caching is enabled.
138
+ - "all": cache the mamba state of all tokens at position i * block_size.
139
- "align": only cache the mamba state of the last token of each scheduler step and
141
- when the token is at position i * block_size.
140
+ when the token is at position i * block_size. This is the default when prefix
141
+ caching is enabled.
142
"""
143
replayssm_buffer_len: int = Field(default=16, gt=0)
144
"""ReplaySSM history buffer length B: with use_replayssm, standard decode
vllm/engine/arg_utils.py
+1
-5
@@ -2587,11 +2587,7 @@ class EngineArgs:
2587
self, model_config: ModelConfig
2588
) -> None:
2589
default_chunked_prefill = model_config.is_chunked_prefill_supported
2590
- # Hybrid models support prefix caching but keep it opt-in for now
2591
- # while the feature matures.
2592
- default_prefix_caching = (
2593
- model_config.is_prefix_caching_supported and not model_config.is_hybrid
2594
- )
2590
+ default_prefix_caching = model_config.is_prefix_caching_supported
2591
2592
if self.enable_chunked_prefill is None:
2593
self.enable_chunked_prefill = default_chunked_prefill
vllm/model_executor/models/config.py
+2
-11
@@ -595,10 +595,8 @@ class MambaModelConfig(VerifyAndUpdateConfig):
595
596
if cache_config.enable_prefix_caching:
597
if cache_config.mamba_cache_mode == "none":
598
- cache_config.mamba_cache_mode = (
599
- "all" if model_config.supports_mamba_prefix_caching else "align"
600
- )
601
- logger.warning(
598
+ cache_config.mamba_cache_mode = "align"
599
+ logger.info(
600
"Mamba cache mode is set to '%s' for %s by default "
601
"when prefix caching is enabled",
602
cache_config.mamba_cache_mode,
@@ -618,13 +616,6 @@ class MambaModelConfig(VerifyAndUpdateConfig):
616
assert vllm_config.scheduler_config.enable_chunked_prefill, (
617
"Chunked prefill is required for mamba cache mode 'align'."
618
)
621
- logger.info(
622
- "Warning: Prefix caching in Mamba cache '%s' "
623
- "mode is currently enabled. "
624
- "Its support for Mamba layers is experimental. "
625
- "Please report any issues you may observe.",
626
- cache_config.mamba_cache_mode,
627
- )
619
# By default, mamba block size will be set to max_model_len (see
620
# below). When enabling prefix caching, we align mamba block size
621
# to the block size as the basic granularity for prefix caching.