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 8, 07:30 UTC
[Hybrid] Enable spec decoding in mamba cache align mode (#33705)
Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
Harry Huang committed
Feb 14, 2026 at 05:02 UTC
c027541eaf05c6ca1e9a544804afe28caef671fc
2 files changed
+10
-8
tests/v1/e2e/test_mamba_prefix_cache.py
+10
-4
@@ -11,8 +11,10 @@ import datasets
11
import pytest
12
import torch
13
14
+from tests.utils import create_new_process_for_each_test
15
from vllm import LLM, SamplingParams, TokensPrompt
16
from vllm.config import CacheConfig
17
+from vllm.distributed import cleanup_dist_env_and_memory
18
from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
19
from vllm.sequence import IntermediateTensors
20
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
@@ -103,6 +105,7 @@ def get_fake_propose_draft_token_ids_fn():
105
aux_hidden_states: list[torch.Tensor] | None,
106
spec_decode_metadata: SpecDecodeMetadata | None,
107
common_attn_metadata: CommonAttentionMetadata,
108
+ slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None,
109
) -> list[list[int]]:
110
num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
111
num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
@@ -401,6 +404,9 @@ def _run_ref_mamba_state_worker():
404
}
405
torch.save(cpu_state_ref, "mamba_kv_cache_dict_ref.pth")
406
mamba_kv_cache_dict.clear()
407
+ del engine
408
+ torch.cuda.empty_cache()
409
+ cleanup_dist_env_and_memory()
410
except Exception:
411
traceback.print_exc()
412
raise
@@ -473,10 +479,7 @@ def apply_patch(monkeypatch: pytest.MonkeyPatch):
479
monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn)
480
481
476
-@pytest.mark.skip(
477
- reason="Skipping test_mamba_prefix_cache because it is based on spec "
478
- "decode which is not allowed now."
479
-)
482
+@create_new_process_for_each_test()
483
def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
484
run_ref_mamba_state_in_subprocess()
485
apply_patch(monkeypatch)
@@ -762,3 +765,6 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
765
mamba_state_ref = torch.load("mamba_kv_cache_dict_ref.pth")
766
check_mamba_state_equal(mamba_state_ref, mamba_kv_cache_dict, keys_to_check)
767
mamba_kv_cache_dict.clear()
768
+ del engine
769
+ torch.cuda.empty_cache()
770
+ cleanup_dist_env_and_memory()
vllm/model_executor/models/config.py
-4
@@ -354,10 +354,6 @@ class MambaModelConfig(VerifyAndUpdateConfig):
354
assert vllm_config.scheduler_config.enable_chunked_prefill, (
355
"Chunked prefill is required for mamba cache mode 'align'."
356
)
357
- assert not vllm_config.speculative_config, (
358
- "Mamba cache mode 'align' is currently not compatible "
359
- "with speculative decoding."
360
- )
357
logger.info(
358
"Warning: Prefix caching in Mamba cache '%s' "
359
"mode is currently enabled. "