diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 717461261135..855ec8c056cd 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -48,8 +48,9 @@ steps: - tests/models/test_terratorch.py - tests/models/transformers/test_backend.py - tests/models/test_registry.py + - tests/models/test_deepseek_v4_vl_rocm.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py mirror: amd: label: ":amd: (MI300) Basic Models (Other)" @@ -59,9 +60,10 @@ steps: depends_on: - image-build-amd source_file_dependencies: + - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_hyv4_rocm.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 1c389a111733..f0ac4d838511 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -336,6 +336,101 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None: torch.testing.assert_close(actual_lens, expected_lens) +@torch.inference_mode() +def test_combine_topk_swa_indices_adds_image_visibility() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + num_tokens = 8 + topk_indices = torch.full((num_tokens, 1), -1, dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device) + gather_lens = seq_lens.clone() + left_visible = torch.tensor( + [0, 0, 0, 1, 2, 3, 4, 0], dtype=torch.int32, device=device + ) + right_visible = torch.tensor( + [0, 0, 4, 3, 2, 1, 0, 0], dtype=torch.int32, device=device + ) + + indices, lens = combine_topk_swa_indices( + topk_indices, + query_start_loc, + seq_lens, + gather_lens, + window_size=4, + compress_ratio=1, + topk=0, + M=16, + N=0, + max_image_tokens=5, + left_visible=left_visible, + right_visible=right_visible, + ) + + expected_rows = [ + [0], + [0, 1], + [0, 1, 2, 3, 4, 5, 6], + [0, 1, 2, 3, 4, 5, 6], + [1, 2, 3, 4, 5, 6], + [2, 3, 4, 5, 6], + [2, 3, 4, 5, 6], + [4, 5, 6, 7], + ] + for token_idx, expected in enumerate(expected_rows): + actual = indices[token_idx, : lens[token_idx]].cpu().tolist() + assert actual == expected + + +@torch.inference_mode() +def test_combine_topk_swa_indices_apc_hit_inside_image() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + indices, lens = combine_topk_swa_indices( + torch.full((2, 1), -1, dtype=torch.int32, device=device), + torch.tensor([0, 2], dtype=torch.int32, device=device), + torch.tensor([10], dtype=torch.int32, device=device), + # Only positions [5, 10) exist in the gathered SWA workspace. + torch.tensor([5], dtype=torch.int32, device=device), + window_size=4, + compress_ratio=1, + topk=0, + M=10, + N=0, + max_image_tokens=10, + left_visible=torch.tensor([8, 9], dtype=torch.int32, device=device), + # Deliberately extends beyond seq_len to exercise the upper clamp too. + right_visible=torch.tensor([5, 5], dtype=torch.int32, device=device), + ) + + assert lens.cpu().tolist() == [5, 5] + assert indices[:, :5].cpu().tolist() == [list(range(5)), list(range(5))] + + +@torch.inference_mode() +def test_combine_topk_swa_indices_keeps_vision_row_width_without_images() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + indices, lens = combine_topk_swa_indices( + torch.full((1, 1), -1, dtype=torch.int32, device=device), + torch.tensor([0, 1], dtype=torch.int32, device=device), + torch.tensor([1], dtype=torch.int32, device=device), + torch.tensor([1], dtype=torch.int32, device=device), + window_size=120, + compress_ratio=1, + topk=0, + M=256, + N=0, + max_image_tokens=16, + ) + + assert indices.shape == (1, 256) + assert lens.item() == 1 + + def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None: from vllm.models.deepseek_v4.amd import rocm as mod diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 0ba18e9f0315..e1fb0cef7439 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -164,10 +164,11 @@ def test_model_tensor_schema(model_id: str): "Kimi-K2.5's offline inference has issues about vision chunks. Fix later." ) - if model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" and not ( - current_platform.is_cuda() + if ( + model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" + and not current_platform.is_cuda_alike() ): - pytest.skip("Deepseek V4 is only supported on CUDA") + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") if model_id == "zai-org/GLM-5.3-Flash" and (current_platform.is_xpu()): pytest.skip("GLM-5.3-Flash is not supported on XPU") diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py new file mode 100644 index 000000000000..750927972ef3 --- /dev/null +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -0,0 +1,458 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.model_executor.models.utils import WeightsMapper +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + + +def _preshuffled_fp8_linear( + holder: str = "quant_method", weight_shape: tuple[int, int] = (256, 128) +) -> nn.Module: + pytest.importorskip("aiter") + from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( + AiterPreshuffledFp8BlockScaledMMKernel, + ) + from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + FP8ScaledMMLinearLayerConfig, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + ) + + config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8Static128BlockSym, + activation_quant_key=kFp8Dynamic128Sym, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + weight_shape=weight_shape, + ) + layer = nn.Module() + layer.weight = nn.Parameter( + torch.randn(*weight_shape, device="cuda").to(current_platform.fp8_dtype()), + requires_grad=False, + ) + layer.register_parameter( + "weight_scale" if holder == "scheme" else "weight_scale_inv", + nn.Parameter( + torch.ones(*(dim // 128 for dim in weight_shape), device="cuda"), + requires_grad=False, + ), + ) + setattr( + layer, + holder, + SimpleNamespace(fp8_linear=AiterPreshuffledFp8BlockScaledMMKernel(config)), + ) + return layer + + +@pytest.mark.parametrize("holder", ["quant_method", "scheme"]) +def test_rocm_wo_a_keeps_row_major_weights(holder: str, default_vllm_config) -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 + + layer = _preshuffled_fp8_linear(holder) + layer.is_bmm = True + original = layer.weight.detach().clone() + kernel = getattr(layer, holder).fp8_linear + kernel.process_weights_after_loading(layer) + + actual = _get_cached_wo_a_bf16(layer, 2, 128, 128) + torch.testing.assert_close(actual, original.to(torch.bfloat16).view(2, 128, 128)) + + +@pytest.mark.parametrize("holder", ["quant_method", "scheme"]) +def test_rocm_gateup_shuffles_once_and_down_proj_keeps_preshuffled_backend( + holder: str, monkeypatch: pytest.MonkeyPatch, default_vllm_config +) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.models.deepseek_v4.amd import model as rocm_model + + gate_up, down = _preshuffled_fp8_linear(holder), _preshuffled_fp8_linear(holder) + original_gate_up = gate_up.weight.detach().clone() + original_down = down.weight.detach().clone() + monkeypatch.setattr( + rocm_model, "MergedColumnParallelLinear", lambda *a, **k: gate_up + ) + monkeypatch.setattr(rocm_model, "RowParallelLinear", lambda *a, **k: down) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + model = rocm_model.DeepseekV4MLP(128, 128, "silu") + for linear in (gate_up, down): + getattr(linear, holder).fp8_linear.process_weights_after_loading(linear) + model.prepare_gateup_preshuffle() + + for linear, original in ((gate_up, original_gate_up), (down, original_down)): + expected = rocm_aiter_ops.shuffle_weight(original, layout=(16, 16)) + torch.testing.assert_close( + linear.weight.view(torch.uint8), expected.view(torch.uint8) + ) + + +def test_rocm_fused_qkv_quant_matches_preshuffled_gemm_scale_layout( + monkeypatch: pytest.MonkeyPatch, default_vllm_config +) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLAAttention, + apply_pre_quantized_block_scaled_mm, + ) + + if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(1024, 256): + pytest.skip("Requires a tuned AITER preshuffled FP8 GEMM") + linear = _preshuffled_fp8_linear(weight_shape=(1024, 256)) + original = linear.weight.detach().clone() + linear.quant_method.fp8_linear.process_weights_after_loading(linear) + attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention) + nn.Module.__init__(attention) + attention.wq_b = linear + attention.indexer = None + attention.q_lora_rank, attention.head_dim = 256, 128 + attention.eps = 1e-5 + attention.q_norm = SimpleNamespace( + weight=torch.ones(256, dtype=torch.bfloat16, device="cuda") + ) + attention.kv_norm = SimpleNamespace( + weight=torch.ones(128, dtype=torch.bfloat16, device="cuda") + ) + monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: True) + # Multiple rows and groups with different ranges expose byte-order errors. + inputs = torch.randn(16, 384, dtype=torch.bfloat16, device="cuda") + inputs[:, :128] *= 4 + quantized, scales, _ = attention._split_qkv_and_norm(inputs) + assert scales is not None + output = apply_pre_quantized_block_scaled_mm(linear, quantized, scales) + q = inputs[:, :256].float() + normalized = q * torch.rsqrt(q.square().mean(dim=-1, keepdim=True) + attention.eps) + reference_scales = normalized.view(16, 2, 128).abs().amax(dim=-1) + reference_scales /= torch.finfo(current_platform.fp8_dtype()).max + dequantized = quantized.float() * reference_scales.repeat_interleave(128, dim=1) + expected = dequantized @ original.float().T + torch.testing.assert_close(output.float(), expected, atol=0.125, rtol=0.01) + + +def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None: + from vllm.config import CacheConfig + from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype + + cache_config = CacheConfig() + + resolved_dtype, torch_dtype = _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout=True, + kv_cache_dtype=cache_config.cache_dtype, + cache_config=cache_config, + ) + + assert resolved_dtype == "fp8_ds_mla" + assert torch_dtype is torch.uint8 + assert cache_config.cache_dtype == "fp8_ds_mla" + + +def test_rocm_packed_kv_cache_rejects_unquantized_dtype() -> None: + from vllm.config import CacheConfig + from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype + + cache_config = CacheConfig(cache_dtype="bfloat16") + + with pytest.raises(ValueError, match="only supports fp8 kv-cache"): + _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout=True, + kv_cache_dtype=cache_config.cache_dtype, + cache_config=cache_config, + ) + + +def test_vl_mapper_preserves_rocm_weight_mapping() -> None: + from vllm.models.deepseek_v4.amd.model import _make_deepseek_v4_weights_mapper + from vllm.models.deepseek_v4.common.vl_model import ( + _make_deepseek_v4_vl_weights_mapper, + ) + + text_mapper = _make_deepseek_v4_weights_mapper("fp4", fuse_shared_experts=True) + mapper = _make_deepseek_v4_vl_weights_mapper(text_mapper, image_enabled=True) + + assert mapper._map_name("layers.3.attn.wq_a.input_scale") == ( + "language_model.model.layers.3.attn.wq_a.input_scale_2" + ) + assert mapper._map_name("layers.3.ffn.shared_experts.w2.weight") == ( + "language_model.model.layers.3.ffn.shared_experts.w2.weight" + ) + assert mapper._map_name("head.weight") == "language_model.lm_head.weight" + + +def test_rocm_moe_wires_vision_routing_on_hash_and_regular_layers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.models.deepseek_v4.amd import model as rocm_model + from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID + + captured: list[dict] = [] + + class FakeGate(nn.Module): + def __init__(self, **kwargs) -> None: + super().__init__() + + def fake_factory(**kwargs): + captured.append(kwargs) + return nn.Identity() + + monkeypatch.setattr(rocm_model, "GateLinear", FakeGate) + monkeypatch.setattr(rocm_model, "FusedMoEFactory", fake_factory) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_rank", lambda: 0) + + config = SimpleNamespace( + hidden_size=16, + n_routed_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + swiglu_limit=None, + norm_topk_prob=True, + scoring_func="sqrtsoftplus", + num_hash_layers=1, + vocab_size=32, + topk_method="noaux_tc", + vision_n_layers=1, + n_shared_experts=None, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), quant_config=None + ) + + hash_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.0.ffn") + regular_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.1.ffn") + + assert hash_moe.gate.tid2eid is not None + assert regular_moe.gate.tid2eid is None + for moe, factory_kwargs in zip((hash_moe, regular_moe), captured, strict=True): + assert moe.gate.e_score_correction_bias is not None + assert moe.gate.bias_vl is not None + assert factory_kwargs["bias_vl"] is moe.gate.bias_vl + assert factory_kwargs["image_sentinel_lo"] == IMAGE_SENTINEL_BASE_ID + + +def test_rocm_mtp_forwards_input_ids_for_vision_routing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.models.deepseek_v4.amd import mtp as rocm_mtp + + hidden_size = 4 + hc_mult = 2 + + class FakeNorm(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = 1e-6 + + class FakeMTPBlock(nn.Module): + use_fused_mhc = False + + def __init__(self) -> None: + super().__init__() + self.seen_input_ids: torch.Tensor | None = None + + def forward( + self, + *, + positions: torch.Tensor, + x: torch.Tensor, + input_ids: torch.Tensor | None, + ): + self.seen_input_ids = input_ids + return x, None, None, None + + def passthrough_mtp_input( + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + *args, + ) -> tuple[torch.Tensor, torch.Tensor]: + return inputs_embeds, previous_hidden_states + + monkeypatch.setattr( + rocm_mtp, "_FUSED_MTP_INPUT_RMSNORM_KERNEL", passthrough_mtp_input + ) + + layer = object.__new__(rocm_mtp.DeepSeekV4MultiTokenPredictorLayer) + nn.Module.__init__(layer) + layer.config = SimpleNamespace(hidden_size=hidden_size) + layer.hc_mult = hc_mult + layer.enorm = FakeNorm() + layer.hnorm = FakeNorm() + layer.e_proj = nn.Identity() + layer.h_proj = nn.Identity() + layer.mtp_block = FakeMTPBlock() + + input_ids = torch.tensor([11, 12]) + positions = torch.tensor([3, 4]) + inputs_embeds = torch.arange(8, dtype=torch.float32).view(2, hidden_size) + previous_hidden_states = torch.arange(16, dtype=torch.float32).view(2, -1) + + output = layer( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + ) + + assert layer.mtp_block.seen_input_ids is input_ids + expected = previous_hidden_states.view(2, hc_mult, hidden_size) + expected = expected + inputs_embeds.unsqueeze(-2) + torch.testing.assert_close(output, expected.flatten(1)) + + +def test_rocm_compute_logits_local_skips_gather() -> None: + from vllm.models.deepseek_v4.amd.model import DeepseekV4ForCausalLM + + calls: list[tuple[nn.Module, torch.Tensor, bool]] = [] + + def logits_processor( + lm_head: nn.Module, hidden_states: torch.Tensor, *, skip_gather: bool = False + ) -> torch.Tensor: + calls.append((lm_head, hidden_states, skip_gather)) + return hidden_states + 1 + + model = object.__new__(DeepseekV4ForCausalLM) + nn.Module.__init__(model) + model.lm_head = nn.Identity() + model.logits_processor = logits_processor + hidden_states = torch.tensor([4.0]) + + result = model.compute_logits_local(hidden_states) + + assert torch.equal(result, torch.tensor([5.0])) + assert calls == [(model.lm_head, hidden_states, True)] + + +class _FakeLanguageModel(nn.Module): + finalizes_weights_during_load = False + + def __init__(self) -> None: + super().__init__() + self.tensor_a = nn.Parameter(torch.zeros(1)) + self.tensor_c = nn.Parameter(torch.zeros(1)) + self.finalized_values: list[tuple[float, float]] = [] + + def process_weights_after_loading(self) -> None: + self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item())) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + 1 + + +def test_vl_wrapper_streams_then_delegates_finalization() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeLanguageModel() + model.vision = nn.Module() + model.vision.tensor_b = nn.Parameter(torch.zeros(1)) + model.hf_to_vllm_mapper = WeightsMapper() + + def interleaved_weights(): + yield "language_model.tensor_a", torch.tensor([1.0]) + assert model.language_model.tensor_a.item() == 1.0 + yield "vision.tensor_b", torch.tensor([2.0]) + assert model.vision.tensor_b.item() == 2.0 + yield "language_model.tensor_c", torch.tensor([3.0]) + + loaded = model.load_weights(interleaved_weights()) + + assert loaded == { + "language_model.tensor_a", + "vision.tensor_b", + "language_model.tensor_c", + } + assert model.language_model.finalized_values == [] + + model.process_weights_after_loading() + + assert model.language_model.finalized_values == [(1.0, 3.0)] + assert torch.equal( + model.compute_logits_local(torch.tensor([4.0])), torch.tensor([5.0]) + ) + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(1.0, 3.0)] + + +class _FakeFinalizingLanguageModel(_FakeLanguageModel): + finalizes_weights_during_load = True + + def __init__(self) -> None: + super().__init__() + self.load_calls = 0 + + def load_weights(self, weights) -> set[str]: + self.load_calls += 1 + loaded = set() + for name, value in weights: + getattr(self, name).data.copy_(value) + loaded.add(name) + self.process_weights_after_loading() + return loaded + + +def test_vl_wrapper_groups_child_that_finalizes_during_load() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeFinalizingLanguageModel() + model.vision = nn.Module() + model.vision.tensor_b = nn.Parameter(torch.zeros(1)) + model.hf_to_vllm_mapper = WeightsMapper() + + loaded = model.load_weights( + iter( + ( + ("language_model.tensor_a", torch.tensor([1.0])), + ("vision.tensor_b", torch.tensor([2.0])), + ("language_model.tensor_c", torch.tensor([3.0])), + ) + ) + ) + + assert loaded == { + "language_model.tensor_a", + "vision.tensor_b", + "language_model.tensor_c", + } + assert model.language_model.load_calls == 1 + assert model.language_model.finalized_values == [(1.0, 3.0)] + + # The framework's later model-level hook must not double-finalize a child + # which already completed this work in load_weights. + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(1.0, 3.0)] + + +def test_vl_wrapper_dummy_load_delegates_finalization() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeFinalizingLanguageModel() + + # DummyModelLoader bypasses model.load_weights(), so no finalized marker + # exists and the framework-level hook must still delegate to the child. + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(0.0, 0.0)] + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(0.0, 0.0)] diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 3b68f424788f..cc982f699686 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -135,8 +135,8 @@ def _initialize_kv_caches_v1(self, vllm_config): if model_arch == "DeepseekV4ForConditionalGeneration": from vllm.platforms import current_platform - if not current_platform.is_cuda(): - pytest.skip("Deepseek V4 is only supported on CUDA") + if not (current_platform.is_cuda() or current_platform.is_rocm()): + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") with ( patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1), diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 14d85fc09d84..5f6bda8e2949 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -64,10 +64,11 @@ def test_registry_imports(model_arch): ): pytest.skip("HY V4 is only supported on CUDA and ROCm") - if model_arch == "DeepseekV4ForConditionalGeneration" and not ( - current_platform.is_cuda() + if ( + model_arch == "DeepseekV4ForConditionalGeneration" + and not current_platform.is_cuda_alike() ): - pytest.skip("Deepseek V4 is only supported on CUDA") + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") # Ensure all model classes can be imported successfully model_cls = ModelRegistry._try_load_model_cls(model_arch) diff --git a/tests/test_config.py b/tests/test_config.py index 3982507b6b63..2524968cce0d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -57,6 +57,45 @@ def _write_json(path: Path, value: object) -> None: path.write_text(json.dumps(value), encoding="utf-8") +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") +@pytest.mark.parametrize( + ("is_mm_prefix_lm", "is_multimodal_model", "expected"), + [ + pytest.param(True, True, True, id="multimodal-prefix-lm"), + pytest.param(False, True, False, id="multimodal-causal"), + pytest.param(True, False, False, id="text-prefix-lm"), + pytest.param(None, True, False, id="missing-model-config"), + ], +) +def test_rocm_mm_prefix_lm_disables_chunked_mm_input( + is_mm_prefix_lm: bool | None, + is_multimodal_model: bool, + expected: bool, +) -> None: + from vllm.platforms.rocm import RocmPlatform + + config = SimpleNamespace( + compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE), + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + worker_cls="test-worker", + ), + model_config=( + None + if is_mm_prefix_lm is None + else SimpleNamespace(is_mm_prefix_lm=is_mm_prefix_lm) + ), + scheduler_config=SimpleNamespace( + is_multimodal_model=is_multimodal_model, + disable_chunked_mm_input=False, + ), + ) + + RocmPlatform.check_and_update_config(config) + + assert config.scheduler_config.disable_chunked_mm_input is expected + + def test_kda_recoverssm_derivation_is_revalidated(): config = SimpleNamespace( cache_config=SimpleNamespace( diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index b14fc126bbe9..116f4a7c3612 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -42,6 +42,7 @@ def make_manager( ("mode", "piecewise_capture_available", "expected"), [ ("FULL_DECODE_ONLY", True, "FULL_DECODE_ONLY"), + ("FULL_DECODE_ONLY", False, "FULL_DECODE_ONLY"), ("FULL", True, "FULL_AND_PIECEWISE"), ("FULL", False, "FULL_DECODE_ONLY"), ("FULL_AND_PIECEWISE", True, "FULL_AND_PIECEWISE"), diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index 4170a39042ed..6bbd6e0f29b5 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -20,7 +20,7 @@ ) from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP - from .vl_stub import ( # type: ignore[assignment] + from .common.vl_model import ( # type: ignore[assignment] DeepseekV4ForConditionalGeneration, ) elif current_platform.is_xpu(): @@ -38,14 +38,14 @@ DeepseekV4ForConditionalGeneration, ) else: + from .common.vl_model import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, ) from .nvidia.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .nvidia.mtp import DeepSeekV4MTP # type: ignore[assignment] - from .nvidia.vl_model import ( # type: ignore[assignment] - DeepseekV4ForConditionalGeneration, - ) __all__ = [ "DSparkDeepseekV4ForCausalLM", diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index e9ff6cd80f71..349ff5267981 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -79,6 +79,8 @@ from vllm.platforms.rocm import on_gfx950 from vllm.sequence import IntermediateTensors +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID + logger = init_logger(__name__) @@ -549,6 +551,10 @@ def __init__( self.gate.e_score_correction_bias = None self.gate.tid2eid = None + self.gate.bias_vl = None + self.image_sentinel_lo = ( + IMAGE_SENTINEL_BASE_ID if getattr(config, "vision_n_layers", 0) > 0 else 0 + ) is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers self.hash_indices_dtype = torch.int32 if is_hash_moe: @@ -564,12 +570,20 @@ def __init__( ), requires_grad=False, ) - elif getattr(config, "topk_method", None) == "noaux_tc": + if getattr(config, "topk_method", None) == "noaux_tc" and ( + not is_hash_moe or getattr(config, "vision_n_layers", 0) > 0 + ): self.gate.e_score_correction_bias = nn.Parameter( torch.empty(config.n_routed_experts, dtype=torch.float32), requires_grad=False, ) + if getattr(config, "vision_n_layers", 0) > 0: + self.gate.bias_vl = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + self.n_shared_experts = config.n_shared_experts # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` @@ -638,6 +652,8 @@ def __init__( routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=self.gate.e_score_correction_bias, hash_indices_table=self.gate.tid2eid, + bias_vl=self.gate.bias_vl, + image_sentinel_lo=self.image_sentinel_lo, swiglu_limit=self.swiglu_limit, router_logits_dtype=torch.float32, routed_experts_cls=( @@ -660,6 +676,8 @@ def forward( ) -> torch.Tensor: if self.gate.tid2eid is not None and input_ids is None: raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + if self.gate.bias_vl is not None and input_ids is None: + raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.") org_shape = hidden_states.shape final_hidden_states = self.experts( @@ -1351,6 +1369,7 @@ def _make_deepseek_v4_weights_mapper( class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3): model_cls = DeepseekV4Model + finalizes_weights_during_load = False # Default mapper assumes the original FP4-expert checkpoint layout. # Overridden per-instance in __init__ when expert_dtype != "fp4". @@ -1397,6 +1416,9 @@ def compute_logits( logits = self.logits_processor(self.lm_head, hidden_states) return logits + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True) + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index cafb6bb0ed24..d94ff5c2709e 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -160,7 +160,7 @@ def forward( inputs_embeds ).unsqueeze(-2) hidden_states, residual, post_mix, res_mix = self.mtp_block( - positions=positions, x=hidden_states, input_ids=None + positions=positions, x=hidden_states, input_ids=input_ids ) if self.mtp_block.use_fused_mhc: hidden_states = self.mtp_block.hc_post( diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index a798053a4974..b39ad953e75d 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -67,8 +67,13 @@ def weight_already_preshuffled(linear: torch.nn.Module) -> bool: The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped for those, since shuffle_weight is a permutation rather than an involution. """ - kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) - return bool(getattr(kernel, "preshuffles_weight", False)) + return any( + getattr(getattr(method, "fp8_linear", None), "preshuffles_weight", False) + for method in ( + getattr(linear, "quant_method", None), + getattr(linear, "scheme", None), + ) + ) def apply_pre_quantized_block_scaled_mm( @@ -116,13 +121,18 @@ def _combine_topk_swa_indices_kernel( query_start_loc_ptr, seq_lens_ptr, gather_lens_ptr, + left_visible_ptr, + right_visible_ptr, M, N, TOP_K: tl.constexpr, COMPRESS_RATIO: tl.constexpr, WINDOW_SIZE: tl.constexpr, + SWA_WIDTH: tl.constexpr, TOPK_WIDTH: tl.constexpr, PADDED_TOP_K: tl.constexpr, + PADDED_SWA_WIDTH: tl.constexpr, + HAS_IMAGE: tl.constexpr, ): batch_idx = tl.program_id(0) worker_id = tl.program_id(1) @@ -141,7 +151,20 @@ def _combine_topk_swa_indices_kernel( token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + if HAS_IMAGE: + left = tl.load(left_visible_ptr + token_idx) + right = tl.load(right_visible_ptr + token_idx) + else: + left = 0 + right = 0 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) + # Prefix caching can resume inside an image span. Do not generate + # indices outside the SWA rows present in the gathered workspace. + swa_start = tl.maximum( + tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0), gather_start + ) + swa_end = tl.minimum(pos + right + 1, seq_len) + swa_len = tl.maximum(swa_end - swa_start, 0) topk_offset = tl.arange(0, PADDED_TOP_K) topk_mask = topk_offset < topk_len @@ -159,14 +182,14 @@ def _combine_topk_swa_indices_kernel( mask=topk_mask, ) - swa_offset = tl.arange(0, WINDOW_SIZE) + swa_offset = tl.arange(0, PADDED_SWA_WIDTH) tl.store( combined_indices_ptr + token_idx * combined_indices_stride + topk_len + swa_offset, - M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, - mask=swa_offset < swa_len, + M * batch_idx + N + swa_offset + swa_start - gather_start, + mask=(swa_offset < swa_len) & (swa_offset < SWA_WIDTH), ) tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) @@ -182,12 +205,21 @@ def combine_topk_swa_indices( topk: int, M: int, N: int, + max_image_tokens: int = 0, + left_visible: torch.Tensor | None = None, + right_visible: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + if (left_visible is None) != (right_visible is None): + raise ValueError("left_visible and right_visible must be provided together") topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() num_tokens = topk_indices.shape[0] num_reqs = seq_lens.shape[0] + has_image = left_visible is not None + # Keep the row shape fixed for a vision model even when a particular batch + # has no image. + swa_width = window_size + max_image_tokens combined_topk = ( - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + (topk + swa_width + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) // _SPARSE_PREFILL_TOPK_ALIGNMENT * _SPARSE_PREFILL_TOPK_ALIGNMENT ) @@ -211,13 +243,18 @@ def combine_topk_swa_indices( query_start_loc, seq_lens, gather_lens, + left_visible if left_visible is not None else topk_indices, + right_visible if right_visible is not None else topk_indices, M, N, TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, + SWA_WIDTH=swa_width, TOPK_WIDTH=topk_indices.shape[-1], PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), + PADDED_SWA_WIDTH=triton.next_power_of_2(swa_width), + HAS_IMAGE=has_image, ) return combined_indices, combined_lens @@ -948,6 +985,12 @@ def _forward_prefill( assert query_start_loc_cpu is not None assert query_start_loc is not None prefill_token_base = query_start_loc_cpu[num_decodes] + left_visible = swa_metadata.prefill_left_visible + right_visible = swa_metadata.prefill_right_visible + if left_visible is not None: + left_visible = left_visible[num_decode_tokens:] + assert right_visible is not None + right_visible = right_visible[num_decode_tokens:] if not swa_only: if self.compress_ratio == 4: @@ -1026,6 +1069,17 @@ def _forward_prefill( top_k, M, N, + max_image_tokens=self.max_image_tokens, + left_visible=( + left_visible[query_start:query_end] + if left_visible is not None + else None + ), + right_visible=( + right_visible[query_start:query_end] + if right_visible is not None + else None + ), ) rocm_sparse_attn_prefill( q=q[query_start:query_end], diff --git a/vllm/models/deepseek_v4/common/vl_model.py b/vllm/models/deepseek_v4/common/vl_model.py new file mode 100644 index 000000000000..4cccb85a8302 --- /dev/null +++ b/vllm/models/deepseek_v4/common/vl_model.py @@ -0,0 +1,331 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). + +Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: + +- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE + sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / + ``image_newline`` / ``image_end``) fill the remaining sentinel positions. +- Image placeholders (``<|deepseek_image|>``) are expanded by the processor + in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved + in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` + (see ``common/vision.py`` for the tower itself). +- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before + its hyper-connection stream expansion. Raw ``input_ids`` still flow into + the model so the MoE router can apply ``bias_vl`` to image tokens + (``requires_raw_input_tokens``). +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .mm_preprocess import ( + IMAGE_PLACEHOLDER, + IMAGE_SENTINEL_BASE_ID, + DeepseekV4VLDummyInputsBuilder, + DeepseekV4VLMultiModalProcessor, + DeepseekV4VLProcessingInfo, + image_sentinel_mask, +) +from .vision import DeepseekV4Aligner, DeepseekV4ViT + + +def _make_deepseek_v4_vl_weights_mapper( + text_mapper: WeightsMapper, image_enabled: bool +) -> WeightsMapper: + """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" + orig_to_new_prefix = { + src: None if dst is None else f"language_model.{dst}" + for src, dst in text_mapper.orig_to_new_prefix.items() + } + if not image_enabled: + orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) + return WeightsMapper( + orig_to_new_renaming=text_mapper.orig_to_new_renaming, + orig_to_new_prefix=orig_to_new_prefix, + orig_to_new_regex=text_mapper.orig_to_new_regex, + orig_to_new_stacked=text_mapper.orig_to_new_stacked, + orig_to_new_suffix={ + **text_mapper.orig_to_new_suffix, + "head.weight": "language_model.lm_head.weight", + }, + orig_to_new_substr={ + **text_mapper.orig_to_new_substr, + # Draft models load the checkpoint's MTP weights separately from + # the target model. + "mtp.": None, + }, + ) + + +@MULTIMODAL_REGISTRY.register_processor( + DeepseekV4VLMultiModalProcessor, + info=DeepseekV4VLProcessingInfo, + dummy_inputs=DeepseekV4VLDummyInputsBuilder, +) +class DeepseekV4ForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 +): + """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. + + ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) + delegates through ``language_model`` via the protocol defaults. + """ + + # The MoE router needs raw token ids to detect image sentinel tokens + # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. + requires_raw_input_tokens = True + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return IMAGE_PLACEHOLDER + raise ValueError(f"Unsupported modality: {modality!r}") + + def __init__(self, *, vllm_config, prefix: str = "") -> None: + super().__init__() + model_config = vllm_config.model_config + config = model_config.hf_config + self.config = config + self.multimodal_config = model_config.multimodal_config + assert self.multimodal_config is not None + + image_enabled = ( + config.vision_n_layers > 0 + and self.multimodal_config.get_limit_per_prompt("image") > 0 + ) + with self._mark_tower_model(vllm_config, {"image"}): + self.vision: DeepseekV4ViT | None = None + self.aligner: DeepseekV4Aligner | None = None + self.image_start: nn.Parameter | None = None + self.image_end: nn.Parameter | None = None + self.image_newline: nn.Parameter | None = None + self.image_pad: nn.Parameter | None = None + if image_enabled: + self.vision = DeepseekV4ViT(config) + self.aligner = DeepseekV4Aligner(config) + for name in ( + "image_start", + "image_end", + "image_newline", + "image_pad", + ): + setattr( + self, + name, + nn.Parameter( + torch.empty(config.hidden_size, dtype=torch.float32) + ), + ) + self.vision.to(dtype=model_config.dtype) + self.aligner.to(dtype=model_config.dtype) + + with self._mark_language_model(vllm_config): + # The arch convertor routes any config with a vision tower to + # this wrapper class; mark the copy handed to the text backbone + # so it resolves to DeepseekV4ForCausalLM instead of recursing + # (with_hf_config deepcopies the config, the marker survives). + config._dsv4_vl_inner = True # type: ignore[attr-defined] + try: + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["DeepseekV4ForCausalLM"], + ) + finally: + del config._dsv4_vl_inner # type: ignore[attr-defined] + # The outer mapper (see load_weights) fully resolves HF names into + # this wrapper's namespace before AutoWeightsLoader strips the + # "language_model." prefix and delegates to the child's load_weights, + # so the child's own mapper must be a no-op. Its suffix rules are not + # idempotent (e.g. "lm_head.weight".endswith("head.weight") would + # re-fire "head.weight" -> "lm_head.weight"). + text_mapper = self.language_model.hf_to_vllm_mapper + self.language_model.hf_to_vllm_mapper = WeightsMapper() + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + + self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( + text_mapper, image_enabled + ) + self._weights_finalized = False + + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: + patches = kwargs.pop("patches", None) + if patches is None: + return None + vit_grid = kwargs.pop("vit_grid", None) + llm_grid = kwargs.pop("llm_grid", None) + perm = kwargs.pop("perm", None) + assert vit_grid is not None and llm_grid is not None and perm is not None + return { + "patches": patches, + "vit_grid": vit_grid, + "llm_grid": llm_grid, + "perm": perm, + } + + def _encode_image( + self, + patches: torch.Tensor, + n_vit_h: int, + n_vit_w: int, + perm: torch.Tensor, + ) -> torch.Tensor: + assert self.vision is not None and self.aligner is not None + image_embeds = self.aligner( + self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w + ) + # Reorder into the N-layout block order used in the prompt. + return image_embeds[perm.to(image_embeds.device)] + + def _process_image_input( + self, + patches: torch.Tensor, + vit_grid: torch.Tensor, + llm_grid: torch.Tensor, + perm: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + assert self.vision is not None and self.aligner is not None + patches = patches.to(self.aligner.w1.weight.dtype) + + embeds: list[torch.Tensor] = [] + vit_offset = 0 + llm_offset = 0 + for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( + vit_grid.tolist(), llm_grid.tolist(), strict=True + ): + n_vit = n_vit_h * n_vit_w + n_llm = n_llm_h * n_llm_w + embeds.append( + self._encode_image( + patches[vit_offset : vit_offset + n_vit], + n_vit_h, + n_vit_w, + perm[llm_offset : llm_offset + n_llm], + ) + ) + vit_offset += n_vit + llm_offset += n_llm + return tuple(embeds) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None or self.vision is None: + return [] + return self._process_image_input( + image_input["patches"], + image_input["vit_grid"], + image_input["llm_grid"], + image_input["perm"], + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + # All ids are in-vocab here: image-block sentinels are borrowed + # reserved tokens (their embedding rows are always overwritten below). + inputs_embeds = self.language_model.embed_input_ids(input_ids) + + if self.image_start is not None: + # Branch-free sentinel overwrite: safe inside compiled/captured + # regions (no data-dependent control flow). + sentinel_mask = image_sentinel_mask(input_ids) + if is_multimodal is not None: + # IMAGE positions get vision embeddings via the merge below. + sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) + table = torch.stack( + [ + self.image_start, + self.image_pad, + self.image_pad, + self.image_newline, + self.image_end, + ] + ).to(inputs_embeds.dtype) + idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) + inputs_embeds = torch.where( + sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.language_model.compute_logits_local(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.language_model.get_expert_mapping() + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" + return self.language_model.get_mtp_target_hidden_states() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + child_finalizes = getattr( + self.language_model, "finalizes_weights_during_load", True + ) + mapped = self.hf_to_vllm_mapper.apply(weights) + if child_finalizes: + # A child which finalizes inside load_weights must see all of its + # weights in one contiguous delegation from AutoWeightsLoader. + mapped = iter(sorted(mapped, key=lambda x: x[0])) + loader = AutoWeightsLoader(self) + loaded_params = loader.load_weights(mapped) + self._weights_finalized = child_finalizes + return loaded_params + + def process_weights_after_loading(self) -> None: + # Backbones such as the ROCm implementation require this to run only + # after the loader's generic per-layer quantization finalization. + if getattr(self, "_weights_finalized", False): + return + self.language_model.process_weights_after_loading() + self._weights_finalized = True diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 021001b9ce01..97e7e617fe9c 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1790,6 +1790,7 @@ class DeepseekV4ForCausalLM( SupportsLoRA, DeepseekV4MixtureOfExperts, ): + finalizes_weights_during_load = True model_cls = DeepseekV4Model # Default mapper assumes the original FP4-expert checkpoint layout. diff --git a/vllm/models/deepseek_v4/nvidia/vl_model.py b/vllm/models/deepseek_v4/nvidia/vl_model.py index bd40b7bce9e7..a66217f2fe65 100644 --- a/vllm/models/deepseek_v4/nvidia/vl_model.py +++ b/vllm/models/deepseek_v4/nvidia/vl_model.py @@ -1,333 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). +"""Compatibility imports for the platform-neutral DeepSeek-V4 vision model.""" -Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: - -- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE - sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / - ``image_newline`` / ``image_end``) fill the remaining sentinel positions. -- Image placeholders (``<|deepseek_image|>``) are expanded by the processor - in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved - in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` - (see ``common/vision.py`` for the tower itself). -- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before - its hyper-connection stream expansion. Raw ``input_ids`` still flow into - the model so the MoE router can apply ``bias_vl`` to image tokens - (``requires_raw_input_tokens``). -""" - -from collections.abc import Iterable - -import torch -from torch import nn - -from vllm.model_executor.models.interfaces import ( - MultiModalEmbeddings, - SupportsEagle3, - SupportsMultiModal, - SupportsPP, +from ..common.vl_model import ( + DeepseekV4ForConditionalGeneration, ) -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - WeightsMapper, - init_vllm_registered_model, - maybe_prefix, +from ..common.vl_model import ( + _make_deepseek_v4_vl_weights_mapper as _make_common_vl_weights_mapper, ) -from vllm.multimodal import MULTIMODAL_REGISTRY - -from ..common.mm_preprocess import ( - IMAGE_PLACEHOLDER, - IMAGE_SENTINEL_BASE_ID, - DeepseekV4VLDummyInputsBuilder, - DeepseekV4VLMultiModalProcessor, - DeepseekV4VLProcessingInfo, - image_sentinel_mask, -) -from ..common.vision import DeepseekV4Aligner, DeepseekV4ViT from .model import _make_deepseek_v4_weights_mapper -def _make_deepseek_v4_vl_weights_mapper( - expert_dtype: str, image_enabled: bool -) -> WeightsMapper: - """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" - base = _make_deepseek_v4_weights_mapper(expert_dtype) - orig_to_new_prefix: dict[str, str | None] = { - "layers.": "language_model.model.layers.", - "embed.": "language_model.model.embed.", - "norm.": "language_model.model.norm.", - "hc_head": "language_model.model.hc_head", - "mtp.": "language_model.model.mtp.", - } - if not image_enabled: - orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) - return WeightsMapper( - orig_to_new_prefix=orig_to_new_prefix, - orig_to_new_regex=base.orig_to_new_regex, - orig_to_new_suffix={ - "head.weight": "language_model.lm_head.weight", - "embed.weight": "embed_tokens.weight", - ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", - }, - orig_to_new_substr={ - ".shared_experts.w2": ".shared_experts.down_proj", - # The MTP/DSpark draft heads are not supported for the vision - # variant; drop their weights. - "mtp.": None, - }, +def _make_deepseek_v4_vl_weights_mapper(expert_dtype: str, image_enabled: bool): + """Retain the original NVIDIA helper signature for downstream imports.""" + return _make_common_vl_weights_mapper( + _make_deepseek_v4_weights_mapper(expert_dtype), image_enabled ) -@MULTIMODAL_REGISTRY.register_processor( - DeepseekV4VLMultiModalProcessor, - info=DeepseekV4VLProcessingInfo, - dummy_inputs=DeepseekV4VLDummyInputsBuilder, -) -class DeepseekV4ForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 -): - """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. - - ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) - delegates through ``language_model`` via the protocol defaults. - """ - - # The MoE router needs raw token ids to detect image sentinel tokens - # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. - requires_raw_input_tokens = True - - @classmethod - def get_placeholder_str(cls, modality: str, i: int) -> str | None: - if modality == "image": - return IMAGE_PLACEHOLDER - raise ValueError(f"Unsupported modality: {modality!r}") - - def __init__(self, *, vllm_config, prefix: str = "") -> None: - super().__init__() - model_config = vllm_config.model_config - config = model_config.hf_config - self.config = config - self.multimodal_config = model_config.multimodal_config - assert self.multimodal_config is not None - - image_enabled = ( - config.vision_n_layers > 0 - and self.multimodal_config.get_limit_per_prompt("image") > 0 - ) - with self._mark_tower_model(vllm_config, {"image"}): - self.vision: DeepseekV4ViT | None = None - self.aligner: DeepseekV4Aligner | None = None - self.image_start: nn.Parameter | None = None - self.image_end: nn.Parameter | None = None - self.image_newline: nn.Parameter | None = None - self.image_pad: nn.Parameter | None = None - if image_enabled: - self.vision = DeepseekV4ViT(config) - self.aligner = DeepseekV4Aligner(config) - for name in ( - "image_start", - "image_end", - "image_newline", - "image_pad", - ): - setattr( - self, - name, - nn.Parameter( - torch.empty(config.hidden_size, dtype=torch.float32) - ), - ) - self.vision.to(dtype=model_config.dtype) - self.aligner.to(dtype=model_config.dtype) - - with self._mark_language_model(vllm_config): - # The arch convertor routes any config with a vision tower to - # this wrapper class; mark the copy handed to the text backbone - # so it resolves to DeepseekV4ForCausalLM instead of recursing - # (with_hf_config deepcopies the config, the marker survives). - config._dsv4_vl_inner = True # type: ignore[attr-defined] - try: - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - hf_config=config, - prefix=maybe_prefix(prefix, "language_model"), - architectures=["DeepseekV4ForCausalLM"], - ) - finally: - del config._dsv4_vl_inner # type: ignore[attr-defined] - # The outer mapper (see load_weights) fully resolves HF names into - # this wrapper's namespace before AutoWeightsLoader strips the - # "language_model." prefix and delegates to the child's load_weights, - # so the child's own mapper must be a no-op. Its suffix rules are not - # idempotent (e.g. "lm_head.weight".endswith("head.weight") would - # re-fire "head.weight" -> "lm_head.weight"). - self.language_model.hf_to_vllm_mapper = WeightsMapper() - self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] - self.language_model.make_empty_intermediate_tensors - ) - - expert_dtype = getattr(config, "expert_dtype", "fp4") - self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( - expert_dtype, image_enabled - ) - - def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: - patches = kwargs.pop("patches", None) - if patches is None: - return None - vit_grid = kwargs.pop("vit_grid", None) - llm_grid = kwargs.pop("llm_grid", None) - perm = kwargs.pop("perm", None) - assert vit_grid is not None and llm_grid is not None and perm is not None - return { - "patches": patches, - "vit_grid": vit_grid, - "llm_grid": llm_grid, - "perm": perm, - } - - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - perm: torch.Tensor, - ) -> torch.Tensor: - assert self.vision is not None and self.aligner is not None - image_embeds = self.aligner( - self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w - ) - # Reorder into the N-layout block order used in the prompt. - return image_embeds[perm.to(image_embeds.device)] - - def _process_image_input( - self, - patches: torch.Tensor, - vit_grid: torch.Tensor, - llm_grid: torch.Tensor, - perm: torch.Tensor, - ) -> tuple[torch.Tensor, ...]: - assert self.vision is not None and self.aligner is not None - patches = patches.to(self.aligner.w1.weight.dtype) - - embeds: list[torch.Tensor] = [] - vit_offset = 0 - llm_offset = 0 - for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( - vit_grid.tolist(), llm_grid.tolist(), strict=True - ): - n_vit = n_vit_h * n_vit_w - n_llm = n_llm_h * n_llm_w - embeds.append( - self._encode_image( - patches[vit_offset : vit_offset + n_vit], - n_vit_h, - n_vit_w, - perm[llm_offset : llm_offset + n_llm], - ) - ) - vit_offset += n_vit - llm_offset += n_llm - return tuple(embeds) - - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: - image_input = self._parse_and_validate_image_input(**kwargs) - if image_input is None or self.vision is None: - return [] - return self._process_image_input( - image_input["patches"], - image_input["vit_grid"], - image_input["llm_grid"], - image_input["perm"], - ) - - def embed_input_ids( - self, - input_ids: torch.Tensor, - multimodal_embeddings: MultiModalEmbeddings | None = None, - *, - is_multimodal: torch.Tensor | None = None, - ) -> torch.Tensor: - from vllm.model_executor.models.utils import _merge_multimodal_embeddings - - # All ids are in-vocab here: image-block sentinels are borrowed - # reserved tokens (their embedding rows are always overwritten below). - inputs_embeds = self.language_model.embed_input_ids(input_ids) - - if self.image_start is not None: - # Branch-free sentinel overwrite: safe inside compiled/captured - # regions (no data-dependent control flow). - sentinel_mask = image_sentinel_mask(input_ids) - if is_multimodal is not None: - # IMAGE positions get vision embeddings via the merge below. - sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) - table = torch.stack( - [ - self.image_start, - self.image_pad, - self.image_pad, - self.image_newline, - self.image_end, - ] - ).to(inputs_embeds.dtype) - idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) - inputs_embeds = torch.where( - sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds - ) - - if multimodal_embeddings is None or len(multimodal_embeddings) == 0: - return inputs_embeds - - assert is_multimodal is not None - return _merge_multimodal_embeddings( - inputs_embeds=inputs_embeds, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors=None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - return self.language_model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: - return self.language_model.compute_logits(hidden_states) - - def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.language_model.compute_logits_local(hidden_states) - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - return self.language_model.get_expert_mapping() - - def get_mtp_target_hidden_states(self) -> torch.Tensor | None: - """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" - return self.language_model.get_mtp_target_hidden_states() - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Map HF names into this wrapper's namespace up front and sort, so - # the "language_model." group reaches the child loader as one - # contiguous block (AutoWeightsLoader delegates per contiguous group, - # and the child's load_weights finalizes fused expert weights, which - # must not run on a partially loaded model). - mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0]) - loader = AutoWeightsLoader(self) - loaded_params = loader.load_weights(mapped) - # The child's load_weights already ran its post-load finalization. - self._weights_finalized = True - return loaded_params - - def process_weights_after_loading(self) -> None: - # Model-level post-load hook (called by the loader after any load - # format). Under DummyModelLoader the child's load_weights — and - # hence its finalize step — is bypassed, so run it here instead. - if getattr(self, "_weights_finalized", False): - return - self.language_model.process_weights_after_loading() +__all__ = ["DeepseekV4ForConditionalGeneration"] diff --git a/vllm/models/deepseek_v4/vl_stub.py b/vllm/models/deepseek_v4/vl_stub.py index 0f15a82f0190..4ade7b17d387 100644 --- a/vllm/models/deepseek_v4/vl_stub.py +++ b/vllm/models/deepseek_v4/vl_stub.py @@ -9,6 +9,6 @@ class DeepseekV4ForConditionalGeneration(nn.Module): def __init__(self, *, vllm_config, prefix: str = ""): super().__init__() raise NotImplementedError( - "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is only " - "supported on NVIDIA GPUs for now." + "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is not " + "supported on the current platform; use CUDA or ROCm." ) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 297efc632f92..0723075137ae 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -953,6 +953,21 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: if parallel_config.worker_cls == "auto": parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker" + model_config = vllm_config.model_config + scheduler_config = vllm_config.scheduler_config + # Note: model_config may be None during testing + if ( + model_config is not None + and model_config.is_mm_prefix_lm + and scheduler_config.is_multimodal_model + and not scheduler_config.disable_chunked_mm_input + ): + logger.warning_once( + "Forcing --disable_chunked_mm_input for models " + "with multimodal-bidirectional attention." + ) + scheduler_config.disable_chunked_mm_input = True + @classmethod def verify_model_arch(cls, model_arch: str) -> None: if model_arch in _ROCM_UNSUPPORTED_MODELS: