[BugFix][Model] Align DeepSeek V4 Vision runtime contracts

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

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,
+        )

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"

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

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