[Bugfix][ROCm] DeepSeek-V4/V4.1: Hybrid FP8/MXFP4 Quant Routing Guard, Fused Shared Experts, Engram Scale Translation, and VLM Generator Streaming

tests/models/deepseek_v41/test_fse_routing.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for DeepSeek-V4.1 Fused Shared Experts (FSE) routing, padding,
+and expert mapping."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from vllm.models.deepseek_v41.amd.model import DeepseekV4Model, DeepseekV4MoEBase
+
+
+def test_fse_weight_name_redirection():
+    """Verify shared_experts weights are redirected to virtual expert slot
+    when FSE is enabled."""
+    n_routed = 256
+    model = MagicMock(spec=DeepseekV4Model)
+    model.config = SimpleNamespace(n_routed_experts=n_routed, n_shared_experts=1)
+    model.quant_config = None
+    model.use_sequence_parallel = False
+
+    # Mock MoE module on layer 0 with FSE enabled
+    moe_layer0 = MagicMock(spec=DeepseekV4MoEBase)
+    moe_layer0.is_fused_shared_expert_enabled = True
+
+    # Layer 1 has FSE disabled
+    moe_layer1 = MagicMock(spec=DeepseekV4MoEBase)
+    moe_layer1.is_fused_shared_expert_enabled = False
+
+    model.named_modules.return_value = [
+        ("model.layers.0.mlp", moe_layer0),
+        ("model.layers.1.mlp", moe_layer1),
+    ]
+
+    # Test names across layer 0 (FSE enabled) and layer 1 (FSE disabled)
+    dummy_tensor = torch.zeros(1)
+    weights = [
+        ("model.layers.0.mlp.shared_experts.gate_proj.weight", dummy_tensor),
+        ("model.layers.0.mlp.shared_experts.down_proj.weight", dummy_tensor),
+        ("model.layers.0.mlp.shared_experts.up_proj.weight", dummy_tensor),
+        ("model.layers.0.mlp.shared_experts.w1.weight", dummy_tensor),
+        ("model.layers.0.mlp.shared_experts.w2.weight", dummy_tensor),
+        ("model.layers.0.mlp.shared_experts.w3.weight", dummy_tensor),
+        ("model.layers.0.ffn.shared_experts.gate_proj.weight", dummy_tensor),
+        ("model.layers.1.mlp.shared_experts.gate_proj.weight", dummy_tensor),
+        ("model.layers.1.mlp.shared_experts.down_proj.weight", dummy_tensor),
+    ]
+
+    # Intercept names processed in the weight iteration loop
+    processed_names = []
+
+    def mock_pad_shared(quant_config, name, loaded_weight):
+        return loaded_weight
+
+    model._pad_shared_expert_weight = mock_pad_shared
+
+    # Simulate load_weights loop name transformation logic
+    from vllm.model_executor.models.utils import extract_layer_index
+
+    fuse_by_layer = {
+        extract_layer_index(mod_name): getattr(
+            mod, "is_fused_shared_expert_enabled", False
+        )
+        for mod_name, mod in model.named_modules()
+        if isinstance(mod, DeepseekV4MoEBase)
+    }
+
+    for name, loaded_weight in weights:
+        is_shared = ".ffn.shared_experts." in name or ".shared_experts." in name
+        is_fse_redirect = is_shared and fuse_by_layer.get(
+            extract_layer_index(name), False
+        )
+        if is_fse_redirect:
+            name = name.replace(".shared_experts.down_proj", f".experts.{n_routed}.w2")
+            name = name.replace(".shared_experts.gate_proj", f".experts.{n_routed}.w1")
+            name = name.replace(".shared_experts.up_proj", f".experts.{n_routed}.w3")
+            name = name.replace(".shared_experts.w", f".experts.{n_routed}.w")
+        processed_names.append(name)
+
+    # Layer 0 must redirect to .experts.256.w{1,2,3}
+    assert processed_names[0] == f"model.layers.0.mlp.experts.{n_routed}.w1.weight"
+    assert processed_names[1] == f"model.layers.0.mlp.experts.{n_routed}.w2.weight"
+    assert processed_names[2] == f"model.layers.0.mlp.experts.{n_routed}.w3.weight"
+    assert processed_names[3] == f"model.layers.0.mlp.experts.{n_routed}.w1.weight"
+    assert processed_names[4] == f"model.layers.0.mlp.experts.{n_routed}.w2.weight"
+    assert processed_names[5] == f"model.layers.0.mlp.experts.{n_routed}.w3.weight"
+    assert processed_names[6] == f"model.layers.0.ffn.experts.{n_routed}.w1.weight"
+
+    # Layer 1 (FSE disabled) must preserve original .shared_experts. names
+    assert processed_names[7] == "model.layers.1.mlp.shared_experts.gate_proj.weight"
+    assert processed_names[8] == "model.layers.1.mlp.shared_experts.down_proj.weight"
+
+
+def test_pad_shared_expert_weight_dim1_for_w2_and_down_proj():
+    """Verify down_proj and w2 weights/scales are padded along dimension 1
+    (intermediate axis)."""
+    with patch(
+        "vllm.models.deepseek_v41.amd.model.get_tensor_model_parallel_world_size",
+        return_value=2,
+    ):
+        quant_cfg = SimpleNamespace(weight_block_size=[128, 128])
+        # mult = tp_size (2) * block_size (128) = 256
+        # Down proj shape [H, I] = [1024, 2000]. 2000 -> padded to 2048 along dim 1.
+        w_down = torch.ones(1024, 2000)
+
+        # 1. Test .down_proj.
+        padded_down = DeepseekV4Model._pad_shared_expert_weight(
+            quant_cfg, "model.layers.0.mlp.shared_experts.down_proj.weight", w_down
+        )
+        assert padded_down.shape == (1024, 2048)
+        assert torch.all(padded_down[:, :2000] == 1.0)
+        assert torch.all(padded_down[:, 2000:] == 0.0)
+
+        # 2. Test .w2. (FSE redirected name)
+        padded_w2 = DeepseekV4Model._pad_shared_expert_weight(
+            quant_cfg, "model.layers.0.mlp.experts.256.w2.weight", w_down
+        )
+        assert padded_w2.shape == (1024, 2048)
+        assert torch.all(padded_w2[:, :2000] == 1.0)
+        assert torch.all(padded_w2[:, 2000:] == 0.0)
+
+        # 3. Test scale for .w2. (step = 1 instead of 128)
+        # Scale shape in blocks: [1024 // 128, 15] = [8, 15]. mult = 2 * 1 = 2.
+        # 15 -> padded to 16 along dim 1.
+        scale_w2 = torch.ones(8, 15)
+        padded_scale = DeepseekV4Model._pad_shared_expert_weight(
+            quant_cfg, "model.layers.0.mlp.experts.256.w2.weight_scale", scale_w2
+        )
+        assert padded_scale.shape == (8, 16)
+        assert torch.all(padded_scale[:, :15] == 1.0)
+        assert torch.all(padded_scale[:, 15:] == 0.0)
+
+
+def test_pad_shared_expert_weight_dim0_for_w1_w3_and_gate_up():
+    """Verify gate/up and w1/w3 weights/scales are padded along dimension 0."""
+    with patch(
+        "vllm.models.deepseek_v41.amd.model.get_tensor_model_parallel_world_size",
+        return_value=2,
+    ):
+        quant_cfg = SimpleNamespace(weight_block_size=[128, 128])
+        # Gate/Up shape [I, H] = [2000, 1024]. 2000 -> padded to 2048 along dim 0.
+        w_gate = torch.ones(2000, 1024)
+
+        # 1. Test .gate_proj.
+        padded_gate = DeepseekV4Model._pad_shared_expert_weight(
+            quant_cfg, "model.layers.0.mlp.shared_experts.gate_proj.weight", w_gate
+        )
+        assert padded_gate.shape == (2048, 1024)
+        assert torch.all(padded_gate[:2000, :] == 1.0)
+        assert torch.all(padded_gate[2000:, :] == 0.0)
+
+        # 2. Test .w1. (FSE redirected name)
+        padded_w1 = DeepseekV4Model._pad_shared_expert_weight(
+            quant_cfg, "model.layers.0.mlp.experts.256.w1.weight", w_gate
+        )
+        assert padded_w1.shape == (2048, 1024)
+
+        # 3. Test .w3. (FSE redirected name)
+        padded_w3 = DeepseekV4Model._pad_shared_expert_weight(
+            quant_cfg, "model.layers.0.mlp.experts.256.w3.weight", w_gate
+        )
+        assert padded_w3.shape == (2048, 1024)
+
+        # 4. Test scale for .w1. (step = 1)
+        # Scale shape: [15, 8]. mult = 2 * 1 = 2. 15 -> padded to 16 along dim 0.
+        scale_w1 = torch.ones(15, 8)
+        padded_scale = DeepseekV4Model._pad_shared_expert_weight(
+            quant_cfg, "model.layers.0.mlp.experts.256.w1.weight_scale", scale_w1
+        )
+        assert padded_scale.shape == (16, 8)
+
+
+def test_pad_shared_expert_weight_no_op_when_aligned():
+    """Verify already-aligned tensor is returned unchanged without copying or
+    allocating pad."""
+    with patch(
+        "vllm.models.deepseek_v41.amd.model.get_tensor_model_parallel_world_size",
+        return_value=2,
+    ):
+        quant_cfg = SimpleNamespace(weight_block_size=[128, 128])
+        # mult = 2 * 128 = 256. 2048 is already a multiple of 256.
+        aligned_w = torch.ones(1024, 2048)
+        result = DeepseekV4Model._pad_shared_expert_weight(
+            quant_cfg, "model.layers.0.mlp.experts.256.w2.weight", aligned_w
+        )
+        assert result is aligned_w  # Exact identity preserved
+
+
+def test_get_expert_mapping_expansion_under_fse():
+    """Verify get_expert_mapping expands num_experts by n_shared_experts
+    when FSE is enabled."""
+    model = MagicMock(spec=DeepseekV4Model)
+    model.config = SimpleNamespace(n_routed_experts=256, n_shared_experts=1)
+
+    moe_module = MagicMock(spec=DeepseekV4MoEBase)
+    moe_module.is_fused_shared_expert_enabled = True
+    model.modules.return_value = [moe_module]
+
+    mapping_fn = DeepseekV4Model.get_expert_mapping.__get__(model, DeepseekV4Model)
+
+    with patch(
+        "vllm.models.deepseek_v41.amd.model.fused_moe_make_expert_params_mapping"
+    ) as mock_mapping_gen:
+        mapping_fn()
+        mock_mapping_gen.assert_called_once_with(
+            model,
+            ckpt_gate_proj_name="w1",
+            ckpt_down_proj_name="w2",
+            ckpt_up_proj_name="w3",
+            num_experts=257,
+        )
+
+    # When FSE is disabled, num_experts must remain exactly n_routed_experts (256)
+    moe_module.is_fused_shared_expert_enabled = False
+    with patch(
+        "vllm.models.deepseek_v41.amd.model.fused_moe_make_expert_params_mapping"
+    ) as mock_mapping_gen:
+        mapping_fn()
+        mock_mapping_gen.assert_called_once_with(
+            model,
+            ckpt_gate_proj_name="w1",
+            ckpt_down_proj_name="w2",
+            ckpt_up_proj_name="w3",
+            num_experts=256,
+        )

tests/models/deepseek_v41/test_quant_config.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for DeepSeek-V4.1 quant method routing and shared experts exclusion."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.config import set_current_vllm_config
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod
+from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config
+
+
+class DummyLinear(LinearBase):
+    """Dummy linear layer for testing quant method resolution."""
+
+    def __init__(self):
+        super().__init__(input_size=128, output_size=128, tp_rank=0, tp_size=1)
+
+
+@pytest.fixture
+def vllm_config_context():
+    """Mock vllm_config context required for quant method instantiation."""
+    mock_config = SimpleNamespace(
+        model_config=SimpleNamespace(dtype=torch.bfloat16, is_multimodal_model=False),
+        kernel_config=SimpleNamespace(linear_backend="auto"),
+    )
+    with set_current_vllm_config(mock_config):
+        yield
+
+
+def test_shared_experts_excluded_from_modelopt_linear_method(vllm_config_context):
+    """Verify shared_experts falls back to Fp8LinearMethod under [32, 32] block size."""
+    cfg = DeepseekV4FP8Config(
+        is_checkpoint_fp8_serialized=True,
+        activation_scheme="dynamic",
+        weight_block_size=[32, 32],
+    )
+    # Set resolved expert dtype to fp4 to trigger is_scale_e8m0 = True
+    cfg._resolved_expert_dtype = "fp4"
+
+    layer = DummyLinear()
+    method_gate = cfg.get_quant_method(
+        layer, prefix="model.layers.0.mlp.shared_experts.gate_proj"
+    )
+    method_down = cfg.get_quant_method(
+        layer, prefix="model.layers.0.mlp.shared_experts.down_proj"
+    )
+
+    # Shared experts must NOT use ModelOptLinearMethod
+    assert not isinstance(method_gate, ModelOptLinearMethod)
+    assert isinstance(method_gate, Fp8LinearMethod)
+
+    assert not isinstance(method_down, ModelOptLinearMethod)
+    assert isinstance(method_down, Fp8LinearMethod)
+
+
+def test_dense_and_attn_linear_use_modelopt_linear_method(vllm_config_context):
+    """Verify attention and dense linear layers use ModelOptLinearMethod
+    under [32, 32]."""
+    cfg = DeepseekV4FP8Config(
+        is_checkpoint_fp8_serialized=True,
+        activation_scheme="dynamic",
+        weight_block_size=[32, 32],
+    )
+    cfg._resolved_expert_dtype = "fp4"
+
+    layer = DummyLinear()
+    method_q = cfg.get_quant_method(layer, prefix="model.layers.0.self_attn.q_proj")
+    method_o = cfg.get_quant_method(layer, prefix="model.layers.0.self_attn.o_proj")
+
+    assert isinstance(method_q, ModelOptLinearMethod)
+    assert isinstance(method_o, ModelOptLinearMethod)
+
+
+def test_quant_method_fallback_for_128x128_block_size(vllm_config_context):
+    """Verify fallback to standard Fp8LinearMethod for standard [128, 128]
+    block size."""
+    cfg = DeepseekV4FP8Config(
+        is_checkpoint_fp8_serialized=True,
+        activation_scheme="dynamic",
+        weight_block_size=[128, 128],
+    )
+    cfg._resolved_expert_dtype = "fp4"
+
+    layer = DummyLinear()
+    method_attn = cfg.get_quant_method(layer, prefix="model.layers.0.self_attn.q_proj")
+    method_shared = cfg.get_quant_method(
+        layer, prefix="model.layers.0.mlp.shared_experts.gate_proj"
+    )
+
+    # When block size is not [32, 32], standard Fp8LinearMethod is used for both
+    assert isinstance(method_attn, Fp8LinearMethod)
+    assert isinstance(method_shared, Fp8LinearMethod)
+
+
+def test_deepseek_v4_quark_block_size_discovery():
+    """Verify deepseek_v4 dynamic block size extraction from Quark configs."""
+    from vllm.models.deepseek_v4.quant_config import (
+        DeepseekV4FP8Config as DeepseekV4CoreFP8Config,
+    )
+
+    # 1. Top-level weight_block_size
+    cfg1 = DeepseekV4CoreFP8Config.from_config(
+        {
+            "quant_method": "quark",
+            "weight_block_size": [32, 32],
+        }
+    )
+    assert cfg1.weight_block_size == [32, 32]
+
+    # 2. Nested in layer_quant_config
+    cfg2 = DeepseekV4CoreFP8Config.from_config(
+        {
+            "quant_method": "quark",
+            "layer_quant_config": {
+                "model.layers.0.self_attn.q_proj": {"weight": {"block_size": [32, 32]}}
+            },
+        }
+    )
+    assert cfg2.weight_block_size == [32, 32]
+
+    # 3. Default fallback to [128, 128]
+    cfg3 = DeepseekV4CoreFP8Config.from_config(
+        {
+            "quant_method": "quark",
+        }
+    )
+    assert cfg3.weight_block_size == [128, 128]