Fix TurboMind Qwen3 VL family vision support

Author: web-flow

lmdeploy/turbomind/deploy/source_model/qwen.py

@@ -153,6 +153,17 @@ def moe_ffn_shared_gate(self, i):
         return self.params.get(f'{self.attn_layer_prefix}.{i}.mlp.shared_expert_gate.weight')
 
 
+def _configure_nested_language_model_prefix(reader):
+    """Handle VL checkpoints whose text weights live under
+    ``model.language_model``."""
+    if any(k.startswith('model.language_model.') for k in reader.params.keys()):
+        reader.attn_layer_prefix = 'model.language_model.layers'
+        reader.tok_embeddings_key = 'model.language_model.embed_tokens.weight'
+        reader.norm_weight_key = 'model.language_model.norm.weight'
+        if reader.model_cfg.get('tie_word_embeddings', False):
+            reader.output_weight_key = reader.tok_embeddings_key
+
+
 @INPUT_MODELS.register_module(name='qwen2-moe')
 class Qwen2MoeModel(LlamaModel):
 
@@ -172,6 +183,11 @@ def model_info(self):
 
 
 class Qwen3Reader(LlamaReader):
+    attn_layer_patten = r'(?:model\.language_model\.|model\.)layers\.([0-9]+)\.'
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        _configure_nested_language_model_prefix(self)
 
     def qk_norm(self, i: int):
         result = []
@@ -193,6 +209,11 @@ def model_info(self):
 
 
 class Qwen3MoeReader(Qwen2MoeReader):
+    attn_layer_patten = r'(?:model\.language_model\.|model\.)layers\.([0-9]+)\.'
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        _configure_nested_language_model_prefix(self)
 
     def qk_norm(self, i: int):
         result = []
@@ -236,13 +257,7 @@ class Qwen3_5ReaderMixin:
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        if any(k.startswith('model.language_model.') for k in self.params.keys()):
-            self.attn_layer_prefix = 'model.language_model.layers'
-            self.tok_embeddings_key = 'model.language_model.embed_tokens.weight'
-            self.norm_weight_key = 'model.language_model.norm.weight'
-        tie_word_embeddings = self.model_cfg.get('tie_word_embeddings', False)
-        if tie_word_embeddings:
-            self.output_weight_key = self.tok_embeddings_key
+        _configure_nested_language_model_prefix(self)
 
     # ---- zero-centered RMSNorm: add 1 to weights during export ----
     def attn_norm(self, i: int):

lmdeploy/turbomind/supported_models.py

@@ -32,6 +32,8 @@
     # Qwen3
     Qwen3ForCausalLM='qwen3',
     Qwen3MoeForCausalLM='qwen3-moe',
+    Qwen3VLForConditionalGeneration='qwen3',
+    Qwen3VLMoeForConditionalGeneration='qwen3-moe',
     # Qwen 3.5
     Qwen3_5ForConditionalGeneration='qwen3_5',
     Qwen3_5MoeForConditionalGeneration='qwen3_5-moe',

lmdeploy/vl/model/qwen3.py

@@ -256,7 +256,13 @@ def proc_messages(self, messages, chat_template, sequence_start, chat_template_k
         else:
             prompt_messages = messages
         prompt = chat_template.messages2prompt(prompt_messages, sequence_start, **chat_template_kwargs)
-        return prompt, None
+        return prompt, self.image_token
+
+    def _ensure_turbomind_image_only(self, inputs: list[dict]):
+        """TurboMind split vision currently supports image inputs only."""
+        has_video = self.contains_video_input or any('video_grid_thw' in item for item in inputs)
+        if has_video:
+            raise NotImplementedError('TurboMind split vision for the Qwen3 VL family currently supports images only.')
 
     def to_pytorch_aux_video(self, messages, prompt, VIDEO_TOKEN, tokenizer, sequence_start):
         """Pack the video input to the compatible format with pytorch
@@ -332,6 +338,7 @@ def forward(self, messages: list[dict], max_batch_size: int = 1) -> list[dict]:
         """Run vision encoder for TurboMind split path (shared Qwen3 VL
         family)."""
         inputs = [x['content'] for x in messages if x['role'] == 'preprocess'][0]
+        self._ensure_turbomind_image_only(inputs)
         dtype = torch.half
         device = next(self.model.visual.parameters()).device
         outputs = []
@@ -381,8 +388,9 @@ def to_turbomind(self,
                      chat_template_kwargs: dict | None = None,
                      **kwargs):
         prompt, IMAGE_TOKEN = self.proc_messages(messages, chat_template, sequence_start, chat_template_kwargs)
-        info = super().to_turbomind_aux(messages, prompt, IMAGE_TOKEN, tokenizer, sequence_start)
         inputs = [x['content'] for x in messages if x['role'] == 'preprocess'][0]
+        self._ensure_turbomind_image_only(inputs)
+        info = super().to_turbomind_aux(messages, prompt, IMAGE_TOKEN, tokenizer, sequence_start)
         grid_thws = [x['image_grid_thw'].tolist()[0] for x in inputs]
         seq_len = len(info['input_ids'])
         ranges = info['input_embedding_ranges']

tests/test_lmdeploy/test_vl/test_qwen_vl_family.py

@@ -1,11 +1,12 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import os
+from types import SimpleNamespace
 
 import pytest
+import torch
 
-from lmdeploy.messages import PytorchEngineConfig
-from lmdeploy.serve.core.vl_async_engine import adjust_vl_backend_config_for_prefix_caching
-from lmdeploy.vl.model.qwen3 import resolve_qwen_vl_family_automodel
+from lmdeploy.archs import get_task
+from lmdeploy.messages import TurbomindEngineConfig
+from lmdeploy.vl.model.qwen3 import Qwen3VLModel, resolve_qwen_vl_family_automodel
 
 
 @pytest.mark.parametrize('arch,expected_block', [
@@ -25,44 +26,81 @@ def test_resolve_unknown_arch_raises():
         resolve_qwen_vl_family_automodel('NotAModel')
 
 
-def test_adjust_prefix_caching_leaves_qwen3vl_untouched(monkeypatch):
-    monkeypatch.setattr(
-        'lmdeploy.archs.get_model_arch',
-        lambda _path: ('Qwen3VLForConditionalGeneration', None),
-    )
-    cfg = PytorchEngineConfig(enable_prefix_caching=True)
-    out = adjust_vl_backend_config_for_prefix_caching('/fake', cfg)
-    assert out is cfg
-    assert cfg.enable_prefix_caching is True
-
-
-def test_adjust_prefix_caching_disables_for_other_vl_without_mutation(monkeypatch):
-    monkeypatch.setattr(
-        'lmdeploy.archs.get_model_arch',
-        lambda _path: ('InternVLChatModel', None),
-    )
-    cfg = PytorchEngineConfig(enable_prefix_caching=True)
-    out = adjust_vl_backend_config_for_prefix_caching('/fake', cfg)
-    assert out is not cfg
-    assert out.enable_prefix_caching is False
-    assert cfg.enable_prefix_caching is True
-
-
-_HUB_SMOKE_MODEL_IDS = (
-    'Qwen/Qwen3-VL-4B-Instruct',
-    'Qwen/Qwen3-VL-8B-Instruct',
-)
-
-
-def test_hub_config_resolves_qwen_vl_family_smoke():
-    """Optional: verify published configs use an architecture handled by the shared loader."""
-    if os.environ.get('LMDEPLOY_RUN_VL_LOAD_TESTS') != '1':
-        pytest.skip('HF hub config fetch; enable with LMDEPLOY_RUN_VL_LOAD_TESTS=1')
-    pytest.importorskip('transformers')
-    from transformers import AutoConfig
-
-    for model_id in _HUB_SMOKE_MODEL_IDS:
-        cfg = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
-        cls, no_split = resolve_qwen_vl_family_automodel(cfg.architectures[0])
-        assert cls is not None
-        assert len(no_split) >= 1
+def test_get_task_routes_qwen3_vl_to_vl_engine(monkeypatch):
+    cfg = SimpleNamespace(to_dict=lambda: {'architectures': ['Qwen3VLForConditionalGeneration']})
+    monkeypatch.setattr('lmdeploy.archs.get_model_arch', lambda _path: ('Qwen3VLForConditionalGeneration', cfg))
+
+    task, pipeline_class = get_task('/fake-model', TurbomindEngineConfig())
+    assert task == 'vlm'
+    assert pipeline_class.__name__ == 'VLAsyncEngine'
+
+
+class _DummyChatTemplate:
+
+    def __init__(self, prompt):
+        self.prompt = prompt
+
+    def messages2prompt(self, messages, sequence_start, **kwargs):
+        return self.prompt
+
+
+class _DummyTokenizer:
+
+    def encode(self, text, add_bos=False):
+        tokens = [] if not text else [len(text)]
+        if add_bos:
+            return [0] + tokens
+        return tokens
+
+
+def _build_qwen3_vl_stub():
+    model = Qwen3VLModel.__new__(Qwen3VLModel)
+    model.image_token = '<|image_pad|>'
+    model.image_token_id = 151655
+    model.contains_video_input = False
+    return model
+
+
+def test_qwen3_vl_to_turbomind_uses_image_token_placeholder():
+    model = _build_qwen3_vl_stub()
+    tokenizer = _DummyTokenizer()
+    prompt = 'prefix<|vision_start|><|image_pad|><|vision_end|>suffix'
+    chat_template = _DummyChatTemplate(prompt)
+    image_grid_thw = torch.tensor([[1, 2, 2]])
+    image_embed = torch.randn(1, 4)
+    messages = [{
+        'role': 'user',
+        'content': [{
+            'type': 'image',
+            'data': object(),
+        }],
+    }, {
+        'role': 'preprocess',
+        'content': [{
+            'image_grid_thw': image_grid_thw,
+        }],
+    }, {
+        'role': 'forward',
+        'content': [image_embed],
+    }]
+
+    info = model.to_turbomind(messages, chat_template, tokenizer, sequence_start=True)
+
+    begin = len(tokenizer.encode('prefix<|vision_start|>', add_bos=True))
+    assert info['input_embedding_ranges'] == [[begin, begin + image_embed.shape[0]]]
+    assert len(info['input_embeddings']) == 1
+    assert info['input_meta']['mrope_position_ids'].shape[1] == len(info['input_ids'])
+
+
+def test_qwen3_vl_to_turbomind_rejects_video():
+    model = _build_qwen3_vl_stub()
+    model.contains_video_input = True
+    messages = [{
+        'role': 'preprocess',
+        'content': [{
+            'video_grid_thw': torch.tensor([[1, 2, 2]]),
+        }],
+    }]
+
+    with pytest.raises(NotImplementedError, match='supports images only'):
+        model.to_turbomind(messages, _DummyChatTemplate(''), _DummyTokenizer(), sequence_start=True)