From a0f4947bbb72942806a8ac3fcdb2772a0ef70ec7 Mon Sep 17 00:00:00 2001 From: GDzhu01 <116337067+GDzhu01@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:41:45 +0800 Subject: [PATCH 1/2] [Feature][Model] Support DeepSeek V4 Flash Vision Add the DeepSeek V4 multimodal preprocessing and vision runtime, Ascend quantization and MoE routing integration, bidirectional vision attention, and DSpark support. Isolate tokenizer backends per preprocessing thread to make concurrent multimodal requests safe. Signed-off-by: GDzhu01 <116337067+GDzhu01@users.noreply.github.com> Co-authored-by: MengLong Chen <71744434+dragondream-chen@users.noreply.github.com> Co-authored-by: RenYuKai <184603735+pgzddxx@users.noreply.github.com> --- tests/ut/attention/test_dsa_v1.py | 41 ++ tests/ut/models/test_deepseek_v4_moe.py | 90 +++ tests/ut/models/test_deepseek_v4_vision.py | 20 + .../test_deepseek_v4_vision_preprocess.py | 160 ++++++ .../ut/patch/platform/test_patch_fused_moe.py | 34 ++ .../test_patch_speculative_config_dspark.py | 55 ++ .../ut/quantization/test_modelslim_config.py | 21 + .../ut/spec_decode/test_llm_base_proposer.py | 9 + vllm_ascend/attention/dsa_v1.py | 123 ++++ vllm_ascend/models/__init__.py | 4 + vllm_ascend/models/deepseek_v4/dspark.py | 9 + .../models/deepseek_v4/mm_preprocess.py | 540 ++++++++++++++++++ vllm_ascend/models/deepseek_v4/model.py | 18 +- vllm_ascend/models/deepseek_v4/vision.py | 132 +++++ vllm_ascend/models/deepseek_v4/vl_model.py | 267 +++++++++ .../ops/fused_moe/router/fused_topk_router.py | 78 ++- .../ops/fused_moe/router/router_factory.py | 4 + vllm_ascend/patch/__init__.py | 57 ++ vllm_ascend/patch/platform/__init__.py | 1 + .../platform/patch_deepseek_v4_vision.py | 53 ++ vllm_ascend/patch/platform/patch_fused_moe.py | 4 + .../platform/patch_speculative_config.py | 37 ++ vllm_ascend/platform.py | 9 + .../quantization/configs/modelslim_config.py | 8 + vllm_ascend/spec_decode/llm_base_proposer.py | 12 +- vllm_ascend/worker/model_runner_v1.py | 69 ++- 26 files changed, 1831 insertions(+), 24 deletions(-) create mode 100644 tests/ut/models/test_deepseek_v4_vision.py create mode 100644 tests/ut/models/test_deepseek_v4_vision_preprocess.py create mode 100644 vllm_ascend/models/deepseek_v4/mm_preprocess.py create mode 100644 vllm_ascend/models/deepseek_v4/vision.py create mode 100644 vllm_ascend/models/deepseek_v4/vl_model.py create mode 100644 vllm_ascend/patch/platform/patch_deepseek_v4_vision.py 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/vllm_ascend/attention/dsa_v1.py b/vllm_ascend/attention/dsa_v1.py index f6bbdf261cde..feaf353a6093 100644 --- a/vllm_ascend/attention/dsa_v1.py +++ b/vllm_ascend/attention/dsa_v1.py @@ -284,6 +284,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 @@ -400,6 +401,91 @@ 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] + 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) + 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 @@ -827,6 +913,7 @@ def build_req_metadata( 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 +933,35 @@ 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: + actual_reqs = num_reqs if num_actual_reqs is None else num_actual_reqs + vision_swa_indices, _ = build_vision_bidirectional_swa_indices( + block_table=self.block_table[: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[: actual_reqs + 1], + seq_lens=seq_lens[: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, @@ -947,6 +1063,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,6 +2044,12 @@ def compressor_tail_fn() -> CompressorForwardOutput: layout_kv=_dsa_layout_kv(self.vllm_config), ) + # 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 self.compress_ratio <= 1: if swa_req_metadata.dspark_swa_indices is not None: attn_kwargs["ori_sparse_indices"] = swa_req_metadata.dspark_swa_indices 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/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..3059fb1a6f7f 100644 --- a/vllm_ascend/models/deepseek_v4/model.py +++ b/vllm_ascend/models/deepseek_v4/model.py @@ -320,6 +320,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 +362,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=129257, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, @@ -1123,7 +1134,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..43d33580f322 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -2355,10 +2355,15 @@ 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. + # Encoder-decoder models and raw-token multimodal models can only + # compile pure decode steps where no encoder inputs are present. The + # DeepSeek-V4 vision router needs raw sentinel ids during image + # prefill, so keep that pass eager. num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) - has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0 + has_encoder_input = num_encoder_reqs > 0 and ( + self.model_config.is_encoder_decoder + or self.model_config.requires_raw_input_tokens + ) # Run forward pass defer_kv_connector_finalize = self.speculative_config is not None and ( @@ -3236,6 +3241,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: + image_doc_ranges = [] + req_state = self.requests[req_id] + for mm_feature in req_state.mm_features: + if mm_feature.modality == "audio": + 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 +3319,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: @@ -3406,18 +3447,7 @@ 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 req_doc_ranges is not None: if isinstance(attn_metadata, list): for ub_metadata in attn_metadata: for _metadata in ub_metadata.values(): @@ -3664,7 +3694,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: From 3eb808938821597cf9bec07d0b174e255c9b4b25 Mon Sep 17 00:00:00 2001 From: QwertyJack <7554089+QwertyJack@users.noreply.github.com> Date: Fri, 4 Sep 2026 03:17:15 +0000 Subject: [PATCH 2/2] fix(deepseek_v4): align vision runtime contracts Signed-off-by: QwertyJack <7554089+QwertyJack@users.noreply.github.com> --- tests/ut/attention/test_attention_v1.py | 2 + .../attention/test_dsa_compressor_seqused.py | 101 ++++++++++++++++++ tests/ut/worker/test_model_runner_v1.py | 55 +++++++++- vllm_ascend/attention/dsa_v1.py | 43 ++++++-- vllm_ascend/attention/utils.py | 9 ++ vllm_ascend/models/deepseek_v4/compressor.py | 2 +- vllm_ascend/models/deepseek_v4/model.py | 3 +- vllm_ascend/worker/model_runner_v1.py | 51 ++++----- 8 files changed, 233 insertions(+), 33 deletions(-) create mode 100644 tests/ut/attention/test_dsa_compressor_seqused.py 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/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 feaf353a6093..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 @@ -328,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) @@ -423,6 +438,8 @@ def build_vision_bidirectional_swa_indices( 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], @@ -448,6 +465,10 @@ def build_vision_bidirectional_swa_indices( 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( @@ -555,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 ) @@ -909,6 +933,12 @@ 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 @@ -951,14 +981,13 @@ def build_req_metadata( else 0 ) if has_prefill and max_image_tokens > 0 and mm_ranges: - actual_reqs = num_reqs if num_actual_reqs is None else num_actual_reqs vision_swa_indices, _ = build_vision_bidirectional_swa_indices( - block_table=self.block_table[:actual_reqs], + 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[: actual_reqs + 1], - seq_lens=seq_lens[:actual_reqs], + 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, ) @@ -1055,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, @@ -2050,10 +2080,11 @@ def compressor_tail_fn() -> CompressorForwardOutput: if swa_req_metadata.vision_swa_indices is not None: attn_kwargs["ori_sparse_indices"] = swa_req_metadata.vision_swa_indices - if self.compress_ratio <= 1: + 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/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/model.py b/vllm_ascend/models/deepseek_v4/model.py index 3059fb1a6f7f..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 @@ -363,7 +364,7 @@ def __init__( swiglu_limit=self.swiglu_limit, e_score_correction_bias=self.gate.e_score_correction_bias, bias_vl=self.gate.bias_vl, - image_sentinel_lo=129257, + 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, diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 43d33580f322..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,15 +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 and raw-token multimodal models can only - # compile pure decode steps where no encoder inputs are present. The - # DeepSeek-V4 vision router needs raw sentinel ids during image - # prefill, so keep that pass eager. - num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) - has_encoder_input = num_encoder_reqs > 0 and ( - self.model_config.is_encoder_decoder - or self.model_config.requires_raw_input_tokens - ) + # 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 ( @@ -2382,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, ), @@ -3132,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] @@ -3250,11 +3260,11 @@ def _get_block_table_and_slot_mapping( "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: + 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: - if mm_feature.modality == "audio": + 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: @@ -3356,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, ) @@ -3447,15 +3457,6 @@ def _build_attn_group_metadata( cm, common_ratio_to_sas_metadata, ) - if req_doc_ranges is not None: - 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 ): @@ -3502,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() @@ -3681,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,