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

docs/design/feature/attention_backend_selection.md

@@ -76,6 +76,34 @@ preserve common `self` or `cross` policies. A model also decides whether its
 path is compatible with automatic TRTLLM selection, because only the model
 knows whether masking and packing satisfy the kernel contract.
 
+### Model-owned implementation specializations
+
+Use `Attention(..., impl_overrides={backend_name: implementation_class})`
+when a model needs different computation within an existing backend contract.
+The model defines the implementation in its own directory and passes it when
+constructing the layer:
+
+```python
+# MyModelVSAImpl is a model-owned subclass of FastVideoVSAImpl.
+self.attention = Attention(
+    num_heads=num_heads,
+    head_size=head_size,
+    softmax_scale=head_size**-0.5,
+    causal=False,
+    role="self",
+    impl_overrides={"FASTVIDEO_VSA": MyModelVSAImpl},
+)
+```
+
+The override is applied after role/platform selection, keyed by the selected
+backend's `get_name()`. It must subclass the **selected implementation**;
+incompatible platform implementations raise `TypeError`. Unselected overrides
+have no effect. Backend capabilities, constructor options, and shared
+parallel dispatch remain unchanged and must be honored by the specialization.
+
+Use a new backend for a new provider/capability contract. `custom_attention`
+instead owns communication and requires `skip_sequence_parallel=True`.
+
 ## Adding or changing a backend
 
 An implementation change is complete only when it:
@@ -91,8 +119,9 @@ An implementation change is complete only when it:
 
 ## FastVideo VSA model metadata
 
-FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in
-the original latent video grid. Wan integrations therefore attach the
+FastWan's text-to-video and image-to-video modes both use `Wan22Pipeline`.
+It operates on a flattened DiT sequence but partitions tokens in the original
+latent video grid. Wan integrations therefore attach the
 post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The
 backend validates that `T * H * W` equals the sequence length, derives the
 runtime block count, and selects the configured top-k key/value blocks for
@@ -104,3 +133,8 @@ the checkpoint contains those weights. Checkpoints without the projection do
 not allocate or execute it. When top-k selects every block, native checkpoints
 route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path
 to retain checkpoint semantics.
+
+MiniMax-H3 supplies `MiniMaxH3VSAImpl` through `impl_overrides`.
+`models/minimax_h3/attention/fastvideo_h3.py` owns its prefix layout, video tiling, and
+learned compression gate. Pooling, block-map construction, and tile64 provider
+calls are shared through `attention/ops/block_sparse.py`.

docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md

@@ -1,24 +1,33 @@
 # FastVideo VSA
 
-FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's
-self-attention by partitioning the post-patch latent grid into spatiotemporal
+FastVideo Variable Sparse Attention (VSA) partitions the post-patch latent grid into spatiotemporal
 blocks. For every query block, VSA scores the key/value blocks and computes
 attention only against the selected top-k blocks.
 
-The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
+## Supported models
 
-MiniMax-H3 served with a FastH3 VSA adapter uses a second, H3-specific route
-through the same backend. See the
-[MiniMax-H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for its geometry,
-supported topologies, and commands.
+| Model / checkpoint | Required adapter | Tasks | Sequence parallelism |
+| --- | --- | --- | --- |
+| `FastVideo/FastWan2.2-TI2V-5B-Diffusers` | None | T2V, I2V through `Wan22Pipeline` | Disabled |
+| `MiniMaxAI/MiniMax-H3` | [FastH3 VSA (recipe)](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter) | T2VA | Disabled or pure Ulysses |
+| Wan I2V-14B, S2V, VACE | — | Unsupported | — |
 
-VSA is a CUDA-only, explicitly selected backend. It requires the
-`fastvideo-kernel` package and currently supports non-causal self-attention
-with equal query and key/value sequence lengths. Unsupported shapes, masks,
-dtypes, sequence-parallel execution, or kernel failures fall back to
-`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
-sequence-parallel context is one of those fallbacks; the H3 route supports pure
-Ulysses and rejects ring or all-gather sequence parallelism at startup.
+## Installation
+
+Install vLLM-Omni from this checkout with the `vsa` extra:
+
+```bash
+uv pip install -e '.[vsa]'
+```
+
+The extra installs the tested kernel dependency automatically. Prebuilt kernels require
+Linux, Python 3.12, and glibc 2.34 or newer (x86-64 or aarch64). The full
+FastVideo framework and provider environment variables are not required.
+
+`FASTVIDEO_VSA` selects the attention algorithm. On SM120, the H3 integration
+currently executes FastVideo's 64-token Triton block-sparse kernel; it does not
+dispatch to FlashInfer. The `vsa` extra installs `fastvideo-kernel==0.3.4` for
+this execution path.
 
 ## Enable the backend
 
@@ -27,24 +36,18 @@ flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained
 for every query block:
 
 ```bash
-vllm-omni serve <model> \
+vllm serve <model> --omni \
   --diffusion-attention-backend FASTVIDEO_VSA \
   --fastvideo-vsa-topk 64
 ```
 
-The backwards-compatible environment variable selects the backend with the
-default `topk=64`:
+<details markdown="1">
+<summary>Alternative configuration formats</summary>
+
+Equivalent structured configuration:
 
 ```bash
-export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA
-vllm-omni serve <model>
-```
-
-To tune top-k, pass the CLI backend and top-k flags together as shown above,
-or use the equivalent structured configuration:
-
-```bash
-vllm-omni serve <model> \
+vllm serve <model> --omni \
   --diffusion-attention-config \
   '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}'
 ```
@@ -73,8 +76,15 @@ stages:
         fastvideo_vsa_topk: 64
 ```
 
+</details>
+
 ## Choose top-k
 
+H3 uses 64-token video blocks and keeps
+all prefix blocks; see the [FastH3 VSA recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter).
+
+### Wan top-k behavior
+
 At runtime the backend logs the sequence shape and derived block count:
 
 ```text
@@ -117,8 +127,16 @@ selecting VSA does not turn a native checkpoint into a distilled model.
 
 ## Verify routing and fallback
 
-Check the startup and first-forward logs instead of assuming that selecting
-the backend guarantees sparse execution:
+VSA is a CUDA-only, explicitly selected backend. It requires the
+`fastvideo-kernel` package and currently supports non-causal self-attention
+with equal query and key/value sequence lengths. Unsupported shapes, masks,
+dtypes, sequence-parallel execution, or kernel failures fall back to
+`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
+sequence-parallel context is one of those fallbacks; the H3 route supports pure
+Ulysses and rejects ring or all-gather sequence parallelism at startup.
+H3 accelerator faults propagate instead of attempting dense recovery.
+
+On the Wan route, check the startup and first-forward logs:
 
 - `route=VSA` means top-k block selection is active.
 - `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all
@@ -126,6 +144,10 @@ the backend guarantees sparse execution:
 - `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
   executed; the warning includes the reason.
 
+For H3, check `FastH3 adapter active` at startup and
+`FASTVIDEO_VSA H3 routing` during DiT execution, as described in the
+[model recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter).
+
 The Wan route requires CUDA tensors in FP16 or BF16, 256-token blocks,
 standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
 and no active sequence-parallel context. The MiniMax-H3 route uses 64-token

pyproject.toml

@@ -123,6 +123,11 @@ fa4 = [
     "flash-attn-4[cu13]==4.0.0b18",
 ]
 
+# FastVideo VSA kernels for the supported Linux/Python 3.12 wheel environment.
+vsa = [
+    "fastvideo-kernel==0.3.4",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",

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

tests/helpers/runtime.py

@@ -23,6 +23,7 @@ from typing import Any, NamedTuple
 import numpy as np
 import psutil
 import yaml
+from filelock import FileLock, Timeout
 from vllm import TextPrompt
 from vllm.logger import init_logger
 
@@ -170,9 +171,54 @@ class OmniServer:
         self.use_omni = use_omni
         self.proc: subprocess.Popen | None = None
         self.host = "127.0.0.1"
+        self._auto_port = port is None
+        self._port_lock: FileLock | None = None
         self.port = get_open_port() if port is None else port
 
+    def _reserve_port(self) -> None:
+        # bind(0)/close does not reserve a port across concurrent pytest workers.
+        # Keep an advisory lock until teardown, including the long import phase
+        # before the subprocess binds its HTTP socket. Never unlink lock files:
+        # another worker may already have the same inode open.
+        for attempt in range(128):
+            if attempt:
+                self.port = get_open_port(self.host)
+            lock_path = Path(tempfile.gettempdir()) / f"vllm-omni-test-port-{os.getuid()}-{self.port}.lock"
+            lock = FileLock(lock_path)
+            try:
+                lock.acquire(timeout=0)
+            except Timeout as error:
+                if not self._auto_port:
+                    raise RuntimeError(f"HTTP test port {self.port} is already reserved") from error
+                continue
+            try:
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
+                    probe.bind((self.host, self.port))
+            except OSError as error:
+                lock.release()
+                if not self._auto_port or error.errno != errno.EADDRINUSE:
+                    raise
+                continue
+            self._port_lock = lock
+            return
+        raise RuntimeError("Could not reserve an HTTP test port after 128 attempts")
+
+    def _owns_listening_port(self) -> bool:
+        assert self.proc is not None
+        try:
+            parent = psutil.Process(self.proc.pid)
+            processes = [parent, *parent.children(recursive=True)]
+            return any(
+                conn.status == psutil.CONN_LISTEN and conn.laddr.port == self.port
+                for process in processes
+                for conn in process.net_connections(kind="tcp")
+            )
+        except psutil.NoSuchProcess:
+            # A child can exit between enumeration and inspecting its sockets.
+            return False
+
     def _start_server(self) -> None:
+        self._reserve_port()
         env = os.environ.copy()
         if self.env_dict is not None:
             env.update(self.env_dict)
@@ -208,7 +254,10 @@ class OmniServer:
                 raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.")
             with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
                 sock.settimeout(1)
-                if sock.connect_ex((self.host, self.port)) == 0:
+                if sock.connect_ex((self.host, self.port)) == 0 and self._owns_listening_port():
+                    ret = self.proc.poll()
+                    if ret is not None:
+                        raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.")
                     startup_s = time.perf_counter() - startup_t0
                     if self.log_stats:
                         print(
@@ -368,9 +417,14 @@ class OmniServer:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.proc:
-            self._kill_process_tree(self.proc.pid)
-        cleanup_test_environment()
+        try:
+            if self.proc:
+                self._kill_process_tree(self.proc.pid)
+        finally:
+            if self._port_lock is not None:
+                self._port_lock.release()
+                self._port_lock = None
+            cleanup_test_environment()
 
 
 class OmniServerStageCli(OmniServer):

tests/helpers/tests/test_runtime_server_ports.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""A listening port alone must not make a different test's server ready."""
+
+import socket
+import subprocess
+import sys
+
+import pytest
+from filelock import FileLock
+
+from tests.helpers import runtime
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_server_does_not_accept_another_process_listener(monkeypatch):
+    monkeypatch.setattr(runtime, "cleanup_test_environment", lambda: None)
+    server = runtime.OmniServer("fake-model", [])
+    popen = subprocess.Popen
+
+    # Occupy the chosen port after allocation, while the requested subprocess
+    # is still importing. The old TCP-only probe returns ready for this socket.
+    with socket.socket() as foreign_listener:
+
+        def launch(*args, **kwargs):
+            foreign_listener.bind((server.host, server.port))
+            foreign_listener.listen()
+            return popen([sys.executable, "-c", "import time; time.sleep(0.2); raise SystemExit(7)"])
+
+        monkeypatch.setattr(runtime.subprocess, "Popen", launch)
+        with pytest.raises(RuntimeError, match="exited with code 7"):
+            with server:
+                pass
+    # Failed __enter__ must also release the lease for a later server.
+    retry = runtime.OmniServer("retry-model", [], port=server.port)
+    try:
+        retry._reserve_port()
+    finally:
+        retry.__exit__(None, None, None)
+
+
+def test_auto_ports_are_reserved_before_subprocesses_bind(monkeypatch, tmp_path):
+    monkeypatch.setattr(runtime, "cleanup_test_environment", lambda: None)
+    monkeypatch.setattr(runtime.tempfile, "gettempdir", lambda: str(tmp_path))
+    # Both constructors see the same free port before either server has bound.
+    first = runtime.OmniServer("first-model", [])
+    candidate = first.port
+    monkeypatch.setattr(runtime, "get_open_port", lambda *args: candidate)
+    second = runtime.OmniServer("second-model", [])
+    first._reserve_port()
+    with socket.socket() as occupied, socket.socket() as probe:
+        occupied.bind((first.host, candidate))
+        probe.bind((first.host, 0))
+        alternative = probe.getsockname()[1]
+    monkeypatch.setattr(runtime, "get_open_port", lambda *args: alternative)
+    try:
+        second._reserve_port()
+        assert first.port == candidate
+        assert second.port == alternative
+        # Check OS-level lock exclusion, not only the Python object's state.
+        lock_path = next(tmp_path.glob(f"*-{candidate}.lock"))
+        contender = FileLock(lock_path)
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                "import sys\nfrom filelock import FileLock, Timeout\n"
+                "try:\n FileLock(sys.argv[1]).acquire(timeout=0)\n"
+                "except Timeout:\n sys.exit(23)\n",
+                str(lock_path),
+            ],
+            timeout=10,
+            check=False,
+        )
+        assert result.returncode == 23
+        explicit = runtime.OmniServer("explicit-model", [], port=candidate)
+        with pytest.raises(RuntimeError, match="already reserved"):
+            with explicit:
+                pass
+    finally:
+        first.__exit__(None, None, None)
+        second.__exit__(None, None, None)
+    with contender.acquire(timeout=0):
+        pass
+
+
+@pytest.mark.parametrize("in_child", [False, True])
+def test_server_accepts_its_own_listener(monkeypatch, in_child):
+    monkeypatch.setattr(runtime, "cleanup_test_environment", lambda: None)
+    server = runtime.OmniServer("fake-model", [])
+    popen = subprocess.Popen
+
+    def launch(*args, **kwargs):
+        code = (
+            "import socket,time; s=socket.socket(); "
+            f"s.bind(({server.host!r}, {server.port})); s.listen(); time.sleep(30)"
+        )
+        if in_child:
+            code = f"import subprocess,sys; child=subprocess.Popen([sys.executable,'-c',{code!r}]); child.wait()"
+        return popen([sys.executable, "-c", code])
+
+    monkeypatch.setattr(runtime.subprocess, "Popen", launch)
+    with server:
+        assert server.proc is not None
+        assert server.proc.poll() is None
+        with socket.create_connection((server.host, server.port), timeout=1):
+            pass

vllm_omni/diffusion/attention/backends/fastvideo_vsa.py

@@ -6,7 +6,6 @@ from __future__ import annotations
 import functools
 import importlib.util
 import math
-import os
 from collections.abc import Mapping
 from typing import Any
 
@@ -88,77 +87,6 @@ if not hasattr(torch.ops.vllm_omni, "fastvideo_vsa_bshd"):
 _fastvideo_vsa_bshd_op = torch.ops.vllm_omni.fastvideo_vsa_bshd
 
 
-# H3 needs an explicit per-query block map: prefix queries are dense, while
-# video queries select prefix + top-k video tiles. The generic
-# video_sparse_attn() entry point cannot express that contract.
-if not hasattr(torch.ops.vllm_omni, "fastvideo_h3_vsa_bhsd"):
-
-    @torch.library.custom_op("vllm_omni::fastvideo_h3_vsa_bhsd", mutates_args=())
-    def _fastvideo_h3_vsa_bhsd_op(
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        block_map: torch.Tensor,
-        variable_block_sizes: torch.Tensor,
-        logical_blocks: int,
-    ) -> torch.Tensor:
-        q = query.transpose(1, 2).contiguous()
-        k = key.transpose(1, 2).contiguous()
-        v = value.transpose(1, 2).contiguous()
-
-        # Official FastVideo's measured FastH3 route opts into this native
-        # Blackwell forward. Wheels without the extension (including SM103
-        # builds today) retain the corrected explicit-mask Triton route.
-        if os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1":
-            try:
-                from fastvideo_kernel import block_sparse_attn_sm100a
-                from fastvideo_kernel.triton_kernels.index import map_to_index
-
-                if block_sparse_attn_sm100a.is_supported(q, variable_block_sizes):
-                    q2k_idx, q2k_num = map_to_index(block_map)
-                    out, _ = block_sparse_attn_sm100a.block_sparse_attn_sm100a(
-                        q,
-                        k,
-                        v,
-                        q2k_idx.to(torch.int32).contiguous(),
-                        q2k_num.to(torch.int32).contiguous(),
-                        variable_block_sizes.to(torch.int32).contiguous(),
-                        need_lse=False,
-                    )
-                    return out.transpose(1, 2).contiguous()
-            except (ImportError, RuntimeError) as exc:
-                # Opting in explicitly and then silently getting a different
-                # numeric path is worse than the slower route it lands on.
-                logger.warning_once(
-                    "FASTVIDEO_VSA_SM100A=1 requested but the native Blackwell forward is "
-                    "unavailable (%s); using the Triton block-sparse route instead.",
-                    exc,
-                )
-
-        from fastvideo_kernel.block_sparse_attn import block_sparse_attn
-
-        logical_len = logical_blocks * 64
-        out, _ = block_sparse_attn(
-            q[:, :, :logical_len].contiguous(),
-            k[:, :, :logical_len].contiguous(),
-            v[:, :, :logical_len].contiguous(),
-            block_map[..., :logical_blocks, :logical_blocks].contiguous(),
-            variable_block_sizes[:logical_blocks].to(torch.int32).contiguous(),
-        )
-        out = out.transpose(1, 2).contiguous()
-        if out.shape[1] != query.shape[1]:
-            out = torch.nn.functional.pad(out, (0, 0, 0, 0, 0, query.shape[1] - out.shape[1]))
-        return out
-
-    @_fastvideo_h3_vsa_bhsd_op.register_fake
-    def _(query, key, value, block_map, variable_block_sizes, logical_blocks):
-        del key, value, block_map, variable_block_sizes, logical_blocks
-        return torch.empty_like(query)
-
-
-_fastvideo_h3_vsa_bhsd_op = torch.ops.vllm_omni.fastvideo_h3_vsa_bhsd
-
-
 @functools.lru_cache(maxsize=32)
 def _get_tile_partition_indices(
     dit_seq_shape: tuple[int, int, int],
@@ -226,81 +154,6 @@ def _get_tile_metadata(
     return tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index
 
 
-@functools.lru_cache(maxsize=32)
-def _get_h3_tile_metadata(
-    prefix_segments: tuple[int, ...],
-    video_shape: tuple[int, int, int],
-    device: torch.device,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
-    """Official FastVideo H3 geometry: pure prefix chunks + 3-D video tiles."""
-    block_size = (4, 4, 4)
-    block_elements = 64
-    prefix_len = sum(prefix_segments)
-    prefix_sizes: list[int] = []
-    for segment in prefix_segments:
-        full, remainder = divmod(segment, block_elements)
-        prefix_sizes.extend([block_elements] * full)
-        if remainder:
-            prefix_sizes.append(remainder)
-
-    video_indices = _get_tile_partition_indices(video_shape, block_size, device) + prefix_len
-    video_sizes = _construct_variable_block_sizes(video_shape, block_size, device)
-    partition = torch.cat([torch.arange(prefix_len, device=device, dtype=torch.long), video_indices])
-    sizes = torch.cat([torch.tensor(prefix_sizes, device=device, dtype=torch.int32), video_sizes.to(torch.int32)])
-    non_pad = _get_non_pad_index(sizes, block_elements)
-    untile = non_pad[torch.argsort(partition)]
-    total = prefix_len + math.prod(video_shape)
-    if int(sizes.sum()) != total or untile.numel() != total:
-        raise ValueError(
-            f"invalid H3 VSA geometry: prefix={prefix_segments}, video={video_shape}, "
-            f"sizes_sum={int(sizes.sum())}, total={total}"
-        )
-    return partition, sizes, non_pad, untile, len(prefix_sizes), int(video_sizes.numel())
-
-
-def _get_h3_layout(
-    attn_metadata: AttentionMetadata | None,
-) -> tuple[tuple[int, ...], tuple[int, int, int], int] | None:
-    if attn_metadata is None or attn_metadata.video_layout is None:
-        return None
-    prefix = attn_metadata.extra.get("vsa_h3_prefix_segments")
-    if not isinstance(prefix, (tuple, list)):
-        return None
-    target = next(
-        (span for span in reversed(attn_metadata.video_layout.video_spans) if span.role == "target"),
-        None,
-    )
-    if target is None:
-        return None
-    return tuple(int(x) for x in prefix if int(x) > 0), target.latent_grid, target.start
-
-
-def _pool_h3_tiles(x: torch.Tensor, sizes: torch.Tensor) -> torch.Tensor:
-    batch, seq_len, heads, dim = x.shape
-    blocks = seq_len // 64
-    pooled = x.view(batch, blocks, 64, heads, dim).sum(dim=2, dtype=torch.float32)
-    pooled = pooled / sizes.view(1, -1, 1, 1).clamp_min(1)
-    return pooled.permute(0, 2, 1, 3)
-
-
-def _build_h3_block_map(
-    scores: torch.Tensor,
-    num_prefix_blocks: int,
-    num_video_blocks: int,
-    topk: int,
-) -> torch.Tensor:
-    """Prefix K/V are exempt and prefix queries stay dense, as in FastVideo."""
-    keep_video = min(topk, num_video_blocks)
-    if keep_video == num_video_blocks:
-        return torch.ones_like(scores, dtype=torch.bool)
-    block_map = torch.zeros_like(scores, dtype=torch.bool)
-    indices = scores[..., num_prefix_blocks:].topk(keep_video, dim=-1).indices + num_prefix_blocks
-    block_map.scatter_(-1, indices, True)
-    block_map[..., :num_prefix_blocks] = True
-    block_map[:, :, :num_prefix_blocks, :] = True
-    return block_map
-
-
 def _get_vsa_dit_seq_shape(attn_metadata: AttentionMetadata | None) -> tuple[int, int, int] | None:
     if attn_metadata is None:
         return None
@@ -343,7 +196,9 @@ class FastVideoVSABackend(AttentionBackend):
         if importlib.util.find_spec("fastvideo_kernel") is None:
             raise ImportError(
                 "FASTVIDEO_VSA requires the optional fastvideo-kernel package "
-                "(the installation must provide the 'fastvideo_kernel' Python module)."
+                "included in vllm-omni[vsa]. Install with `uv pip install 'vllm-omni[vsa]'` "
+                "(from source: `uv pip install -e '.[vsa]'`). "
+                "Prebuilt kernels require Linux, Python 3.12 and glibc >= 2.34."
             )
 
     @staticmethod
@@ -504,93 +359,6 @@ class FastVideoVSAImpl(AttentionImpl):
                 return "sequence parallel context is active"
         return None
 
-    def _forward_h3(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-    ) -> torch.Tensor:
-        layout = _get_h3_layout(attn_metadata)
-        if layout is None:
-            raise ValueError("incomplete VSA-H3 layout metadata")
-        prefix_segments, video_shape, target_start = layout
-        if sum(prefix_segments) != target_start:
-            raise ValueError(f"VSA-H3 prefix segments sum to {sum(prefix_segments)}, target starts at {target_start}")
-        expected = target_start + math.prod(video_shape)
-        if query.shape[1] != expected:
-            raise ValueError(f"VSA-H3 layout has {expected} rows but attention received {query.shape[1]}")
-        if query.dtype not in (torch.float16, torch.bfloat16):
-            raise ValueError(f"VSA-H3 requires fp16/bf16 tensors, got {query.dtype}")
-        gate = _get_gate_compress(attn_metadata)
-        if gate is not None:
-            # H3 pads the packed document to 64 rows, while VSA operates on
-            # the valid prefix. Validate/slice before launching attention so a
-            # metadata error can never trigger an unsafe dense fallback after
-            # an asynchronous custom kernel.
-            if gate.shape[0] != query.shape[0] or gate.shape[2:] != query.shape[2:] or gate.shape[1] < query.shape[1]:
-                raise ValueError(f"gate_compress shape {gate.shape} cannot cover query shape {query.shape}")
-            gate = gate[:, : query.shape[1]]
-
-        partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
-            prefix_segments, video_shape, query.device
-        )
-        logical_blocks = int(sizes.numel())
-        # The native sm100a kernel assigns pairs of query blocks to CTAs. Its
-        # contract requires an even block count; the synthetic partner is
-        # transport-only and is removed before returning.
-        pair_pad = logical_blocks % 2
-        kernel_blocks = logical_blocks + pair_pad
-        target_shape = (query.shape[0], kernel_blocks * 64, query.shape[2], query.shape[3])
-        q_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype)
-        k_tiled = torch.zeros_like(q_tiled)
-        v_tiled = torch.zeros_like(q_tiled)
-        q_tiled[:, non_pad] = query[:, partition]
-        k_tiled[:, non_pad] = key[:, partition]
-        v_tiled[:, non_pad] = value[:, partition]
-
-        q_pool = _pool_h3_tiles(q_tiled[:, : logical_blocks * 64], sizes)
-        k_pool = _pool_h3_tiles(k_tiled[:, : logical_blocks * 64], sizes)
-        scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) * self.softmax_scale
-        block_map = _build_h3_block_map(scores, prefix_blocks, video_blocks, self.topk)
-        kernel_sizes = sizes
-        if pair_pad:
-            block_map = torch.nn.functional.pad(block_map, (0, 1, 0, 1), value=False)
-            kernel_sizes = torch.nn.functional.pad(sizes, (0, 1), value=0)
-
-        logger.info_once(
-            "FASTVIDEO_VSA H3 routing: seq_len=%d, prefix_segments=%s, video_shape=%s, "
-            "prefix_blocks=%d, video_blocks=%d, topk=%d, kernel_blocks=%d",
-            query.shape[1],
-            prefix_segments,
-            video_shape,
-            prefix_blocks,
-            video_blocks,
-            min(self.topk, video_blocks),
-            kernel_blocks,
-        )
-        output = _fastvideo_h3_vsa_bhsd_op(
-            q_tiled.contiguous(),
-            k_tiled.contiguous(),
-            v_tiled.contiguous(),
-            block_map.contiguous(),
-            kernel_sizes.contiguous(),
-            logical_blocks,
-        )[:, : logical_blocks * 64]
-
-        if gate is not None:
-            gate_tiled = torch.zeros_like(q_tiled[:, : logical_blocks * 64])
-            gate_tiled[:, non_pad] = gate[:, partition]
-            v_pool = _pool_h3_tiles(v_tiled[:, : logical_blocks * 64], sizes)
-            compressed = torch.matmul(torch.softmax(scores, dim=-1), v_pool)
-            compressed = compressed.permute(0, 2, 1, 3).to(output.dtype)
-            output = (
-                output.view(output.shape[0], logical_blocks, 64, output.shape[2], output.shape[3])
-                + compressed.unsqueeze(2)
-                * gate_tiled.view(gate_tiled.shape[0], logical_blocks, 64, gate_tiled.shape[2], gate_tiled.shape[3])
-            ).view_as(output)
-        return output[:, untile].contiguous()
-
     def forward_cuda(
         self,
         query: torch.Tensor,
@@ -611,26 +379,6 @@ class FastVideoVSAImpl(AttentionImpl):
             key = key[:, :valid_seq_len]
             value = value[:, :valid_seq_len]
 
-        if attn_metadata is not None and _get_h3_layout(attn_metadata) is not None:
-            try:
-                output = self._forward_h3(query, key, value, attn_metadata)
-                if valid_seq_len == original_seq_len:
-                    return output
-                restored = torch.zeros_like(original_query)
-                restored[:, :valid_seq_len] = output
-                return restored
-            except Exception as exc:
-                # A CUDA fault poisons the process context; attempting SDPA
-                # afterwards obscures the original kernel failure and cannot
-                # recover the request.
-                if isinstance(exc, torch.AcceleratorError):
-                    raise
-                if not self.fallback_on_error:
-                    raise
-                return self._fallback(
-                    original_query, original_key, original_value, attn_metadata, f"VSA-H3 kernel failed: {exc}"
-                )
-
         reason = self._fallback_reason(query, key, value, attn_metadata)
         if reason is not None:
             return self._fallback(original_query, original_key, original_value, attn_metadata, reason)

vllm_omni/diffusion/attention/layer.py

@@ -7,7 +7,9 @@
 # https://github.com/feifeibear/long-context-attention/blob/main/yunchang/attention/layer.py
 
 
+from collections.abc import Mapping
 from dataclasses import replace
+from typing import cast
 
 import torch
 import torch.nn as nn
@@ -16,7 +18,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.utils import extract_layer_index
 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
 from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
 from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
@@ -83,6 +85,10 @@ class Attention(nn.Module):
         custom_attention: nn.Module | None = None,
         # Preserve dense FP32 inference for models opting into CUDA auto fallback.
         allow_fp32_fallback: bool = False,
+        # Model-owned implementation subclasses, keyed by backend name.
+        # Keep the platform-selected backend and its capabilities unchanged.
+        # Each override must preserve the selected implementation's contract.
+        impl_overrides: Mapping[str, type[AttentionImpl]] | None = None,
     ):
         super().__init__()
 
@@ -181,8 +187,18 @@ class Attention(nn.Module):
                 self.backend_pref = attn_backend_cls.get_name()
                 logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref)
 
-            self.attn_backend = attn_backend_cls
+            self.attn_backend: type[AttentionBackend] | None = attn_backend_cls
             self.attn_impl_cls = self.attn_backend.get_impl_cls()
+            if impl_overrides is not None:
+                override = impl_overrides.get(attn_backend_cls.get_name())
+                if override is not None:
+                    if not issubclass(override, self.attn_impl_cls):
+                        raise TypeError(
+                            f"Attention implementation override {override.__qualname__} must subclass "
+                            f"the selected implementation {self.attn_impl_cls.__qualname__} "
+                            f"for backend {attn_backend_cls.__qualname__}"
+                        )
+                    self.attn_impl_cls = override
             self.attention = self.attn_impl_cls(
                 num_heads=num_heads,
                 head_size=head_size,
@@ -197,7 +213,7 @@ class Attention(nn.Module):
             )
             # Compatibility kernels run inside shared dispatch, between the
             # parallel strategy's input preparation and output restoration.
-            self.sdpa_fallback = SDPABackend.get_impl_cls()(
+            self.sdpa_fallback: AttentionImpl | None = SDPABackend.get_impl_cls()(
                 num_heads=num_heads,
                 head_size=head_size,
                 softmax_scale=softmax_scale,
@@ -299,6 +315,7 @@ class Attention(nn.Module):
     def _init_kv_cache_quantization(self, config) -> None:
         if config is None or self._has_custom_attention:
             return
+        assert self.attn_backend is not None
         dtype = getattr(config, "diffusion_kv_cache_dtype", None)
         if dtype == "auto":
             dtype = None
@@ -400,6 +417,7 @@ class Attention(nn.Module):
             )
         use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
         if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
+            assert self.attn_backend is not None
             raise NotImplementedError(
                 f"Diffusion paged KV requires an Omni backend with paged support; "
                 f"selected {self.attn_backend.get_name()}"
@@ -496,7 +514,7 @@ class Attention(nn.Module):
 
     def _run_local_attention(self, query, key, value, attn_metadata):
         if self._has_custom_attention:
-            return self.attention(query, key, value, attn_metadata)
+            return cast(nn.Module, self.attention)(query, key, value, attn_metadata)
 
         self._assert_metadata_compatible(attn_metadata)
 
@@ -506,7 +524,7 @@ class Attention(nn.Module):
             and query.dtype == torch.float32
             and query.ndim == 4
             and not self.backend_explicit
-            and self.attn_backend.get_name() == "FLASH_ATTN"
+            and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN"
             and (
                 attn_metadata is None
                 or (
@@ -519,7 +537,7 @@ class Attention(nn.Module):
             )
         ):
             logger.warning_once("Using SDPA for this layer's FP32 input with automatic CUDA FlashAttention selection.")
-            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+            return cast(AttentionImpl, self.sdpa_fallback).forward(query, key, value, attn_metadata)
 
         in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
         # The startup KV-capacity profile needs tensor shapes, not a paged
@@ -531,14 +549,14 @@ class Attention(nn.Module):
             self._scheduler_paged_kv
             and self.paged_kv_cache_role is not None
             and in_kv_memory_profile
-            and self.attn_backend.get_name() == "FLASH_ATTN"
+            and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN"
             and not current_omni_platform.supports_diffusion_dense_flash_attention()
         ):
             logger.warning_once(
                 "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. "
                 "Formal paged requests still use the platform-native paged attention backend."
             )
-            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+            return cast(AttentionImpl, self.sdpa_fallback).forward(query, key, value, attn_metadata)
 
         return self.attention.forward(query, key, value, attn_metadata)
 

vllm_omni/diffusion/attention/ops/block_sparse.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shared pooling, routing and provider calls for explicit block-sparse attention."""
+
+import os
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+if not hasattr(torch.ops.vllm_omni, "fastvideo_block_sparse_attn_bshd"):
+
+    @torch.library.custom_op("vllm_omni::fastvideo_block_sparse_attn_bshd", mutates_args=())
+    def fastvideo_block_sparse_attn_bshd(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        block_map: torch.Tensor,
+        variable_block_sizes: torch.Tensor,
+        logical_blocks: int,
+    ) -> torch.Tensor:
+        """Run explicit block-sparse attention on BSHD tensors in 64-row blocks.
+
+        Inputs and the block map may include transport-only blocks after
+        ``logical_blocks``. The Triton provider sees only logical blocks and
+        the wrapper restores the original row count. Native provider selection
+        and its supported-input checks remain owned by FastVideo.
+        """
+        q = query.transpose(1, 2).contiguous()
+        k = key.transpose(1, 2).contiguous()
+        v = value.transpose(1, 2).contiguous()
+
+        # Prefer the explicitly selected native provider when it supports
+        # these tensors; retain the existing Triton provider otherwise.
+        if os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1":
+            try:
+                from fastvideo_kernel import block_sparse_attn_sm100a
+                from fastvideo_kernel.triton_kernels.index import map_to_index
+
+                if block_sparse_attn_sm100a.is_supported(q, variable_block_sizes):
+                    q2k_idx, q2k_num = map_to_index(block_map)
+                    out, _ = block_sparse_attn_sm100a.block_sparse_attn_sm100a(
+                        q,
+                        k,
+                        v,
+                        q2k_idx.to(torch.int32).contiguous(),
+                        q2k_num.to(torch.int32).contiguous(),
+                        variable_block_sizes.to(torch.int32).contiguous(),
+                        need_lse=False,
+                    )
+                    return out.transpose(1, 2).contiguous()
+            except torch.AcceleratorError:
+                # A device fault can poison the context; do not launch another kernel.
+                raise
+            except (ImportError, RuntimeError) as exc:
+                logger.warning_once(
+                    "FASTVIDEO_VSA_SM100A=1 requested but the native Blackwell forward is "
+                    "unavailable (%s); using the Triton block-sparse route instead.",
+                    exc,
+                )
+
+        from fastvideo_kernel.block_sparse_attn import block_sparse_attn
+
+        logical_len = logical_blocks * 64
+        out, _ = block_sparse_attn(
+            q[:, :, :logical_len].contiguous(),
+            k[:, :, :logical_len].contiguous(),
+            v[:, :, :logical_len].contiguous(),
+            block_map[..., :logical_blocks, :logical_blocks].contiguous(),
+            variable_block_sizes[:logical_blocks].to(torch.int32).contiguous(),
+        )
+        out = out.transpose(1, 2).contiguous()
+        if out.shape[1] != query.shape[1]:
+            out = torch.nn.functional.pad(out, (0, 0, 0, 0, 0, query.shape[1] - out.shape[1]))
+        return out
+
+    @fastvideo_block_sparse_attn_bshd.register_fake
+    def _(query, key, value, block_map, variable_block_sizes, logical_blocks):
+        del key, value, block_map, variable_block_sizes, logical_blocks
+        return torch.empty_like(query)
+
+
+fastvideo_block_sparse_attn_bshd = torch.ops.vllm_omni.fastvideo_block_sparse_attn_bshd
+
+
+def mean_pool_tiles(x: torch.Tensor, sizes: torch.Tensor, block_size: int) -> torch.Tensor:
+    """Mean-pool zero-padded BSHD tiles with FP32 accumulation.
+
+    ``sizes`` gives the number of valid rows in each tile.
+    """
+    batch, seq_len, heads, dim = x.shape
+    blocks = seq_len // block_size
+    pooled = x.view(batch, blocks, block_size, heads, dim).sum(dim=2, dtype=torch.float32)
+    pooled = pooled / sizes.view(1, -1, 1, 1).clamp_min(1)
+    return pooled.permute(0, 2, 1, 3)
+
+
+def build_prefix_dense_block_map(
+    scores: torch.Tensor,
+    num_prefix_blocks: int,
+    num_sparse_blocks: int,
+    topk: int,
+) -> torch.Tensor:
+    """Keep prefix queries/keys dense and select top-k non-prefix keys per row."""
+    keep_sparse = min(topk, num_sparse_blocks)
+    if keep_sparse == num_sparse_blocks:
+        return torch.ones_like(scores, dtype=torch.bool)
+    block_map = torch.zeros_like(scores, dtype=torch.bool)
+    indices = scores[..., num_prefix_blocks:].topk(keep_sparse, dim=-1).indices + num_prefix_blocks
+    block_map.scatter_(-1, indices, True)
+    block_map[..., :num_prefix_blocks] = True
+    block_map[:, :, :num_prefix_blocks, :] = True
+    return block_map

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1262,7 +1262,7 @@ def extract_minimax_h3_context(
     preprocessing via ``_embed``, cacheable ``blocks`` loop, and
     ``final_layer`` postprocessing with row selection and update masks.
     """
-    from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import indexed_scale_shift_
+    from vllm_omni.diffusion.layers.indexed_modulation import indexed_scale_shift_
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
         _FORWARD_SUPPORTED_KWARGS,
         MINIMAX_H3_ADALN_MODALITY_NUM,

vllm_omni/diffusion/layers/indexed_modulation.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Fused MiniMax H3 modulation with FP32 accumulation."""
+"""Fused indexed modulation with FP32 accumulation."""
 
 from __future__ import annotations
 
@@ -244,7 +244,7 @@ def rms_norm_indexed_scale_shift(
     indices: torch.Tensor,
     eps: float,
 ) -> torch.Tensor:
-    """Fuse H3 RMSNorm with its indexed AdaLN affine transform."""
+    """Fuse RMSNorm with an indexed affine transform."""
     if x.is_cpu:
         input_dtype = x.dtype
         normalized = x.float()

vllm_omni/diffusion/models/minimax_h3/attention/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

vllm_omni/diffusion/models/minimax_h3/attention/fastvideo_h3.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MiniMax-H3 VSA prefix routing, tile layout, and learned compression gate."""
+
+from __future__ import annotations
+
+import functools
+import math
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
+    FastVideoVSAImpl,
+    _construct_variable_block_sizes,
+    _get_gate_compress,
+    _get_non_pad_index,
+    _get_tile_partition_indices,
+)
+from vllm_omni.diffusion.attention.ops.block_sparse import (
+    build_prefix_dense_block_map,
+    fastvideo_block_sparse_attn_bshd,
+    mean_pool_tiles,
+)
+
+logger = init_logger(__name__)
+
+
+@functools.lru_cache(maxsize=32)
+def _get_h3_tile_metadata(
+    prefix_segments: tuple[int, ...],
+    video_shape: tuple[int, int, int],
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+    """Official FastVideo H3 geometry: pure prefix chunks + 3-D video tiles."""
+    block_size = (4, 4, 4)
+    block_elements = 64
+    prefix_len = sum(prefix_segments)
+    prefix_sizes: list[int] = []
+    for segment in prefix_segments:
+        full, remainder = divmod(segment, block_elements)
+        prefix_sizes.extend([block_elements] * full)
+        if remainder:
+            prefix_sizes.append(remainder)
+
+    video_indices = _get_tile_partition_indices(video_shape, block_size, device) + prefix_len
+    video_sizes = _construct_variable_block_sizes(video_shape, block_size, device)
+    partition = torch.cat([torch.arange(prefix_len, device=device, dtype=torch.long), video_indices])
+    sizes = torch.cat([torch.tensor(prefix_sizes, device=device, dtype=torch.int32), video_sizes.to(torch.int32)])
+    non_pad = _get_non_pad_index(sizes, block_elements)
+    untile = non_pad[torch.argsort(partition)]
+    total = prefix_len + math.prod(video_shape)
+    if int(sizes.sum()) != total or untile.numel() != total:
+        raise ValueError(
+            f"invalid H3 VSA geometry: prefix={prefix_segments}, video={video_shape}, "
+            f"sizes_sum={int(sizes.sum())}, total={total}"
+        )
+    return partition, sizes, non_pad, untile, len(prefix_sizes), int(video_sizes.numel())
+
+
+def _get_h3_layout(
+    attn_metadata: AttentionMetadata | None,
+) -> tuple[tuple[int, ...], tuple[int, int, int], int] | None:
+    if attn_metadata is None or attn_metadata.video_layout is None:
+        return None
+    prefix = attn_metadata.extra.get("vsa_h3_prefix_segments")
+    if not isinstance(prefix, (tuple, list)):
+        return None
+    target = next(
+        (span for span in reversed(attn_metadata.video_layout.video_spans) if span.role == "target"),
+        None,
+    )
+    if target is None:
+        return None
+    return tuple(int(x) for x in prefix if int(x) > 0), target.latent_grid, target.start
+
+
+class MiniMaxH3VSAImpl(FastVideoVSAImpl):
+    """Apply H3 tile64 routing after shared parallel attention dispatch."""
+
+    def _forward_h3(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        layout = _get_h3_layout(attn_metadata)
+        if layout is None:
+            raise ValueError("incomplete VSA-H3 layout metadata")
+        prefix_segments, video_shape, target_start = layout
+        if sum(prefix_segments) != target_start:
+            raise ValueError(f"VSA-H3 prefix segments sum to {sum(prefix_segments)}, target starts at {target_start}")
+        expected = target_start + math.prod(video_shape)
+        if query.shape[1] != expected:
+            raise ValueError(f"VSA-H3 layout has {expected} rows but attention received {query.shape[1]}")
+        if query.dtype not in (torch.float16, torch.bfloat16):
+            raise ValueError(f"VSA-H3 requires fp16/bf16 tensors, got {query.dtype}")
+        gate = _get_gate_compress(attn_metadata)
+        if gate is not None:
+            # H3 pads the packed document to 64 rows, while VSA operates on
+            # the valid prefix. Validate/slice before launching attention so a
+            # metadata error can never trigger an unsafe dense fallback after
+            # an asynchronous custom kernel.
+            if gate.shape[0] != query.shape[0] or gate.shape[2:] != query.shape[2:] or gate.shape[1] < query.shape[1]:
+                raise ValueError(f"gate_compress shape {gate.shape} cannot cover query shape {query.shape}")
+            gate = gate[:, : query.shape[1]]
+
+        partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
+            prefix_segments, video_shape, query.device
+        )
+        logical_blocks = int(sizes.numel())
+        # The native sm100a kernel assigns pairs of query blocks to CTAs. Its
+        # contract requires an even block count; the synthetic partner is
+        # transport-only and is removed before returning.
+        pair_pad = logical_blocks % 2
+        kernel_blocks = logical_blocks + pair_pad
+        target_shape = (query.shape[0], kernel_blocks * 64, query.shape[2], query.shape[3])
+        q_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype)
+        k_tiled = torch.zeros_like(q_tiled)
+        v_tiled = torch.zeros_like(q_tiled)
+        q_tiled[:, non_pad] = query[:, partition]
+        k_tiled[:, non_pad] = key[:, partition]
+        v_tiled[:, non_pad] = value[:, partition]
+
+        q_pool = mean_pool_tiles(q_tiled[:, : logical_blocks * 64], sizes, block_size=64)
+        k_pool = mean_pool_tiles(k_tiled[:, : logical_blocks * 64], sizes, block_size=64)
+        scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) * self.softmax_scale
+        block_map = build_prefix_dense_block_map(scores, prefix_blocks, video_blocks, self.topk)
+        kernel_sizes = sizes
+        if pair_pad:
+            block_map = torch.nn.functional.pad(block_map, (0, 1, 0, 1), value=False)
+            kernel_sizes = torch.nn.functional.pad(sizes, (0, 1), value=0)
+
+        logger.info_once(
+            "FASTVIDEO_VSA H3 routing: seq_len=%d, prefix_segments=%s, video_shape=%s, "
+            "prefix_blocks=%d, video_blocks=%d, topk=%d, kernel_blocks=%d",
+            query.shape[1],
+            prefix_segments,
+            video_shape,
+            prefix_blocks,
+            video_blocks,
+            min(self.topk, video_blocks),
+            kernel_blocks,
+        )
+        output = fastvideo_block_sparse_attn_bshd(
+            q_tiled.contiguous(),
+            k_tiled.contiguous(),
+            v_tiled.contiguous(),
+            block_map.contiguous(),
+            kernel_sizes.contiguous(),
+            logical_blocks,
+        )[:, : logical_blocks * 64]
+
+        if gate is not None:
+            gate_tiled = torch.zeros_like(q_tiled[:, : logical_blocks * 64])
+            gate_tiled[:, non_pad] = gate[:, partition]
+            v_pool = mean_pool_tiles(v_tiled[:, : logical_blocks * 64], sizes, block_size=64)
+            compressed = torch.matmul(torch.softmax(scores, dim=-1), v_pool)
+            compressed = compressed.permute(0, 2, 1, 3).to(output.dtype)
+            output = (
+                output.view(output.shape[0], logical_blocks, 64, output.shape[2], output.shape[3])
+                + compressed.unsqueeze(2)
+                * gate_tiled.view(gate_tiled.shape[0], logical_blocks, 64, gate_tiled.shape[2], gate_tiled.shape[3])
+            ).view_as(output)
+        return output[:, untile].contiguous()
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        if _get_h3_layout(attn_metadata) is None:
+            return super().forward_cuda(query, key, value, attn_metadata)
+
+        original_query, original_key, original_value = query, key, value
+        original_seq_len = query.shape[1]
+        valid_seq_len = original_seq_len
+        if attn_metadata is not None and attn_metadata.packed_padding is not None:
+            valid_seq_len = attn_metadata.packed_padding.q_length
+            if attn_metadata.packed_padding.kv_length != valid_seq_len:
+                return self._fallback(
+                    original_query, original_key, original_value, attn_metadata, "packed Q/KV lengths must match"
+                )
+            query = query[:, :valid_seq_len]
+            key = key[:, :valid_seq_len]
+            value = value[:, :valid_seq_len]
+
+        try:
+            output = self._forward_h3(query, key, value, attn_metadata)
+            if valid_seq_len == original_seq_len:
+                return output
+            restored = torch.zeros_like(original_query)
+            restored[:, :valid_seq_len] = output
+            return restored
+        except Exception as exc:
+            # A CUDA fault poisons the process context; attempting SDPA
+            # afterwards obscures the original kernel failure and cannot
+            # recover the request.
+            if isinstance(exc, torch.AcceleratorError):
+                raise
+            if not self.fallback_on_error:
+                raise
+            return self._fallback(
+                original_query, original_key, original_value, attn_metadata, f"VSA-H3 kernel failed: {exc}"
+            )

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -33,12 +33,6 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     VideoTokenLayout,
 )
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
-    indexed_gate,
-    indexed_gate_rms_norm_scale_shift,
-    indexed_scale_shift_,
-    rms_norm_indexed_scale_shift,
-)
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -50,6 +44,12 @@ from vllm_omni.diffusion.forward_context import (
 )
 from vllm_omni.diffusion.layers.activation import SiluAndMul
 from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
+from vllm_omni.diffusion.layers.indexed_modulation import (
+    indexed_gate,
+    indexed_gate_rms_norm_scale_shift,
+    indexed_scale_shift_,
+    rms_norm_indexed_scale_shift,
+)
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
@@ -442,6 +442,8 @@ class MiniMaxH3Attention(nn.Module):
         self._gate_hidden_size = arch.hidden_size
         self._gate_quant_config = quant_config
         self._gate_prefix = f"{prefix}.to_gate_compress"
+        from .attention.fastvideo_h3 import MiniMaxH3VSAImpl
+
         self.attention = Attention(
             num_heads=self.num_heads,
             num_kv_heads=self.num_kv_heads,
@@ -454,6 +456,7 @@ class MiniMaxH3Attention(nn.Module):
             role_category=role_category,
             skip_sequence_parallel=skip_sequence_parallel,
             prefix=prefix,
+            impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl},
         )
 
     def enable_vsa_gate(self) -> None: