[MRV2] Support attention-free models (#52374)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Aug 14, 2026 at 16:54 UTC
d6f17f3d530efbd83a730d3d73e9f483726dfde1
3 files changed
+7
-2
tests/v1/e2e/general/test_mamba_prefix_cache.py
+4
@@ -873,6 +873,10 @@ def get_mamba_prefix_cache_step_configs(
873
def _run_mamba_prefix_cache_mrv1(
874
monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
875
):
876
+ # This test patches the V1 model runner, so pin V1 explicitly: MoE/hybrid
877
+ # models like Qwen3-Next now default to the V2 runner.
878
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
879
+ envs.disable_envs_cache()
880
global async_scheduling_mode
881
async_scheduling_mode = async_scheduling
882
run_ref_mamba_state_in_subprocess()
vllm/v1/worker/gpu/model_runner.py
+2
-1
@@ -372,7 +372,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
372
373
self.model_memory_usage = m.consumed_memory
374
logger.info(
375
- "Model loading took %s GiB and %.6f seconds",
375
+ "Model loading took %s GiB memory and %.6f seconds",
376
format_gib(m.consumed_memory),
377
time_after_load - time_before_load,
378
)
@@ -1889,6 +1889,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1889
"""Release GPU tensors (model weights, KV caches, workspace) so that
1890
memory is reclaimable when running in the same process."""
1891
torch.accelerator.synchronize()
1892
+ self.cudagraph_manager = None
1893
if hasattr(self, "kv_caches"):
1894
self.kv_caches.clear()
1895
if hasattr(self, "attn_groups"):
vllm/v1/worker/gpu/model_states/__init__.py
+1
-1
@@ -33,7 +33,7 @@ def init_model_state(
33
34
return EncoderOnlyModelState(vllm_config, model, encoder_cache, device)
35
36
- if vllm_config.model_config.is_hybrid:
36
+ if vllm_config.model_config.is_hybrid or vllm_config.model_config.is_attention_free:
37
from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
38
39
return MambaHybridModelState(vllm_config, model, encoder_cache, device)