diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 0e753d8ace4a..8acf89ab927b 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -15,11 +15,14 @@ import re from types import SimpleNamespace -from typing import Dict, List, Literal +from typing import TYPE_CHECKING, Dict, List, Literal import torch from transformers import PretrainedConfig +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger from tensorrt_llm.quantization import QuantAlgo @@ -32,6 +35,7 @@ support_multimodal_disaggregated, ) from ..pyexecutor.config_utils import get_qwen3_hybrid_layer_types +from ..utils import is_nvfp4_marlin_supported_sm from .checkpoints.base_weight_mapper import BaseWeightMapper from .checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper from .modeling_qwen3_next import Qwen3NextForCausalLM @@ -54,6 +58,28 @@ } +def _get_qwen35_moe_model_defaults(llm_args: "TorchLlmArgs") -> dict: + """Return Marlin defaults for Qwen3.5 MoE with NVFP4 experts on Ada/Hopper.""" + defaults = Qwen3NextForCausalLM.get_model_defaults(llm_args) + quant_config = getattr(llm_args, "quant_config", None) + if getattr(quant_config, "quant_algo", None) in ( + QuantAlgo.NVFP4, + QuantAlgo.MIXED_PRECISION, + ) and is_nvfp4_marlin_supported_sm(get_sm_version()): + # CUTLASS W4A4 requires Blackwell; use Marlin's W4A16 path instead. + defaults.update( + { + "moe_config": { + "backend": "MARLIN", + }, + "nvfp4_gemm_config": { + "allowed_backends": ["marlin"], + }, + } + ) + return defaults + + def _translate_mtp_pattern(name, n_hidden_layers): """Translate an HF ``mtp.*`` exclude pattern to a TRT-LLM module path. @@ -614,6 +640,10 @@ class Qwen3_5MoeForCausalLM(Qwen3NextForCausalLM): class that serves the vanilla Qwen3NextForCausalLM architecture. """ + @classmethod + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: + return _get_qwen35_moe_model_defaults(llm_args) + def __init__(self, model_config): keep_lm_head_quant = _lm_head_nvfp4_enabled(model_config) _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=keep_lm_head_quant) @@ -745,6 +775,10 @@ def load_weights( class Qwen3_5MoeVLModel(_Qwen3_5VLModel): """VLM wrapper composing Qwen3 vision encoder with Qwen3.5 MoE text decoder.""" + @classmethod + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: + return _get_qwen35_moe_model_defaults(llm_args) + # TODO(TRTLLM-13417): Add tests for disaggregated support. @support_multimodal_disaggregated diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py index e715c53a3ac2..9aa8e7e11a06 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import List, Optional +import pytest import torch import transformers from test_modeling_multimodal import MultimodalScenario, TestModelingMultimodal @@ -15,7 +16,7 @@ from utils.util import skip_pre_hopper from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models import Qwen3_5MoeVLModel +from tensorrt_llm._torch.models import Qwen3_5MoeForCausalLM, Qwen3_5MoeVLModel from tensorrt_llm._torch.models.checkpoints.auto_mapper import AutoCheckpointMapper from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM @@ -27,6 +28,10 @@ from tensorrt_llm._torch.pyexecutor.model_loader import validate_and_set_mamba_ssm_cache_dtype from tensorrt_llm.inputs import ContentFormat from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo def _write_qwen35_moe_vl_config(tmp_path: Path) -> Path: @@ -156,6 +161,39 @@ def test_qwen35_moe_vl_resolves_model_and_mapper(tmp_path: Path) -> None: ) +@pytest.mark.parametrize( + ("quant_algo", "sm_version", "use_marlin"), + [ + pytest.param(QuantAlgo.NVFP4, 90, True, id="hopper-nvfp4"), + pytest.param(QuantAlgo.MIXED_PRECISION, 90, True, id="hopper-mixed-precision"), + pytest.param(QuantAlgo.NVFP4, 100, False, id="blackwell-nvfp4"), + ], +) +def test_qwen35_moe_model_defaults( + monkeypatch: pytest.MonkeyPatch, + quant_algo: QuantAlgo, + sm_version: int, + use_marlin: bool, +) -> None: + monkeypatch.setattr( + "tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", + lambda: sm_version, + ) + + expected_moe_backend = "MARLIN" if use_marlin else "AUTO" + expected_gemm_backends = ["marlin"] if use_marlin else ["cutlass", "cublaslt", "cuda_core"] + for model_cls in (Qwen3_5MoeForCausalLM, Qwen3_5MoeVLModel): + llm_args = TorchLlmArgs(model="/tmp/dummy_model") + llm_args.quant_config = QuantConfig(quant_algo=quant_algo) + defaults = model_cls.get_model_defaults(llm_args) + apply_model_defaults_to_llm_args(llm_args, defaults) + + assert llm_args.kv_cache_config.enable_block_reuse is False + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + assert llm_args.moe_config.backend == expected_moe_backend + assert llm_args.nvfp4_gemm_config.allowed_backends == expected_gemm_backends + + def test_qwen35_moe_vl_placeholder_metadata_registered() -> None: metadata = MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_metadata("qwen3_5_moe")