[Feature][MiniMax-H3] Complete VSA and Ulysses support (#6909)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>

tests/diffusion/attention/test_fastvideo_vsa.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+import math
 import sys
 import types
 from typing import Any
@@ -8,9 +9,17 @@ from typing import Any
 import pytest
 import torch
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+    VideoTokenLayout,
+    VideoTokenSpan,
+)
 from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
+    FastVideoVSABackend,
     FastVideoVSAImpl,
+    _build_h3_block_map,
+    _get_h3_tile_metadata,
 )
 from vllm_omni.diffusion.attention.backends.registry import (
     DiffusionAttentionBackendEnum,
@@ -23,6 +32,12 @@ def test_fastvideo_vsa_backend_is_registered():
     assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
 
 
+def test_fastvideo_vsa_reports_missing_optional_kernel(monkeypatch):
+    monkeypatch.setattr("importlib.util.find_spec", lambda name: None)
+    with pytest.raises(ImportError, match="fastvideo-kernel"):
+        FastVideoVSABackend.validate_available()
+
+
 def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
     calls: dict[str, Any] = {}
     fake_module = types.ModuleType("fastvideo_kernel")
@@ -201,3 +216,190 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
     # CPU/float32 is rejected later, but k=N itself must not trigger fallback.
     assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
     assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
+
+
+def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
+    extra: dict[str, Any] = {
+        "vsa_h3_prefix_segments": prefix_segments,
+    }
+    if gate is not None:
+        extra["gate_compress"] = gate
+    return AttentionMetadata(
+        packed_padding=packed_padding,
+        video_layout=VideoTokenLayout(
+            used_len=sum(prefix_segments) + math.prod(video_shape),
+            video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
+        ),
+        extra=extra,
+    )
+
+
+def _fake_block_sparse_kernel(monkeypatch, calls):
+    """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
+
+    def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
+        calls["block_map"] = block_map.clone()
+        calls["variable_block_sizes"] = variable_block_sizes.tolist()
+        calls["q_shape"] = tuple(q.shape)
+        blocks = block_map.shape[-1]
+        assert blocks * 64 == q.shape[2]
+        # Expand the per-block map to a per-token mask so the stand-in honours
+        # the sparsity contract the real kernel implements, and drop the rows
+        # each partial block pads with, which the real kernel skips through
+        # variable_block_sizes.
+        token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
+        within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
+        token_mask = token_mask & within.reshape(1, 1, 1, -1)
+        out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
+        return out.to(q.dtype), None
+
+    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    setattr(module, "block_sparse_attn", block_sparse_attn)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
+
+
+def _h3_impl(**backend_kwargs):
+    kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
+    kwargs.update(backend_kwargs)
+    return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
+
+
+def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5, 70), (2, 4, 4)
+    seq_len = sum(prefix_segments) + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    assert output.shape == query.shape
+    # topk covers every video block, so the block map is dense and the tiled
+    # kernel must reproduce plain attention over the same rows.
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
+    # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
+    assert calls["variable_block_sizes"] == [5, 64, 6, 32]
+
+
+def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=1)
+    prefix_segments, video_shape = (5,), (8, 4, 4)
+    seq_len = 5 + 8 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    # 1 prefix block + 2 video blocks. The odd count is padded to an even one
+    # for the paired-CTA kernel contract, and the Triton route slices that
+    # transport-only partner off again before launching.
+    block_map = calls["block_map"]
+    assert block_map.shape[-1] == 3
+    assert calls["variable_block_sizes"] == [5, 64, 64]
+    assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
+    assert block_map[..., :1].all(), "prefix keys are exempt from selection"
+    assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
+
+    # The same odd geometry must still untile back onto the original rows.
+    dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
+
+
+def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    seq_len = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    zero_gate = torch.zeros_like(query)
+    with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
+    live_gate = torch.full_like(query, 0.5)
+    with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
+
+    torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
+    assert not torch.allclose(with_live_gate.float(), without_gate.float())
+
+
+def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    valid = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
+    padding = PackedPaddingMetadata(
+        q_length=valid,
+        kv_length=valid,
+        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
+        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
+    )
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
+
+    assert output.shape == query.shape
+    assert torch.count_nonzero(output[:, valid:]) == 0
+
+
+def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
+    # The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
+    # building the padding mask. Every fallback out of the VSA path must honour
+    # that contract itself; SDPA reads attn_mask and nothing else.
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        backend_kwargs={"topk": 1, "min_seq_len": 4096},
+    )
+    valid = 40
+    torch.manual_seed(0)
+    query = torch.randn(1, 64, 2, 8)
+    padding = PackedPaddingMetadata(
+        q_length=valid,
+        kv_length=valid,
+        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
+        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
+    )
+
+    baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
+    perturbed_input = query.clone()
+    perturbed_input[:, valid:] += 100.0
+    perturbed = impl.forward_cuda(
+        perturbed_input, perturbed_input, perturbed_input, AttentionMetadata(packed_padding=padding, extra={})
+    )
+
+    assert baseline.shape == query.shape
+    torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
+    assert torch.count_nonzero(baseline[:, valid:]) == 0
+
+
+def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
+    partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
+        (5, 70, 9), (5, 6, 6), torch.device("cpu")
+    )
+    assert sizes[:4].tolist() == [5, 64, 6, 9]
+    assert prefix_blocks == 4
+    assert video_blocks == 8
+    assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
+    assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
+
+
+def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
+    scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
+    block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
+    assert block_map[:, :, :2].all()
+    assert block_map[..., :2].all()
+    assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()

tests/diffusion/attention/test_ulysses_uaa.py

@@ -322,6 +322,50 @@ def test_advanced_uaa_hybrid_rejects_non_gqa_shapes(monkeypatch: pytest.MonkeyPa
         strategy.pre_attention(query, key, value, None)
 
 
+def test_strict_ulysses_reshards_vsa_gate_with_qkv(monkeypatch) -> None:
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
+        lambda default: "strict",
+    )
+    calls: list[torch.Tensor] = []
+
+    def fake_all_to_all(pg, tensor, scatter_idx, gather_idx, use_sync):
+        del pg, scatter_idx, gather_idx, use_sync
+        calls.append(tensor)
+        return tensor.reshape(tensor.shape[0], tensor.shape[1] * 2, tensor.shape[2] // 2, tensor.shape[3])
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.parallel.ulysses.SeqAllToAll4D.apply",
+        fake_all_to_all,
+    )
+    strategy = UlyssesParallelAttention(
+        sp_group=SimpleNamespace(
+            ulysses_group=object(),
+            ulysses_world_size=2,
+            ulysses_rank=0,
+            ring_world_size=1,
+        ),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+    query = torch.arange(24).reshape(1, 3, 4, 2)
+    key = query + 10
+    value = query + 20
+    gate = query + 100
+    metadata = AttentionMetadata(extra={"gate_compress": gate})
+
+    query_out, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
+
+    assert all(actual is expected for actual, expected in zip(calls, (query, key, value, gate), strict=True))
+    assert query_out.shape == (1, 6, 2, 2)
+    assert metadata_out is metadata
+    assert metadata.extra["gate_compress"].shape == query_out.shape
+    torch.testing.assert_close(metadata.extra["gate_compress"], gate.reshape_as(query_out))
+
+
 def _run_attention_case(
     local_rank: int,
     world_size: int,

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -795,6 +795,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
 
 
 def test_packed_attention_skips_mask_for_packed_mask_free_backend():
+    from vllm_omni.diffusion.attention.backends.abstract import (
+        VideoTokenLayout,
+        VideoTokenSpan,
+    )
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
         MiniMaxH3Attention,
     )
@@ -821,6 +825,10 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
     torch.nn.Module.__init__(attention)
     attention.attention = FakeAttention()
     q = torch.randn(8, 2, 4)
+    video_layout = VideoTokenLayout(
+        used_len=5,
+        video_spans=(VideoTokenSpan(start=3, latent_grid=(1, 1, 2), role="target"),),
+    )
 
     attention._run_packed_attention(
         q,
@@ -829,6 +837,9 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
         cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
         max_seqlen=5,
         packed_total=8,
+        video_layout=video_layout,
+        vsa_prefix_segments=(1, 2),
+        gate_compress=torch.ones_like(q),
     )
 
     metadata = attention.attention.metadata
@@ -836,6 +847,11 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
     assert metadata.attn_mask is None
     assert metadata.extra["valid_kv_length"] == 5
     assert metadata.extra["npu_attn_varlen"] is True
+    assert metadata.video_layout is video_layout
+    assert metadata.extra["vsa_h3_prefix_segments"] == (1, 2)
+    assert "vsa_h3_video_shape" not in metadata.extra
+    assert "vsa_h3_target_start" not in metadata.extra
+    assert metadata.extra["gate_compress"].shape == (1, *q.shape)
     packed_padding = metadata.packed_padding
     assert packed_padding is not None
     assert packed_padding.q_length == 5

tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py

@@ -437,6 +437,83 @@ def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path):
         _fusion(tmp_path).check_serving_contract(partition="ref2va", **contract)
 
 
+def test_a_vsa_variant_accepts_pure_ulysses_and_refuses_ring(tmp_path):
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    od_config = SimpleNamespace(
+        diffusion_attention_backend="FASTVIDEO_VSA",
+        parallel_config=SimpleNamespace(sequence_parallel_size=8, ulysses_degree=8, ring_degree=1),
+    )
+    fusion = _load(sparse.parent)
+    fusion.check_serving_contract(partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0)
+
+    od_config.parallel_config.ring_degree = 2
+    with pytest.raises(ValueError, match="pure Ulysses"):
+        fusion.check_serving_contract(partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0)
+
+
+def test_a_vsa_variant_accepts_normalized_attention_config(tmp_path):
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    od_config = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(default=SimpleNamespace(backend="FASTVIDEO_VSA")),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    _load(sparse.parent).check_serving_contract(
+        partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0
+    )
+
+
+def test_a_vsa_variant_reads_the_backend_the_dit_will_actually_resolve(tmp_path):
+    # The 50 DiT blocks carry attention role "self", so a per_role entry - not
+    # the default - decides whether the compression gates ever reach a sparse
+    # kernel.
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    fusion = _load(sparse.parent)
+    contract = {"partition": "fl2va", "video_shift": 12.0, "audio_shift": 3.0}
+
+    dense_dit = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(
+            default=SimpleNamespace(backend="FASTVIDEO_VSA"),
+            per_role={"self": SimpleNamespace(backend="TRTLLM_ATTN")},
+        ),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    with pytest.raises(ValueError, match="TRTLLM_ATTN"):
+        fusion.check_serving_contract(od_config=dense_dit, **contract)
+
+    sparse_dit = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(
+            default=None,
+            per_role={"self": SimpleNamespace(backend="FASTVIDEO_VSA")},
+        ),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    fusion.check_serving_contract(od_config=sparse_dit, **contract)
+
+
+def test_a_gate_that_never_reached_the_model_is_refused(tmp_path):
+    # load_weights only logs a skip for a parameter the model does not have, so
+    # yielding the gate is not evidence it arrived.
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    fusion = _load(sparse.parent)
+    checkpoint = [
+        ("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN))),
+        ("blocks.0.attn.out_proj.weight", torch.zeros((_HIDDEN, _INNER))),
+        ("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))),
+        ("blocks.0.mlp.fc2.weight", torch.zeros((_HIDDEN, _FFN))),
+    ]
+    gate = "blocks.0.attn.to_gate_compress.weight"
+    streamed = {name for name, _ in fusion.apply(checkpoint)}
+    assert gate in streamed
+
+    with pytest.raises(FastH3AdapterError, match="never reached the model"):
+        fusion.validate_fully_applied(streamed - {gate})
+    fusion.validate_fully_applied(streamed)
+
+
 def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path):
     # Only low-rank factors are placed into H3's fused QKV and gate/up layouts,
     # so a .diff aimed at one would otherwise be added transposed.

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -67,6 +67,33 @@ def test_packed_attention_is_a_regional_compile_boundary():
     assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
 
 
+def test_denoise_branch_keeps_fast_h3_prefix_geometry_model_local():
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
+        MiniMaxH3DenoiseBranch,
+    )
+
+    packed = {
+        "seq_len": torch.tensor(9),
+        "img_pos": torch.arange(4),
+        "audio_pos": torch.empty(0, dtype=torch.long),
+        "text_pos": torch.arange(2),
+        "update_mask": torch.ones(4, dtype=torch.bool),
+        "cu_seqlens": torch.tensor([0, 9, 9], dtype=torch.int32),
+        "img_position_ids": torch.zeros(9, 3, dtype=torch.long),
+        "video_spans": ({"start": 5, "latent_grid": (1, 2, 2), "role": "target"},),
+    }
+    branch = MiniMaxH3DenoiseBranch(
+        packed=packed,
+        text_embeddings=torch.zeros(2, 4),
+        token_tags=torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2]),
+        device=torch.device("cpu"),
+    )
+
+    layout = branch.static_kwargs["video_token_layout"]
+    assert not hasattr(layout, "prefix_segments")
+    assert branch.static_kwargs["packed_seq_params"]["vsa_prefix_segments"] == (2, 3)
+
+
 def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
     from vllm_omni.diffusion.layers.rope import RotaryEmbedding
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (