diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py index f079a1758..b7cedcf43 100644 --- a/tests/diffusion/attention/test_rainfusion_plan.py +++ b/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 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py index efe53f826..53348f579 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 9493e305a..c0b9a738e 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py index 28b032712..06e5483b0 100644 --- a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py +++ b/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) diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index f32a037c1..af6feb963 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/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, diff --git a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py index a57eae700..0dcbd19b8 100644 --- a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py +++ b/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), } diff --git a/vllm_omni/platforms/npu/profiler.py b/vllm_omni/platforms/npu/profiler.py index 65c9c92d9..27f5a647a 100644 --- a/vllm_omni/platforms/npu/profiler.py +++ b/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,