[Bugfix][MammothModa2] Stop allocating AR KV cache for the replaced language model (#8095)

Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_ar_kv_layers.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The AR stage must register attention layers for its own language model only.
+
+The Qwen-VL parent ``__init__`` builds a stock language model before MammothModa2
+replaces it. Its attention layers stay in ``static_forward_context`` unless they are
+removed, and the KV cache is then sized and allocated for both models. The parent
+``__init__`` and the replacement are stubbed here so only the registry is exercised.
+"""
+
+import gc
+import weakref
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from vllm.model_executor.models.utils import StageMissingLayer
+
+import vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2 as mm2
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_LAYERS = 3
+VISION_LAYER = "ar.visual.blocks.0.attn"
+AR_CLASSES = [
+    (mm2.MammothModa2ARForConditionalGeneration, mm2.Qwen2_5_VLForConditionalGeneration),
+    (mm2.MammothModa2Qwen3ARForConditionalGeneration, mm2.Qwen3VLForConditionalGeneration),
+]
+AR_IDS = ["preview_qwen2_5_vl", "dev_qwen3_vl"]
+
+
+class _FakeAttention(nn.Module):
+    def __init__(self, registry: dict, name: str):
+        super().__init__()
+        registry[name] = self
+
+
+class _FakeLM(nn.Module):
+    def __init__(self, registry: dict, prefix: str):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            _FakeAttention(registry, f"{prefix}.layers.{i}.self_attn.attn") for i in range(NUM_LAYERS)
+        )
+        self.make_empty_intermediate_tensors = lambda *args, **kwargs: None
+
+
+class _FakeCompiledLM(_FakeLM):
+    cleanups = 0
+
+    def cleanup(self):
+        type(self).cleanups += 1
+
+
+def _build(
+    monkeypatch,
+    ar_class,
+    parent,
+    *,
+    encoder_only=False,
+    parent_lm=_FakeLM,
+    parent_registers=True,
+    pp_rank=(1, True),
+    stray_parent_layer=False,
+):
+    registry: dict = {}
+    built = {}
+
+    def parent_init(self, *, vllm_config, prefix=""):
+        nn.Module.__init__(self)
+        self.visual = _FakeAttention(registry, VISION_LAYER)
+        stock = nn.Module()
+        stock.model = parent_lm(registry if parent_registers else {}, f"{prefix}.language_model.model")
+        built["parent_lm"] = weakref.ref(stock)
+        if stray_parent_layer:
+            # A parent-prefixed layer that is not part of the module being dropped.
+            registry[f"{prefix}.language_model.model.layers.9.self_attn.attn"] = nn.Module()
+        self.language_model = StageMissingLayer("language_model", stock) if encoder_only else stock
+
+    # The model module only imports get_current_vllm_config with this fix; keep the
+    # reverted source testable so the registry assertion is what fails there.
+    monkeypatch.setattr(
+        mm2,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(compilation_config=SimpleNamespace(static_forward_context=registry)),
+        raising=False,
+    )
+    monkeypatch.setattr(parent, "__init__", parent_init)
+    monkeypatch.setattr(mm2, "init_vllm_registered_model", lambda *, prefix, **kwargs: _FakeLM(registry, prefix))
+    monkeypatch.setattr(mm2, "MammothModa2Qwen3ForCausalLM", lambda *, vllm_config, prefix: _FakeLM(registry, prefix))
+    monkeypatch.setattr(mm2, "TorchCompileWithNoGuardsWrapper", _FakeCompiledLM, raising=False)
+    monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
+    world_size, is_first_rank = pp_rank
+    monkeypatch.setattr(
+        mm2, "get_pp_group", lambda: SimpleNamespace(world_size=world_size, is_first_rank=is_first_rank)
+    )
+
+    hf_config = SimpleNamespace(llm_config=SimpleNamespace(text_config=SimpleNamespace()))
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=hf_config, architectures=["MammothModa2Model"])
+    )
+    vllm_config.with_hf_config = lambda hf, architectures=None: vllm_config
+    model = ar_class(vllm_config=vllm_config, prefix="ar")
+    return model, registry, built
+
+
+@pytest.mark.parametrize("encoder_only", [False, True], ids=["default", "mm_encoder_only"])
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_ar_stage_registers_only_the_replacement_language_model(monkeypatch, ar_class, parent, encoder_only):
+    model, registry, built = _build(monkeypatch, ar_class, parent, encoder_only=encoder_only)
+
+    assert set(registry) == {f"ar.language_model.layers.{i}.self_attn.attn" for i in range(NUM_LAYERS)} | {VISION_LAYER}
+    gc.collect()
+    assert built["parent_lm"]() is None
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_compiled_parent_language_model_releases_its_hook(monkeypatch, ar_class, parent):
+    _FakeCompiledLM.cleanups = 0
+    _build(monkeypatch, ar_class, parent, parent_lm=_FakeCompiledLM)
+    assert _FakeCompiledLM.cleanups == 1
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_unexpected_parent_registry_fails_loudly(monkeypatch, ar_class, parent):
+    with pytest.raises(RuntimeError, match="Unexpected attention registry"):
+        _build(monkeypatch, ar_class, parent, parent_registers=False)
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_later_pipeline_rank_without_parent_layers_starts(monkeypatch, ar_class, parent):
+    _, registry, _ = _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, False))
+    assert set(registry) == {f"ar.language_model.layers.{i}.self_attn.attn" for i in range(NUM_LAYERS)} | {VISION_LAYER}
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_first_pipeline_rank_without_parent_layers_fails(monkeypatch, ar_class, parent):
+    with pytest.raises(RuntimeError, match="Unexpected attention registry"):
+        _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, True))
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_later_pipeline_rank_still_rejects_a_leftover_parent_layer(monkeypatch, ar_class, parent):
+    with pytest.raises(RuntimeError, match="Unexpected attention registry"):
+        _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, False), stray_parent_layer=True)

vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py

@@ -1,5 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
+import gc
 from collections.abc import Callable, Iterable, Mapping
 from itertools import islice
 from typing import Any
@@ -8,8 +12,10 @@ import torch
 from torch import nn
 from transformers import Qwen2Config, Qwen3VLProcessor
 from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor
-from vllm.config import CacheConfig, VllmConfig
+from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
 from vllm.distributed import get_pp_group
+from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -40,6 +46,7 @@ from vllm.model_executor.models.qwen3_vl import (
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    StageMissingLayer,
     WeightsMapper,
     init_vllm_registered_model,
     is_pp_missing_parameter,
@@ -57,6 +64,8 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, is_interleaved
 from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
 
+logger = init_logger(__name__)
+
 
 def _runtime_meta(runtime_info: dict[str, Any]) -> dict[str, Any]:
     """Accept both legacy flat metadata and the canonical nested payload."""
@@ -623,6 +632,42 @@ class MammothModa2Qwen3ForCausalLM(MammothModa2Qwen2ForCausalLM):
         )
 
 
+def _drop_parent_language_model(model: nn.Module, prefix: str) -> None:
+    """Unregister and release the language model built by the Qwen-VL parent ``__init__``.
+
+    The parent builds a stock Qwen LM whose attention layers register in the static
+    forward context as ``{prefix}.language_model.model.layers.*``. MammothModa2 replaces
+    it with its own LM (``{prefix}.language_model.layers.*``); without this, both sets of
+    layers are counted when the KV cache is sized and each gets its own KV blocks.
+    """
+    old = model.language_model
+    if isinstance(old, StageMissingLayer):  # mm_encoder_only
+        old = old.__dict__["module"]
+    ctx = get_current_vllm_config().compilation_config.static_forward_context
+    owned = set(old.modules())
+    dropped = [name for name, layer in ctx.items() if layer in owned]
+    for name in dropped:
+        del ctx[name]
+    stale = [name for name in ctx if name.startswith(maybe_prefix(prefix, "language_model."))]
+    # A later pipeline rank may hold none of the parent's layers, so only the first
+    # rank has to find some; no rank may leave any behind.
+    pp_group = get_pp_group()
+    must_drop = pp_group.world_size == 1 or pp_group.is_first_rank
+    if (must_drop and not dropped) or stale:
+        raise RuntimeError(
+            f"Unexpected attention registry for the parent language model (dropped {len(dropped)}, left {stale}); "
+            "check the Qwen-VL __init__."
+        )
+    # With torch.compile, the compile wrapper's bytecode hook would keep the module alive.
+    for module in owned:
+        if isinstance(module, TorchCompileWithNoGuardsWrapper):
+            module.cleanup()
+    logger.debug("Dropped %d attention layers of the parent language model.", len(dropped))
+    del model.language_model, old, owned
+    gc.collect()
+    torch.accelerator.empty_cache()
+
+
 @MULTIMODAL_REGISTRY.register_processor(
     MammothModa2ARMultiModalProcessor,
     info=MammothModa2ARProcessingInfo,
@@ -656,6 +701,7 @@ class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration)
         ar_vllm_config = vllm_config.with_hf_config(ar_hf_config, architectures=vllm_config.model_config.architectures)
         # Initialize multi-modal components like the vision tower first.
         super().__init__(vllm_config=ar_vllm_config, prefix=prefix)
+        _drop_parent_language_model(self, prefix)
         # Replace with the custom MoE language model.
         lm_hf_config = getattr(
             ar_vllm_config.model_config.hf_config, "text_config", ar_vllm_config.model_config.hf_config
@@ -824,6 +870,7 @@ class MammothModa2Qwen3ARForConditionalGeneration(Qwen3VLForConditionalGeneratio
             architectures=vllm_config.model_config.architectures,
         )
         super().__init__(vllm_config=ar_vllm_config, prefix=prefix)
+        _drop_parent_language_model(self, prefix)
         text_config = ar_hf_config.text_config
         self.language_model = MammothModa2Qwen3ForCausalLM(
             vllm_config=ar_vllm_config.with_hf_config(text_config),