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

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

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 (