From 507cb1d834db66f2dbb776602e82c62d58e92992 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:47:31 +0800 Subject: [PATCH] [Attention] [P1] Keep H3 VSA policy in the model and share sparse primitives (#7535) Signed-off-by: lishunyang12 Co-authored-by: lishunyang12 --- .../feature/attention_backend_selection.md | 38 ++- .../attention_backends/fastvideo_vsa.md | 78 ++++-- pyproject.toml | 5 + .../attention/test_block_sparse_ops.py | 126 +++++++++ .../diffusion/attention/test_fastvideo_vsa.py | 164 +---------- .../attention/test_impl_overrides.py | 136 +++++++++ .../cache/test_teacache_extractors.py | 10 +- .../minimax_h3/test_minimax_h3_modulation.py | 2 +- .../models/minimax_h3/test_vsa_layout.py | 221 +++++++++++++++ tests/helpers/runtime.py | 62 ++++- .../tests/test_runtime_server_ports.py | 109 ++++++++ .../attention/backends/fastvideo_vsa.py | 258 +----------------- vllm_omni/diffusion/attention/layer.py | 34 ++- .../diffusion/attention/ops/block_sparse.py | 115 ++++++++ .../diffusion/cache/teacache/extractors.py | 2 +- .../indexed_modulation.py} | 4 +- .../models/minimax_h3/attention/__init__.py | 2 + .../minimax_h3/attention/fastvideo_h3.py | 209 ++++++++++++++ .../minimax_h3/minimax_h3_transformer.py | 15 +- 19 files changed, 1117 insertions(+), 473 deletions(-) create mode 100644 tests/diffusion/attention/test_block_sparse_ops.py create mode 100644 tests/diffusion/attention/test_impl_overrides.py create mode 100644 tests/diffusion/models/minimax_h3/test_vsa_layout.py create mode 100644 tests/helpers/tests/test_runtime_server_ports.py create mode 100644 vllm_omni/diffusion/attention/ops/block_sparse.py rename vllm_omni/diffusion/{attention/ops/minimax_h3_modulation.py => layers/indexed_modulation.py} (98%) create mode 100644 vllm_omni/diffusion/models/minimax_h3/attention/__init__.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/attention/fastvideo_h3.py diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md index 41cbaf3ff..3ae1115b1 100644 --- a/docs/design/feature/attention_backend_selection.md +++ b/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`. diff --git a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md index c03b9ffde..df6e34b0c 100644 --- a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md +++ b/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 \ +vllm serve --omni \ --diffusion-attention-backend FASTVIDEO_VSA \ --fastvideo-vsa-topk 64 ``` -The backwards-compatible environment variable selects the backend with the -default `topk=64`: +
+Alternative configuration formats + +Equivalent structured configuration: ```bash -export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA -vllm-omni serve -``` - -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 \ +vllm serve --omni \ --diffusion-attention-config \ '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}' ``` @@ -73,8 +76,15 @@ stages: fastvideo_vsa_topk: 64 ``` +
+ ## 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 diff --git a/pyproject.toml b/pyproject.toml index fc1d5573b..434c4305c 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/diffusion/attention/test_block_sparse_ops.py b/tests/diffusion/attention/test_block_sparse_ops.py new file mode 100644 index 000000000..845f09b46 --- /dev/null +++ b/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() diff --git a/tests/diffusion/attention/test_fastvideo_vsa.py b/tests/diffusion/attention/test_fastvideo_vsa.py index 1bd7192ad..e5c854d8d 100644 --- a/tests/diffusion/attention/test_fastvideo_vsa.py +++ b/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() diff --git a/tests/diffusion/attention/test_impl_overrides.py b/tests/diffusion/attention/test_impl_overrides.py new file mode 100644 index 000000000..f7755de5c --- /dev/null +++ b/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}, + ) diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index d821541d0..55b92b820 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/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]) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py index e10fbe56a..7fde01354 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/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, diff --git a/tests/diffusion/models/minimax_h3/test_vsa_layout.py b/tests/diffusion/models/minimax_h3/test_vsa_layout.py new file mode 100644 index 000000000..c556543d0 --- /dev/null +++ b/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))) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 66e0cefab..0632db55d 100644 --- a/tests/helpers/runtime.py +++ b/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): diff --git a/tests/helpers/tests/test_runtime_server_ports.py b/tests/helpers/tests/test_runtime_server_ports.py new file mode 100644 index 000000000..1d8c0b2f1 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py index 465912d39..cbab82c9a 100644 --- a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py +++ b/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) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4c588af0c..2e0b2604d 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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) diff --git a/vllm_omni/diffusion/attention/ops/block_sparse.py b/vllm_omni/diffusion/attention/ops/block_sparse.py new file mode 100644 index 000000000..aeec14699 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 3aca638ec..a7539b34a 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/layers/indexed_modulation.py similarity index 98% rename from vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py rename to vllm_omni/diffusion/layers/indexed_modulation.py index 20327b7e9..385ce1a1b 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/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() diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/__init__.py b/vllm_omni/diffusion/models/minimax_h3/attention/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/fastvideo_h3.py b/vllm_omni/diffusion/models/minimax_h3/attention/fastvideo_h3.py new file mode 100644 index 000000000..35b87a58b --- /dev/null +++ b/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}" + ) diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index f0839429b..e26493931 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/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: