[https://nvbugs/6434512][fix] Select Marlin for Qwen3.5 MoE on Hopper

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

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")