diff --git a/tests/ut/attention/test_attention_v1.py b/tests/ut/attention/test_attention_v1.py index b1033844f560..0c863e3c59a9 100644 --- a/tests/ut/attention/test_attention_v1.py +++ b/tests/ut/attention/test_attention_v1.py @@ -166,6 +166,7 @@ def test_unpadded_preserves_internal_seq_lens_cpu(self): block_table_tensor=torch.zeros((3, 1), dtype=torch.int32), slot_mapping=torch.arange(9, dtype=torch.int32), causal=True, + mm_req_doc_ranges={0: [(1, 2)], 1: [(3, 4)], 2: [(5, 6)]}, actual_seq_lengths_q=[2, 3, 4], positions=torch.arange(9), attn_state=AscendAttentionState.ChunkedPrefill, @@ -176,6 +177,7 @@ def test_unpadded_preserves_internal_seq_lens_cpu(self): self.assertTrue(torch.equal(unpadded_metadata._seq_lens_cpu, internal_seq_lens_cpu[:2])) self.assertIsNone(unpadded_metadata.seq_lens_cpu) + self.assertEqual(unpadded_metadata.mm_req_doc_ranges, {0: [(1, 2)], 1: [(3, 4)]}) @patch.object(AscendAttentionMetadataBuilder, "metadata_cls") def test_build(self, mock_ascend_metadata): diff --git a/tests/ut/attention/test_dsa_compressor_seqused.py b/tests/ut/attention/test_dsa_compressor_seqused.py new file mode 100644 index 000000000000..7538999f3070 --- /dev/null +++ b/tests/ut/attention/test_dsa_compressor_seqused.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_ascend.attention.dsa_v1 import ( + _update_compressor_seqused, + build_vision_bidirectional_swa_indices, +) + + +def _logical_slots( + indices: torch.Tensor, + block_table: torch.Tensor, + block_size: int, +) -> list[int]: + slot_to_pos = { + block_id * block_size + block_offset: block_num * block_size + block_offset + for block_num, block_id in enumerate(block_table.tolist()) + for block_offset in range(block_size) + } + return [slot_to_pos[slot] for slot in indices.tolist() if slot >= 0] + + +def test_compressor_seqused_masks_graph_padding_and_clears_stale_rows(): + buffer = torch.empty(6, dtype=torch.int32) + graph_cu_seqlens = torch.arange(7, dtype=torch.int32) + + first = _update_compressor_seqused(buffer, graph_cu_seqlens, num_reqs=6, num_actual_reqs=4) + assert first.tolist() == [1, 1, 1, 1, 0, 0] + + replay = _update_compressor_seqused(buffer, graph_cu_seqlens, num_reqs=6, num_actual_reqs=1) + assert replay.tolist() == [1, 0, 0, 0, 0, 0] + + idle = _update_compressor_seqused(buffer, graph_cu_seqlens, num_reqs=6, num_actual_reqs=0) + assert idle.tolist() == [0, 0, 0, 0, 0, 0] + + +def test_compressor_seqused_preserves_active_ragged_lengths(): + buffer = torch.empty(4, dtype=torch.int32) + cu_seqlens = torch.tensor([0, 2, 5, 6, 7], dtype=torch.int32) + + result = _update_compressor_seqused(buffer, cu_seqlens, num_reqs=4, num_actual_reqs=2) + + assert result.tolist() == [2, 3, 0, 0] + + +def test_image_queries_see_complete_block_and_text_stays_causal(): + block_size = 4 + block_table = torch.tensor([[7, 2, 9]], dtype=torch.int32) + + indices, _ = build_vision_bidirectional_swa_indices( + block_table=block_table, + window_size=3, + max_image_tokens=5, + block_size=block_size, + query_start_loc=torch.tensor([0, 12], dtype=torch.int32), + seq_lens=torch.tensor([12], dtype=torch.int32), + mm_prefix_ranges={0: [(4, 8)]}, + num_tokens=12, + ) + + assert _logical_slots(indices[3, 0], block_table[0], block_size) == [1, 2, 3] + assert _logical_slots(indices[4, 0], block_table[0], block_size) == list(range(2, 9)) + assert _logical_slots(indices[6, 0], block_table[0], block_size) == list(range(4, 9)) + assert _logical_slots(indices[8, 0], block_table[0], block_size) == list(range(4, 9)) + assert _logical_slots(indices[9, 0], block_table[0], block_size) == [7, 8, 9] + + +def test_multiple_images_and_batches_use_their_own_physical_blocks(): + block_size = 4 + block_table = torch.tensor([[3, 8], [11, 5]], dtype=torch.int32) + + indices, _ = build_vision_bidirectional_swa_indices( + block_table=block_table, + window_size=2, + max_image_tokens=4, + block_size=block_size, + query_start_loc=torch.tensor([0, 8, 16], dtype=torch.int32), + seq_lens=torch.tensor([8, 8], dtype=torch.int32), + mm_prefix_ranges={0: [(1, 2), (5, 7)], 1: [(3, 6)]}, + num_tokens=16, + ) + + assert _logical_slots(indices[5, 0], block_table[0], block_size) == [4, 5, 6, 7] + assert _logical_slots(indices[11, 0], block_table[1], block_size) == [2, 3, 4, 5, 6] + + +def test_image_span_must_fit_current_prefill_chunk(): + with pytest.raises(ValueError, match="single prefill chunk"): + build_vision_bidirectional_swa_indices( + block_table=torch.tensor([[1, 2]], dtype=torch.int32), + window_size=2, + max_image_tokens=4, + block_size=4, + query_start_loc=torch.tensor([0, 4], dtype=torch.int32), + seq_lens=torch.tensor([4], dtype=torch.int32), + mm_prefix_ranges={0: [(2, 5)]}, + num_tokens=4, + ) diff --git a/tests/ut/attention/test_dsa_v1.py b/tests/ut/attention/test_dsa_v1.py index 1166210bc108..37b0f20b39c4 100644 --- a/tests/ut/attention/test_dsa_v1.py +++ b/tests/ut/attention/test_dsa_v1.py @@ -43,6 +43,7 @@ AscendDSAMetadataBuilder, AscendDSAReqMetadata, build_compressor_metadata_out, + build_vision_bidirectional_swa_indices, ) from vllm_ascend.attention.utils import AscendCommonAttentionMetadata from vllm_ascend.device.device_op import DeviceOperator @@ -61,6 +62,46 @@ ) +def test_build_vision_bidirectional_swa_indices(): + indices, lengths = build_vision_bidirectional_swa_indices( + block_table=torch.tensor([[10, 11]], dtype=torch.int32), + window_size=2, + max_image_tokens=4, + block_size=4, + query_start_loc=torch.tensor([0, 8], dtype=torch.int32), + seq_lens=torch.tensor([8], dtype=torch.int32), + mm_prefix_ranges={0: [(2, 5)]}, + num_tokens=8, + ) + + assert indices.shape == (8, 1, 6) + assert lengths.tolist() == [1, 2, 5, 4, 4, 4, 2, 2] + assert indices[:, 0].tolist() == [ + [40, -1, -1, -1, -1, -1], + [40, 41, -1, -1, -1, -1], + [41, 42, 43, 44, 45, -1], + [42, 43, 44, 45, -1, -1], + [42, 43, 44, 45, -1, -1], + [42, 43, 44, 45, -1, -1], + [45, 46, -1, -1, -1, -1], + [46, 47, -1, -1, -1, -1], + ] + + +def test_build_vision_bidirectional_swa_indices_rejects_oversized_span(): + with pytest.raises(ValueError, match="exceeds vision_max_n_token"): + build_vision_bidirectional_swa_indices( + block_table=torch.tensor([[0, 1]], dtype=torch.int32), + window_size=2, + max_image_tokens=3, + block_size=4, + query_start_loc=torch.tensor([0, 4], dtype=torch.int32), + seq_lens=torch.tensor([4], dtype=torch.int32), + mm_prefix_ranges={0: [(0, 3)]}, + num_tokens=4, + ) + + def _mock_dsa_kv_plan(**method_returns) -> MagicMock: plan = MagicMock() plan.layout_kv = "PA_ND" diff --git a/tests/ut/models/test_deepseek_v4_moe.py b/tests/ut/models/test_deepseek_v4_moe.py index 24d2a4eec392..d4cc06bb7097 100644 --- a/tests/ut/models/test_deepseek_v4_moe.py +++ b/tests/ut/models/test_deepseek_v4_moe.py @@ -15,6 +15,9 @@ ) from vllm_ascend.models.deepseek_v4 import model as deepseek_v4_module +from vllm_ascend.ops.fused_moe.router.fused_topk_router import ( + select_deepseek_v4_vision_experts, +) class _FakeGate(nn.Module): @@ -113,3 +116,90 @@ def build_runner(**kwargs): with pytest.raises(ValueError, match="hash MoE routing requires input_ids"): moe(torch.randn(2, config.hidden_size)) + + +def test_deepseek_v4_vision_router_keeps_text_hash_and_applies_bias_vl(): + router_logits = torch.tensor( + [ + [0.0, 1.0, 2.0, 3.0], + [3.0, 2.0, 1.0, 0.0], + [0.0, 1.0, 2.0, 3.0], + ] + ) + input_ids = torch.tensor([7, 129257, 129259]) + tid2eid = torch.zeros(32, 2, dtype=torch.long) + tid2eid[7] = torch.tensor([1, 3]) + bias_vl = torch.tensor([0.0, 0.0, 10.0, 0.0]) + + weights, expert_ids = select_deepseek_v4_vision_experts( + router_logits=router_logits, + input_ids=input_ids, + tid2eid=tid2eid, + bias_vl=bias_vl, + text_bias=None, + top_k=2, + renormalize=True, + ) + + # In-vocabulary text keeps the checkpoint's deterministic hash route. + assert expert_ids[0].tolist() == [1, 3] + # All five sentinel ids use the dynamic vision route. + assert expert_ids[1].tolist() == [2, 0] + assert expert_ids[2].tolist() == [2, 3] + assert torch.allclose(weights.sum(dim=-1), torch.ones(3)) + + +def test_deepseek_v4_hash_vision_layer_exposes_bias_vl(monkeypatch): + gate = _FakeGate() + fused_moe = MagicMock(return_value=_FakeMoERunner(MagicMock())) + ep_group = SimpleNamespace( + device_group=SimpleNamespace(size=lambda: 1), + rank_in_group=0, + ) + config = SimpleNamespace( + hidden_act="silu", + hidden_size=8, + moe_intermediate_size=16, + n_routed_experts=4, + n_shared_experts=None, + norm_topk_prob=True, + num_experts_per_tok=2, + num_hash_layers=1, + routed_scaling_factor=1.5, + scoring_func="sqrtsoftplus", + swiglu_limit=10.0, + vision_n_layers=2, + vocab_size=32, + ) + parallel_config = SimpleNamespace( + enable_eplb=False, + eplb_config=SimpleNamespace(num_redundant_experts=0), + use_sequence_parallel_moe=False, + ) + + monkeypatch.setattr(deepseek_v4_module, "FusedMoEFactory", fused_moe) + monkeypatch.setattr(deepseek_v4_module, "ReplicatedLinear", lambda *args, **kwargs: gate) + monkeypatch.setattr(deepseek_v4_module, "get_ep_group", lambda: ep_group) + monkeypatch.setattr(deepseek_v4_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(deepseek_v4_module, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr( + deepseek_v4_module, + "get_ascend_config", + lambda: SimpleNamespace(mix_placement=False), + ) + monkeypatch.setattr(deepseek_v4_module.rocm_aiter_ops, "is_fused_moe_enabled", lambda: False) + monkeypatch.setattr( + deepseek_v4_module.rocm_aiter_ops, + "is_fusion_moe_shared_experts_enabled", + lambda: False, + ) + + moe = deepseek_v4_module.DeepseekV4MoE( + config=config, + parallel_config=parallel_config, + prefix="model.layers.0.mlp", + ) + + assert moe.gate.bias_vl.shape == (config.n_routed_experts,) + assert fused_moe.call_args.kwargs["bias_vl"] is moe.gate.bias_vl + assert fused_moe.call_args.kwargs["e_score_correction_bias"] is None diff --git a/tests/ut/models/test_deepseek_v4_vision.py b/tests/ut/models/test_deepseek_v4_vision.py new file mode 100644 index 000000000000..3f7fab0c495b --- /dev/null +++ b/tests/ut/models/test_deepseek_v4_vision.py @@ -0,0 +1,20 @@ +from unittest.mock import MagicMock + +from torch import nn +from vllm.model_executor.models.interfaces import supports_eagle3 + +from vllm_ascend.models.deepseek_v4.vl_model import ( + AscendDeepseekV4ForConditionalGeneration, +) + + +def test_vision_wrapper_exposes_dspark_aux_hidden_state_interface(): + model = AscendDeepseekV4ForConditionalGeneration.__new__(AscendDeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + language_model = MagicMock() + model.language_model = language_model + + assert supports_eagle3(model) + + model.set_aux_hidden_state_layers((41, 42, 43)) + language_model.set_aux_hidden_state_layers.assert_called_once_with((41, 42, 43)) diff --git a/tests/ut/models/test_deepseek_v4_vision_preprocess.py b/tests/ut/models/test_deepseek_v4_vision_preprocess.py new file mode 100644 index 000000000000..9a4b7d8e782f --- /dev/null +++ b/tests/ut/models/test_deepseek_v4_vision_preprocess.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project + +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import torch +from PIL import Image +from transformers import BatchFeature +from vllm.multimodal.parse import MultiModalDataParser +from vllm.multimodal.processing import PromptReplacement, PromptUpdateDetails + +from vllm_ascend.models.deepseek_v4 import mm_preprocess as prep +from vllm_ascend.models.deepseek_v4.mm_preprocess import ( + COMPRESS_PAD_TO, + IMAGE, + IMAGE_PAD, + DeepseekV4VLMultiModalProcessor, + DeepseekV4VLProcessor, +) + + +class _StubInfo: + def get_data_parser(self): + return MultiModalDataParser() + + def get_tokenizer(self): + return None + + +class _NonThreadSafeTokenizer: + def __init__(self, entered=None, release=None): + self._state_lock = threading.Lock() + self._active = False + self._entered = entered + self._release = release + + def __deepcopy__(self, memo): + return _NonThreadSafeTokenizer() + + def __call__(self, prompt, return_tensors, **kwargs): + with self._state_lock: + if self._active: + raise RuntimeError("Already borrowed") + self._active = True + try: + if self._entered is not None: + self._entered.set() + if self._release is not None: + self._release.wait(timeout=1) + time.sleep(0.01) + return {"input_ids": torch.tensor([[1]])} + finally: + with self._state_lock: + self._active = False + + +class _ConcurrentStubInfo(_StubInfo): + def __init__(self): + self.tokenizer = _NonThreadSafeTokenizer() + + def get_hf_processor(self, **kwargs): + return lambda **processor_kwargs: BatchFeature({}) + + def get_tokenizer(self): + return self.tokenizer + + +def test_local_image_processor_builds_vit_and_llm_inputs(): + config = SimpleNamespace( + vision_patch_size=14, + vision_downsample_ratio=3, + vision_max_n_token=384, + vision_min_pixels=147456, + vision_max_wh_ratio=8, + ) + output = DeepseekV4VLProcessor(config)(images=[Image.new("RGB", (128, 96), color=(10, 20, 30))]) + + assert output["patches"].dtype == torch.bfloat16 + assert output["patches"].shape[1:] == (3, 14, 14) + assert output["vit_grid"].shape == (1, 2) + assert output["llm_grid"].shape == (1, 2) + assert output["perm"].numel() == output["llm_grid"].prod().item() + + +def test_v027_prompt_updates_add_position_dependent_compress_pad(): + base = prep.IMAGE_SENTINEL_BASE_ID + image_token_id = 7 + n_llm_h, n_llm_w = 3, 2 + processor = DeepseekV4VLMultiModalProcessor(_StubInfo(), None) + + types, _ = prep.build_image_block_pad_free(n_llm_h, n_llm_w) + full = (base + types).tolist() + update = PromptReplacement( + modality="image", + target=[image_token_id], + replacement=PromptUpdateDetails.select_token_id(full, base + IMAGE), + ) + prompt = [11, 12, image_token_id, 13, 14, 15, image_token_id, 16] + mm_prompt_updates = {"image": [[update.resolve(0)], [update.resolve(1)]]} + + token_ids, placeholders = processor._apply_prompt_updates( + prompt, + mm_prompt_updates, + ) + + expected = [11, 12] + first_types, _ = prep.build_image_block( + n_llm_h, + n_llm_w, + len(expected), + ) + expected += (base + first_types).tolist() + expected += [13, 14, 15] + second_types, _ = prep.build_image_block( + n_llm_h, + n_llm_w, + len(expected), + ) + expected += (base + second_types).tolist() + expected += [16] + assert token_ids == expected + + for placeholder in placeholders["image"]: + pad = COMPRESS_PAD_TO - 1 - placeholder.start_idx % COMPRESS_PAD_TO + assert placeholder.tokens[:pad] == [base + IMAGE_PAD] * pad + assert placeholder.tokens[pad:] == full + assert placeholder.is_embed is not None + assert placeholder.is_embed.tolist() == [token == base + IMAGE for token in placeholder.tokens] + + +def test_hf_tokenizer_call_is_thread_safe(): + entered = threading.Event() + release = threading.Event() + info = _ConcurrentStubInfo() + info.tokenizer = _NonThreadSafeTokenizer(entered, release) + processor = DeepseekV4VLMultiModalProcessor(info, None) + + def process(): + return processor._call_hf_processor( + "prompt", + {"images": []}, + {}, + {}, + )["input_ids"] + + with ThreadPoolExecutor(max_workers=8) as pool: + competing_call = pool.submit( + info.tokenizer, + "chat template", + "pt", + ) + assert entered.wait(timeout=1) + outputs = list(pool.map(lambda _: process(), range(16))) + release.set() + competing_call.result() + + assert all(torch.equal(output, torch.tensor([[1]])) for output in outputs) diff --git a/tests/ut/patch/platform/test_patch_fused_moe.py b/tests/ut/patch/platform/test_patch_fused_moe.py index 3a3e284a3500..402a8b384f5b 100644 --- a/tests/ut/patch/platform/test_patch_fused_moe.py +++ b/tests/ut/patch/platform/test_patch_fused_moe.py @@ -150,3 +150,37 @@ def test_factory_shares_upstream_hash_table_with_legacy_ascend_routing(): kwargs = original_factory.call_args.kwargs assert kwargs["hash_indices_table"] is hash_indices_table assert kwargs["routed_experts_args"]["tid2eid"] is hash_indices_table + + +def test_factory_keeps_vision_bias_in_ascend_router_only(): + bias_vl = torch.arange(8, dtype=torch.float32) + router = _Router() + runner = SimpleNamespace(router=router) + original_factory = MagicMock(return_value=runner) + router_factory = MagicMock(return_value=router) + ascend_config = SimpleNamespace( + eplb_config=SimpleNamespace( + dynamic_eplb=False, + expert_map_path=None, + num_redundant_experts=0, + ) + ) + + with ( + patch.object(patch_fused_moe, "_original_FusedMoE", original_factory), + patch.object(patch_fused_moe, "create_ascend_fused_moe_router", router_factory), + patch.object(patch_fused_moe, "get_ascend_config", return_value=ascend_config), + ): + patch_fused_moe._ascend_FusedMoE( + num_experts=8, + top_k=2, + bias_vl=bias_vl, + image_sentinel_lo=129257, + ) + + router_kwargs = router_factory.call_args.kwargs + assert router_kwargs["bias_vl"] is bias_vl + assert router_kwargs["image_sentinel_lo"] == 129257 + factory_kwargs = original_factory.call_args.kwargs + assert "bias_vl" not in factory_kwargs + assert "image_sentinel_lo" not in factory_kwargs diff --git a/tests/ut/patch/platform/test_patch_speculative_config_dspark.py b/tests/ut/patch/platform/test_patch_speculative_config_dspark.py index aabc1d694e3f..31a7ebf7ae06 100644 --- a/tests/ut/patch/platform/test_patch_speculative_config_dspark.py +++ b/tests/ut/patch/platform/test_patch_speculative_config_dspark.py @@ -1,7 +1,14 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + from transformers import Qwen3Config +from vllm.config.model_arch import ModelArchitectureConfig from vllm.config.speculative import SpeculativeConfig import vllm_ascend.patch.platform.patch_speculative_config # noqa: F401 +from vllm_ascend.patch.platform.patch_speculative_config import ( + _normalize_deepseek_v4_dspark_draft, +) def test_legacy_qwen3_dspark_config_uses_qwen3_loader(): @@ -21,3 +28,51 @@ def test_legacy_qwen3_dspark_config_uses_qwen3_loader(): assert normalized.mask_token_id == 163824 assert normalized.target_layer_ids == [7, 23, 51, 67, 83] assert normalized.block_size == 7 + + +def test_deepseek_v4_vision_dspark_restores_draft_architecture(): + hf_config = SimpleNamespace( + model_type="deepseek_v4", + architectures=["DeepseekV4ForConditionalGeneration"], + dspark_target_layer_ids=[40, 41, 42], + ) + hf_config.update = lambda values: hf_config.__dict__.update(values) + model_arch_config = ModelArchitectureConfig( + architectures=["DeepseekV4ForConditionalGeneration"], + model_type="deepseek_v4", + text_model_type=None, + hidden_size=128, + total_num_hidden_layers=43, + total_num_attention_heads=8, + head_size=16, + vocab_size=1024, + total_num_kv_heads=8, + num_experts=256, + num_experts_per_token=8, + quantization_config=None, + is_deepseek_mla=True, + is_mm_prefix_lm=True, + rswa_window=128, + derived_max_model_len_and_key=(8192, "max_position_embeddings"), + ) + registry = MagicMock() + registry.inspect_model_cls.return_value = ("model-info", "DSparkDraftModel") + + class DraftModelConfig(SimpleNamespace): + @property + def architectures(self): + return self.model_arch_config.architectures + + draft_model_config = DraftModelConfig( + hf_config=hf_config, + model_arch_config=model_arch_config, + registry=registry, + ) + + _normalize_deepseek_v4_dspark_draft(draft_model_config) + + assert hf_config.architectures == ["DSparkDraftModel"] + assert draft_model_config.architectures == ["DSparkDraftModel"] + assert draft_model_config.model_arch_config.is_mm_prefix_lm is False + assert draft_model_config._architecture == "DSparkDraftModel" + registry.inspect_model_cls.assert_called_once_with(["DSparkDraftModel"], draft_model_config) diff --git a/tests/ut/quantization/test_modelslim_config.py b/tests/ut/quantization/test_modelslim_config.py index 39761ef374d4..2aed199a86e1 100644 --- a/tests/ut/quantization/test_modelslim_config.py +++ b/tests/ut/quantization/test_modelslim_config.py @@ -537,6 +537,27 @@ def test_apply_mapper_preserves_optional_metadata(self): class TestQuantPrefixMapper(TestBase): + def test_deepseek_v4_vision_maps_wrapped_language_model_prefixes(self): + config = AscendModelSlimConfig( + { + "model.embed_tokens.weight": "W8A8_DYNAMIC", + "model.layers.0.self_attn.q_proj.weight": "W8A8_DYNAMIC", + "lm_head.weight": "FLOAT", + } + ) + + cases = { + "language_model.model.embed_tokens": "model.embed_tokens", + "language_model.model.layers.0.self_attn.q_proj": ("model.layers.0.self_attn.q_proj"), + "language_model.lm_head": "lm_head", + } + for prefix, expected in cases.items(): + with self.subTest(prefix=prefix): + self.assertEqual( + config.quant_prefix_mapper("deepseek_v4", prefix), + expected, + ) + def test_qwen3_5_text_backbones_use_packed_module_mappings(self): dense_mapping = get_packed_modules_mapping("qwen3_5_text") moe_mapping = get_packed_modules_mapping("qwen3_5_moe_text") diff --git a/tests/ut/spec_decode/test_llm_base_proposer.py b/tests/ut/spec_decode/test_llm_base_proposer.py index 145e1ca9a133..50ff1bd1211d 100644 --- a/tests/ut/spec_decode/test_llm_base_proposer.py +++ b/tests/ut/spec_decode/test_llm_base_proposer.py @@ -109,6 +109,15 @@ def test_default_uses_image_token_index(self): assert image_token_index == 456 + def test_model_with_multiple_image_sentinels_needs_no_single_index(self): + config = SimpleNamespace() + + image_token_index = AscendSpecDecodeBaseProposer._get_multimodal_image_token_index( + "AscendDeepseekV4ForConditionalGeneration", config + ) + + assert image_token_index is None + class TestMtpSharesTheTargetLmHead: """``_maybe_share_lm_head`` for the MTP branch. diff --git a/tests/ut/worker/test_model_runner_v1.py b/tests/ut/worker/test_model_runner_v1.py index 9ea49bd0fe60..bf08572a2ada 100644 --- a/tests/ut/worker/test_model_runner_v1.py +++ b/tests/ut/worker/test_model_runner_v1.py @@ -173,6 +173,59 @@ def forward_context(*args, **kwargs): self.assertEqual(events, ["context_enter", "forward", "context_exit", "release"]) +class TestMultimodalPrefillCompilationGuard(unittest.TestCase): + def _build_runner(self, *, mm_features, num_computed_tokens, num_prompt_tokens): + runner = NPUModelRunner.__new__(NPUModelRunner) + runner.requests = { + "req": SimpleNamespace(mm_features=mm_features), + } + runner.input_batch = SimpleNamespace( + req_id_to_index={"req": 0}, + num_computed_tokens_cpu=np.array([num_computed_tokens], dtype=np.int32), + num_prompt_tokens=np.array([num_prompt_tokens], dtype=np.int32), + ) + return runner + + def test_scheduled_encoder_input_skips_compiled_backbone(self): + runner = self._build_runner(mm_features=[], num_computed_tokens=0, num_prompt_tokens=0) + scheduler_output = SimpleNamespace( + scheduled_encoder_inputs={"req": [0]}, + num_scheduled_tokens={}, + ) + + self.assertTrue(runner._should_skip_compiled_for_encoder_input(scheduler_output)) + + def test_initial_multimodal_prefill_skips_compiled_backbone(self): + runner = self._build_runner(mm_features=[object()], num_computed_tokens=0, num_prompt_tokens=237) + scheduler_output = SimpleNamespace(scheduled_encoder_inputs={}, num_scheduled_tokens={"req": 237}) + + self.assertTrue(runner._should_skip_compiled_for_encoder_input(scheduler_output)) + + def test_cache_hit_multimodal_prefill_skips_compiled_backbone(self): + runner = self._build_runner(mm_features=[object()], num_computed_tokens=224, num_prompt_tokens=237) + scheduler_output = SimpleNamespace(scheduled_encoder_inputs={}, num_scheduled_tokens={"req": 13}) + + self.assertTrue(runner._should_skip_compiled_for_encoder_input(scheduler_output)) + + def test_pure_text_prefill_keeps_compiled_backbone(self): + runner = self._build_runner(mm_features=[], num_computed_tokens=0, num_prompt_tokens=88) + scheduler_output = SimpleNamespace(scheduled_encoder_inputs={}, num_scheduled_tokens={"req": 88}) + + self.assertFalse(runner._should_skip_compiled_for_encoder_input(scheduler_output)) + + def test_multimodal_decode_keeps_compiled_backbone(self): + runner = self._build_runner(mm_features=[object()], num_computed_tokens=237, num_prompt_tokens=237) + scheduler_output = SimpleNamespace(scheduled_encoder_inputs={}, num_scheduled_tokens={"req": 1}) + + self.assertFalse(runner._should_skip_compiled_for_encoder_input(scheduler_output)) + + def test_unscheduled_multimodal_prefill_keeps_compiled_backbone(self): + runner = self._build_runner(mm_features=[object()], num_computed_tokens=0, num_prompt_tokens=237) + scheduler_output = SimpleNamespace(scheduled_encoder_inputs={}, num_scheduled_tokens={}) + + self.assertFalse(runner._should_skip_compiled_for_encoder_input(scheduler_output)) + + class TestDSparkAuxCaptureMode(unittest.TestCase): def _build_runner( self, @@ -1485,7 +1538,7 @@ def test_execute_model_skips_dump_start_for_dp_dummy_run( runner.execute_model(scheduler_output) - runner._dummy_run.assert_called_once_with(1) + runner._dummy_run.assert_called_once_with(1, num_actual_reqs=0) runner._start_dump_data.assert_not_called() @patch("vllm_ascend.worker.model_runner_v1.has_kv_transfer_group", return_value=False) diff --git a/vllm_ascend/attention/dsa_v1.py b/vllm_ascend/attention/dsa_v1.py index f6bbdf261cde..0b9e419d5879 100644 --- a/vllm_ascend/attention/dsa_v1.py +++ b/vllm_ascend/attention/dsa_v1.py @@ -274,6 +274,7 @@ class AscendDSAReqMetadata: full_compress_sin: torch.Tensor = None full_compress_cos: torch.Tensor = None start_pos: torch.Tensor | None = None + seqused: torch.Tensor | None = None num_actual_reqs: int | None = None sas_metadata: torch.Tensor = None qli_metadata: torch.Tensor = None @@ -284,6 +285,7 @@ class AscendDSAReqMetadata: ori_win_left: int | None = None ori_win_right: int | None = None dspark_swa_indices: torch.Tensor | None = None + vision_swa_indices: torch.Tensor | None = None @dataclass @@ -327,6 +329,20 @@ def _require_req_metadata(metadata: AscendDSAMetadata) -> AscendDSAReqMetadata: return metadata.req_metadata +def _update_compressor_seqused( + buffer: torch.Tensor, + query_start_loc: torch.Tensor, + num_reqs: int, + num_actual_reqs: int | None, +) -> torch.Tensor: + """Update the stable per-batch compressor token-count buffer.""" + num_actual_reqs = num_reqs if num_actual_reqs is None else min(num_actual_reqs, num_reqs) + buffer[:num_reqs].fill_(0) + if num_actual_reqs > 0: + buffer[:num_actual_reqs].copy_(query_start_loc[1 : num_actual_reqs + 1] - query_start_loc[:num_actual_reqs]) + return buffer[:num_reqs] + + def get_dspark_sparse_sas_window(vllm_config: Any) -> tuple[int, int]: hf_config = vllm_config.model_config.hf_config window_size = int(hf_config.sliding_window) @@ -400,6 +416,97 @@ def build_dspark_swa_indices( return per_token_slots, per_token_lens +def build_vision_bidirectional_swa_indices( + block_table: torch.Tensor, + window_size: int, + max_image_tokens: int, + block_size: int, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + mm_prefix_ranges: dict[int, list[tuple[int, int]]], + num_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build paged original-KV indices with bidirectional image spans. + + Ranges are inclusive absolute token positions. Tokens outside an image + keep the normal causal sliding window. A token inside an image sees the + union of that causal window and its complete image span. The fixed output + width is ``window_size + max_image_tokens`` so it is graph- and + operator-workspace friendly. + """ + if max_image_tokens <= 0: + raise ValueError("max_image_tokens must be positive for vision SWA") + + query_lens = query_start_loc[1:] - query_start_loc[:-1] + if int(query_lens.sum().item()) != num_tokens: + raise ValueError("Vision SWA must contain one query position per token") + req_ids = torch.repeat_interleave( + torch.arange( + query_lens.shape[0], + device=query_start_loc.device, + dtype=torch.long, + ), + query_lens, + output_size=num_tokens, + ) + token_offsets = torch.arange(num_tokens, device=query_start_loc.device) - query_start_loc[req_ids] + positions = seq_lens[req_ids] - query_lens[req_ids] + token_offsets + start_positions = (positions - int(window_size) + 1).clamp_min(0) + end_positions = positions.clone() + + for req_idx, ranges in mm_prefix_ranges.items(): + if req_idx >= query_lens.shape[0]: + continue + for span_start, span_end in ranges: + if span_end < span_start: + raise ValueError(f"Invalid image span [{span_start}, {span_end}]") + if span_end - span_start + 1 > max_image_tokens: + raise ValueError( + f"Image span exceeds vision_max_n_token: span=[{span_start}, {span_end}], max={max_image_tokens}" + ) + in_span = (req_ids == req_idx) & (positions >= span_start) & (positions <= span_end) + if bool(in_span.any().item()) and span_end >= int(seq_lens[req_idx].item()): + raise ValueError( + "Image spans must be scheduled in a single prefill chunk before bidirectional attention is built" + ) + start_positions = torch.where( + in_span, + torch.minimum( + start_positions, + start_positions.new_full((), span_start), + ), + start_positions, + ) + end_positions = torch.where( + in_span, + torch.maximum( + end_positions, + end_positions.new_full((), span_end), + ), + end_positions, + ) + + visible_lens = end_positions - start_positions + 1 + index_width = int(window_size) + int(max_image_tokens) + columns = torch.arange(index_width, device=block_table.device) + visible = columns.unsqueeze(0) < visible_lens.unsqueeze(1) + visible_positions = start_positions.unsqueeze(1) + columns.unsqueeze(0) + block_numbers = visible_positions // int(block_size) + safe_block_numbers = block_numbers.clamp( + min=0, + max=block_table.shape[1] - 1, + ) + request_block_tables = block_table[req_ids] + block_ids = torch.gather( + request_block_tables, + 1, + safe_block_numbers, + ) + slot_ids = (block_ids * int(block_size) + visible_positions % int(block_size)).to(torch.int32) + slot_ids = slot_ids.where(visible, torch.full_like(slot_ids, -1)) + return slot_ids.unsqueeze(1), visible_lens.to(torch.int32) + + class AscendDSAMetadataBuilder(AttentionMetadataBuilder[AscendDSAMetadata]): """ NOTE: Please read the comment at the top of the file before trying to @@ -469,6 +576,9 @@ def __init__( self.start_pos_prefill: torch.Tensor = torch.zeros( scheduler_config.max_num_seqs, dtype=torch.int32, device=self.device ) + self.compressor_seqused: torch.Tensor = torch.zeros( + scheduler_config.max_num_seqs, dtype=torch.int32, device=self.device + ) self.sas_metadata_buffer: torch.Tensor = torch.zeros( DSA_METADATA_BUFFER_SIZE, dtype=torch.int32, device=self.device ) @@ -823,10 +933,17 @@ def build_req_metadata( if num_actual_reqs < num_reqs: self.start_pos_prefill[num_actual_reqs:num_reqs].fill_(0) self.block_table[num_actual_reqs:num_reqs, ...].fill_(0) + compressor_seqused = _update_compressor_seqused( + self.compressor_seqused, + query_start_loc, + num_reqs, + num_actual_reqs, + ) layer_name = f"c{self.compressor_ratio}" cu_seqlens_ori_kv = None cu_seqlens_cmp_kv = None dspark_swa_indices = None + vision_swa_indices = None ori_win_left, ori_win_right = self.model_config.hf_config.sliding_window - 1, 0 if not has_prefill and not common_attn_metadata.causal: # DSpark non-causal parallel drafting: every draft query attends to @@ -846,6 +963,34 @@ def build_req_metadata( ) dspark_swa_indices = dspark_swa_indices[: self.num_decode_tokens] ori_win_left, ori_win_right = get_dspark_sparse_sas_window(self.vllm_config) + # Text-only requests and lightweight metadata fixtures do not carry + # multimodal document ranges. Treat those as having no vision spans. + mm_ranges = getattr(common_attn_metadata, "mm_req_doc_ranges", None) + max_image_tokens = ( + getattr( + self.model_config.hf_config, + "vision_max_n_token", + 0, + ) + if getattr( + self.model_config.hf_config, + "vision_n_layers", + 0, + ) + > 0 + else 0 + ) + if has_prefill and max_image_tokens > 0 and mm_ranges: + vision_swa_indices, _ = build_vision_bidirectional_swa_indices( + block_table=self.block_table[:num_actual_reqs], + window_size=self.model_config.hf_config.sliding_window, + max_image_tokens=max_image_tokens, + block_size=self.storage_block_size, + query_start_loc=query_start_loc[: num_actual_reqs + 1], + seq_lens=seq_lens[:num_actual_reqs], + mm_prefix_ranges=mm_ranges, + num_tokens=self.num_actual_tokens, + ) if not has_prefill and self.common_ratio_to_sas_metadata.get(layer_name) is None: cu_seqlens_ori_kv = DeviceOperator.get_dsa_decode_cu_seqlens_ori_kv( self.common_ratio_to_sas_metadata, @@ -939,6 +1084,7 @@ def build_qli_metadata() -> None: full_compress_sin=full_compress_sin, full_compress_cos=full_compress_cos, start_pos=self.start_pos_prefill[:num_reqs], + seqused=compressor_seqused, num_actual_reqs=num_actual_reqs, sas_metadata=sas_metadata, qli_metadata=qli_metadata, @@ -947,6 +1093,7 @@ def build_qli_metadata() -> None: ori_win_left=ori_win_left, ori_win_right=ori_win_right, dspark_swa_indices=dspark_swa_indices, + vision_swa_indices=vision_swa_indices, ) if self._device_metadata_enabled and self.compressor_metadata_buffers is not None: assert num_compressed_tokens is not None @@ -1927,10 +2074,17 @@ def compressor_tail_fn() -> CompressorForwardOutput: layout_kv=_dsa_layout_kv(self.vllm_config), ) - if self.compress_ratio <= 1: + # Vision prefill uses explicit original-KV indices so tokens inside an + # image span can see the complete span bidirectionally. Compressed KV + # selection remains active and is supplied independently below. + if swa_req_metadata.vision_swa_indices is not None: + attn_kwargs["ori_sparse_indices"] = swa_req_metadata.vision_swa_indices + + if swa_req_metadata.vision_swa_indices is None: if swa_req_metadata.dspark_swa_indices is not None: attn_kwargs["ori_sparse_indices"] = swa_req_metadata.dspark_swa_indices - else: + + if self.compress_ratio > 1: assert compressor_metadata is not None attn_kwargs.update( cmp_kv=compress_kv_cache, diff --git a/vllm_ascend/attention/utils.py b/vllm_ascend/attention/utils.py index 294b21c9afbf..86ac067964d5 100644 --- a/vllm_ascend/attention/utils.py +++ b/vllm_ascend/attention/utils.py @@ -330,6 +330,15 @@ def _slice_reqs(x): block_table_tensor=self.block_table_tensor, slot_mapping=self.slot_mapping, causal=self.causal, + mm_req_doc_ranges=( + { + req_idx: doc_ranges + for req_idx, doc_ranges in self.mm_req_doc_ranges.items() + if req_idx < num_actual_reqs + } + if self.mm_req_doc_ranges is not None + else None + ), actual_seq_lengths_q=self.actual_seq_lengths_q[:num_actual_tokens], positions=self.positions, positions_cpu=self.positions_cpu, diff --git a/vllm_ascend/models/__init__.py b/vllm_ascend/models/__init__.py index 5b33f6489dc9..f0a0414d7958 100644 --- a/vllm_ascend/models/__init__.py +++ b/vllm_ascend/models/__init__.py @@ -31,6 +31,10 @@ def register_model(): ModelRegistry.register_model( "DeepseekV4ForCausalLM", "vllm_ascend.models.deepseek_v4.model:AscendDeepseekV4ForCausalLM" ) + ModelRegistry.register_model( + "DeepseekV4ForConditionalGeneration", + "vllm_ascend.models.deepseek_v4.vl_model:AscendDeepseekV4ForConditionalGeneration", + ) ModelRegistry.register_model( "MiniMaxM3SparseForCausalLM", "vllm_ascend.models.minimax_m3:MiniMaxM3SparseForCausalLM", diff --git a/vllm_ascend/models/deepseek_v4/compressor.py b/vllm_ascend/models/deepseek_v4/compressor.py index aad940685f6d..8d07bac198d9 100644 --- a/vllm_ascend/models/deepseek_v4/compressor.py +++ b/vllm_ascend/models/deepseek_v4/compressor.py @@ -234,7 +234,7 @@ def forward( compress_cos.view(-1, compress_cos.shape[-1]), state_block_table=state_metadata.block_table, cu_seqlens=compressor_metadata.query_start_loc, - seqused=None, + seqused=getattr(compressor_metadata, "seqused", None), start_pos=compressor_metadata.start_pos, rope_head_dim=self.rope_head_dim, cmp_ratio=self.compress_ratio, diff --git a/vllm_ascend/models/deepseek_v4/dspark.py b/vllm_ascend/models/deepseek_v4/dspark.py index 57f83c8b1ab1..4f0e39b1da46 100644 --- a/vllm_ascend/models/deepseek_v4/dspark.py +++ b/vllm_ascend/models/deepseek_v4/dspark.py @@ -456,6 +456,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if name.endswith(".scale"): name = name.replace(".scale", ".weight_scale") + # The multimodal checkpoint also contains one vision-router bias + # for each MTP/DSpark layer. DSpark runs only during text decode, + # so draft MoE gates intentionally do not expose ``bias_vl``. + # Do not alias it to the text correction bias: that would change + # text routing whenever speculative decoding is enabled. + if name.endswith(".e_score_correction_bias_vl") and name not in params_dict: + logger.info_once("Ignoring vision-only router bias while loading the text-only DSpark drafter") + continue + if ".experts." in name: for param_name, weight_name, expert_id, shard_id in expert_mapping: if weight_name not in name: diff --git a/vllm_ascend/models/deepseek_v4/mm_preprocess.py b/vllm_ascend/models/deepseek_v4/mm_preprocess.py new file mode 100644 index 000000000000..84da0a16fa41 --- /dev/null +++ b/vllm_ascend/models/deepseek_v4/mm_preprocess.py @@ -0,0 +1,540 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Multimodal preprocessing for the DeepSeek-V4 vision variants +(DeepSeek-V4-Flash-Vision-Exp). + +The image transform and sentinel-block construction are ported from the +official repository's ``image_processor.py`` so that token counts bit-match +the reference. Each ``<｜deepseek_image｜>`` placeholder in the prompt expands +to a variable-length block of sentinel tokens; only positions with +``type == IMAGE`` receive vision embeddings, the other sentinels are looked +up from learned embedding vectors in the model. + +Unlike the reference (which uses out-of-vocab ids ``vocab_size + type``), +the sentinel block borrows five consecutive reserved tokenizer tokens +(``<|place_holder_mm_span_0431|>`` .. ``_0435|>``): they are special tokens +the tokenizer never emits from plain text, so the ids stay in-vocabulary and +work with stock token validation, logprobs and detokenization. The ids are +pure markers — every sentinel position's embedding is overwritten with +vision/sentinel vectors before the decoder layers see it, exactly like the +reference's out-of-vocab scheme. +""" + +import copy +import math +import threading +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import numpy as np +import torch +from PIL import Image, ImageOps +from transformers import BatchFeature +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import ImageSize, MultiModalDataItems +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.multimodal.processing.processor import ( + MultiModalPromptUpdates, + PlaceholderFeaturesInfo, +) +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config + +IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5) +COMPRESS_PAD_TO = 4 + +IMAGE_PLACEHOLDER = "<｜deepseek_image｜>" + +# Sentinel roles borrow five consecutive ``<|place_holder_mm_span_XXXX|>`` +# tokens (reserved special tokens, never emitted from plain text). The order +# must match IMAGE_START..IMAGE_END above so that +# ``id == IMAGE_SENTINEL_BASE_ID + type``. +IMAGE_SENTINEL_BASE_ID = 129257 +IMAGE_SENTINEL_TOKEN_NAMES = ( + "<|place_holder_mm_span_0431|>", # IMAGE_START + "<|place_holder_mm_span_0432|>", # IMAGE_PAD + "<|place_holder_mm_span_0433|>", # IMAGE + "<|place_holder_mm_span_0434|>", # IMAGE_NEW_LINE + "<|place_holder_mm_span_0435|>", # IMAGE_END +) + +# Fast tokenizers temporarily mutate truncation and padding state during a +# call. Multimodal preprocessing runs in a thread pool while the renderer can +# use the same tokenizer concurrently, so keep an independent tokenizer +# backend per preprocessing thread. +_TOKENIZER_THREAD_LOCAL = threading.local() + + +def _get_thread_local_tokenizer(tokenizer): + cached = getattr(_TOKENIZER_THREAD_LOCAL, "tokenizer", None) + source_id = getattr(_TOKENIZER_THREAD_LOCAL, "source_id", None) + if cached is None or source_id != id(tokenizer): + cached = copy.deepcopy(tokenizer) + _TOKENIZER_THREAD_LOCAL.tokenizer = cached + _TOKENIZER_THREAD_LOCAL.source_id = id(tokenizer) + return cached + + +def image_sentinel_mask(token_ids: torch.Tensor) -> torch.Tensor: + """Boolean mask for image-block sentinel positions (in-vocab ids).""" + return (token_ids >= IMAGE_SENTINEL_BASE_ID) & ( + token_ids < IMAGE_SENTINEL_BASE_ID + len(IMAGE_SENTINEL_TOKEN_NAMES) + ) + + +def validate_image_sentinel_ids(tokenizer) -> None: + """Check the borrowed sentinel ids against the tokenizer.""" + for i, name in enumerate(IMAGE_SENTINEL_TOKEN_NAMES): + token_id = tokenizer.convert_tokens_to_ids(name) + if token_id != IMAGE_SENTINEL_BASE_ID + i: + raise ValueError( + f"Image sentinel token {name!r} has id {token_id}, expected " + f"{IMAGE_SENTINEL_BASE_ID + i}; the DeepSeek-V4 vision " + "sentinel block requires these consecutive reserved ids." + ) + + +def grid_tokens(best_height, best_width, patch_size, downsample_ratio): + """Number of LLM tokens the aligner grid occupies (N-layout, including + row/align padding).""" + n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio) + n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio) + num_tokens = n_llm_h * (n_llm_w + 1) + 2 + if n_llm_h % 2 == 1: + num_tokens += n_llm_w + 1 + num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 + return n_llm_h, n_llm_w, num_tokens + + +def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token): + r = height / width + max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5 + max_h_float = max_w_float * r + if max_w_float < 1.0: + max_w = 1 + max_h = (max_n_token - 2) // (max_w + 1) + if max_h % 2 == 1: + max_h -= 1 + best_width = max_w * patch_size * downsample_ratio + best_height = max_h * patch_size * downsample_ratio + elif max_h_float < 2.0: + max_h = 2 + max_w = ((max_n_token - 2) // max_h) - 1 + assert max_w > 1 + best_width = max_w * patch_size * downsample_ratio + best_height = max_h * patch_size * downsample_ratio + else: + max_w = math.floor(max_w_float) + max_h = math.floor(max_h_float) + if max_h % 2 == 1: + max_h -= 1 + beta = min( + max_w * patch_size * downsample_ratio / width, + max_h * patch_size * downsample_ratio / height, + ) + best_width = math.floor(width * beta / patch_size) * patch_size + best_height = math.floor(height * beta / patch_size) * patch_size + n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio) + return n_llm_h, n_llm_w, best_height, best_width, num_tokens + + +def safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token): + max_n_token -= COMPRESS_PAD_TO - 1 + n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio) + budget = max_n_token + while num_tokens > max_n_token: + n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio( + height, width, patch_size, downsample_ratio, budget + ) + budget -= 1 + return n_llm_h, n_llm_w, best_height, best_width + + +def load_image( + image: Image.Image, + *, + patch_size: int, + downsample_ratio: int, + max_n_token: int, + min_pixels: int, + max_wh_ratio: float | None, +): + """Transform one PIL image into ViT patches. + + Same math as the reference ``load_image``, except the image is already + decoded (vLLM supplies PIL images instead of a record dict). + """ + p = patch_size + image = image.convert("RGB") + width, height = image.size + if max_wh_ratio is not None and width > height * max_wh_ratio: + width = height * max_wh_ratio + if 0 < width * height < min_pixels: + ratio = (min_pixels / (width * height)) ** 0.5 + width = int(width * ratio) + height = int(height * ratio) + best_width = math.ceil(width / p) * p + best_height = math.ceil(height / p) * p + n_llm_h, n_llm_w, best_height, best_width = safe_resize( + height, width, best_height, best_width, p, downsample_ratio, max_n_token + ) + n_vit_h, n_vit_w = best_height // p, best_width // p + if max_wh_ratio is not None and image.width >= max_wh_ratio * image.height: + image = image.resize((best_width, best_height)) + else: + image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127)) + x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255 + x = ((x - 0.5) / 0.5).to(torch.bfloat16) + patches = x.reshape(3, n_vit_h, p, n_vit_w, p).permute(1, 3, 0, 2, 4).reshape(n_vit_h * n_vit_w, 3, p, p) + return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w + + +def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int): + """Builds the N-layout token types (final order) and the aligner-row order + for IMAGE slots.""" + compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO + pad_h = n_llm_h % 2 + rows = n_llm_h + pad_h + row_len = n_llm_w + 1 + pad_last = rows // 2 * row_len % 2 * 2 + types = torch.tensor( + ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + [IMAGE_PAD] * (row_len * pad_h), + dtype=torch.int64, + ) + order = torch.arange(rows * row_len).view(rows // 2, 2, row_len) + order = order.transpose(1, 2).reshape(-1) + image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64) + image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange(n_llm_h * n_llm_w).view(n_llm_h, n_llm_w) + perm = image_idx[order] + perm = perm[perm >= 0] + types = torch.cat( + [ + torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64), + torch.tensor([IMAGE_START]), + types[order], + torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64), + torch.tensor([IMAGE_END]), + ] + ) + return types, perm + + +def build_image_block_pad_free(n_llm_h: int, n_llm_w: int): + """``build_image_block`` without the position-dependent compressor pad. + + ``start_pos`` is chosen so that ``compress_pad == 0``; the pad is instead + prepended when the block is spliced into the final prompt, where its + position is known. + """ + return build_image_block(n_llm_h, n_llm_w, COMPRESS_PAD_TO - 1) + + +class DeepseekV4VLImageProcessor: + """Per-image transform (the PIL-input equivalent of the reference + ``load_image``).""" + + def __init__(self, config: DeepseekV4Config) -> None: + super().__init__() + self.patch_size = config.vision_patch_size + self.downsample_ratio = config.vision_downsample_ratio + self.max_n_token = config.vision_max_n_token + self.min_pixels = config.vision_min_pixels + self.max_wh_ratio = config.vision_max_wh_ratio + + def __call__(self, image: Image.Image): + return load_image( + image, + patch_size=self.patch_size, + downsample_ratio=self.downsample_ratio, + max_n_token=self.max_n_token, + min_pixels=self.min_pixels, + max_wh_ratio=self.max_wh_ratio, + ) + + +class DeepseekV4VLProcessor: + """Minimal stand-in for the HF processor of DeepSeek-V4 vision models. + + The official repository ships image preprocessing as plain functions in + ``image_processor.py`` (no ``auto_map`` processor), so this class wraps + their ports directly and the model loads without ``--trust-remote-code``. + + ``__call__`` returns a ``BatchFeature`` with one entry per image + (flattened across images): + + - ``patches``: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` bf16 ViT patches. + - ``vit_grid``: ``(num_images, 2)`` int64 ``[n_vit_h, n_vit_w]``. + - ``llm_grid``: ``(num_images, 2)`` int64 ``[n_llm_h, n_llm_w]``. + - ``perm``: concatenated per-image ``(n_llm_h * n_llm_w,)`` int64 index + selecting aligner outputs into the final N-layout order. + - ``types``: concatenated per-image pad-free sentinel block types; + ``block ids = IMAGE_SENTINEL_BASE_ID + types``. + """ + + def __init__(self, config: DeepseekV4Config) -> None: + super().__init__() + self.config = config + self.image_processor = DeepseekV4VLImageProcessor(config) + + def __call__( + self, + text: str | None = None, + images: Sequence[Image.Image] | None = None, + return_tensors: str | None = None, + **kwargs: Any, + ) -> BatchFeature: + patches_list = [] + vit_grid = [] + llm_grid = [] + perm_list = [] + types_list = [] + for image in images or []: + patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = self.image_processor(image) + types, perm = build_image_block_pad_free(n_llm_h, n_llm_w) + patches_list.append(patches) + vit_grid.append((n_vit_h, n_vit_w)) + llm_grid.append((n_llm_h, n_llm_w)) + perm_list.append(perm) + types_list.append(types) + + if not patches_list: + return BatchFeature({}) + + return BatchFeature( + { + "patches": torch.cat(patches_list), + "vit_grid": torch.tensor(vit_grid, dtype=torch.int64), + "llm_grid": torch.tensor(llm_grid, dtype=torch.int64), + "perm": torch.cat(perm_list), + "types": torch.cat(types_list), + } + ) + + +class DeepseekV4VLProcessingInfo(BaseProcessingInfo): + def get_hf_config(self) -> DeepseekV4Config: + return self.ctx.get_hf_config(DeepseekV4Config) + + def get_hf_processor(self, **kwargs: object) -> DeepseekV4VLProcessor: + if kwargs: + raise ValueError(f"Unexpected processor kwargs: {sorted(kwargs)}") + return DeepseekV4VLProcessor(self.get_hf_config()) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None} + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + # ``safe_resize`` reserves COMPRESS_PAD_TO - 1 tokens of + # vision_max_n_token for the compressor-alignment pad, so the full + # sentinel block is bounded by vision_max_n_token; the margin is kept + # in case that reservation changes. + return { + "image": self.get_hf_config().vision_max_n_token + COMPRESS_PAD_TO - 1, + } + + def get_image_placeholder_token_id(self) -> int: + token_id = self.get_tokenizer().convert_tokens_to_ids(IMAGE_PLACEHOLDER) + if token_id is None: + raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}") + return token_id + + def get_image_size_with_most_features(self) -> ImageSize: + hf_config = self.get_hf_config() + patch_size = hf_config.vision_patch_size + downsample_ratio = hf_config.vision_downsample_ratio + # A square maximizes the ViT patch count (area) within the token + # budget; solve the budget-derived size directly to keep the dummy + # image small. + budget = hf_config.vision_max_n_token - (COMPRESS_PAD_TO - 1) + side = budget * patch_size * downsample_ratio + _, _, best_h, best_w, _ = solve_resize_ratio(side, side, patch_size, downsample_ratio, budget) + return ImageSize(width=best_w, height=best_h) + + +class DeepseekV4VLDummyInputsBuilder(BaseDummyInputsBuilder[DeepseekV4VLProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return IMAGE_PLACEHOLDER * mm_counts.get("image", 0) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + size = self.info.get_image_size_with_most_features() + return { + "image": self._get_dummy_images( + width=size.width, + height=size.height, + num_images=mm_counts.get("image", 0), + overrides=cast(ImageDummyOptions | None, mm_options.get("image")), + ), + } + + +class DeepseekV4VLMultiModalProcessor(BaseMultiModalProcessor[DeepseekV4VLProcessingInfo]): + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + """Combine the local image transform with v0.27 tokenization.""" + processor = self.info.get_hf_processor(**mm_kwargs) + processed = processor( + text=prompt, + images=cast(Sequence[Image.Image] | None, mm_data.get("images")), + return_tensors="pt", + ) + tokenizer = _get_thread_local_tokenizer(self.info.get_tokenizer()) + tokenizer_outputs = tokenizer( + prompt, + return_tensors="pt", + **tok_kwargs, + ) + processed["input_ids"] = tokenizer_outputs["input_ids"] + return processed + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + del prompt_text, mm_items, hf_processor_mm_kwargs, tokenization_kwargs + # The local processor transforms images only; vLLM performs the + # placeholder replacement after tokenization. + return False + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + vit_grid = hf_inputs.get("vit_grid") + llm_grid = hf_inputs.get("llm_grid") + + if vit_grid is None or llm_grid is None: + empty = torch.empty(0, dtype=torch.long) + patch_sizes = perm_sizes = types_sizes = empty + else: + patch_sizes = vit_grid.prod(-1) + perm_sizes = llm_grid.prod(-1) + n_llm_h, n_llm_w = llm_grid[:, 0], llm_grid[:, 1] + # Pad-free block length; same formula as ``grid_tokens`` given + # the LLM grid. + types_sizes = ( + n_llm_h * (n_llm_w + 1) + 2 + (n_llm_h % 2) * (n_llm_w + 1) + (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 + ) + + return { + "patches": MultiModalFieldConfig.flat_from_sizes("image", patch_sizes), + "vit_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True), + "llm_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True), + "perm": MultiModalFieldConfig.flat_from_sizes("image", perm_sizes), + "types": MultiModalFieldConfig.flat_from_sizes("image", types_sizes, keep_on_cpu=True), + } + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + image_token_id = self.info.get_image_placeholder_token_id() + validate_image_sentinel_ids(self.info.get_tokenizer()) + image_embed_id = IMAGE_SENTINEL_BASE_ID + IMAGE + + def get_image_replacement(item_idx: int) -> PromptUpdateDetails: + types: torch.Tensor = out_mm_kwargs["image"][item_idx]["types"].data + full = (IMAGE_SENTINEL_BASE_ID + types).tolist() + return PromptUpdateDetails.select_token_id(full, image_embed_id) + + return [ + PromptReplacement( + modality="image", + target=[image_token_id], + replacement=get_image_replacement, + ), + ] + + def _apply_prompt_updates( + self, + token_ids: list[int], + mm_prompt_updates: MultiModalPromptUpdates, + ) -> tuple[ + list[int], + Mapping[str, list[PlaceholderFeaturesInfo]], + ]: + """Apply v0.27 prompt updates and inject position-dependent padding. + + vLLM main plans prompt replacements before rendering and exposes that + plan to model processors. v0.27 does not expose that private API, so + first use its supported base implementation, then prepend each image + block's compressor-alignment padding and rebuild placeholder offsets. + """ + new_token_ids, base_placeholders = super()._apply_prompt_updates( + token_ids, + mm_prompt_updates, + ) + placeholders: dict[str, list[PlaceholderFeaturesInfo]] = {modality: [] for modality in base_placeholders} + pad_id = IMAGE_SENTINEL_BASE_ID + IMAGE_PAD + ordered = sorted( + ( + placeholder.start_idx, + modality, + placeholder, + ) + for modality, items in base_placeholders.items() + for placeholder in items + ) + inserted = 0 + for _, modality, placeholder in ordered: + start_idx = placeholder.start_idx + inserted + tokens = list(placeholder.tokens) + is_embed = placeholder.is_embed + if modality == "image": + compress_pad = COMPRESS_PAD_TO - 1 - start_idx % COMPRESS_PAD_TO + new_token_ids[start_idx:start_idx] = [pad_id] * compress_pad + tokens = [pad_id] * compress_pad + tokens + original_mask = ( + is_embed + if is_embed is not None + else torch.ones( + len(placeholder.tokens), + dtype=torch.bool, + ) + ) + is_embed = torch.cat( + [ + torch.zeros(compress_pad, dtype=torch.bool), + original_mask, + ] + ) + inserted += compress_pad + + placeholders[modality].append( + PlaceholderFeaturesInfo( + modality=modality, + item_idx=placeholder.item_idx, + start_idx=start_idx, + tokens=tokens, + is_embed=is_embed, + ) + ) + + return new_token_ids, placeholders diff --git a/vllm_ascend/models/deepseek_v4/model.py b/vllm_ascend/models/deepseek_v4/model.py index 16f559757f19..06a16611ff64 100644 --- a/vllm_ascend/models/deepseek_v4/model.py +++ b/vllm_ascend/models/deepseek_v4/model.py @@ -80,6 +80,7 @@ from vllm_ascend.core.kv_cache_interface import AscendSlidingWindowMLASpec from vllm_ascend.models.deepseek_v4.compressor import Compressor from vllm_ascend.models.deepseek_v4.indexer import DeepseekV4Indexer +from vllm_ascend.models.deepseek_v4.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm_ascend.ops.dsa import AscendDeepseekSparseAttention, DSAModules from vllm_ascend.ops.rope_dsv4 import ComplexExpRotaryEmbedding from vllm_ascend.ops.triton.mul_add import muls_add_triton @@ -320,6 +321,15 @@ def __init__( ) self.hash = layer_idx < config.num_hash_layers and not is_draft_layer + self.gate.bias_vl = None + 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, + ) if self.hash: # Use zeros instead of empty to avoid garbage values causing # invalid memory access in dummy mode (--load-format="dummy") @@ -353,6 +363,8 @@ def __init__( routed_scaling_factor=self.routed_scaling_factor, swiglu_limit=self.swiglu_limit, e_score_correction_bias=self.gate.e_score_correction_bias, + bias_vl=self.gate.bias_vl, + image_sentinel_lo=IMAGE_SENTINEL_BASE_ID, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, @@ -1123,7 +1135,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if "rotary_emb.inv_freq" in name: continue - if ".gate.bias" in name: + if ".gate.bias_vl" in name: + # The parameter keeps the checkpoint name on Ascend. It is + # passed to the hash router as its vision-only correction + # bias, while text rows continue to use tid2eid. + pass + elif ".gate.bias" in name: name = name.replace(".gate.bias", ".gate.e_score_correction_bias") if "sink" in name: diff --git a/vllm_ascend/models/deepseek_v4/vision.py b/vllm_ascend/models/deepseek_v4/vision.py new file mode 100644 index 000000000000..2e93646a9ff2 --- /dev/null +++ b/vllm_ascend/models/deepseek_v4/vision.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 vision tower (ViT + aligner), replicated (no TP/DP sharding). + +Ported from the official reference implementation +(deepseek-ai/DeepSeek-V4-Flash-Vision-Exp). Weight names match the HF +checkpoint so no renaming is needed at load time. +""" + +from functools import lru_cache + +import torch +import torch.nn.functional as F +from torch import nn + + +@lru_cache(8) +def get_vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w) + wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w) + freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() + freqs = (freqs * inv_freq).flatten(1) + return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) + + +def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x1, x2 = x.float().chunk(2, dim=-1) + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype) + + +class DeepseekV4RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) + return (self.weight * x).to(dtype) + + +class DeepseekV4PatchEmbed(nn.Module): + def __init__(self, config): + super().__init__() + self.proj = nn.Linear(3 * config.vision_patch_size**2, config.vision_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x.flatten(1)) + + +class DeepseekV4VisionAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.n_heads = config.vision_n_heads + self.head_dim = config.vision_dim // config.vision_n_heads + self.wqkv = nn.Linear(config.vision_dim, 3 * config.vision_dim) + self.wo = nn.Linear(config.vision_dim, config.vision_dim) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + n = x.size(0) + q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in self.wqkv(x).chunk(3, dim=-1)) + q = apply_rotary(q, cos, sin) + k = apply_rotary(k, cos, sin) + o = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1)) + return self.wo(o.transpose(0, 1).reshape(n, -1)) + + +class DeepseekV4VisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.w1 = nn.Linear(config.vision_dim, 2 * config.vision_inter_dim, bias=False) + self.w2 = nn.Linear(config.vision_inter_dim, config.vision_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, up = self.w1(x).chunk(2, dim=-1) + return self.w2(F.silu(gate) * up) + + +class DeepseekV4VisionBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.norm1 = DeepseekV4RMSNorm(config.vision_dim) + self.attn = DeepseekV4VisionAttention(config) + self.norm2 = DeepseekV4RMSNorm(config.vision_dim) + self.mlp = DeepseekV4VisionMLP(config) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.norm1(x), cos, sin) + return x + self.mlp(self.norm2(x)) + + +class DeepseekV4ViT(nn.Module): + """DeepSeek-V4 ViT: full bidirectional attention per image, 2D RoPE.""" + + def __init__(self, config): + super().__init__() + self.rope_dim = config.vision_dim // config.vision_n_heads // 2 + self.rope_theta = config.vision_rope_theta + self.patch_embed = DeepseekV4PatchEmbed(config) + self.blocks = nn.ModuleList([DeepseekV4VisionBlock(config) for _ in range(config.vision_n_layers)]) + self.norm = DeepseekV4RMSNorm(config.vision_dim) + + def forward(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: + x = self.patch_embed(patches) + cos, sin = get_vision_cos_sin(n_vit_h, n_vit_w, self.rope_dim, self.rope_theta) + cos = cos.to(device=x.device) + sin = sin.to(device=x.device) + for block in self.blocks: + x = block(x, cos, sin) + return self.norm(x) + + +class DeepseekV4Aligner(nn.Module): + """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.""" + + def __init__(self, config): + super().__init__() + self.downsample_ratio = config.vision_downsample_ratio + in_dim = config.vision_dim * self.downsample_ratio**2 + self.w1 = nn.Linear(in_dim, config.hidden_size) + self.w2 = nn.Linear(config.hidden_size, config.hidden_size) + + def forward(self, x: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: + r = self.downsample_ratio + x = x.view(n_vit_h, n_vit_w, -1).permute(2, 0, 1) + x = F.pad(x, (0, -n_vit_w % r, 0, -n_vit_h % r)) + x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1) + return self.w2(F.gelu(self.w1(x))) diff --git a/vllm_ascend/models/deepseek_v4/vl_model.py b/vllm_ascend/models/deepseek_v4/vl_model.py new file mode 100644 index 000000000000..b9844dd51aae --- /dev/null +++ b/vllm_ascend/models/deepseek_v4/vl_model.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Ascend wrapper for DeepSeek-V4-Flash-Vision-Exp. + +The processor, sentinel layout, ViT, and aligner are shared with the upstream +vLLM implementation from vllm-project/vllm#54566. This module contains only +the Ascend language-backbone integration and weight-loading boundary. +""" + +from collections.abc import Iterable, Iterator + +import torch +from torch import nn +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY + +from vllm_ascend.models.deepseek_v4.mm_preprocess import ( + IMAGE_PLACEHOLDER, + IMAGE_SENTINEL_BASE_ID, + DeepseekV4VLDummyInputsBuilder, + DeepseekV4VLMultiModalProcessor, + DeepseekV4VLProcessingInfo, + image_sentinel_mask, +) +from vllm_ascend.models.deepseek_v4.model import AscendDeepseekV4ForCausalLM +from vllm_ascend.models.deepseek_v4.vision import ( + DeepseekV4Aligner, + DeepseekV4ViT, +) + + +def _vision_parameter_name(name: str) -> str | None: + """Map a checkpoint vision tensor to the wrapper parameter namespace.""" + if name.startswith("model."): + name = name.removeprefix("model.") + if name.startswith(("vision.", "aligner.", "image_")): + return name + return None + + +@MULTIMODAL_REGISTRY.register_processor( + DeepseekV4VLMultiModalProcessor, + info=DeepseekV4VLProcessingInfo, + dummy_inputs=DeepseekV4VLDummyInputsBuilder, +) +class AscendDeepseekV4ForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsEagle3, +): + """DeepSeek-V4 vision entry point using the Ascend text backbone.""" + + requires_raw_input_tokens = True + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + del i + 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 + if getattr(config, "vision_n_layers", 0) > 0: + config.mm_prefix_clamp_sliding_window = True + config.mm_prefix_span_leading_pad_modulus = 4 + 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): + self.language_model = AscendDeepseekV4ForCausalLM( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + ) + self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + + 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) + if vit_grid is None or llm_grid is None or perm is None: + raise ValueError("DeepSeek-V4 vision input requires patches, vit_grid, llm_grid, and perm.") + return { + "patches": patches, + "vit_grid": vit_grid, + "llm_grid": llm_grid, + "perm": perm, + } + + 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 + image_embeds = self.aligner( + self.vision( + patches[vit_offset : vit_offset + n_vit], + n_vit_h, + n_vit_w, + ), + n_vit_h, + n_vit_w, + ) + item_perm = perm[llm_offset : llm_offset + n_llm].to(image_embeds.device) + embeds.append(image_embeds[item_perm]) + 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, + ) + + inputs_embeds = self.language_model.embed_input_ids(input_ids) + if self.image_start is not None: + sentinel_mask = image_sentinel_mask(input_ids) + if is_multimodal is not None: + 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 + if is_multimodal is None: + raise ValueError("is_multimodal is required when merging image embeddings.") + 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: + del kwargs + 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 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: + return self.language_model.get_mtp_target_hidden_states() + + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + self.language_model.set_aux_hidden_state_layers(layers) + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + params = dict(self.named_parameters()) + loaded_vision: set[str] = set() + + def language_weights() -> Iterator[tuple[str, torch.Tensor]]: + for name, loaded_weight in weights: + vision_name = _vision_parameter_name(name) + if vision_name is None: + yield name, loaded_weight + continue + if vision_name not in params: + raise KeyError(f"Vision weight {name!r} has no parameter {vision_name!r}.") + param = params[vision_name] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded_vision.add(vision_name) + + loaded_language = self.language_model.load_weights(language_weights()) + return loaded_vision | {f"language_model.{name}" for name in loaded_language} + + def process_weights_after_loading(self) -> None: + hook = getattr( + self.language_model, + "process_weights_after_loading", + None, + ) + if hook is not None: + hook() diff --git a/vllm_ascend/ops/fused_moe/router/fused_topk_router.py b/vllm_ascend/ops/fused_moe/router/fused_topk_router.py index 5d70b53f9ac7..c09e002a5453 100644 --- a/vllm_ascend/ops/fused_moe/router/fused_topk_router.py +++ b/vllm_ascend/ops/fused_moe/router/fused_topk_router.py @@ -24,6 +24,57 @@ from vllm_ascend.device.device_op import DeviceOperator from vllm_ascend.ops.fused_moe.router.grouped_topk_router import AscendGroupedTopKRouter +DEEPSEEK_V4_IMAGE_SENTINEL_BASE_ID = 129257 +DEEPSEEK_V4_IMAGE_SENTINEL_COUNT = 5 + + +def select_deepseek_v4_vision_experts( + router_logits: torch.Tensor, + input_ids: torch.Tensor, + tid2eid: torch.Tensor | None, + bias_vl: torch.Tensor, + text_bias: torch.Tensor | None, + top_k: int, + renormalize: bool, + routed_scaling_factor: float = 1.0, + image_sentinel_lo: int = DEEPSEEK_V4_IMAGE_SENTINEL_BASE_ID, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select text experts and apply the vision route to image rows. + + DeepSeek-V4 vision checkpoints borrow five consecutive in-vocabulary + sentinel ids for IMAGE_START..IMAGE_END. Text rows retain the deterministic + ``tid2eid`` lookup used by the text-only model, while image rows use the + checkpoint's ``bias_vl`` with the sqrt-softplus router scores. + """ + scores = torch.nn.functional.softplus(router_logits).sqrt() + image_hi = image_sentinel_lo + DEEPSEEK_V4_IMAGE_SENTINEL_COUNT + image_mask = (input_ids >= image_sentinel_lo) & (input_ids < image_hi) + row_bias = torch.where( + image_mask.unsqueeze(-1), + bias_vl.to(scores.dtype).unsqueeze(0), + (text_bias.to(scores.dtype).unsqueeze(0) if text_bias is not None else torch.zeros_like(scores)), + ) + dynamic_ids = torch.topk( + scores + row_bias, + k=top_k, + dim=-1, + sorted=True, + ).indices + if tid2eid is None: + topk_ids = dynamic_ids + else: + lookup_ids = torch.where(image_mask, 0, input_ids) + text_ids = tid2eid[lookup_ids].to(torch.int64) + topk_ids = torch.where(image_mask.unsqueeze(-1), dynamic_ids, text_ids) + topk_weights = scores.gather(1, topk_ids) + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True).clamp_min( + torch.finfo(topk_weights.dtype).tiny + ) + if routed_scaling_factor != 1.0: + topk_weights = topk_weights * routed_scaling_factor + return topk_weights, topk_ids + class AscendFusedTopKRouter(AscendGroupedTopKRouter): """Router adapter that uses Ascend's existing expert-selection path.""" @@ -43,6 +94,8 @@ def __init__( eplb_state: EplbLayerState | None = None, num_logical_experts: int | None = None, tid2eid: torch.Tensor | None = None, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = DEEPSEEK_V4_IMAGE_SENTINEL_BASE_ID, select_experts_fn: Callable[..., tuple[torch.Tensor, torch.Tensor]] | None = None, ): super().__init__( @@ -62,6 +115,8 @@ def __init__( self.e_score_correction_bias = e_score_correction_bias self.num_logical_experts = num_logical_experts if num_logical_experts is not None else global_num_experts self.tid2eid = tid2eid + self.bias_vl = bias_vl + self.image_sentinel_lo = image_sentinel_lo def is_fused_supported( self, @@ -89,7 +144,7 @@ def _compute_routing( *, input_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - if not self.is_fused_supported(hidden_states): + if self.bias_vl is None and not self.is_fused_supported(hidden_states): return super()._compute_routing( hidden_states=hidden_states, router_logits=router_logits, @@ -101,11 +156,11 @@ def _compute_routing( num_expert_group = self.num_expert_group if self.num_expert_group is not None else 1 renorm = int(self.renormalize) if self.scoring_func == "sqrtsoftplus": - if self.tid2eid is not None: + if self.tid2eid is not None or self.bias_vl is not None: if input_ids is None: - raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + raise ValueError("DeepSeek V4 vision/hash MoE routing requires input_ids.") input_ids = input_ids.to(torch.int64) - tid2eid_ones = self.tid2eid.to(torch.int32) + tid2eid_ones = self.tid2eid.to(torch.int32) if self.tid2eid is not None else None if _EXTRA_CTX.moe_comm_type == MoECommType.ALLGATHER: prepare_finalize = _EXTRA_CTX.moe_comm_method.prepare_finalize input_ids = prepare_finalize.all_gather_input_id_with_dp_group(input_ids) @@ -121,6 +176,21 @@ def _compute_routing( else: input_ids = None tid2eid_ones = None + if self.bias_vl is not None and input_ids is not None: + topk_weights, topk_ids = select_deepseek_v4_vision_experts( + router_logits=router_logits, + input_ids=input_ids, + tid2eid=tid2eid_ones, + bias_vl=self.bias_vl, + text_bias=self.e_score_correction_bias, + top_k=self.top_k, + renormalize=self.renormalize, + routed_scaling_factor=self.routed_scaling_factor, + image_sentinel_lo=self.image_sentinel_lo, + ) + return topk_weights.to(torch.float32), topk_ids.to( + torch.int32 if indices_type is None else indices_type + ) topk_weights, topk_ids, _ = torch.ops._C_ascend.moe_gating_top_k_hash( x=router_logits, k=self.top_k, diff --git a/vllm_ascend/ops/fused_moe/router/router_factory.py b/vllm_ascend/ops/fused_moe/router/router_factory.py index e336be51efc0..37342a0e54eb 100644 --- a/vllm_ascend/ops/fused_moe/router/router_factory.py +++ b/vllm_ascend/ops/fused_moe/router/router_factory.py @@ -50,6 +50,8 @@ def create_ascend_fused_moe_router( num_logical_experts: int | None = None, hash_indices_table: torch.Tensor | None = None, tid2eid: torch.Tensor | None = None, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 129257, ) -> FusedMoERouter: if custom_routing_function is not None: return CustomRoutingRouter( @@ -98,6 +100,8 @@ def create_ascend_fused_moe_router( e_score_correction_bias=e_score_correction_bias, num_logical_experts=num_logical_experts, tid2eid=tid2eid, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) return AscendGroupedTopKRouter( top_k=top_k, diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 169d4c719c24..780f2da6f4e7 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -48,6 +48,28 @@ # Future Plan: # Remove this patch when vLLM merge the PR. # +# ** 2. File: platform/patch_deepseek_v4_vision.py** +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# 1. `vllm.transformers_utils.model_arch_config_convertor.MODEL_ARCH_CONFIG_CONVERTORS` +# Why: +# The supported vLLM revision has the generic DeepSeek-V4 text config +# conversion but does not identify a checkpoint with `vision_n_layers` +# as the multimodal conditional-generation architecture. Without this +# distinction, vllm-ascend cannot select its DeepSeek-V4 vision wrapper +# or enable bidirectional attention over the image prefix. +# How: +# Register an Ascend DeepSeek-V4 config conversion handler. For vision checkpoints +# it selects `DeepseekV4ForConditionalGeneration`, enables multimodal +# prefix-LM attention, and records the prefix-padding constraints used by +# the Ascend DSA path. Text-only DeepSeek-V4 behavior is unchanged. +# Related PR (if no, explain why): +# https://github.com/vllm-project/vllm/pull/54566 +# Future Plan: +# Remove this patch once the supported vLLM revision natively maps +# DeepSeek-V4 vision checkpoints to the conditional-generation model and +# exposes the required multimodal prefix-LM and padding metadata without +# replacing `MODEL_ARCH_CONFIG_CONVERTORS["deepseek_v4"]`. +# # ** 3. File: platform/patch_distributed.py** # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 1. `torch.distributed.all_reduce`, `torch.distributed.broadcast` @@ -148,6 +170,23 @@ # Remove this patch once upstream exposes a backend dispatch / plugin hook # for selecting the MoE runner implementation. # +# 2. `vllm.model_executor.layers.fused_moe.FusedMoEFactory` +# Why: +# DeepSeek-V4 vision routing supplies `bias_vl` and +# `image_sentinel_lo` through the upstream MoE factory. The Ascend +# replacement factory must preserve those arguments so image tokens use +# the checkpoint's vision-specific expert-routing bias. +# How: +# Accept the two DeepSeek-V4 vision arguments in `_ascend_FusedMoE` and +# pass them to the Ascend router while leaving every other model's +# defaults unchanged. +# Related PR (if no, explain why): +# https://github.com/vllm-project/vllm/pull/54566 +# Future Plan: +# Remove this DeepSeek-V4-specific argument bridge once upstream exposes +# a backend-neutral router configuration object or MoE factory extension +# hook that carries vision routing metadata into the Ascend runner. +# # ** 7a. File: platform/patch_glm5next_config.py** # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 1. `vllm.transformers_utils.config._CONFIG_REGISTRY` @@ -533,6 +572,24 @@ # models without a custom `hf_config_override`, or exposes a plugin hook # for MTP model_type/architecture remapping. # +# 2. `vllm.config.speculative.SpeculativeConfig.__post_init__` +# Why: +# DeepSeek-V4 Vision uses the target checkpoint for its DSpark drafter. +# Multimodal model-architecture conversion can overwrite the draft's +# `DSparkDraftModel` architecture with the full conditional-generation +# architecture, which constructs a second target model and produces +# duplicate attention-layer registrations. +# How: +# After upstream speculative-config initialization, restore the draft's +# `DSparkDraftModel` architecture in both Hugging Face and normalized +# model configs, then refresh the cached registry inspection result. +# Related PR (if no, explain why): +# https://github.com/vllm-project/vllm/pull/54566 +# Future Plan: +# Remove this normalization once upstream performs multimodal conversion +# before DSpark draft selection, or otherwise guarantees that rebuilding +# `model_arch_config` preserves the selected draft architecture. +# # ** 19. File: platform/patch_structured_output.py** # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 1. `vllm.sampling_params.SamplingParams._validate_structured_outputs` diff --git a/vllm_ascend/patch/platform/__init__.py b/vllm_ascend/patch/platform/__init__.py index 4e6ab5dc5385..ee33bb772bfc 100644 --- a/vllm_ascend/patch/platform/__init__.py +++ b/vllm_ascend/patch/platform/__init__.py @@ -16,6 +16,7 @@ import os +import vllm_ascend.patch.platform.patch_deepseek_v4_vision # noqa import vllm_ascend.patch.platform.patch_distributed # noqa import vllm_ascend.patch.platform.patch_kv_cache_utils # noqa import vllm_ascend.patch.platform.patch_mamba_block_aligned_split # noqa diff --git a/vllm_ascend/patch/platform/patch_deepseek_v4_vision.py b/vllm_ascend/patch/platform/patch_deepseek_v4_vision.py new file mode 100644 index 000000000000..f5db062e8ddd --- /dev/null +++ b/vllm_ascend/patch/platform/patch_deepseek_v4_vision.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""vLLM v0.27 model-config compatibility for DeepSeek-V4 vision.""" + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from transformers import PretrainedConfig + +_REGISTERED = False + + +def register_deepseek_v4_vision_config_convertor() -> None: + """Route vision checkpoints to the Ascend multimodal wrapper. + + Keep vLLM config imports out of module scope. Global patches are imported + while spawned engine processes unpickle their state, which can happen while + ``model_arch_config_convertor`` itself is only partially initialized. + """ + global _REGISTERED + if _REGISTERED: + return + + from vllm.transformers_utils.model_arch_config_convertor import ( + MODEL_ARCH_CONFIG_CONVERTORS, + ModelArchConfigConvertorBase, + ) + + from vllm_ascend.utils import vllm_version_is + + class AscendDeepseekV4ModelArchConfigConvertor(ModelArchConfigConvertorBase): + """Route vision checkpoints to the Ascend multimodal wrapper.""" + + def __init__( + self, + hf_config: "PretrainedConfig", + hf_text_config: "PretrainedConfig", + revision: str | None = None, + ) -> None: + if getattr(hf_config, "vision_n_layers", 0) > 0: + hf_config.architectures = ["DeepseekV4ForConditionalGeneration"] + hf_config.mm_prefix_clamp_sliding_window = True + hf_config.mm_prefix_span_leading_pad_modulus = 4 + if vllm_version_is("0.27.1"): + super().__init__(hf_config, hf_text_config) + else: + super().__init__(hf_config, hf_text_config, revision) + + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: + return supports_multimodal and (getattr(self.hf_config, "vision_n_layers", 0) > 0) + + MODEL_ARCH_CONFIG_CONVERTORS["deepseek_v4"] = AscendDeepseekV4ModelArchConfigConvertor + _REGISTERED = True diff --git a/vllm_ascend/patch/platform/patch_fused_moe.py b/vllm_ascend/patch/platform/patch_fused_moe.py index 0cf547d12d3a..4f43fc7acea9 100644 --- a/vllm_ascend/patch/platform/patch_fused_moe.py +++ b/vllm_ascend/patch/platform/patch_fused_moe.py @@ -112,6 +112,8 @@ def _ascend_FusedMoE( routed_experts_args: dict[str, Any] | None = None, hash: Any | None = None, tid2eid: torch.Tensor | None = None, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 129257, **kwargs, ): # RoutedExperts allocates its parameters before AscendMoERunner is @@ -147,6 +149,8 @@ def _ascend_FusedMoE( num_logical_experts=num_experts, hash_indices_table=hash_indices_table, tid2eid=hash_indices_table_for_legacy_path, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, eplb_state=AscendEplbLayerState() if enable_router_eplb else None, ) routed_experts_args = dict(routed_experts_args) if routed_experts_args is not None else {} diff --git a/vllm_ascend/patch/platform/patch_speculative_config.py b/vllm_ascend/patch/platform/patch_speculative_config.py index 0e35e889c3e2..8f936dafde2c 100644 --- a/vllm_ascend/patch/platform/patch_speculative_config.py +++ b/vllm_ascend/patch/platform/patch_speculative_config.py @@ -1,3 +1,5 @@ +from dataclasses import replace + from transformers import DeepseekV2Config, PretrainedConfig from vllm.config.speculative import SpeculativeConfig @@ -37,11 +39,46 @@ def _normalize_legacy_qwen3_dspark_config(hf_config: PretrainedConfig) -> Pretra return hf_config +def _normalize_deepseek_v4_dspark_draft(draft_model_config) -> None: + """Restore the DSpark draft architecture after VL config conversion. + + DeepSeek-V4-Vision uses the same checkpoint for the target and DSpark + drafter. vLLM first rewrites that checkpoint to ``DSparkDraftModel``, but + rebuilding ``model_arch_config`` with multimodal detection can restore the + top-level ``*ForConditionalGeneration`` architecture. The drafter would + then instantiate a second full VL target and register duplicate attention + layer names. Update both config representations without re-running the + multimodal architecture conversion. + """ + hf_config = getattr(draft_model_config, "hf_config", None) + if ( + hf_config is None + or getattr(hf_config, "model_type", None) != "deepseek_v4" + or getattr(hf_config, "dspark_target_layer_ids", None) is None + ): + return + + hf_config.update({"architectures": ["DSparkDraftModel"]}) + draft_model_config.model_arch_config = replace( + draft_model_config.model_arch_config, + architectures=["DSparkDraftModel"], + model_type="deepseek_v4", + is_mm_prefix_lm=False, + ) + model_info, architecture = draft_model_config.registry.inspect_model_cls( + draft_model_config.architectures, + draft_model_config, + ) + draft_model_config._model_info = model_info + draft_model_config._architecture = architecture + + def _dspark_post_init(self): _orig_post_init(self) if self.use_dspark(): draft_model_config = getattr(self, "draft_model_config", None) draft_hf_config = getattr(draft_model_config, "hf_config", None) + _normalize_deepseek_v4_dspark_draft(draft_model_config) # deepseek v4 dspark if getattr(draft_hf_config, "ptd_token_id", None) is None: # type: ignore draft_hf_config.ptd_token_id = getattr(draft_hf_config, "dspark_noise_token_id", None) # type: ignore diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index 42e2db4ecb5d..bebea11e5ace 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -296,6 +296,15 @@ def pre_register_and_update(cls, parser: FlexibleArgumentParser | None = None) - adapt_patch(is_global_patch=True) + # Registration imports vLLM's model config converter and therefore must + # happen after the global patch package has finished importing. Keeping + # it out of patch module scope also makes multiprocessing spawn safe. + from vllm_ascend.patch.platform.patch_deepseek_v4_vision import ( + register_deepseek_v4_vision_config_convertor, + ) + + register_deepseek_v4_vision_config_convertor() + # For online serving, "ascend" quantization method is not a choice natively, # so we need to add "ascend" quantization method to quantization methods list # and the user can enable quantization using "vllm serve --quantization ascend". diff --git a/vllm_ascend/quantization/configs/modelslim_config.py b/vllm_ascend/quantization/configs/modelslim_config.py index a0f512c8b48c..c5cdd28a8d0c 100644 --- a/vllm_ascend/quantization/configs/modelslim_config.py +++ b/vllm_ascend/quantization/configs/modelslim_config.py @@ -382,6 +382,14 @@ def modelslim_moe_weight_loader( QUANT_MODEL_PREFIX_MAPPINGS = { "deepseek_v4": { + # Multimodal checkpoints keep the text-backbone quantization keys in + # the original causal-LM namespace (``model.*`` / ``lm_head.*``), + # while the vLLM wrapper nests the module under ``language_model``. + "language_model.model.": "model.", + "language_model.lm_head.": "lm_head.", + # ``WeightsMapper`` treats the trailing dot literally, while the + # quantization lookup also queries the module prefix itself. + "language_model.lm_head": "lm_head", "layers.": "model.layers.", "embed.": "model.embed_tokens.", "head.": "lm_head.", diff --git a/vllm_ascend/spec_decode/llm_base_proposer.py b/vllm_ascend/spec_decode/llm_base_proposer.py index 3f4f74e23029..095a2ac4b752 100644 --- a/vllm_ascend/spec_decode/llm_base_proposer.py +++ b/vllm_ascend/spec_decode/llm_base_proposer.py @@ -114,7 +114,7 @@ class AscendSpecDecodeBaseProposer(SpecDecodeBaseProposer): _runnable: ACLGraphWrapper | Callable @staticmethod - def _get_multimodal_image_token_index(model_name: str, config: Any) -> int: + def _get_multimodal_image_token_index(model_name: str, config: Any) -> int | None: if model_name in [ "Qwen2_5_VLForConditionalGeneration", "Qwen3VLForConditionalGeneration", @@ -135,7 +135,11 @@ def _get_multimodal_image_token_index(model_name: str, config: Any) -> int: "AscendKimiK3ForConditionalGeneration", }: return config.media_placeholder_token_id - return config.image_token_index + # Some models (for example DeepSeek-V4 Vision) use multiple + # position-dependent image sentinel tokens instead of one placeholder + # token. Their text-only drafter does not need a synthetic image token + # index during decode. + return getattr(config, "image_token_index", None) def __init__(self, vllm_config: VllmConfig, device: torch.device, pass_hidden_states_to_model: bool, runner=None): super().__init__(vllm_config, device, pass_hidden_states_to_model, runner=runner) @@ -405,7 +409,9 @@ def load_model(self, model: nn.Module) -> None: if supports_multimodal(model): # handle multimodality model_name = self.get_model_name(model) - self.model.config.image_token_index = self._get_multimodal_image_token_index(model_name, model.config) + image_token_index = self._get_multimodal_image_token_index(model_name, model.config) + if image_token_index is not None: + self.model.config.image_token_index = image_token_index target_language_model = model.get_language_model() else: target_language_model = model diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index c1dc496ae6ed..a78f744802d2 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -853,6 +853,21 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None self._track_tmp_encoder_cache_refs(scheduler_output) return sampling_metadata + def _should_skip_compiled_for_encoder_input(self, scheduler_output: "SchedulerOutput") -> bool: + if scheduler_output.scheduled_encoder_inputs: + return True + for req_id in scheduler_output.num_scheduled_tokens: + req_state = self.requests.get(req_id) + if req_state is None or not req_state.mm_features: + continue + req_idx = self.input_batch.req_id_to_index[req_id] + if ( + self.input_batch.num_computed_tokens_cpu[req_idx] + < self.input_batch.num_prompt_tokens[req_idx] + ): + return True + return False + def _update_states_after_model_execute( self, output_token_ids: torch.Tensor, scheduler_output: "SchedulerOutput" ) -> None: @@ -2150,7 +2165,7 @@ def execute_model( # returns True. before returning early here we call # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. - self._dummy_run(1) + self._dummy_run(1, num_actual_reqs=0) if not has_kv_transfer_group(): # Return empty ModelRunnerOutput if no work to do. return EMPTY_MODEL_RUNNER_OUTPUT @@ -2355,10 +2370,9 @@ def execute_model( cudagraph_mode = CUDAGraphMode.NONE # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False # type: ignore[has-type] - # Encoder-decoder models can only compile the pure decode steps where no - # encoder inputs are present. Use eager for the first pass. - num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) - has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0 + # Encoder inputs and multimodal prompt embeddings are not safe to feed + # through the compiled backbone. Pure text and decode remain compiled. + skip_compiled = self._should_skip_compiled_for_encoder_input(scheduler_output) # Run forward pass defer_kv_connector_finalize = self.speculative_config is not None and ( @@ -2377,7 +2391,7 @@ def execute_model( num_actual_tokens=scheduler_output.total_num_scheduled_tokens, model_instance=self.model, device_metadata_executor=active_device_metadata_executor, - skip_compiled=has_encoder_input, + skip_compiled=skip_compiled, has_sinks=self._has_sinks, eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ), @@ -3127,6 +3141,7 @@ def _build_attention_metadata( cascade_attn_prefix_lens: list[list[int]] | None = None, cudagraph_runtime_mode: CUDAGraphMode | None = None, batch_descriptor: BatchDescriptor | None = None, + num_actual_reqs: int | None = None, ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]: """ :return: tuple[attn_metadata, spec_decode_common_attn_metadata] @@ -3236,6 +3251,41 @@ def _get_block_table_and_slot_mapping( seq_lens_cpu = None num_computed_tokens_cpu = None + req_doc_ranges = None + if self.is_mm_prefix_lm: + req_doc_ranges = {} + hf_text_config = self.model_config.hf_text_config + span_pad_modulus = getattr( + hf_text_config, + "mm_prefix_span_leading_pad_modulus", + 4 if getattr(hf_text_config, "vision_n_layers", 0) > 0 else 0, + ) + for req_id in self.input_batch.req_ids[:num_reqs]: + image_doc_ranges = [] + req_state = self.requests[req_id] + for mm_feature in req_state.mm_features or (): + if mm_feature.modality not in ("image", "video"): + continue + pos_info = mm_feature.mm_position + if span_pad_modulus: + leading_pad = ( + span_pad_modulus + - 1 + - pos_info.offset % span_pad_modulus + ) + image_doc_ranges.append( + ( + pos_info.offset + leading_pad, + pos_info.offset + pos_info.length - 1, + ) + ) + else: + image_doc_ranges.extend( + pos_info.extract_embeds_range() + ) + req_idx = self.input_batch.req_id_to_index[req_id] + req_doc_ranges[req_idx] = image_doc_ranges + cm_base = AscendCommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], @@ -3279,6 +3329,7 @@ def _get_block_table_and_slot_mapping( if self._offload_token_to_req is not None else None ), + mm_req_doc_ranges=req_doc_ranges, ) if logits_indices is not None and self.cache_config.kv_sharing_fast_prefill: @@ -3315,7 +3366,7 @@ def _build_attn_group_metadata( if for_cudagraph_capture: common_ratio_to_sas_metadata = {} extra_attn_metadata_args = dict( - num_actual_reqs=num_reqs, + num_actual_reqs=num_reqs if num_actual_reqs is None else min(num_actual_reqs, num_reqs), common_ratio_to_sas_metadata=common_ratio_to_sas_metadata, full_graph_mode=cudagraph_runtime_mode == CUDAGraphMode.FULL, ) @@ -3406,26 +3457,6 @@ def _build_attn_group_metadata( cm, common_ratio_to_sas_metadata, ) - if self.is_mm_prefix_lm: - req_doc_ranges = {} - for req_id in self.input_batch.req_ids: - image_doc_ranges = [] - req_state = self.requests[req_id] - for mm_feature in req_state.mm_features: - pos_info = mm_feature.mm_position - img_doc_range = pos_info.extract_embeds_range() - image_doc_ranges.extend(img_doc_range) - req_idx = self.input_batch.req_id_to_index[req_id] - req_doc_ranges[req_idx] = image_doc_ranges - - if isinstance(attn_metadata, list): - for ub_metadata in attn_metadata: - for _metadata in ub_metadata.values(): - _metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined] - else: - for _metadata in attn_metadata.values(): - _metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined] - if spec_decode_common_attn_metadata is not None and ( num_reqs != num_reqs_padded or num_tokens != num_tokens_padded ): @@ -3472,6 +3503,7 @@ def _dummy_run( num_active_loras: int = 0, profile_seq_lens: int | None = None, profile_cpp: bool = False, + num_actual_reqs: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() @@ -3651,6 +3683,7 @@ def _dummy_run( num_scheduled_tokens_np=num_scheduled_tokens, cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, + num_actual_reqs=num_actual_reqs, ) with self.maybe_dummy_run_with_lora( self.lora_config, @@ -3664,7 +3697,14 @@ def _dummy_run( ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens - if self.supports_mm_inputs and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds: + if ( + ( + self.supports_mm_inputs + and not self.model_config.is_encoder_decoder + and not self.model_config.requires_raw_input_tokens + ) + or self.enable_prompt_embeds + ): input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: