[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>

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