[NPU][Quantization] Add RainFusion attention and INT8 online quantization for MiniMax H3 (#5706)

Signed-off-by: Zijun Huang <zijunhua@usc.edu>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>

docs/design/module/dit_module.md

@@ -539,7 +539,7 @@ def get_attn_backend_for_role(role, head_size, attention_config=None, role_categ
 
 4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
 
-A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today the only such field is `skip_softmax` (TRTLLM_ATTN), so e.g. `target_sparsity` can be set without changing model code.
+A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today those fields are `skip_softmax` (TRTLLM_ATTN) and `block_sparse` (any block-sparse backend, currently RAINFUSION_ATTN), so e.g. `target_sparsity` or `sparsity` can be set without changing model code.
 
 For the user-facing CLI surface, see [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../contributing/model/adding_diffusion_model.md).
 

docs/user_guide/diffusion/attention_backends.md

@@ -23,6 +23,7 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
 | `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
 | `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
+| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |
 
 
 ## Configuration
@@ -69,8 +70,9 @@ vllm-omni serve <model> \
     --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
 ```
 
-A backend that needs configuration exposes it as a typed field on the spec. Today the only
-one is `TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)):
+A backend that needs configuration exposes it as a typed field on the spec:
+`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
+`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
 
 ```bash
 --diffusion-attention-config.default.backend TRTLLM_ATTN \
@@ -204,11 +206,80 @@ AttentionConfig(
 )
 ```
 
-
 **Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
 availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
 (B300); `int8` QK kernels are compiled for **SM100 only**.
 
+## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
+
+`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
+128-token block of keys, ranks them per query block, and computes attention against only the
+top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
+spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
+
+Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
+first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
+nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
+exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
+delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
+
+Enable it through the typed `block_sparse` block on the attention spec, shared by every
+block-sparse backend:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
+| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
+| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+  --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
+      "block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
+```
+
+Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="RAINFUSION_ATTN",
+        block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
+    ),
+)
+```
+
+**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
+cheapest way to recover structure, because the early high-noise steps decide global layout while the
+later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
+`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
+
+Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
+incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
+blocks — use Ulysses SP (`ring_degree=1`).
+
+### Which geometries run sparse
+
+Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
+is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
+tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
+that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
+attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
+speedup, so pick an aligned geometry when you want one.
+
+Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
+spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
+patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
+are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
+distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
+to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
+
+"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
+resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
+which you can buy back with `start_step`.
+
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 
 Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -50,7 +50,8 @@ uv pip install -e .
 ```
 
 Install the **mindie-sd** operator library to enable Ascend-optimized fused
-operators (`adalayernorm`, etc.):
+operators (`adalayernorm`, etc.) and the RainFusion `rf_v2` block-sparse
+attention kernel:
 
 ```bash
 git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
@@ -111,6 +112,28 @@ Do not add `--enforce-eager`. The first request includes regional
 compilation; warm the server once before measuring steady-state latency.
 H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
 
+### Optional optimizations
+
+Two independent optimizations may be enabled on top of the configuration
+above. Both are validated for T2VA only.
+
+**RainFusion block-sparse attention** — switch the attention backend, keeping
+every other flag unchanged:
+
+```bash
+  --diffusion-attention-backend RAINFUSION_ATTN
+```
+
+**INT8 online quantization** — add one flag to either command above:
+
+```bash
+  --quantization int8
+```
+
+Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
+over the whole sequence, so ring parallelism would split away the keys it
+needs. Scale with `--usp` instead.
+
 To serve Ref2VA, stop the FL2VA server and restart with:
 
 ```bash
@@ -154,9 +177,16 @@ throughput guarantee.
   the same HTTP request contract.
 - VAE patch parallelism requires size 1 or the full DiT group size and
   supports the H3 native `tile` mode only.
+- RainFusion block-sparse attention and INT8 quantization are validated for
+  T2VA only; use the BF16 dense configuration for FL2VA and Ref2VA.
+- Online quantization cannot be combined with distributed layerwise offload
+  while AllGather is enabled; pass `--dlo-no-use-allgather` in that case.
 
 ## Additional resources
 
 - [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
+- [Attention backends § RAINFUSION_ATTN](../../docs/user_guide/diffusion/attention_backends.md#rainfusion_attn-backend-and-block-sparse-video-attention)
+  — RainFusion knobs and tuning
+- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
 - [Supported models](../../docs/models/supported_models.md)
 - [Video API](../../docs/serving/videos_api.md)

tests/diffusion/attention/test_attention_config.py

@@ -90,6 +90,35 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match=match):
             AttentionSpec(**spec)
 
+    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}
+
+    def test_block_sparse_skip_layers_selector_expanded(self):
+        spec = AttentionSpec(
+            backend="RAINFUSION_ATTN",
+            block_sparse={"sparsity": 0.9, "start_step": 12, "skip_layers": "0-2,38"},
+        )
+        assert spec.block_sparse.skip_layer_indices == {0, 1, 2, 38}
+        assert spec.backend_kwargs() == {
+            "sparsity": 0.9,
+            "start_step": 12,
+            "skip_layers": [0, 1, 2, 38],
+        }
+
+    def test_block_sparse_rejected_on_dense_backend(self):
+        with pytest.raises(ValueError, match="block_sparse is only supported by"):
+            AttentionSpec(backend="FLASH_ATTN", block_sparse={"sparsity": 0.8})
+
+    @pytest.mark.parametrize(
+        "block_sparse",
+        [{"sparsity": 1.5}, {"start_step": -1}],
+    )
+    def test_block_sparse_invalid_values(self, block_sparse):
+        with pytest.raises(ValueError):
+            AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse)
+
 
 class TestAttentionConfig:
     def test_empty_config(self):

tests/diffusion/attention/test_rainfusion_plan.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2.
+
+The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
+1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
+first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
+mask from the video and prefix segments separately, so the two only agree on the
+block count when the video segment ends on a boundary.
+"""
+
+import dataclasses
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
+from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
+    _BLOCK_SIZE,
+    RainFusionAttentionImpl,
+    RainFusionPlan,
+)
+from vllm_omni.platforms import current_omni_platform
+
+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
+
+
+def make_impl(**backend_kwargs):
+    return RainFusionAttentionImpl(
+        num_heads=8,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        prefix="transformer_blocks.0.attn",
+        qkv_layout="BSND",
+        backend_kwargs={"sparsity": 0.8, **backend_kwargs},
+    )
+
+
+def make_metadata(grid, prefix_len=PREFIX_ROWS, max_seqlen_q=None):
+    video_rows = grid[0] * grid[1] * grid[2]
+    return AttentionMetadata(
+        extra={"max_seqlen_q": prefix_len + video_rows if max_seqlen_q is None else max_seqlen_q},
+        video_layout=VideoTokenLayout(prefix_len=prefix_len, latent_grid=grid),
+    )
+
+
+def test_block_aligned_video_segment_runs_sparse():
+    plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID))
+
+    assert plan is not None
+    assert plan.prefix_len == PREFIX_ROWS
+    assert plan.used_len == PREFIX_ROWS + 59520
+    assert plan.latent_shape == list(ALIGNED_GRID)
+
+
+def test_misaligned_video_segment_stays_dense():
+    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
+    # pad keys would take a share of every real query's softmax denominator.
+    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+
+
+@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
+def test_prefix_length_does_not_affect_alignment(prefix_len):
+    # Only the video segment has to land on a block boundary; rf_v2 pools the
+    # prefix separately and keeps every one of its blocks.
+    plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID, prefix_len=prefix_len))
+
+    assert plan is not None
+    assert plan.prefix_len == prefix_len
+
+
+def test_sparsity_zero_never_resolves_a_plan():
+    assert make_impl(sparsity=0.0)._resolve_plan(make_metadata(ALIGNED_GRID)) is None
+
+
+def test_missing_video_layout_falls_back_to_dense():
+    assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 60230})) is None
+
+
+def test_missing_max_seqlen_falls_back_to_dense():
+    metadata = AttentionMetadata(
+        video_layout=VideoTokenLayout(prefix_len=PREFIX_ROWS, latent_grid=ALIGNED_GRID),
+    )
+
+    assert make_impl()._resolve_plan(metadata) is None
+
+
+def test_video_segment_must_be_the_tail_of_packed_document_zero():
+    metadata = make_metadata(ALIGNED_GRID, max_seqlen_q=PREFIX_ROWS + 59520 + 128)
+
+    assert make_impl()._resolve_plan(metadata) is None
+
+
+@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
+
+
+@pytest.mark.parametrize("qkv_layout", ["BNSD", "BSH"])
+def test_explicitly_wrong_layout_is_rejected(qkv_layout):
+    with pytest.raises(ValueError, match="BSND"):
+        RainFusionAttentionImpl(
+            num_heads=8,
+            head_size=128,
+            softmax_scale=128**-0.5,
+            qkv_layout=qkv_layout,
+            backend_kwargs={"sparsity": 0.8},
+        )
+
+
+def test_undeclared_layout_stays_dense():
+    # An undeclared layout is not an error -- the layer keeps working, just
+    # densely, and the fallback sees the same absent layout plain FLASH_ATTN would.
+    impl = RainFusionAttentionImpl(
+        num_heads=8,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        prefix="transformer_blocks.0.attn",
+        backend_kwargs={"sparsity": 0.8},
+    )
+
+    assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None
+    assert impl.dense_fallback.qkv_layout is None
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
+def test_fully_populated_mask_reproduces_dense_attention():
+    """At sparsity=0 every key block is kept, so rf_v2 must match dense attention.
+
+    This is the guard on the geometry the backend hands the kernel: a wrong
+    prefix length or latent grid still produces plausible output under sparse
+    selection, but shows up here as a mismatch against SDPA.
+    """
+    torch.manual_seed(0)
+    prefix_len, grid = 384, (4, 16, 64)  # 4096 video rows = 32 blocks
+    video_len = grid[0] * grid[1] * grid[2]
+    used = prefix_len + video_len
+    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()
+    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)
+
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        k.transpose(1, 2),
+        v.transpose(1, 2),
+        scale=impl.softmax_scale,
+    ).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"

tests/diffusion/quantization/test_int8_config.py

@@ -147,6 +147,130 @@ def test_get_npu_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyP
     assert isinstance(method, UnquantizedLinearMethod)
 
 
+def test_fused_layer_can_be_ignored_by_its_prefix(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture):
+    """A layer opts out of quantization under the prefix its module was built with.
+
+    MiniMax H3 stores qkv as one checkpoint tensor, so the fused name is the
+    only handle a user has on it.
+    """
+    from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+    config = build_quant_config("int8", ignored_layers=["blocks.0.attn.qkv_proj"])
+    layer = mocker.Mock(spec=LinearBase)
+
+    assert isinstance(config.get_quant_method(layer, "blocks.0.attn.qkv_proj"), UnquantizedLinearMethod)
+    assert isinstance(config.get_quant_method(layer, "blocks.1.attn.qkv_proj"), NPUInt8OnlineLinearMethod)
+    assert isinstance(config.get_quant_method(layer, "blocks.0.mlp.fc1"), NPUInt8OnlineLinearMethod)
+
+
+class TestNPUQuantMatmulShapeFallback:
+    """npu_quant_matmul rejects wide outputs, so those layers stay unquantized.
+
+    The limit applies to the per-rank shard, which is what the kernel sees. The
+    same layer can therefore be quantized at a higher TP degree.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _mock_tp(self, mocker):
+        # The fallback delegates to UnquantizedLinearMethod, which reads the TP
+        # group; stand one in so the shape gating can be tested without a
+        # distributed init.
+        mock_group = mocker.Mock()
+        mock_group.rank_in_group = 0
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+        mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+        mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
+    @pytest.fixture
+    def quant_config(self):
+        from vllm_omni.quantization.int8_config import DiffusionInt8Config
+
+        return DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+
+    def create_weights(self, method, layer, output_partition_sizes):
+        # Narrow input dim: the fallback path allocates for real, and only the
+        # output dimension is under test.
+        method.create_weights(
+            layer,
+            input_size_per_partition=8,
+            output_partition_sizes=output_partition_sizes,
+            input_size=8,
+            output_size=sum(output_partition_sizes),
+            params_dtype=torch.bfloat16,
+            weight_loader=lambda param, loaded_weight, *args, **kwargs: None,
+        )
+
+    @pytest.mark.parametrize(
+        "method_name",
+        ["NPUInt8LinearMethod", "NPUInt8OnlineLinearMethod"],
+    )
+    def test_over_wide_layer_falls_back_to_unquantized(self, quant_config, method_name):
+        import vllm_omni.quantization.int8_config as int8_config
+
+        method = getattr(int8_config, method_name)(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [int8_config.NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1])
+
+        assert isinstance(layer.quant_method, UnquantizedLinearMethod)
+        assert layer.weight.dtype == torch.bfloat16
+
+    def test_layer_within_the_limit_keeps_int8(self, quant_config):
+        from vllm_omni.quantization.int8_config import (
+            NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+            NPUInt8OnlineLinearMethod,
+        )
+
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [NPU_QUANT_MATMUL_MAX_OUT_FEATURES])
+
+        assert layer.quant_method is method
+
+    def test_tensor_parallel_sharding_brings_a_wide_layer_back(self, quant_config):
+        from vllm_omni.quantization.int8_config import (
+            NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+            NPUInt8OnlineLinearMethod,
+        )
+
+        global_out_features = 4 * NPU_QUANT_MATMUL_MAX_OUT_FEATURES
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = Module()
+        layer.quant_method = method
+
+        self.create_weights(method, layer, [global_out_features // 4])
+
+        assert layer.quant_method is method
+
+
+class TestOffloadAfterQuant:
+    def test_only_methods_advertising_the_capability_are_asked(self):
+        from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+        from vllm_omni.quantization.int8_config import DiffusionInt8Config, NPUInt8OnlineLinearMethod
+
+        class MetaDeviceOnlyMethod:
+            """Stands in for upstream online FP8: lazy, but no offload hook."""
+
+            uses_meta_device = True
+
+        model = torch.nn.Sequential(Module(), Module())
+        lazy_int8 = NPUInt8OnlineLinearMethod(
+            DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+        )
+        model[0].quant_method = lazy_int8
+        model[1].quant_method = MetaDeviceOnlyMethod()
+
+        marked = DiffusersPipelineLoader._request_offload_after_quant(model)
+
+        assert marked == 1
+        assert lazy_int8._offload_after_quant
+
+
 class TestInt8LinearMethod:
     @pytest.fixture
     def mock_quant_config(self, mocker):

vllm_omni/diffusion/attention/backends/abstract.py

@@ -20,6 +20,21 @@ class AttentionBackend(ABC):
     # avoid a slower masked-attention plan when tail padding is not semantic.
     supports_prefix_kv_slicing: bool = False
 
+    # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
+    # Platform resolution rejects an explicit selection outside this set, so a
+    # hardware-specific backend fails before the model is built.
+    supported_platforms: tuple[str, ...] | None = None
+
+    @classmethod
+    def validate_available(cls) -> None:
+        """Raise if this backend's optional dependencies are missing.
+
+        Called during platform resolution, i.e. before model construction, so a
+        backend that probes its kernel package lazily still reports the problem
+        while the user can still act on it.
+        """
+        return None
+
     @classmethod
     def supports_attention_mask(cls) -> bool:
         return False
@@ -63,6 +78,23 @@ class QueryRange:
     global_start: int
 
 
+@dataclass(frozen=True, slots=True)
+class VideoTokenLayout:
+    """Where the video segment sits inside 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
+    holds everything that is not video (text, visual conditions, audio).
+    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.
+
+    Plain ints, so reading it never forces a device-to-host sync.
+    """
+
+    prefix_len: int
+    latent_grid: tuple[int, int, int]
+
+
 @dataclass
 class AttentionMetadata:
     attn_mask: torch.Tensor | None = None
@@ -95,6 +127,10 @@ class AttentionMetadata:
     full_attn_spans: list[list[tuple[int, int]]] | None = None
     query_ranges: tuple[QueryRange, ...] | None = None
 
+    # Geometry of the video segment for backends that exploit spatiotemporal
+    # locality (block-sparse selection, tiled masks). Dense backends ignore it.
+    video_layout: VideoTokenLayout | None = None
+
 
 T = TypeVar("T", bound=AttentionMetadata)
 

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import extract_layer_index
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
+from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+logger = init_logger(__name__)
+
+# The rf_v2 kernel only implements a 128-token block.
+_BLOCK_SIZE = 128
+
+# Below this many video blocks, the pooling and gather that block selection adds
+# cost more than the QK work it removes, so stay dense.
+_MIN_VIDEO_BLOCKS = 32
+
+# rf_v2's ``input_layout`` describes the caller's tensors, and everything below
+# slices the sequence on dim 1. vLLM-Omni diffusion attention hands the impl
+# [B, S, N, D], so that is the only layout this backend accepts.
+_INPUT_LAYOUT = "BSND"
+
+# rf_v2 precision mode: 0 = high precision, 1 = high performance. Kept at the
+# precise setting because the sparsity knob is the intended perf lever, and on
+# A5 devices the kernel is routed to rf_v3, which overrides this anyway.
+_INNER_PRECISE = 0
+
+_WRONG_PLATFORM = (
+    "RAINFUSION_ATTN runs the MindIE-SD rf_v2 kernel and is available on Ascend NPU only. "
+    "Select FLASH_ATTN or TORCH_SDPA on this platform."
+)
+
+_MISSING_MINDIESD = (
+    "RAINFUSION_ATTN requires MindIE-SD. Please install MindIE-SD to enable RainFusion sparse "
+    "attention on Ascend NPU. For installation details, see https://gitcode.com/Ascend/MindIE-SD "
+    "Otherwise, use FlashAttention by setting DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN"
+)
+
+
+def _try_extract_layer_index(prefix: str) -> int | None:
+    if not prefix:
+        return None
+    try:
+        return extract_layer_index(prefix)
+    except (AssertionError, ValueError):
+        return None
+
+
+@dataclass(frozen=True)
+class RainFusionConfig:
+    """Resolved RainFusion controls for one attention layer.
+
+    ``sparsity`` is the nominal fraction of key blocks dropped per query block.
+    The realized sparsity is lower because rf_v2 always keeps the prefix rows and
+    the first-frame blocks. ``start_step`` and ``skip_layers`` are the accuracy
+    knobs: early denoise steps and specific DiT blocks stay dense.
+    """
+
+    sparsity: float = 0.0
+    start_step: int = 0
+    skip_layers: frozenset[int] = frozenset()
+
+    @classmethod
+    def from_backend_kwargs(cls, backend_kwargs: dict | None) -> RainFusionConfig:
+        bk = backend_kwargs or {}
+        return cls(
+            sparsity=float(bk.get("sparsity", 0.0)),
+            start_step=int(bk.get("start_step", 0)),
+            skip_layers=frozenset(bk.get("skip_layers") or ()),
+        )
+
+    @property
+    def enabled(self) -> bool:
+        return self.sparsity > 0.0
+
+
+@dataclass(frozen=True)
+class RainFusionPlan:
+    """Per-forward geometry handed to the rf_v2 kernel."""
+
+    prefix_len: int
+    used_len: int
+    latent_shape: list[int]
+
+
+class RainFusionAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+    supported_platforms: tuple[str, ...] = ("npu",)
+
+    @classmethod
+    def validate_available(cls) -> None:
+        from importlib.util import find_spec
+
+        if find_spec("mindiesd") is None:
+            raise ValueError(_MISSING_MINDIESD)
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "RAINFUSION_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type[RainFusionAttentionImpl]:
+        return RainFusionAttentionImpl
+
+
+class RainFusionAttentionImpl(AttentionImpl):
+    """Block-sparse video attention via MindIE-SD RainFusion (rf_v2) on Ascend NPU.
+
+    Sparsity applies only to the video segment of a packed multimodal sequence,
+    whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
+    other case — warmup denoise steps, exempt layers, a layer that does not declare
+    ``qkv_layout="BSND"``, sequences without a published video segment, video
+    segments too short to pay for block selection or not aligned to the kernel's
+    block size — delegates to FlashAttention, so a model can select this backend
+    unconditionally.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict[str, Any] | None = None,
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.qkv_layout = qkv_layout
+
+        self.rainfusion = RainFusionConfig.from_backend_kwargs(backend_kwargs)
+        self.layer_idx = _try_extract_layer_index(prefix)
+
+        if self.rainfusion.enabled:
+            self._validate_parallel_config()
+            if causal:
+                raise ValueError(
+                    "RAINFUSION_ATTN does not support causal attention: rf_v2 selects key "
+                    "blocks by pooled relevance and cannot express a causal mask. Select "
+                    "FLASH_ATTN for causal roles."
+                )
+            if qkv_layout is not None and qkv_layout.upper() != _INPUT_LAYOUT:
+                raise ValueError(
+                    f"RAINFUSION_ATTN needs {_INPUT_LAYOUT} tensors to locate the video segment along "
+                    f"the sequence axis, but this layer declares qkv_layout={qkv_layout!r}. Select "
+                    "FLASH_ATTN for this role."
+                )
+
+        self.dense_fallback = FlashAttentionBackend.get_impl_cls()(
+            num_heads=num_heads,
+            head_size=head_size,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            num_kv_heads=num_kv_heads,
+            prefix=prefix,
+            qkv_layout=qkv_layout,
+        )
+
+    def _validate_parallel_config(self) -> None:
+        config = get_current_diffusion_config_or_none()
+        parallel_config = getattr(config, "parallel_config", None)
+        ring_degree = getattr(parallel_config, "ring_degree", 1)
+        if ring_degree > 1:
+            # Ring gives each rank a slice of the sequence, so block selection
+            # would score only local keys and the layer bypasses the backend
+            # entirely (see Attention._run_ring_attention).
+            raise ValueError(
+                "RAINFUSION_ATTN is not compatible with ring sequence parallelism "
+                f"(ring_degree={ring_degree}): rf_v2 needs the whole key sequence to rank "
+                "blocks. Use Ulysses SP (ring_degree=1) instead."
+            )
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        # ROCm and MUSA route through forward_cuda by default, so this covers them too.
+        raise NotImplementedError(_WRONG_PLATFORM)
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError(_WRONG_PLATFORM)
+
+    def forward_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        plan = self._resolve_plan(attn_metadata)
+        if plan is None:
+            return self.dense_fallback.forward_npu(query, key, value, attn_metadata)
+        return self._forward_sparse_npu(query, key, value, plan)
+
+    def _resolve_plan(self, attn_metadata: AttentionMetadata | None) -> RainFusionPlan | None:
+        """Return the rf_v2 geometry, or None when this forward must stay dense."""
+        rf = self.rainfusion
+        if not rf.enabled:
+            return None
+        if self.layer_idx is not None and self.layer_idx in rf.skip_layers:
+            return None
+        if is_forward_context_available():
+            step_idx = get_forward_context().denoise_step_idx
+            if step_idx is not None and step_idx < rf.start_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
+            # own way. Sparsifying on an assumption would put the two paths on
+            # different axes, so an undeclared layout stays dense.
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: this layer does not declare qkv_layout, and rf_v2 "
+                "needs %s to locate the video segment along the sequence axis. Set qkv_layout=%r on "
+                "the Attention layer to enable sparsity.",
+                _INPUT_LAYOUT,
+                _INPUT_LAYOUT,
+            )
+            return None
+
+        if attn_metadata is None:
+            return None
+
+        layout = attn_metadata.video_layout
+        if layout is None:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: this attention role carries no video segment. The "
+                "model must publish AttentionMetadata.video_layout for the sequence to be sparsified."
+            )
+            return None
+        max_seqlen_q = attn_metadata.extra.get("max_seqlen_q")
+        if max_seqlen_q is None:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: attention metadata is missing max_seqlen_q, so the "
+                "video segment cannot be confirmed to be the tail of packed document 0."
+            )
+            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
+        # the packed sequence holds those rows; anything past it is alignment
+        # padding that rf_v2 must not see.
+        video_len = math.prod(latent_shape)
+        used_len = prefix_len + video_len
+
+        if used_len != int(max_seqlen_q):
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: prefix (%d) plus latent grid %s does not fill "
+                "packed document 0 (%d rows). rf_v2 requires the video segment to be its tail.",
+                prefix_len,
+                tuple(latent_shape),
+                int(max_seqlen_q),
+            )
+            return None
+        if video_len < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE:
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: %d video rows is under the %d-row "
+                "(%d block) threshold where sparse selection pays off.",
+                video_len,
+                _MIN_VIDEO_BLOCKS * _BLOCK_SIZE,
+                _MIN_VIDEO_BLOCKS,
+            )
+            return None
+        if video_len % _BLOCK_SIZE != 0:
+            # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
+            # builds its block mask from the two segments pooled apart. The counts
+            # only agree when the video segment ends on a block boundary; otherwise
+            # a prefix block is forced on out of range, and the block straddling the
+            # seam holds real prefix rows that selection may drop. Padding the
+            # sequence to realign it is not an option because rf_v2 ignores
+            # attn_mask, so pad keys would take a share of every softmax denominator.
+            logger.warning_once(
+                "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
+                "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
+                "multiple of %d to enable sparsity.",
+                tuple(latent_shape),
+                video_len,
+                _BLOCK_SIZE,
+                _BLOCK_SIZE,
+            )
+            return None
+
+        logger.info_once(
+            "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
+            "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "
+            "because prefix and first-frame blocks are always kept.",
+            rf.sparsity,
+            rf.start_step,
+            len(rf.skip_layers),
+            tuple(latent_shape),
+            prefix_len,
+            video_len,
+        )
+        return RainFusionPlan(
+            prefix_len=prefix_len,
+            used_len=used_len,
+            latent_shape=latent_shape,
+        )
+
+    def _forward_sparse_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        plan: RainFusionPlan,
+    ) -> torch.Tensor:
+        try:
+            from mindiesd import sparse_attention
+        except ImportError:
+            raise ImportError(_MISSING_MINDIESD)
+
+        used = plan.used_len
+        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,
+        )
+        if used == query.shape[1]:
+            return out
+        padded = torch.zeros_like(query)
+        padded[:, :used] = out
+        return padded

vllm_omni/diffusion/attention/backends/registry.py

@@ -65,6 +65,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
     FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
     TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend"
+    RAINFUSION_ATTN = "vllm_omni.diffusion.attention.backends.rainfusion_attn.RainFusionAttentionBackend"
 
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).

vllm_omni/diffusion/data.py

@@ -1461,6 +1461,33 @@ class AttnQuantSpec:
         return self.dtype_qk is not None or self.dtype_vo is not None
 
 
+# Backends that select key blocks instead of attending densely, and so accept
+# a ``block_sparse`` spec. Each maps the same knobs onto its own kernel.
+BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"})
+
+
+@dataclass
+class BlockSparseSpec:
+    """User-facing controls shared by block-sparse attention backends.
+
+    ``sparsity`` is the nominal fraction of key blocks dropped per query block;
+    ``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``.
+    """
+
+    sparsity: float = 0.8
+    start_step: int = 0
+    skip_layers: str | list[int] | None = None
+    skip_layer_indices: set[int] | None = field(default=None, repr=False)
+
+    def __post_init__(self) -> None:
+        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}.")
+        self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers)
+
+
 @dataclass
 class AttentionSpec:
     """Specifies a backend and its typed backend-specific config for one attention role."""
@@ -1468,6 +1495,7 @@ class AttentionSpec:
     backend: str
     skip_softmax: SkipSoftmaxSpec | None = None
     quant: AttnQuantSpec | None = None
+    block_sparse: BlockSparseSpec | None = None
     skip_calibration: dict | None = field(default=None, repr=False)
 
     def __post_init__(self) -> None:
@@ -1475,6 +1503,7 @@ class AttentionSpec:
             raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}")
         self.skip_softmax = self._coerce(self.skip_softmax, SkipSoftmaxSpec, "skip_softmax")
         self.quant = self._coerce(self.quant, AttnQuantSpec, "quant")
+        self.block_sparse = self._coerce(self.block_sparse, BlockSparseSpec, "block_sparse")
         if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN":
             raise ValueError(
                 f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. "
@@ -1485,6 +1514,15 @@ class AttentionSpec:
                 f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but "
                 f"backend={self.backend!r}. Remove quant or set a supported backend."
             )
+        if self.backend.upper() in BLOCK_SPARSE_BACKENDS:
+            # Selecting the backend is the opt-in; without an explicit block the
+            # defaults apply rather than silently running dense.
+            self.block_sparse = self.block_sparse or BlockSparseSpec()
+        elif self.block_sparse is not None:
+            raise ValueError(
+                f"block_sparse is only supported by the {sorted(BLOCK_SPARSE_BACKENDS)} backends, but "
+                f"backend={self.backend!r}. Remove block_sparse or set a supported backend."
+            )
 
     @staticmethod
     def _coerce(value: Any, cls: type, field_name: str) -> Any:
@@ -1517,6 +1555,12 @@ class AttentionSpec:
             if q.flashinfer_backend is not None:
                 quant_kw["flashinfer_backend"] = q.flashinfer_backend
             kw["quant"] = quant_kw
+        if self.block_sparse is not None:
+            bs = self.block_sparse
+            kw["sparsity"] = bs.sparsity
+            kw["start_step"] = bs.start_step
+            if bs.skip_layer_indices:
+                kw["skip_layers"] = sorted(bs.skip_layer_indices)
         return kw or None
 
 
@@ -1586,7 +1630,7 @@ class AttentionConfig:
             normalized[role] = node
             return
 
-        spec_keys = {"backend", "skip_softmax", "quant"}
+        spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"}
         node_dict = dict(node)
         node_keys = set(node_dict)
         if node_keys & spec_keys:

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -393,6 +393,13 @@ class DiffusersPipelineLoader:
                             "DLO+AllGather active but model does not support mmap: loading weights via regular loader."
                         )
                     logger.debug("Loading weights on %s ...", load_device)
+                    if offload_after_quant:
+                        marked = self._request_offload_after_quant(model)
+                        if marked:
+                            logger.info(
+                                "Online quantization will return each of %d layers to CPU as it is quantized",
+                                marked,
+                            )
                     if load_format == "diffusers":
                         cast(DiffusersAdapterPipeline, model).load_weights()
                     else:
@@ -406,6 +413,29 @@ class DiffusersPipelineLoader:
         self._apply_skip_softmax_calibration(model)
         return model.eval()
 
+    @staticmethod
+    def _request_offload_after_quant(model: nn.Module) -> int:
+        """Ask online-quant layers to return to host memory once quantized.
+
+        The weights only visit the accelerator so the quant kernels can run on
+        them; ``load_model`` sends the model back to the host afterwards either
+        way. Without this the whole transformer accumulates on device until that
+        final move, which for MiniMax H3 is a ~43 GiB peak that no longer fits
+        beside a resident TP-sharded text encoder — even though layer-wise
+        offload means none of it is supposed to be resident at inference time.
+
+        Only quant methods that advertise ``supports_offload_after_quant`` are
+        asked, since the implementation has to know when a layer is finished.
+        Deferring materialization to the ``meta`` device does not imply that.
+        """
+        marked = 0
+        for module in model.modules():
+            quant_method = getattr(module, "quant_method", None)
+            if getattr(quant_method, "supports_offload_after_quant", False):
+                quant_method.enable_offload_after_quant()
+                marked += 1
+        return marked
+
     @staticmethod
     def _has_online_quant(model: nn.Module) -> bool:
         """Whether any layer uses an online-quant method that defers weight
@@ -460,6 +490,13 @@ class DiffusersPipelineLoader:
             if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
                 continue
 
+            # Layers finished during loading would only be staged onto the target
+            # device for a process call that immediately returns. That round trip
+            # is wasted work in general, and undoes the point of the offload for
+            # layers that already went back to the host.
+            if getattr(module, "_already_called_process_weights_after_loading", False):
+                continue
+
             if has_online_quant:
                 # Online quant may leave straggler params on the ``meta`` device.
                 # Move only real (non-meta) params onto the target device for

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -14,6 +14,9 @@ from typing import Any
 
 import torch
 
+from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout
+from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+
 from .scheduling_minimax_h3_euler_ancestral import (
     minimax_h3_euler_eta0_step,
     minimax_h3_rf_v_to_x0,
@@ -87,6 +90,13 @@ class MiniMaxH3DenoiseBranch:
                 "max_seqlen_q": text_len,
             },
         }
+        # 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])),
+        )
 
     def forward_kwargs(
         self,
@@ -191,6 +201,9 @@ def minimax_h3_denoise_loop(
     for step in range(num_steps):
         step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
         with step_cm:
+            # Publish the step index so step-gated attention features (e.g. the
+            # dense warmup of RAINFUSION_ATTN) can see where we are.
+            set_forward_context_denoise_step_idx(step)
             s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
             s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1]
             t_v, t_a = 1.0 - s_v, 1.0 - s_a
@@ -236,6 +249,7 @@ def minimax_h3_denoise_loop(
             if on_step is not None:
                 on_step(step, video_rows, audio_rows)
 
+    set_forward_context_denoise_step_idx(None)
     return video_rows, audio_rows
 
 

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -26,7 +26,7 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
@@ -134,6 +134,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset(
         "img_pos_for_infer_output_info",
         "packed_seq_params",
         "refiner_packed_seq_params",
+        "video_token_layout",
     }
 )
 
@@ -348,7 +349,11 @@ class MiniMaxH3Attention(nn.Module):
             head_size=self.head_dim,
             softmax_scale=self.softmax_scale,
             causal=False,
+            # Packed rows reach the impl as [B, S, N, D].
+            qkv_layout="BSND",
             skip_sequence_parallel=skip_sequence_parallel,
+            role="self",
+            prefix=prefix,
         )
 
     def _apply_rope(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
@@ -374,6 +379,7 @@ class MiniMaxH3Attention(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int,
+        video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
         """Run packed attention as a small eager island.
 
@@ -407,6 +413,7 @@ class MiniMaxH3Attention(nn.Module):
                 "max_seqlen_k": max_seqlen,
                 "valid_kv_length": used,
             },
+            video_layout=video_layout,
         )
         return self.attention(
             q.unsqueeze(0),
@@ -424,6 +431,7 @@ class MiniMaxH3Attention(nn.Module):
         max_seqlen: int,
         packed_total: int | None = None,
         sp_seq_lens: list[int] | None = None,
+        video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
         """x: [T, hidden] packed thd rows -> [T, hidden].
 
@@ -463,6 +471,7 @@ class MiniMaxH3Attention(nn.Module):
             # backend receives the global sequence after all-to-all, so carry
             # its Python length explicitly instead of inferring it from q.
             packed_total=packed_total if packed_total is not None else q.shape[0],
+            video_layout=video_layout,
         )
         out = out.reshape(total, self.num_heads * self.head_dim)
         out, _ = self.out_proj(out)
@@ -642,6 +651,8 @@ class MiniMaxH3DiTBlock(nn.Module):
         super().__init__()
         self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps)
         self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps)
+        # The prefix also carries the block index that block-sparse attention
+        # backends match against their skip_layers selector.
         self.attn = MiniMaxH3Attention(
             arch,
             quant_config,
@@ -672,6 +683,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         max_seqlen: int,
         packed_total: int,
         sp_seq_lens: list[int] | None = None,
+        video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
         """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T]
         (= inverse_indices * modality_num + token_tags.clamp(min=0)).
@@ -699,6 +711,7 @@ class MiniMaxH3DiTBlock(nn.Module):
             max_seqlen=max_seqlen,
             packed_total=packed_total,
             sp_seq_lens=sp_seq_lens,
+            video_layout=video_layout,
         )
         x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE)
 
@@ -831,6 +844,10 @@ class MiniMaxH3DiTModel(nn.Module):
         },
         "sp_gather": SequenceParallelOutput(gather_dim=0, expected_dims=2),
     }
+    # The checkpoint already stores qkv and the MLP gate/up as single tensors
+    # (see the reordering in load_weights), so there are no unfused names for
+    # quantization or LoRA to map onto. Address the fused layers directly, e.g.
+    # ignored_layers=["blocks.0.attn.qkv_proj"].
     packed_modules_mapping = {}
 
     def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None:
@@ -1097,6 +1114,7 @@ class MiniMaxH3DiTModel(nn.Module):
         refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
         refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
         refiner_max = int(self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q"))
+        video_layout = kwargs.get("video_token_layout")
 
         if x.dim() != 3 or x.shape[0] != 1:
             raise ValueError(f"x must be [1, S, C], got {list(x.shape)}")
@@ -1145,6 +1163,7 @@ class MiniMaxH3DiTModel(nn.Module):
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
                 packed_total=seq_len,
+                video_layout=video_layout,
             )
         hidden = self.sp_gather(hidden)
 

vllm_omni/diffusion/models/minimax_h3/packed_sequence.py

@@ -236,6 +236,9 @@ def minimax_h3_packed_sequence(
         "token_tags": token_tags,
         "cu_seqlens": cu,
         "document_id": doc,
+        # 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),
     }
 
 
@@ -554,6 +557,9 @@ def minimax_h3_packed_sequence_ref2va_blocks(
         "token_tags": token_tags,
         "cu_seqlens": cu,
         "document_id": doc,
+        # 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),
     }
 
 

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -51,7 +51,14 @@ def _load_remote_component(
         class_reference,
         component_path,
     )
-    return component_cls.from_pretrained(component_path)
+    # Build on the host regardless of the ambient default device. Online
+    # quantization wraps pipeline construction in a `with torch.device(<accel>)`
+    # block for the DiT's quantized linears, and the checkpoint's own VAE code
+    # builds constants with ops that have no accelerator kernel (BigVGAN's
+    # anti-aliasing filters call torch.kaiser_window). Callers place the module
+    # explicitly right after this returns, so nothing depends on the context.
+    with torch.device("cpu"):
+        return component_cls.from_pretrained(component_path)
 
 
 class _AudioVAEDeterminismContext(AbstractContextManager):

vllm_omni/platforms/cuda/platform.py

@@ -134,6 +134,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 try:
                     importlib.import_module("kernels")

vllm_omni/platforms/interface.py

@@ -127,6 +127,29 @@ class OmniPlatform(Platform):
         """
         raise NotImplementedError
 
+    @classmethod
+    def validate_diffusion_attn_backend(cls, selected_backend: str) -> None:
+        """Reject an explicitly selected backend this platform cannot run.
+
+        Platforms call this from ``get_diffusion_attn_backend_cls`` so that a
+        backend restricted to other hardware, or one whose kernel package is
+        missing, fails during resolution rather than at the first forward.
+        """
+        from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+
+        backend_upper = selected_backend.upper()
+        backend_cls = DiffusionAttentionBackendEnum[backend_upper].get_class()
+
+        supported = backend_cls.supported_platforms
+        platform = cls._omni_enum.value
+        if supported is not None and platform not in supported:
+            raise ValueError(
+                f"The {backend_upper} diffusion attention backend runs on {', '.join(supported)} only, "
+                f"but the current platform is {platform}. Select a backend supported here, "
+                "such as FLASH_ATTN or TORCH_SDPA."
+            )
+        backend_cls.validate_available()
+
     @classmethod
     def supports_torch_inductor(cls) -> bool:
         """Check if the platform supports torch.compile with inductor backend."""

vllm_omni/platforms/musa/platform.py

@@ -80,6 +80,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 logger.warning(
                     "HuggingFace kernels-backed FlashAttention is "

vllm_omni/platforms/npu/platform.py

@@ -130,6 +130,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 logger.warning(
                     "HuggingFace kernels-backed FlashAttention is "

vllm_omni/platforms/rocm/platform.py

@@ -102,6 +102,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 logger.warning(
                     "HuggingFace kernels-backed FlashAttention is "

vllm_omni/platforms/xpu/platform.py

@@ -53,6 +53,7 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            cls.validate_diffusion_attn_backend(backend_upper)
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 logger.warning(
                     "HuggingFace kernels-backed FlashAttention is "

vllm_omni/quantization/int8_config.py

@@ -50,9 +50,55 @@ if TYPE_CHECKING:
 # Dynamic quantization is supported first.
 ACTIVATION_SCHEMES = ["dynamic"]
 
+# Ascend's npu_quant_matmul (QuantBatchMatmulV3) refuses a weight whose last
+# dimension exceeds this, and it fails at the first forward rather than at load
+# time. Layers still wider than the limit after TP sharding stay unquantized.
+NPU_QUANT_MATMUL_MAX_OUT_FEATURES = 65535
+
 logger = init_logger(__name__)
 
 
+def _fell_back_to_unquantized_npu(
+    layer: torch.nn.Module,
+    input_size_per_partition: int,
+    output_partition_sizes: list[int],
+    input_size: int,
+    output_size: int,
+    params_dtype: torch.dtype,
+    **extra_weight_attrs,
+) -> bool:
+    """Swap a layer to unquantized weights when npu_quant_matmul cannot run its shape.
+
+    The check uses the per-partition output size because that is what the kernel
+    actually sees; a layer over the limit on one rank can be within it at a
+    higher TP degree. Returns True when the layer was swapped and its weights
+    created, in which case the caller must not create its own.
+    """
+    output_size_per_partition = sum(output_partition_sizes)
+    if output_size_per_partition <= NPU_QUANT_MATMUL_MAX_OUT_FEATURES:
+        return False
+
+    logger.warning_once(
+        "Keeping a %d-wide linear unquantized: npu_quant_matmul rejects an output dimension past "
+        "%d. Tensor parallelism shrinks this per-rank dimension, so a higher TP degree brings such "
+        "layers back into range.",
+        output_size_per_partition,
+        NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+    )
+    fallback = UnquantizedLinearMethod()
+    layer.quant_method = fallback
+    fallback.create_weights(
+        layer,
+        input_size_per_partition,
+        output_partition_sizes,
+        input_size,
+        output_size,
+        params_dtype,
+        **extra_weight_attrs,
+    )
+    return True
+
+
 def create_weight_parameter(
     output_size_per_partition: int,
     input_size_per_partition: int,
@@ -231,6 +277,21 @@ class LazyWeightMixin:
 
     uses_meta_device: bool = True
 
+    # This mixin knows when a layer's weight is final, so it can hand the layer
+    # back to the host right there. Loaders that intend to offload the whole
+    # model after loading opt in per layer via ``enable_offload_after_quant``.
+    supports_offload_after_quant: bool = True
+    _offload_after_quant: bool = False
+
+    def enable_offload_after_quant(self) -> None:
+        """Return each layer to host memory as soon as it has been quantized.
+
+        Caps the load-time device footprint at one layer instead of the whole
+        model. A quant method instance belongs to a single layer, so this is not
+        a global switch.
+        """
+        self._offload_after_quant = True
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -287,6 +348,11 @@ class LazyWeightMixin:
                 # anything
                 layer._already_called_process_weights_after_loading = True
 
+                # This layer's weight is final, so nothing needs it on the
+                # accelerator until inference.
+                if self._offload_after_quant:
+                    layer.to("cpu")
+
                 # Note that we keep `layer._loaded_numel` around just in case
                 # there is logic added to vllm in the future which calls a
                 # weight loader twice - we do not want to re-initialize in
@@ -354,6 +420,11 @@ class NPUInt8LinearMethod(BaseInt8LinearMethod):
     def __init__(self, quant_config: DiffusionInt8Config):
         super().__init__(quant_config)
 
+    def create_weights(self, layer: torch.nn.Module, *args, **kwargs) -> None:
+        if _fell_back_to_unquantized_npu(layer, *args, **kwargs):
+            return
+        super().create_weights(layer, *args, **kwargs)
+
     def process_weights_after_loading(self, layer: Module) -> None:
         layer.weight.data = layer.weight.data.t().contiguous()
         layer.weight_scale.data = layer.weight_scale.data.squeeze()
@@ -424,6 +495,13 @@ class NPUInt8OnlineLinearMethod(LazyWeightMixin, NPUInt8LinearMethod):
     and quantized the weights during loading.
     """
 
+    def create_weights(self, layer: torch.nn.Module, *args, **kwargs) -> None:
+        # NPUInt8LinearMethod's override is unreachable from here: LazyWeightMixin
+        # comes first in the MRO and does not call super().
+        if _fell_back_to_unquantized_npu(layer, *args, **kwargs):
+            return
+        super().create_weights(layer, *args, **kwargs)
+
     def process_weights_after_loading(self, layer: Module) -> None:
         if getattr(layer, "_already_called_process_weights_after_loading", False):
             return