diff --git a/docs/user_guide/diffusion/attention_backends/rainfusion.md b/docs/user_guide/diffusion/attention_backends/rainfusion.md index 10f7458be..00f4e37a1 100644 --- a/docs/user_guide/diffusion/attention_backends/rainfusion.md +++ b/docs/user_guide/diffusion/attention_backends/rainfusion.md @@ -17,6 +17,8 @@ globally. | --- | --- | --- | | `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity | | `start_step` | integer, `>= 0` | Number of early denoise steps kept dense | +| `end_step` | integer, `>= 0` | Number of final denoise steps kept dense (tail fallback) | +| `precision` | `"bf16"`, `"fp8"`, `"mix"` | Kernel precision mode; default `"bf16"`. Requires MindIE-SD with `sparse_attention(precision=...)` | | `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense | ```bash @@ -52,6 +54,12 @@ platform raises. It is incompatible with ring sequence parallelism because the kernel needs the complete key sequence for ranking; use Ulysses sequence parallelism with `ring_degree=1`. +The `precision` knob requires a MindIE-SD release whose `sparse_attention` +accepts `precision=`; older releases accept it through `**kwargs` but +silently ignore it and run the BF16 path. `RAINFUSION_ATTN` raises at runtime +when a non-`bf16` precision is requested without that support, so pin a +compatible MindIE-SD release when using `precision="mix"` or `"fp8"`. + ## Geometry handling RainFusion handles arbitrary video grids by rearranging video spatially and, diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 7d0b2a5fa..247f80677 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -101,7 +101,12 @@ class TestAttentionSpec: def test_block_sparse_defaults_applied_when_backend_selected(self): spec = AttentionSpec(backend="RAINFUSION_ATTN") assert spec.block_sparse.sparsity == 0.8 - assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0} + assert spec.backend_kwargs() == { + "sparsity": 0.8, + "start_step": 0, + "end_step": 0, + "precision": "bf16", + } def test_block_sparse_skip_layers_selector_expanded(self): spec = AttentionSpec( @@ -112,6 +117,8 @@ class TestAttentionSpec: assert spec.backend_kwargs() == { "sparsity": 0.9, "start_step": 12, + "end_step": 0, + "precision": "bf16", "skip_layers": [0, 1, 2, 38], } @@ -121,12 +128,44 @@ class TestAttentionSpec: @pytest.mark.parametrize( "block_sparse", - [{"sparsity": 1.5}, {"start_step": -1}], + [{"sparsity": 1.5}, {"start_step": -1}, {"end_step": -1}, {"precision": "bogus"}], ) def test_block_sparse_invalid_values(self, block_sparse): with pytest.raises(ValueError): AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse) + def test_block_sparse_end_step_precision_forwarded(self): + spec = AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse={"sparsity": 0.9, "start_step": 5, "end_step": 3, "precision": "fp8"}, + ) + assert spec.backend_kwargs() == { + "sparsity": 0.9, + "start_step": 5, + "end_step": 3, + "precision": "fp8", + } + + def test_block_sparse_precision_enum_input_normalized(self): + """RainFusionPrecision enum members must normalize to their string value. + + Regression: the str,Enum members passed the value check but str() later + produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it). + """ + from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig + from vllm_omni.diffusion.data import RainFusionPrecision + + spec = AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse={"sparsity": 0.8, "precision": RainFusionPrecision.MIX}, + ) + kwargs = spec.backend_kwargs() + assert kwargs["precision"] == "mix" + # The config path is where str(value) used to mangle enum members. + cfg = RainFusionConfig.from_backend_kwargs(kwargs) + assert cfg.precision == "mix" + assert type(cfg.precision) is str + class TestAttentionConfig: def test_empty_config(self): diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py index b7cedcf43..5b5fcc405 100644 --- a/tests/diffusion/attention/test_rainfusion_plan.py +++ b/tests/diffusion/attention/test_rainfusion_plan.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2. @@ -13,10 +13,12 @@ to the always-kept prefix segment before generating the block mask. import dataclasses import sys import types +from unittest import mock import pytest import torch +from vllm_omni.diffusion.attention.backends import rainfusion_attn from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout, VideoTokenSpan from vllm_omni.diffusion.attention.backends.rainfusion_attn import ( _BLOCK_SIZE, @@ -26,6 +28,8 @@ from vllm_omni.diffusion.attention.backends.rainfusion_attn import ( ) from vllm_omni.platforms import current_omni_platform +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu] + PREFIX_ROWS = 710 # 14 text rows + 696 audio rows ALIGNED_GRID = (62, 24, 40) # 1280x768 -> 59520 video rows, 465 blocks MISALIGNED_GRID = (62, 24, 42) # 1344x768 -> 62496 video rows, 488 blocks + 32 rows @@ -228,7 +232,9 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd() heads, head_dim = 4, 128 q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3)) - impl = make_impl() + # Explicit bf16: the dense comparison exercises the plain sparse kernel, + # and older MindIE-SD releases do not accept precision= (gate would raise). + impl = make_impl(precision="bf16") impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0) plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid)) out = impl._forward_sparse_npu(q, k, v, plan) @@ -241,3 +247,90 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd() ).transpose(1, 2) error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean() assert error < 2e-3, f"mean relative error {error:.4%} against dense attention" + + +# end_step tail fallback: the last ``end_step`` denoise steps must stay dense. + + +def test_end_step_tail_window_stays_dense(): + """step_idx inside the last end_step window must resolve no sparse plan.""" + impl = make_impl(end_step=3) + fc = mock.Mock() + fc.denoise_step_idx = 47 + fc.total_denoise_steps = 50 + with ( + mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True), + mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc), + ): + assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None + + +def test_end_step_before_tail_window_still_sparse(): + """step_idx before the tail window must still resolve a sparse plan.""" + impl = make_impl(end_step=3) + fc = mock.Mock() + fc.denoise_step_idx = 46 + fc.total_denoise_steps = 50 + with ( + mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True), + mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc), + ): + assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None + + +def test_end_step_zero_never_triggers_tail_fallback(): + """end_step=0 must not fall back even on the final denoise step.""" + impl = make_impl(end_step=0) + fc = mock.Mock() + fc.denoise_step_idx = 49 + fc.total_denoise_steps = 50 + with ( + mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True), + mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc), + ): + assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None + + +# precision capability gate: a mindiesd without sparse_attention(precision=) +# must fail fast instead of silently running the BF16 path. + + +def _fake_mindiesd_module(): + """Install a minimal stand-in so the import inside _forward_sparse_npu succeeds.""" + import types + + fake = types.ModuleType("mindiesd") + fake.sparse_attention = lambda *args, **kwargs: None + return fake + + +def test_precision_non_bf16_requires_mindiesd_support(): + """precision != bf16 against a mindiesd lacking the kwarg must raise RuntimeError.""" + import sys + + impl = make_impl(precision="mix") + sys.modules["mindiesd"] = _fake_mindiesd_module() + try: + with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=False): + with pytest.raises(RuntimeError, match="requires MindIE-SD"): + impl._forward_sparse_npu(None, None, None, None) + finally: + sys.modules.pop("mindiesd", None) + + +def test_precision_non_bf16_passes_gate_when_supported(): + """precision != bf16 with a capable mindiesd must not raise from the gate.""" + import sys + + impl = make_impl(precision="mix") + sys.modules["mindiesd"] = _fake_mindiesd_module() + try: + with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=True): + # q/k/v shapes: [B, S, N, D]; plan geometry must match S. + q = torch.randn(1, 8, 4, 128) + plan = RainFusionPlan(prefix_len=0, used_len=8, latent_shape=[2, 2, 2]) + # The gate must pass; the fake mindiesd returns None so no crash. + out = impl._forward_sparse_npu(q, q, q, plan) + assert out is None + finally: + sys.modules.pop("mindiesd", None) diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py index 06e5483b0..9bbde02fa 100644 --- a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py +++ b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +import functools import inspect import math from dataclasses import dataclass @@ -57,6 +58,22 @@ _INCOMPATIBLE_MINDIESD = ( ) +# Whether the installed mindiesd ``sparse_attention`` accepts ``precision=``. +# Releases without it accept the kwarg through ``**kwargs`` but silently ignore +# it, so a requested mix/fp8 mode would silently run the BF16 path. Cached +# because forwards run per layer per denoise step. +@functools.cache +def _mindiesd_supports_precision() -> bool: + try: + from inspect import signature + + from mindiesd import sparse_attention + + return "precision" in signature(sparse_attention).parameters + except Exception: + return False + + def _try_extract_layer_index(prefix: str) -> int | None: if not prefix: return None @@ -85,6 +102,8 @@ class RainFusionConfig: sparsity: float = 0.0 start_step: int = 0 + end_step: int = 0 + precision: str = "bf16" skip_layers: frozenset[int] = frozenset() @classmethod @@ -93,6 +112,8 @@ class RainFusionConfig: return cls( sparsity=float(bk.get("sparsity", 0.0)), start_step=int(bk.get("start_step", 0)), + end_step=int(bk.get("end_step", 0)), + precision=str(bk.get("precision", "bf16")), skip_layers=frozenset(bk.get("skip_layers") or ()), ) @@ -252,8 +273,17 @@ class RainFusionAttentionImpl(AttentionImpl): return None if is_forward_context_available(): step_idx = get_forward_context().denoise_step_idx + total_steps = get_forward_context().total_denoise_steps if step_idx is not None and step_idx < rf.start_step: return None + # Tail fallback: keep the last ``end_step`` denoise steps dense. + if ( + rf.end_step > 0 + and step_idx is not None + and total_steps is not None + and step_idx >= total_steps - rf.end_step + ): + return None if self.qkv_layout is None: # The sparse path reads the sequence off dim 1, which the tensors alone # do not establish, and the dense fallback resolves an absent layout its @@ -434,6 +464,13 @@ class RainFusionAttentionImpl(AttentionImpl): from mindiesd import sparse_attention except ImportError: raise ImportError(_MISSING_MINDIESD) + if self.rainfusion.precision != "bf16" and not _mindiesd_supports_precision(): + raise RuntimeError( + f"block_sparse.precision={self.rainfusion.precision!r} requires MindIE-SD " + "with sparse_attention(precision=...) support; the installed mindiesd " + "silently ignores it and would run the BF16 path. Install a compatible " + "MindIE-SD release or use precision='bf16'." + ) used = plan.used_len q, k, v = (tensor[:, :used] for tensor in (query, key, value)) @@ -446,6 +483,7 @@ class RainFusionAttentionImpl(AttentionImpl): "inner_precise": _INNER_PRECISE, "block_size": _BLOCK_SIZE, "sparsity": self.rainfusion.sparsity, + "precision": self.rainfusion.precision, } if plan.video_spans is not None: out = sparse_attention( diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index e89825e0e..590afcd8c 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -7,7 +7,7 @@ import os import random from collections.abc import Callable, Mapping from dataclasses import dataclass, field, fields -from enum import Enum +from enum import Enum, StrEnum from typing import TYPE_CHECKING, Any import diffusers @@ -1675,6 +1675,20 @@ class AttnQuantSpec: BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"}) +class RainFusionPrecision(StrEnum): + """Execution precision for block-sparse RainFusion (rf_v3) attention. + + ``bf16``: no quantization, pure BF16 sparse attention (official baseline). + ``fp8``: BSA FP8 path - Hadamard rotation then full FP8 block quantization + of Q/K/V before the BSA kernel. + ``mix``: EagleQBSA mixed precision - Q/K per-block INT8 + V per-channel FP8. + """ + + BF16 = "bf16" + FP8 = "fp8" + MIX = "mix" + + @dataclass class BlockSparseSpec: """User-facing controls shared by block-sparse attention backends. @@ -1683,10 +1697,13 @@ class BlockSparseSpec: ``start_step`` keeps the first N denoise steps dense and ``skip_layers`` (an index selector such as "0-3,38") exempts individual DiT blocks. Those two are the accuracy knobs to trade back quality at a fixed ``sparsity``. + ``precision`` selects the kernel precision mode (see ``RainFusionPrecision``). """ sparsity: float = 0.8 start_step: int = 0 + end_step: int = 0 + precision: str = RainFusionPrecision.BF16.value skip_layers: str | list[int] | None = None skip_layer_indices: set[int] | None = field(default=None, repr=False) @@ -1694,6 +1711,13 @@ class BlockSparseSpec: self.sparsity = _in_range(self.sparsity, "block_sparse.sparsity", 0.0, 1.0) or 0.0 if self.start_step < 0: raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.") + if self.end_step < 0: + raise ValueError(f"block_sparse.end_step must be >= 0; got {self.end_step!r}.") + if self.precision not in {p.value for p in RainFusionPrecision}: + raise ValueError( + f"block_sparse.precision must be one of {sorted(p.value for p in RainFusionPrecision)}; " + f"got {self.precision!r}." + ) self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers) @@ -1776,6 +1800,8 @@ class AttentionSpec: bs = self.block_sparse kw["sparsity"] = bs.sparsity kw["start_step"] = bs.start_step + kw["end_step"] = bs.end_step + kw["precision"] = bs.precision if bs.skip_layer_indices: kw["skip_layers"] = sorted(bs.skip_layer_indices) diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index f069cee85..b8fddcfe8 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from contextlib import contextmanager @@ -34,6 +37,7 @@ class ForwardContext: split_text_embed_in_sp: bool = False denoise_step_idx: int | None = None denoise_timestep: float | None = None + total_denoise_steps: int | None = None # Per-request reference latent for img2img DiT models (e.g. Ming) ref_latent: torch.Tensor | None = None # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel @@ -255,6 +259,16 @@ def set_forward_context_denoise_timestep(timestep: float | None) -> None: _forward_context.denoise_timestep = None if timestep is None else float(timestep) +def set_forward_context_denoise_total_steps(total_steps: int | None) -> None: + """Set the total denoise step count on the active ForwardContext. + + Denoise loops publish it so tail-fallback gates (e.g. ``end_step`` in + RAINFUSION_ATTN) know when the final denoise steps begin. + """ + if _forward_context is not None: + _forward_context.total_denoise_steps = total_steps + + class DenoiseProgressMixin: def record_denoise_step( self, @@ -262,8 +276,11 @@ class DenoiseProgressMixin: timestep=None, scheduler=None, normalized_timestep: float | None = None, + total_steps: int | None = None, ) -> None: set_forward_context_denoise_step_idx(step_idx) + if _forward_context is not None: + _forward_context.total_denoise_steps = total_steps if _forward_context is None: return if normalized_timestep is not None: diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index f6b92aaf7..653dfce71 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -19,6 +19,7 @@ from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout, Vi from vllm_omni.diffusion.forward_context import ( set_forward_context_denoise_step_idx, set_forward_context_denoise_timestep, + set_forward_context_denoise_total_steps, ) from vllm_omni.platforms import current_omni_platform @@ -36,15 +37,19 @@ MINIMAX_H3_VIDEO_ROW_WIDTH = 96 MINIMAX_H3_AUDIO_ROW_WIDTH = 32 -def minimax_h3_publish_denoise_progress(step: int | None, sigma_video: float | None) -> None: +def minimax_h3_publish_denoise_progress( + step: int | None, sigma_video: float | None, total_steps: int | None = None +) -> None: """Publish denoise progress for step-gated attention features. - Both execution modes must publish the same pair: the step index drives the - dense warmup of RAINFUSION_ATTN, and the normalized descending timestep - drives the TRTLLM_ATTN skip gate, which stays dense while it is unset. + Both execution modes must publish the same trio: the step index drives the + dense warmup of RAINFUSION_ATTN, the normalized descending timestep + drives the TRTLLM_ATTN skip gate (which stays dense while it is unset), + and the total step count enables the ``end_step`` tail fallback. """ set_forward_context_denoise_step_idx(step) set_forward_context_denoise_timestep(sigma_video) + set_forward_context_denoise_total_steps(total_steps) class MiniMaxH3DenoiseBranch: @@ -297,7 +302,7 @@ def minimax_h3_denoise_loop( # warmup of RAINFUSION_ATTN, the timestep gate of TRTLLM_ATTN) can # see it. Gates use the scheduler-style descending timestep, which # for a rectified-flow schedule is the video sigma. - minimax_h3_publish_denoise_progress(step, s_v) + minimax_h3_publish_denoise_progress(step, s_v, num_steps) t_v, t_a = 1.0 - s_v, 1.0 - s_a imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference)) audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference)) @@ -341,7 +346,7 @@ def minimax_h3_denoise_loop( if on_step is not None: on_step(step, video_rows, audio_rows) - minimax_h3_publish_denoise_progress(None, None) + minimax_h3_publish_denoise_progress(None, None, None) return video_rows, audio_rows diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index d14bfd0fd..5117dc02d 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -2567,8 +2567,11 @@ class MiniMaxH3Pipeline( # attention features. Requests can differ in both step index and sigma # schedule, so a batch that is not at one single point has nothing to # publish and those gates stay dense -- which is their safe default. - progress = {(state.step_index, schedule["sigma_video"]) for state, schedule in zip(batch_states, schedules)} - minimax_h3_publish_denoise_progress(*(progress.pop() if len(progress) == 1 else (None, None))) + progress = { + (state.step_index, schedule["sigma_video"], len(state.extra[_STEP_SIGMAS_VIDEO]) - 1) + for state, schedule in zip(batch_states, schedules) + } + minimax_h3_publish_denoise_progress(*(progress.pop() if len(progress) == 1 else (None, None, None))) if len(batch_states) > 1 and (mixed_transformers or not self._packed_batch_supported(transformers[0])): if mixed_transformers: