[Feat] Minimax H3 support ref and tgt video sparse attention on NPU (#6518)

Signed-off-by: fan2956 <zhoufan53@huawei.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>

tests/diffusion/attention/test_rainfusion_plan.py

@@ -11,13 +11,16 @@ to the always-kept prefix segment before generating the block mask.
 """
 
 import dataclasses
+import sys
+import types
 
 import pytest
 import torch
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout, VideoTokenSpan
 from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
     _BLOCK_SIZE,
+    RainFusionAttentionBackend,
     RainFusionAttentionImpl,
     RainFusionPlan,
 )
@@ -97,6 +100,62 @@ def test_video_segment_must_be_the_tail_of_packed_document_zero():
     assert make_impl()._resolve_plan(metadata) is None
 
 
+def test_ref2va_multi_video_spans_resolve():
+    layout = VideoTokenLayout(
+        used_len=12000,
+        video_spans=(
+            VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
+            VideoTokenSpan(start=5000, latent_grid=(4, 16, 64), role="target"),
+        ),
+    )
+    plan = make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout))
+
+    assert plan is not None
+    assert plan.used_len == 12000
+    assert plan.video_spans == [
+        {"start": 128, "latent_shape": [4, 16, 64]},
+        {"start": 5000, "latent_shape": [4, 16, 64]},
+    ]
+
+
+def test_invalid_ref2va_video_spans_fall_back_to_dense():
+    layout = VideoTokenLayout(
+        used_len=12000,
+        video_spans=(
+            VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
+            VideoTokenSpan(start=4000, latent_grid=(4, 16, 64), role="target"),
+        ),
+    )
+    assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None
+
+
+def test_validate_available_rejects_legacy_mindiesd(monkeypatch):
+    mindiesd = types.ModuleType("mindiesd")
+
+    def sparse_attention(query, key, value, **kwargs):
+        return query
+
+    mindiesd.sparse_attention = sparse_attention
+    monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
+    monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
+
+    with pytest.raises(ValueError, match="video_spans"):
+        RainFusionAttentionBackend.validate_available()
+
+
+def test_validate_available_accepts_new_mindiesd(monkeypatch):
+    mindiesd = types.ModuleType("mindiesd")
+
+    def sparse_attention(query, key, value, *, video_spans=None, **kwargs):
+        return query
+
+    mindiesd.sparse_attention = sparse_attention
+    monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
+    monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
+
+    RainFusionAttentionBackend.validate_available()
+
+
 @pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
 def test_short_video_stays_dense(grid):
     assert make_impl()._resolve_plan(make_metadata(grid)) is None

tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py

@@ -74,6 +74,8 @@ def test_t2va_and_fl2va_packing_keep_update_rows_separate():
     assert fl2va["update_mask"].sum().item() == 12
     assert (~fl2va["update_mask"]).sum().item() == 6
     assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
+    assert t2va["video_spans"] == ({"start": 10, "latent_grid": (2, 2, 3), "role": "target"},)
+    assert fl2va["video_spans"] == ({"start": 16, "latent_grid": (2, 2, 3), "role": "target"},)
 
 
 def test_ref2va_packing_tracks_video_and_audio_update_masks():
@@ -101,6 +103,33 @@ def test_ref2va_packing_tracks_video_and_audio_update_masks():
     assert packed["cu_seqlens"].tolist() == [0, 30, 64]
 
 
+def test_ref2va_packing_publishes_physical_video_spans_only():
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence_ref2va_blocks,
+    )
+
+    packed = minimax_h3_packed_sequence_ref2va_blocks(
+        text_len=4,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+        ref_blocks=[
+            {"kind": "image", "latent_h": 4, "latent_w": 4},
+            {"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 6},
+            {"kind": "audio", "ref_audio_t": 1},
+            {"kind": "video", "ref_audio_t": 0, "latent_t": 1, "latent_h": 4, "latent_w": 4},
+        ],
+    )
+
+    # Image and all audio rows intentionally stay out of the sparse spans.
+    assert packed["video_spans"] == (
+        {"start": 12, "latent_grid": (2, 2, 3), "role": "reference"},
+        {"start": 26, "latent_grid": (1, 2, 2), "role": "reference"},
+        {"start": 36, "latent_grid": (2, 2, 3), "role": "target"},
+    )
+
+
 def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
     from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
         minimax_h3_audio_cond_noise_aug_rows,

vllm_omni/diffusion/attention/backends/abstract.py

@@ -3,7 +3,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Literal, TypeVar
 
 import torch
 
@@ -125,9 +125,22 @@ class QueryRange:
     global_start: int
 
 
+@dataclass(frozen=True, slots=True)
+class VideoTokenSpan:
+    """One physical video-grid slice in packed document 0."""
+
+    start: int
+    latent_grid: tuple[int, int, int]
+    role: Literal["reference", "target"]
+
+    @property
+    def length(self) -> int:
+        return self.latent_grid[0] * self.latent_grid[1] * self.latent_grid[2]
+
+
 @dataclass(frozen=True, slots=True)
 class VideoTokenLayout:
-    """Where the video segment sits inside a packed multimodal sequence.
+    """Video-grid slices in a packed multimodal sequence.
 
     A model that packs its sequence as ``[prefix | t*h*w video rows | padding]``
     publishes this so backends can recover spatiotemporal locality; the prefix
@@ -135,11 +148,17 @@ class VideoTokenLayout:
     Publishing it also asserts that any ``attn_mask`` masks only the trailing
     padding, so ``prefix_len + t*h*w`` is the used length of the sequence.
 
+    ``prefix_len``/``latent_grid`` retain the original one-tail contract. A
+    Ref2VA layout instead publishes ``used_len`` and every physical video
+    span, so audio and image rows between videos remain dense.
+
     Plain ints, so reading it never forces a device-to-host sync.
     """
 
-    prefix_len: int
-    latent_grid: tuple[int, int, int]
+    prefix_len: int | None = None
+    latent_grid: tuple[int, int, int] | None = None
+    used_len: int | None = None
+    video_spans: tuple[VideoTokenSpan, ...] = ()
 
 
 @dataclass(frozen=True, slots=True)

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import inspect
 import math
 from dataclasses import dataclass
 from typing import Any
@@ -50,6 +51,11 @@ _MISSING_MINDIESD = (
     "Otherwise, use FlashAttention by setting DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN"
 )
 
+_INCOMPATIBLE_MINDIESD = (
+    "RAINFUSION_ATTN requires a MindIE-SD build whose sparse_attention supports the video_spans "
+    "argument. Please upgrade MindIE-SD or select FLASH_ATTN."
+)
+
 
 def _try_extract_layer_index(prefix: str) -> int | None:
     if not prefix:
@@ -60,6 +66,13 @@ def _try_extract_layer_index(prefix: str) -> int | None:
         return None
 
 
+def _supports_video_spans(sparse_attention: Any) -> bool:
+    try:
+        return "video_spans" in inspect.signature(sparse_attention).parameters
+    except (TypeError, ValueError):
+        return False
+
+
 @dataclass(frozen=True)
 class RainFusionConfig:
     """Resolved RainFusion controls for one attention layer.
@@ -92,9 +105,10 @@ class RainFusionConfig:
 class RainFusionPlan:
     """Per-forward geometry handed to the rf_v2 kernel."""
 
-    prefix_len: int
     used_len: int
-    latent_shape: list[int]
+    prefix_len: int | None = None
+    latent_shape: list[int] | None = None
+    video_spans: list[dict[str, object]] | None = None
 
 
 class RainFusionAttentionBackend(AttentionBackend):
@@ -107,6 +121,12 @@ class RainFusionAttentionBackend(AttentionBackend):
 
         if find_spec("mindiesd") is None:
             raise ValueError(_MISSING_MINDIESD)
+        try:
+            from mindiesd import sparse_attention
+        except ImportError as exc:
+            raise ValueError(_MISSING_MINDIESD) from exc
+        if not _supports_video_spans(sparse_attention):
+            raise ValueError(_INCOMPATIBLE_MINDIESD)
 
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
@@ -266,6 +286,14 @@ class RainFusionAttentionImpl(AttentionImpl):
             )
             return None
 
+        if layout.video_spans:
+            return self._resolve_multi_span_plan(layout, max_seqlen_q)
+
+        if layout.prefix_len is None or layout.latent_grid is None:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: video layout has neither a legacy video tail nor multi-video spans."
+            )
+            return None
         prefix_len = int(layout.prefix_len)
         latent_shape = [int(dim) for dim in layout.latent_grid]
         # rf_v2 splits the sequence as [prefix | t*h*w video rows]. Document 0 of
@@ -304,11 +332,97 @@ class RainFusionAttentionImpl(AttentionImpl):
             video_len,
         )
         return RainFusionPlan(
-            prefix_len=prefix_len,
             used_len=used_len,
+            prefix_len=prefix_len,
             latent_shape=latent_shape,
         )
 
+    def _resolve_multi_span_plan(self, layout, max_seqlen_q: int) -> RainFusionPlan | None:
+        """Validate Ref2VA's non-contiguous video grids before sparse dispatch."""
+        used_len = layout.used_len
+        if used_len is None or int(used_len) != int(max_seqlen_q):
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: multi-video layout used_len=%r does not match packed document 0 (%d).",
+                used_len,
+                int(max_seqlen_q),
+            )
+            return None
+
+        spans: list[dict[str, object]] = []
+        span_summaries: list[str] = []
+        previous_end = 0
+        target_count = 0
+        video_seqlen = 0
+        boundary_dense_seqlen = 0
+        previous_video_length: int | None = None
+        for span in sorted(layout.video_spans, key=lambda item: item.start):
+            grid = tuple(int(dim) for dim in span.latent_grid)
+            length = math.prod(grid)
+            start = int(span.start)
+            if (
+                len(grid) != 3
+                or any(dim <= 0 for dim in grid)
+                or start < previous_end
+                or start + length > int(used_len)
+            ):
+                logger.warning_once(
+                    "RAINFUSION_ATTN staying dense: invalid multi-video span start=%d grid=%s used_len=%d.",
+                    start,
+                    grid,
+                    int(used_len),
+                )
+                return None
+            role = span.role
+            if role not in ("reference", "target"):
+                logger.warning_once("RAINFUSION_ATTN staying dense: unsupported multi-video span role %r.", role)
+                return None
+            if role == "target":
+                target_count += 1
+            if previous_video_length is not None:
+                # rf_v2 works on fixed 128-token blocks. The preceding clip
+                # needs these dense rows to complete its tail block before
+                # this clip begins, otherwise one sparse block would cross a
+                # clip boundary.
+                boundary_dense_seqlen += (-previous_video_length) % _BLOCK_SIZE
+            spans.append({"start": start, "latent_shape": list(grid)})
+            span_summaries.append(f"role={role}, start={start}, seqlen={length}, latent_shape={grid}")
+            previous_end = start + length
+            video_seqlen += length
+            previous_video_length = length
+
+        if target_count != 1:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: Ref2VA layout must contain exactly one target video span, got %d.",
+                target_count,
+            )
+            return None
+        if video_seqlen < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: multi-video seqlen=%d is under the sparse threshold seqlen=%d.",
+                video_seqlen,
+                _MIN_VIDEO_BLOCKS * _BLOCK_SIZE,
+            )
+            return None
+        dense_context_seqlen = int(used_len) - video_seqlen
+        if boundary_dense_seqlen > dense_context_seqlen:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: multi-video spans need dense_context_seqlen=%d "
+                "to isolate clip block boundaries, but this layout has only %d.",
+                boundary_dense_seqlen,
+                dense_context_seqlen,
+            )
+            return None
+        logger.info_once(
+            "RAINFUSION_ATTN multi-video active: sparsity=%.2f, spans=[%s], "
+            "video_seqlen=%d, dense_context_seqlen=%d, valid_packed_seqlen=%d.",
+            self.rainfusion.sparsity,
+            "; ".join(span_summaries),
+            video_seqlen,
+            dense_context_seqlen,
+            int(used_len),
+        )
+        return RainFusionPlan(used_len=int(used_len), video_spans=spans)
+
     def _forward_sparse_npu(
         self,
         query: torch.Tensor,
@@ -325,21 +439,32 @@ class RainFusionAttentionImpl(AttentionImpl):
         q, k, v = (tensor[:, :used] for tensor in (query, key, value))
         # Ulysses has already gathered the full sequence onto this rank and split
         # the heads, so read the head count off the tensor rather than num_heads.
-        out = sparse_attention(
-            q,
-            k,
-            v,
-            scale=self.softmax_scale,
-            head_num=query.shape[-2],
-            input_layout=_INPUT_LAYOUT,
-            inner_precise=_INNER_PRECISE,
-            sparse_type="rf_v2",
-            txt_len=plan.prefix_len,
-            block_size=_BLOCK_SIZE,
-            latent_shape_q=plan.latent_shape,
-            latent_shape_k=plan.latent_shape,
-            sparsity=self.rainfusion.sparsity,
-        )
+        common_kwargs: dict[str, object] = {
+            "scale": self.softmax_scale,
+            "head_num": query.shape[-2],
+            "input_layout": _INPUT_LAYOUT,
+            "inner_precise": _INNER_PRECISE,
+            "block_size": _BLOCK_SIZE,
+            "sparsity": self.rainfusion.sparsity,
+        }
+        if plan.video_spans is not None:
+            out = sparse_attention(
+                q,
+                k,
+                v,
+                sparse_type="rf_v2",
+                video_spans=plan.video_spans,
+                **common_kwargs,
+            )
+        else:
+            assert plan.prefix_len is not None and plan.latent_shape is not None
+            common_kwargs.update(
+                sparse_type="rf_v2",
+                txt_len=plan.prefix_len,
+                latent_shape_q=plan.latent_shape,
+                latent_shape_k=plan.latent_shape,
+            )
+            out = sparse_attention(q, k, v, **common_kwargs)
         if used == query.shape[1]:
             return out
         padded = torch.zeros_like(query)

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -14,7 +14,7 @@ from typing import Any
 
 import torch
 
-from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout
+from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout, VideoTokenSpan
 from vllm_omni.diffusion.forward_context import (
     set_forward_context_denoise_step_idx,
     set_forward_context_denoise_timestep,
@@ -119,11 +119,23 @@ class MiniMaxH3DenoiseBranch:
         }
         # Where the video segment sits in the packed sequence. Resolved to plain
         # ints here so the attention layers never sync on it per step.
-        grid = packed["latent_grid"].tolist()
-        self.static_kwargs["video_token_layout"] = VideoTokenLayout(
-            prefix_len=int(packed["video_row_start"]),
-            latent_grid=(int(grid[0]), int(grid[1]), int(grid[2])),
-        )
+        raw_spans = packed.get("video_spans")
+        if raw_spans:
+            spans = tuple(
+                VideoTokenSpan(
+                    start=int(span["start"]),
+                    latent_grid=tuple(int(dim) for dim in span["latent_grid"]),
+                    role=str(span["role"]),
+                )
+                for span in raw_spans
+            )
+            self.static_kwargs["video_token_layout"] = VideoTokenLayout(used_len=int(cu[1]), video_spans=spans)
+        else:
+            grid = packed["latent_grid"].tolist()
+            self.static_kwargs["video_token_layout"] = VideoTokenLayout(
+                prefix_len=int(packed["video_row_start"]),
+                latent_grid=(int(grid[0]), int(grid[1]), int(grid[2])),
+            )
 
     def forward_kwargs(
         self,

vllm_omni/diffusion/models/minimax_h3/packed_sequence.py

@@ -124,7 +124,7 @@ def minimax_h3_packed_sequence(
     include_keyframe_cond: bool,
     keyframe_frame_indices: list[int] | tuple[int, ...] | None = None,
     frame_count: int | None = None,
-) -> dict[str, torch.Tensor]:
+) -> dict[str, object]:
     """Build the packed-sequence structural fields for one CFG branch.
 
     The used length is padded up to a multiple of 64.
@@ -239,6 +239,14 @@ def minimax_h3_packed_sequence(
         # Extent of the video segment within the packed sequence.
         "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64),
         "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64),
+        # The single target video uses the same span contract as Ref2VA.
+        "video_spans": (
+            {
+                "start": video_sl.start,
+                "latent_grid": (latent_t, ph, pw),
+                "role": "target",
+            },
+        ),
     }
 
 
@@ -300,7 +308,7 @@ def minimax_h3_packed_sequence_ref2va_blocks(
     ref_blocks: Sequence[Mapping[str, object]],
     audio_channel: int = 2,
     seq_len: int | None = None,
-) -> dict[str, torch.Tensor]:
+) -> dict[str, object]:
     """General ref2va-family packed layout.
 
     ``ref_blocks`` are consumed in request/plan order:
@@ -400,6 +408,33 @@ def minimax_h3_packed_sequence_ref2va_blocks(
     video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
     pad_sl = slice(video_sl.stop, seq_len)
 
+    # This is deliberately tensor-free metadata. It is constructed once by
+    # the denoise branch and consumed by every attention layer without a
+    # device-to-host synchronization.
+    video_spans: list[dict[str, object]] = []
+    for item in block_slices:
+        if item["kind"] in ("video", "video_audio"):
+            visual_sl = item["visual_sl"]
+            assert isinstance(visual_sl, slice)
+            video_spans.append(
+                {
+                    "start": visual_sl.start,
+                    "latent_grid": (
+                        int(item["latent_t"]),
+                        int(item["latent_h"]) // _PATCH_H,
+                        int(item["latent_w"]) // _PATCH_W,
+                    ),
+                    "role": "reference",
+                }
+            )
+    video_spans.append(
+        {
+            "start": video_sl.start,
+            "latent_grid": (latent_t, ph, pw),
+            "role": "target",
+        }
+    )
+
     input_ids = torch.full((seq_len,), MINIMAX_H3_PAD_ID, dtype=torch.int64)
     input_ids[text_sl] = MINIMAX_H3_TEXT_ID
     image_mask = torch.zeros(seq_len, dtype=torch.bool)
@@ -560,6 +595,7 @@ def minimax_h3_packed_sequence_ref2va_blocks(
         # Extent of the video segment within the packed sequence.
         "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64),
         "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64),
+        "video_spans": tuple(video_spans),
     }
 
 

vllm_omni/platforms/npu/profiler.py

@@ -53,7 +53,7 @@ class NPUTorchProfilerWrapper(OmniTorchProfilerWrapper):
             export_type=torch_npu.profiler.ExportType.Text,
             profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
             msprof_tx=False,
-            aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
+            aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization,
             l2_cache=False,
             op_attr=False,
             data_simplification=True,