[Attention] [P1] Keep H3 VSA policy in the model and share sparse primitives (#7535)

Signed-off-by: lishunyang12 <lishunyang12@users.noreply.github.com>
Co-authored-by: lishunyang12 <lishunyang12@users.noreply.github.com>

tests/diffusion/attention/test_block_sparse_ops.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import sys
+import types
+
+import pytest
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+from vllm_omni.diffusion.attention.ops.block_sparse import (
+    build_prefix_dense_block_map,
+    fastvideo_block_sparse_attn_bshd,
+    mean_pool_tiles,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("prefix_blocks,topk", [(0, 1), (2, 0), (2, 1), (2, 3), (2, 5)])
+def test_block_map_keeps_prefix_dense_and_selects_non_prefix_keys(prefix_blocks, topk):
+    sparse_blocks = 3
+    blocks = prefix_blocks + sparse_blocks
+    scores = torch.arange(2 * blocks * blocks, dtype=torch.float32).reshape(1, 2, blocks, blocks)
+    block_map = build_prefix_dense_block_map(scores, prefix_blocks, sparse_blocks, topk)
+    expected = torch.zeros_like(scores, dtype=torch.bool)
+    # Scores increase with key index, so top-k always selects the last keys.
+    for query_block in range(blocks):
+        for key_block in range(blocks):
+            if (
+                query_block < prefix_blocks
+                or key_block < prefix_blocks
+                or key_block >= blocks - min(topk, sparse_blocks)
+            ):
+                expected[..., query_block, key_block] = True
+    assert torch.equal(block_map, expected)
+
+
+@pytest.mark.parametrize("block_size", [4, 32, 64])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_pooling_uses_valid_tile_sizes_and_fp32_accumulation(block_size, dtype):
+    sizes = torch.tensor([block_size, block_size - 1, 0], dtype=torch.int32)
+    tiles = torch.arange(2 * 3 * block_size * 2 * 3).reshape(2, 3, block_size, 2, 3).to(dtype)
+    tiles[:, 1, -1] = 0
+    tiles[:, 2] = 0
+    expected = torch.stack(
+        [tiles[:, 0].float().mean(dim=1), tiles[:, 1, :-1].float().mean(dim=1), torch.zeros(2, 2, 3)], dim=2
+    )
+    result = mean_pool_tiles(tiles.flatten(1, 2), sizes, block_size)
+    assert result.dtype == torch.float32
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+def test_provider_wrapper_restores_transport_rows_and_bshd_layout(monkeypatch):
+    monkeypatch.delenv("FASTVIDEO_VSA_SM100A", raising=False)
+    calls = []
+
+    def block_sparse_attn(q, k, v, block_map, sizes):
+        calls.append((q.shape, block_map.shape, sizes.tolist()))
+        assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
+        assert sizes.dtype == torch.int32
+        return q + k + v, None
+
+    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    setattr(module, "block_sparse_attn", block_sparse_attn)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
+    query = torch.randn(1, 4 * 64, 2, 8)
+    block_map = torch.ones(1, 2, 4, 4, dtype=torch.bool)
+    sizes = torch.tensor([64, 17, 64, 0])
+    output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
+    expected = torch.zeros_like(query)
+    expected[:, : 3 * 64] = query[:, : 3 * 64] * 3
+    torch.testing.assert_close(output, expected)
+    assert calls == [(torch.Size([1, 2, 3 * 64, 8]), torch.Size([1, 2, 3, 3]), [64, 17, 64])]
+
+
+def test_custom_op_fake_preserves_metadata_without_loading_provider(monkeypatch):
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", None)
+    with FakeTensorMode():
+        query = torch.empty(1, 4 * 64, 2, 8, dtype=torch.bfloat16)
+        block_map = torch.empty(1, 2, 4, 4, dtype=torch.bool)
+        sizes = torch.empty(4, dtype=torch.int32)
+        output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
+        assert output.shape == query.shape
+        assert output.dtype == query.dtype
+        assert output.device == query.device
+
+
+@pytest.mark.parametrize("error_type", [ImportError, RuntimeError, torch.AcceleratorError])
+def test_native_provider_failure_only_recovers_without_accelerator_fault(monkeypatch, mocker, error_type):
+    monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1")
+    error = error_type("native provider failure")
+    native = mocker.Mock()
+    native.is_supported.return_value = True
+    native.block_sparse_attn_sm100a.side_effect = error
+    provider = types.ModuleType("fastvideo_kernel")
+    setattr(provider, "block_sparse_attn_sm100a", native)
+    index = types.ModuleType("fastvideo_kernel.triton_kernels.index")
+    setattr(
+        index,
+        "map_to_index",
+        mocker.Mock(return_value=(torch.zeros(1, dtype=torch.int32), torch.ones(1, dtype=torch.int32))),
+    )
+    triton = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    fallback = mocker.Mock(side_effect=lambda q, *args: (q.clone(), None))
+    setattr(triton, "block_sparse_attn", fallback)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", provider)
+    monkeypatch.setitem(
+        sys.modules, "fastvideo_kernel.triton_kernels", types.ModuleType("fastvideo_kernel.triton_kernels")
+    )
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.triton_kernels.index", index)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", triton)
+    query = torch.arange(128 * 2 * 8, dtype=torch.float32).reshape(1, 128, 2, 8)
+    block_map = torch.ones(1, 2, 2, 2, dtype=torch.bool)
+    sizes = torch.tensor([64, 64], dtype=torch.int32)
+
+    if error_type is torch.AcceleratorError:
+        with pytest.raises(torch.AcceleratorError, match="native provider failure") as raised:
+            fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
+        assert raised.value is error
+        fallback.assert_not_called()
+    else:
+        output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
+        torch.testing.assert_close(output, query, rtol=0, atol=0)
+        fallback.assert_called_once()

tests/diffusion/attention/test_fastvideo_vsa.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-import math
 import sys
 import types
 from typing import Any
@@ -12,14 +11,10 @@ import torch
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
     PackedPaddingMetadata,
-    VideoTokenLayout,
-    VideoTokenSpan,
 )
 from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
     FastVideoVSABackend,
     FastVideoVSAImpl,
-    _build_h3_block_map,
-    _get_h3_tile_metadata,
 )
 from vllm_omni.diffusion.attention.backends.registry import (
     DiffusionAttentionBackendEnum,
@@ -218,144 +213,8 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
     assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
 
 
-def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
-    extra: dict[str, Any] = {
-        "vsa_h3_prefix_segments": prefix_segments,
-    }
-    if gate is not None:
-        extra["gate_compress"] = gate
-    return AttentionMetadata(
-        packed_padding=packed_padding,
-        video_layout=VideoTokenLayout(
-            used_len=sum(prefix_segments) + math.prod(video_shape),
-            video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
-        ),
-        extra=extra,
-    )
-
-
-def _fake_block_sparse_kernel(monkeypatch, calls):
-    """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
-
-    def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
-        calls["block_map"] = block_map.clone()
-        calls["variable_block_sizes"] = variable_block_sizes.tolist()
-        calls["q_shape"] = tuple(q.shape)
-        blocks = block_map.shape[-1]
-        assert blocks * 64 == q.shape[2]
-        # Expand the per-block map to a per-token mask so the stand-in honours
-        # the sparsity contract the real kernel implements, and drop the rows
-        # each partial block pads with, which the real kernel skips through
-        # variable_block_sizes.
-        token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
-        within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
-        token_mask = token_mask & within.reshape(1, 1, 1, -1)
-        out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
-        return out.to(q.dtype), None
-
-    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
-    setattr(module, "block_sparse_attn", block_sparse_attn)
-    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
-    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
-
-
-def _h3_impl(**backend_kwargs):
-    kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
-    kwargs.update(backend_kwargs)
-    return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
-
-
-def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
-    calls: dict[str, Any] = {}
-    _fake_block_sparse_kernel(monkeypatch, calls)
-    impl = _h3_impl(topk=64)
-    prefix_segments, video_shape = (5, 70), (2, 4, 4)
-    seq_len = sum(prefix_segments) + 2 * 4 * 4
-    torch.manual_seed(0)
-    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
-
-    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
-
-    assert output.shape == query.shape
-    # topk covers every video block, so the block map is dense and the tiled
-    # kernel must reproduce plain attention over the same rows.
-    reference = torch.nn.functional.scaled_dot_product_attention(
-        *(x.float().transpose(1, 2) for x in (query, query, query))
-    ).transpose(1, 2)
-    torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
-    # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
-    assert calls["variable_block_sizes"] == [5, 64, 6, 32]
-
-
-def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
-    calls: dict[str, Any] = {}
-    _fake_block_sparse_kernel(monkeypatch, calls)
-    impl = _h3_impl(topk=1)
-    prefix_segments, video_shape = (5,), (8, 4, 4)
-    seq_len = 5 + 8 * 4 * 4
-    torch.manual_seed(0)
-    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
-
-    impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
-
-    # 1 prefix block + 2 video blocks. The odd count is padded to an even one
-    # for the paired-CTA kernel contract, and the Triton route slices that
-    # transport-only partner off again before launching.
-    block_map = calls["block_map"]
-    assert block_map.shape[-1] == 3
-    assert calls["variable_block_sizes"] == [5, 64, 64]
-    assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
-    assert block_map[..., :1].all(), "prefix keys are exempt from selection"
-    assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
-
-    # The same odd geometry must still untile back onto the original rows.
-    dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
-    reference = torch.nn.functional.scaled_dot_product_attention(
-        *(x.float().transpose(1, 2) for x in (query, query, query))
-    ).transpose(1, 2)
-    torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
-
-
-def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
-    _fake_block_sparse_kernel(monkeypatch, {})
-    impl = _h3_impl(topk=64)
-    prefix_segments, video_shape = (5,), (2, 4, 4)
-    seq_len = 5 + 2 * 4 * 4
-    torch.manual_seed(0)
-    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
-
-    without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
-    zero_gate = torch.zeros_like(query)
-    with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
-    live_gate = torch.full_like(query, 0.5)
-    with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
-
-    torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
-    assert not torch.allclose(with_live_gate.float(), without_gate.float())
-
-
-def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
-    _fake_block_sparse_kernel(monkeypatch, {})
-    impl = _h3_impl(topk=64)
-    prefix_segments, video_shape = (5,), (2, 4, 4)
-    valid = 5 + 2 * 4 * 4
-    torch.manual_seed(0)
-    query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
-    padding = PackedPaddingMetadata(
-        q_length=valid,
-        kv_length=valid,
-        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
-        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
-    )
-
-    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
-
-    assert output.shape == query.shape
-    assert torch.count_nonzero(output[:, valid:]) == 0
-
-
 def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
-    # The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
+    # The backend advertises supports_packed_mask_free, so a model may skip
     # building the padding mask. Every fallback out of the VSA path must honour
     # that contract itself; SDPA reads attn_mask and nothing else.
     impl = FastVideoVSAImpl(
@@ -374,6 +233,8 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
         cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
     )
 
+    monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
+
     baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
     perturbed_input = query.clone()
     perturbed_input[:, valid:] += 100.0
@@ -384,22 +245,3 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
     assert baseline.shape == query.shape
     torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
     assert torch.count_nonzero(baseline[:, valid:]) == 0
-
-
-def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
-    partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
-        (5, 70, 9), (5, 6, 6), torch.device("cpu")
-    )
-    assert sizes[:4].tolist() == [5, 64, 6, 9]
-    assert prefix_blocks == 4
-    assert video_blocks == 8
-    assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
-    assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
-
-
-def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
-    scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
-    block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
-    assert block_map[:, :, :2].all()
-    assert block_map[..., :2].all()
-    assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()

tests/diffusion/attention/test_impl_overrides.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from vllm_omni.diffusion.attention import layer
+from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSABackend, FastVideoVSAImpl
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+from vllm_omni.diffusion.data import AttentionSpec
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _SpecializedImpl(FastVideoVSAImpl):
+    pass
+
+
+@pytest.fixture
+def local_attention(monkeypatch):
+    monkeypatch.setattr(layer, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention())
+
+
+@pytest.mark.parametrize("explicit", [False, True])
+@pytest.mark.parametrize("override", [False, True])
+def test_specialization_preserves_selected_backend_and_options(mocker, local_attention, explicit, override):
+    spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=7) if explicit else None
+    select = mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, spec))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        role="video.self",
+        role_category="self",
+        qkv_layout="BSND",
+        impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl} if override else None,
+    )
+
+    assert attention.attn_backend is FastVideoVSABackend
+    assert type(attention.attention) is (_SpecializedImpl if override else FastVideoVSAImpl)
+    assert attention.attention.topk == (7 if explicit else 64)
+    assert attention.attention.qkv_layout == "BSND"
+    assert attention.backend_pref == "FASTVIDEO_VSA"
+    assert attention.backend_explicit is explicit
+    assert attention.attn_spec is spec
+    assert isinstance(attention.parallel_strategy, NoParallelAttention)
+    select.assert_called_once_with(
+        role="video.self", head_size=8, attention_config=None, role_category="self", allow_trtllm_default=False
+    )
+
+
+def test_unselected_specialization_does_not_replace_dense_backend(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(SDPABackend, None))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+    )
+    assert attention.attn_backend is SDPABackend
+    assert type(attention.attention) is SDPABackend.get_impl_cls()
+
+
+def test_specialization_does_not_bypass_selection_errors(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", side_effect=ImportError("optional kernel unavailable"))
+    with pytest.raises(ImportError, match="optional kernel unavailable"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+        )
+
+
+def test_specialization_must_extend_selected_implementation(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, None))
+    with pytest.raises(TypeError, match="must subclass the selected implementation"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": SDPABackend.get_impl_cls()},
+        )
+
+
+def test_platform_backend_capabilities_survive_model_specialization(mocker, local_attention):
+    from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import MiniMaxH3VSAImpl
+
+    class PlatformBackend(FastVideoVSABackend):
+        @classmethod
+        def supports_packed_mask_free(cls):
+            return False
+
+        @staticmethod
+        def get_supported_head_sizes():
+            return [128]
+
+    # A registry/platform may refine capabilities without changing the
+    # implementation. The model must not replace those capability decisions.
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        causal=False,
+        impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl},
+    )
+    assert attention.attn_backend is PlatformBackend
+    assert not attention.attn_backend.supports_packed_mask_free()
+    assert attention.attn_backend.get_supported_head_sizes() == [128]
+    assert type(attention.attention) is MiniMaxH3VSAImpl
+
+
+def test_incompatible_platform_implementation_is_not_discarded(mocker, local_attention):
+    class PlatformImpl(FastVideoVSAImpl):
+        pass
+
+    class PlatformBackend(FastVideoVSABackend):
+        @staticmethod
+        def get_impl_cls():
+            return PlatformImpl
+
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
+    with pytest.raises(TypeError, match="selected implementation .*PlatformImpl"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+        )

tests/diffusion/cache/test_teacache_extractors.py

@@ -68,12 +68,12 @@ class BaseExtractorTest(ABC):
         pass
 
     @abstractmethod
-    def get_module(self):
+    def get_module(self, module, /):
         """Return model module instance."""
         pass
 
     @abstractmethod
-    def get_sample_inputs(self):
+    def get_sample_inputs(self, sample_inputs, /):
         """Return sample inputs for model."""
         pass
 
@@ -737,10 +737,10 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
         monkeypatch,
     ):
         """Strict SP must not mix local TeaCache state with gathered block rows."""
-        from vllm_omni.diffusion.attention.ops import minimax_h3_modulation
         from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
         from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
         from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.layers import indexed_modulation
 
         seq_len = sample_inputs["x"].shape[1]
         local_len = seq_len // 2
@@ -844,13 +844,13 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
         )
         monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
 
-        original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_
+        original_indexed_scale_shift = indexed_modulation.indexed_scale_shift_
 
         def capture_indexed_scale_shift(hidden, shift, scale, indices):
             captured_modulation_indices.append(indices.clone())
             return original_indexed_scale_shift(hidden, shift, scale, indices)
 
-        monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
+        monkeypatch.setattr(indexed_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
 
         local_img_pos = torch.tensor([0, 1])
         local_audio_pos = torch.tensor([2, 3])

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -6,7 +6,7 @@ import torch
 from vllm.triton_utils import HAS_TRITON
 
 from tests.helpers.mark import hardware_marks
-from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
+from vllm_omni.diffusion.layers.indexed_modulation import (
     _MAX_1D_GRID_SIZE,
     _iter_row_chunks,
     _launch_row_chunks,

tests/diffusion/models/minimax_h3/test_vsa_layout.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import math
+import sys
+import types
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+    VideoTokenLayout,
+    VideoTokenSpan,
+)
+from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import (
+    MiniMaxH3VSAImpl,
+    _get_h3_tile_metadata,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
+    extra: dict[str, Any] = {
+        "vsa_h3_prefix_segments": prefix_segments,
+    }
+    if gate is not None:
+        extra["gate_compress"] = gate
+    return AttentionMetadata(
+        packed_padding=packed_padding,
+        video_layout=VideoTokenLayout(
+            used_len=sum(prefix_segments) + math.prod(video_shape),
+            video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
+        ),
+        extra=extra,
+    )
+
+
+def _fake_block_sparse_kernel(monkeypatch, calls):
+    """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
+
+    def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
+        calls["block_map"] = block_map.clone()
+        calls["variable_block_sizes"] = variable_block_sizes.tolist()
+        calls["q_shape"] = tuple(q.shape)
+        blocks = block_map.shape[-1]
+        assert blocks * 64 == q.shape[2]
+        # Expand the per-block map to a per-token mask so the stand-in honours
+        # the sparsity contract the real kernel implements, and drop the rows
+        # each partial block pads with, which the real kernel skips through
+        # variable_block_sizes.
+        token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
+        within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
+        token_mask = token_mask & within.reshape(1, 1, 1, -1)
+        out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
+        return out.to(q.dtype), None
+
+    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    setattr(module, "block_sparse_attn", block_sparse_attn)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
+
+
+def _h3_impl(**backend_kwargs):
+    kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
+    kwargs.update(backend_kwargs)
+    return MiniMaxH3VSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
+
+
+def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5, 70), (2, 4, 4)
+    seq_len = sum(prefix_segments) + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    assert output.shape == query.shape
+    # topk covers every video block, so the block map is dense and the tiled
+    # kernel must reproduce plain attention over the same rows.
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
+    # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
+    assert calls["variable_block_sizes"] == [5, 64, 6, 32]
+
+
+def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=1)
+    prefix_segments, video_shape = (5,), (8, 4, 4)
+    seq_len = 5 + 8 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    # 1 prefix block + 2 video blocks. The odd count is padded to an even one
+    # for the paired-CTA kernel contract, and the Triton route slices that
+    # transport-only partner off again before launching.
+    block_map = calls["block_map"]
+    assert block_map.shape[-1] == 3
+    assert calls["variable_block_sizes"] == [5, 64, 64]
+    assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
+    assert block_map[..., :1].all(), "prefix keys are exempt from selection"
+    assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
+
+    # The same odd geometry must still untile back onto the original rows.
+    dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
+
+
+def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    seq_len = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    zero_gate = torch.zeros_like(query)
+    with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
+    live_gate = torch.full_like(query, 0.5)
+    with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
+
+    torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
+    assert not torch.allclose(with_live_gate.float(), without_gate.float())
+
+
+def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    valid = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
+    padding = PackedPaddingMetadata(
+        q_length=valid,
+        kv_length=valid,
+        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
+        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
+    )
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
+
+    assert output.shape == query.shape
+    assert torch.count_nonzero(output[:, valid:]) == 0
+
+
+def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
+    partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
+        (5, 70, 9), (5, 6, 6), torch.device("cpu")
+    )
+    assert sizes[:4].tolist() == [5, 64, 6, 9]
+    assert prefix_blocks == 4
+    assert video_blocks == 8
+    assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
+    assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
+
+
+@pytest.mark.parametrize("fallback_on_error", [False, True])
+def test_h3_kernel_failure_preserves_fallback_policy(monkeypatch, fallback_on_error):
+    impl = _h3_impl(fallback_on_error=fallback_on_error)
+    query = torch.randn(1, 37, 2, 8)
+    metadata = _h3_metadata((5,), (2, 4, 4))
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("kernel failure")
+
+    monkeypatch.setattr(impl, "_forward_h3", fail)
+    monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
+    if fallback_on_error:
+        output = impl.forward_cuda(query, query, query, metadata)
+        expected = impl.sdpa_fallback.forward_cuda(query, query, query, metadata)
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+    else:
+        with pytest.raises(RuntimeError, match="kernel failure"):
+            impl.forward_cuda(query, query, query, metadata)
+
+
+def test_h3_accelerator_failure_is_not_swallowed(monkeypatch):
+    impl = _h3_impl(fallback_on_error=True)
+    query = torch.randn(1, 37, 2, 8)
+
+    def fail(*args, **kwargs):
+        raise torch.AcceleratorError("device failure")
+
+    monkeypatch.setattr(impl, "_forward_h3", fail)
+    with pytest.raises(torch.AcceleratorError, match="device failure"):
+        impl.forward_cuda(query, query, query, _h3_metadata((5,), (2, 4, 4)))
+
+
+def test_h3_without_layout_delegates_to_shared_backend(monkeypatch):
+    from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSAImpl
+
+    impl = _h3_impl()
+    query = torch.randn(1, 37, 2, 8)
+    metadata = AttentionMetadata()
+    expected = torch.empty_like(query)
+    calls = []
+
+    def forward(self, q, k, v, attn_metadata):
+        calls.append((self, q, k, v, attn_metadata))
+        return expected
+
+    monkeypatch.setattr(FastVideoVSAImpl, "forward_cuda", forward)
+    assert impl.forward_cuda(query, query, query, metadata) is expected
+    assert len(calls) == 1
+    assert all(actual is wanted for actual, wanted in zip(calls[0], (impl, query, query, query, metadata)))