[NPU][Diffusion] MiniMax-H3 RainFusion end_step tail fallback (#6037)

Signed-off-by: HAAZZZEEEE <wang-taicheng@qq.com>

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,

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):

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)

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(

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)
 

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:

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
 
 

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: