From 9f966f76398ac41731abef8d9291ba61649199b8 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Wed, 30 Sep 2026 09:45:07 -0700 Subject: [PATCH 1/8] [Core] Add shared SubBlock sparse attention with an FA4 adapter Signed-off-by: Rahul Steiger --- .gitignore | 1 + benchmarks/diffusion/fa4_subblock.py | 133 +++ docs/design/feature/fa4_subblock_poc.md | 130 +++ recipes/attention/fa4-subblock.json | 16 + .../attention/minimax-h3-fa4-subblock.json | 16 + .../attention/test_block_selection.py | 247 +++++ .../diffusion/attention/test_block_sparse.py | 879 ++++++++++++++++++ .../attention/test_block_sparse_adapters.py | 87 ++ .../attention/test_block_sparse_owner.py | 89 ++ .../test_diffusers_backend.py | 16 + .../cosmos3/test_cosmos3_transformer.py | 159 +++- .../minimax_h3/test_fasth3_checkpoint.py | 32 + .../minimax_h3/test_minimax_h3_contract.py | 170 ++++ .../minimax_h3/test_minimax_h3_fasth3.py | 18 +- tests/helpers/block_sparse.py | 94 ++ .../diffusion/attention/backends/abstract.py | 92 ++ .../attention/backends/flash_attn.py | 105 ++- .../attention/block_selection/__init__.py | 2 + .../attention/block_selection/abstract.py | 72 ++ .../attention/block_selection/registry.py | 31 + .../block_selection/subblock_topk.py | 405 ++++++++ vllm_omni/diffusion/attention/block_sparse.py | 410 ++++++++ vllm_omni/diffusion/attention/layer.py | 81 +- vllm_omni/diffusion/attention/selector.py | 35 +- vllm_omni/diffusion/data.py | 98 +- .../models/cosmos3/transformer_cosmos3.py | 14 +- .../cosmos3/transformer_cosmos3_edge.py | 4 +- .../pipeline_diffusers_adapter.py | 11 +- .../diffusion/models/minimax_h3/fasth3.py | 29 +- .../models/minimax_h3/fasth3_checkpoint.py | 6 +- .../minimax_h3/minimax_h3_transformer.py | 6 + 31 files changed, 3414 insertions(+), 74 deletions(-) create mode 100644 benchmarks/diffusion/fa4_subblock.py create mode 100644 docs/design/feature/fa4_subblock_poc.md create mode 100644 recipes/attention/fa4-subblock.json create mode 100644 recipes/attention/minimax-h3-fa4-subblock.json create mode 100644 tests/diffusion/attention/test_block_selection.py create mode 100644 tests/diffusion/attention/test_block_sparse.py create mode 100644 tests/diffusion/attention/test_block_sparse_adapters.py create mode 100644 tests/diffusion/attention/test_block_sparse_owner.py create mode 100644 tests/helpers/block_sparse.py create mode 100644 vllm_omni/diffusion/attention/block_selection/__init__.py create mode 100644 vllm_omni/diffusion/attention/block_selection/abstract.py create mode 100644 vllm_omni/diffusion/attention/block_selection/registry.py create mode 100644 vllm_omni/diffusion/attention/block_selection/subblock_topk.py create mode 100644 vllm_omni/diffusion/attention/block_sparse.py diff --git a/.gitignore b/.gitignore index ae59008c010..5d3406844da 100644 --- a/.gitignore +++ b/.gitignore @@ -242,6 +242,7 @@ data/ datasets/ *.csv *.json +!recipes/attention/minimax-h3-fa4-subblock.json !recipes/cosmos3/negative_prompt.json !benchmarks/mammoth_moda2/reference_results.json !tests/dfx/perf/tests/*.json diff --git a/benchmarks/diffusion/fa4_subblock.py b/benchmarks/diffusion/fa4_subblock.py new file mode 100644 index 00000000000..53a65415135 --- /dev/null +++ b/benchmarks/diffusion/fa4_subblock.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Benchmark FA4 SubBlock attention calls on synthetic or captured BSHD inputs.""" + +import argparse +import json +from pathlib import Path + +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl +from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention +from vllm_omni.diffusion.data import BlockSparseAttentionSpec + + +def latency(fn, iterations): + for _ in range(5): + fn() + start, end = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + start.record() + for _ in range(iterations): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +@torch.inference_mode() +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--capture", help="Torch file with q/k/v and optional meta.scale/meta.prefix_len") + parser.add_argument("--q-length", type=int, default=1024) + parser.add_argument("--kv-length", type=int, default=1089) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--kv-heads", type=int, default=2) + parser.add_argument("--block-size", type=int, nargs=2, default=[64, 64], metavar=("BQ", "BKV")) + parser.add_argument("--sparsity", type=float, default=0.75, help="Target fraction of unprotected blocks to drop") + parser.add_argument("--output", type=Path, help="Write the measurement as JSON") + parser.add_argument("--iterations", type=int, default=30) + args = parser.parse_args() + if args.iterations <= 0: + parser.error("--iterations must be positive") + torch.manual_seed(17) + scale, prefix = 128**-0.5, 0 + if args.capture: + capture = torch.load(args.capture, map_location="cpu", weights_only=True) + q, k, v = (capture[name].cuda() for name in ("q", "k", "v")) + scale = capture.get("meta", {}).get("scale", q.shape[-1] ** -0.5) + prefix = capture.get("meta", {}).get("prefix_len", 0) + else: + q = torch.randn(1, args.q_length, args.heads, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, args.kv_length, args.kv_heads, 128, device=q.device, dtype=q.dtype) + v = torch.randn_like(k) + spec = BlockSparseAttentionSpec( + name="block_sparse", + config={ + "block_size": args.block_size, + "selection": {"name": "block_topk", "config": {"target_sparsity": args.sparsity}}, + "backend": {"require": "FLASH_ATTN", "implementation": "auto"}, + }, + ) + sparse = BlockSparseAttention( + q.shape[2], + k.shape[2], + q.shape[-1], + scale, + False, + None, + spec, + adapter=FlashAttentionBackend.get_block_sparse_adapter()(), + ) + dense = FlashAttentionImpl(q.shape[2], q.shape[-1], scale, num_kv_heads=k.shape[2]) + metadata = AttentionMetadata(extra={"protected_kv_prefix": prefix}) + sparse.prepare_request(q, k, v, metadata) + compiled_dense = torch.compile(dense.forward_cuda, fullgraph=True) + compiled = torch.compile(sparse.forward_cuda, fullgraph=True) + torch.testing.assert_close(compiled(q, k, v, metadata), sparse.forward_cuda(q, k, v, metadata), atol=0, rtol=0) + torch.testing.assert_close(compiled_dense(q, k, v), dense.forward_cuda(q, k, v), atol=0, rtol=0) + blocks = (k.shape[1] + args.block_size[1] - 1) // args.block_size[1] + prefix_blocks = (prefix + args.block_size[1] - 1) // args.block_size[1] + # Report the actual selected pattern, rather than duplicating budget arithmetic. + selection = sparse.selector.select(q, k, scale, prefix) + retained_blocks = torch.unique(selection.counts).item() + assert retained_blocks >= prefix_blocks + torch.testing.assert_close( + selection.indices[..., :prefix_blocks], + torch.arange(prefix_blocks, device=q.device, dtype=torch.int32).expand( + *selection.indices.shape[:-1], prefix_blocks + ), + ) + unprotected_blocks = blocks - prefix_blocks + from vllm_omni.diffusion.attention.backends.utils import fa + + report = json.dumps( + { + "benchmark_scope": "attention_call", + "model_role_activation_verified": False, + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "q_shape": list(q.shape), + "kv_shape": list(k.shape), + "protected_prefix": prefix, + "fa4_version": sparse.adapter.dependency_version, + "dense_kernel": getattr(fa.flash_attn_func, "__module__", "varlen fallback"), + "config": spec.config, + "protected_blocks": prefix_blocks, + "target_sparsity_scope": "unprotected_blocks", + "unprotected_blocks": unprotected_blocks, + "retained_unprotected_blocks": retained_blocks - prefix_blocks, + "retained_blocks": retained_blocks, + "overall_block_sparsity": 1 - retained_blocks / blocks, + "unprotected_block_sparsity": ( + 1 - (retained_blocks - prefix_blocks) / unprotected_blocks if unprotected_blocks else None + ), + "iterations": args.iterations, + "warmup_iterations": 5, + "eager_compiled_outputs_equal": True, + "total_blocks": blocks, + "dense_eager_ms": latency(lambda: dense.forward_cuda(q, k, v), args.iterations), + "dense_compiled_ms": latency(lambda: compiled_dense(q, k, v), args.iterations), + "sparse_eager_ms": latency(lambda: sparse.forward_cuda(q, k, v, metadata), args.iterations), + "sparse_compiled_ms": latency(lambda: compiled(q, k, v, metadata), args.iterations), + }, + indent=2, + ) + print(report) + if args.output: + args.output.write_text(report + "\n") + + +if __name__ == "__main__": + main() diff --git a/docs/design/feature/fa4_subblock_poc.md b/docs/design/feature/fa4_subblock_poc.md new file mode 100644 index 00000000000..ff8f1169d75 --- /dev/null +++ b/docs/design/feature/fa4_subblock_poc.md @@ -0,0 +1,130 @@ +# SubBlock attention with FA4: proof of concept + +This PoC adds shared block selection and an FA4 execution adapter on top of +[the attention execution contract (#7379)](https://github.com/vllm-project/vllm-omni/pull/7379). +The selector chooses logical KV blocks; the adapter executes attention over those +blocks. Dense attention remains the default. + +The selector adapts the original +[SGLang MiniMax H3 SubBlock work (#34148)](https://github.com/sgl-project/sglang/pull/34148), +with native GQA routing and protected-prefix extensions. The shared adapter +interface lets other providers consume the same selection; separate drafts cover +[FlashInfer (#8336)](https://github.com/vllm-project/vllm-omni/pull/8336), +[TRTLLM (#8337)](https://github.com/vllm-project/vllm-omni/pull/8337), and +[cuDNN (#8338)](https://github.com/vllm-project/vllm-omni/pull/8338). + +## Configuration + +Pass the JSON contents of a recipe to `--diffusion-attention-config`: + +| Recipe | Sparse role | Dense roles | +| --- | --- | --- | +| `recipes/attention/fa4-subblock.json` | `cosmos3.gen` | Understanding and multi-control | +| `recipes/attention/minimax-h3-fa4-subblock.json` | `minimax_h3.dit` | Token refiner | + +Both recipes use dense `FLASH_ATTN` by default, 64×64 routing blocks, and +`target_sparsity: 0.75`. Sparse execution pins the `FLASH_ATTN` adapter with +`implementation: auto`, which uses FA4's sparse entry point. FA4 exposes no +kernel-ID selector. Provider preference lists are not supported. + +The selected role uses sparse attention on every invocation. There is no layer +or denoising-step schedule. Exact roles take precedence over legacy `self` +overrides; Cosmos3 multi-control requires its own explicit sparse override. + +### What 75% sparsity means + +The target applies to **unprotected candidate blocks**. The retained candidate +budget rounds upward to eight and is capped by the candidate count. Protected +blocks are added to that budget, so overall sparsity can be lower than 75%. + +Cosmos3 supplies its actual understanding KV-prefix length. Every block touching +that prefix is retained in full. For example, 301 blocks with nine protected +blocks retain 80 candidate blocks plus the nine protected blocks: 70.4% overall +block sparsity. + +MiniMax currently supplies a zero protected prefix, so conditioning tokens can +also be sparsified. Protecting them requires the packer to expose their actual +boundary. + +## Execution contract + +Q/K/V use BSHD layout. `BlockSelection` contains contiguous device-local int32 +`indices` shaped `[B, Hq, ceil(Sq / Bq), capacity]` and `counts` shaped +`[B, Hq, ceil(Sq / Bq)]`. Each row's active prefix contains sorted, unique, +in-range KV-block IDs; trailing storage is ignored. + +The adapter may translate that representation, but must preserve the selected +blocks and geometry. It does not select additional blocks, expand K/V, or fall +back to another backend. FA4 supports native MHA, MQA, and GQA through this path. +The selector and installed kernel must both accept the requested geometry; +provider support is checked against the installed implementation rather than a +copied hardware/dtype/head-size allowlist. + +`BlockSparseAttention` owns routing, validation, and request preparation. +The first actual execution of each request signature runs and synchronizes the +kernel before recording support. Capability queries are read-only and report +preparation as required until that succeeds. Provider errors propagate. + +An opaque custom op keeps request preparation and provider execution outside +Dynamo tracing, with a fake implementation for output shape inference. Eager +execution and ordinary `torch.compile`, including `fullgraph=True`, are covered. +This demonstrates the execution contract; it does not implement the entirety of +[the broader capability proposal (#7226)](https://github.com/vllm-project/vllm-omni/issues/7226). + +Typed `PackedPaddingMetadata` supports a single document with suffix padding +at batch size one. Real Q/K/V lengths are trimmed before routing, and padded +output rows are restored as zeros. MiniMax supplies this metadata automatically. + +## Current validation + +The latest combined run used **vLLM-Omni 0.30.0 on GH200**, PyTorch +2.13.0+cu130, and FA4 revision +`e9cf2c1651d2303191eb40a739a3c135fda00999` +(`4.0.0b33.dev13+ge9cf2c1`) with CuTe DSL 4.7.1. + +**669 tests passed and two were skipped** in 96 seconds at `b076b99fd` +(before this documentation-only update): + +```bash +python -m pytest -q \ + tests/diffusion/attention/test_block_selection.py \ + tests/diffusion/attention/test_block_sparse.py \ + tests/diffusion/attention/test_block_sparse_owner.py \ + tests/diffusion/attention/test_block_sparse_adapters.py \ + tests/diffusion/attention/test_attention_config.py \ + tests/diffusion/attention/test_selector.py \ + tests/diffusion/models/cosmos3/test_cosmos3_transformer.py \ + tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py \ + tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py \ + tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py \ + tests/diffusion/diffusion_backend/test_diffusers_backend.py +``` + +Coverage includes independent routing and selected-key attention references, +head mapping, tails, prefix protection, padding, invalid inputs, owner lifetime, +fullgraph/regional compilation, small Cosmos3 and MiniMax model dispatch, and +dense Diffusers end-to-end execution. The two skips concern MiniMax APIs absent +from this branch. + +For an attention-call benchmark, including routing and metadata conversion: + +```bash +python benchmarks/diffusion/fa4_subblock.py --capture /path/to/capture.pt \ + --sparsity 0.75 --block-size 64 64 --iterations 30 --output /tmp/fa4-subblock.json +``` + +Captures contain `q`, `k`, `v`, and optional `meta.scale` and `meta.prefix_len`. +This measures attention calls; it does not establish checkpoint-level role +activation, end-to-end speedup, or video quality. Historical timings are omitted +because they have not been refreshed for the current implementation. + +## Limitations and follow-up + +Arbitrary masks, causal sparse attention, multi-document packing, +paged or quantized KV, and sequence-parallel wrappers are unsupported. Diffusers +rejects sparse configurations before loading weights; sparse execution needs a +native model integration. + +Checkpoint generation quality and end-to-end latency remain unqualified. +A separate attention-strategies RFC will follow for mixing dense and sparse +attention across roles, layers, and denoising steps to improve quality. diff --git a/recipes/attention/fa4-subblock.json b/recipes/attention/fa4-subblock.json new file mode 100644 index 00000000000..96f136410d8 --- /dev/null +++ b/recipes/attention/fa4-subblock.json @@ -0,0 +1,16 @@ +{ + "default": {"backend": "FLASH_ATTN"}, + "per_role": { + "cosmos3.gen": { + "name": "block_sparse", + "config": { + "block_size": [64, 64], + "selection": { + "name": "block_topk", + "config": {"target_sparsity": 0.75} + }, + "backend": {"require": "FLASH_ATTN", "implementation": "auto"} + } + } + } +} diff --git a/recipes/attention/minimax-h3-fa4-subblock.json b/recipes/attention/minimax-h3-fa4-subblock.json new file mode 100644 index 00000000000..9dd1fd13451 --- /dev/null +++ b/recipes/attention/minimax-h3-fa4-subblock.json @@ -0,0 +1,16 @@ +{ + "default": {"backend": "FLASH_ATTN"}, + "per_role": { + "minimax_h3.dit": { + "name": "block_sparse", + "config": { + "block_size": [64, 64], + "selection": { + "name": "block_topk", + "config": {"target_sparsity": 0.75} + }, + "backend": {"require": "FLASH_ATTN", "implementation": "auto"} + } + } + } +} diff --git a/tests/diffusion/attention/test_block_selection.py b/tests/diffusion/attention/test_block_selection.py new file mode 100644 index 00000000000..a779fb173eb --- /dev/null +++ b/tests/diffusion/attention/test_block_selection.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Selector tests require CUDA/Triton, independently of FA4 and Hopper.""" + +import math + +import pytest +import torch + +from tests.helpers.block_sparse import make_attention_inputs + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model] + + +@pytest.fixture +def cuda_selector(): + if not torch.cuda.is_available(): + pytest.skip("Requires CUDA for Triton routing") + torch.compiler.reset() + yield + torch.compiler.reset() + + +@pytest.mark.parametrize("batch,kv_heads", [(1, 4), (2, 2)]) +@pytest.mark.cuda +@torch.inference_mode() +def test_gqa_routing_and_prefix(cuda_selector, batch, kv_heads): + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + torch.manual_seed(17) + q = torch.randn(batch, 129, 4, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(batch, 4226, kv_heads, 128, device="cuda", dtype=q.dtype) + selector = SubBlockTopK({"target_sparsity": 0.75}) + selector.prepare((64, 64), 128, q.device) + expected = selector.select(q, k.repeat_interleave(4 // kv_heads, dim=2), 0.1, 0).indices + torch.testing.assert_close(selector.select(q, k, 0.1, 0).indices, expected, atol=0, rtol=0) + protected = selector.select(q, k, 0.1, 129).indices + assert protected.shape == (batch, 4, 3, 19) + assert torch.all(protected[..., 1:] > protected[..., :-1]) + torch.testing.assert_close( + protected[..., :3], torch.arange(3, device="cuda", dtype=torch.int32).expand_as(protected[..., :3]) + ) + all_protected = selector.select(q, k, 0.1, k.shape[1]) + expected_all = torch.arange(67, dtype=torch.int32, device=q.device).expand(batch, 4, 3, 67) + torch.testing.assert_close(all_protected.indices, expected_all, atol=0, rtol=0) + with pytest.raises(ValueError, match="nonempty keys"): + selector.select(q, k[:, :0], 0.1, 0) + + +@pytest.mark.cuda +@pytest.mark.parametrize("head_size", [8, 64, 96, 128, 192, 256]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@torch.inference_mode() +def test_pooling_and_scores_against_reference(cuda_selector, head_size, dtype): + from vllm_omni.diffusion.attention.block_selection.subblock_topk import _fused_pool, _fused_scores + + torch.manual_seed(23) + q = torch.randn(1, 65, 2, head_size, device="cuda", dtype=dtype) + k = torch.randn(1, 129, 2, head_size, device="cuda", dtype=q.dtype) + + def pool(x, cells, scale): + out = torch.zeros(x.shape[2], cells, head_size, device="cuda", dtype=x.dtype) + for i in range((x.shape[1] + 15) // 16): + out[:, i] = (x[0, i * 16 : (i + 1) * 16].float().mean(0) * scale).to(x.dtype) + return out + + qp, kp = ( + torch.empty(2, 8, head_size, device="cuda", dtype=q.dtype), + torch.empty(2, 12, head_size, device="cuda", dtype=q.dtype), + ) + scale = 0.1 * math.log2(math.e) + _fused_pool(q, 8, 16, qp, scale) + _fused_pool(k, 12, 16, kp) + torch.testing.assert_close(qp, pool(q, 8, scale), atol=0, rtol=0) + torch.testing.assert_close(kp, pool(k, 12, 1), atol=0, rtol=0) + scores = torch.empty(2, 2, 3, device="cuda", dtype=torch.float32) + _fused_scores(qp, kp, scores, n_k=4, n_valid=9, n_q=4, m_valid=5) + dots = qp.float() @ kp.float().transpose(-1, -2) + dots[:, 5:, :] = -float("inf") + dots[:, :, 9:] = -float("inf") + dots = dots.reshape(2, 2, 4, 3, 4).permute(0, 1, 3, 2, 4).reshape(2, 2, 3, 16) + reference = torch.logsumexp(dots * math.log(2), dim=-1) + torch.testing.assert_close(scores, reference, atol=1e-5, rtol=1e-5) + + +@pytest.mark.cuda +@pytest.mark.parametrize( + "sparsity,kv_len,prefix,kept", + [(0.0, 1089, 0, 18), (0.75, 257, 0, 5), (0.75, 1089, 1089, 18), (0.75, 1089, 1025, 18)], +) +@torch.inference_mode() +def test_known_selection_skips_scoring(cuda_selector, monkeypatch, sparsity, kv_len, prefix, kept): + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + q, k, _ = make_attention_inputs(kv_len=kv_len) + selector = SubBlockTopK({"target_sparsity": sparsity}) + selector.prepare((64, 64), 128, q.device) + + def unexpected_scoring(*args, **kwargs): + pytest.fail("A fully determined selection must not compute scores") + + monkeypatch.setattr(SubBlockTopK, "_compute_scores", staticmethod(unexpected_scoring)) + selection = selector.select(q, k, 128**-0.5, prefix) + expected = torch.arange(kept, dtype=torch.int32, device=q.device).expand(2, 4, 3, kept) + torch.testing.assert_close(selection.indices, expected, atol=0, rtol=0) + torch.testing.assert_close(selection.counts, torch.full_like(selection.counts, kept), atol=0, rtol=0) + assert selection.indices.is_contiguous() and selection.counts.is_contiguous() + if sparsity == 0.0: + compiled = torch.compile(selector.select, fullgraph=True) + actual = compiled(q, k, 128**-0.5, prefix) + torch.testing.assert_close(actual.indices, selection.indices, atol=0, rtol=0) + torch.testing.assert_close(actual.counts, selection.counts, atol=0, rtol=0) + + +@pytest.mark.cuda +@pytest.mark.parametrize("prefix", [0, 64, 65, 2048]) +@torch.inference_mode() +def test_prefix_does_not_consume_candidate_budget(cuda_selector, prefix): + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + prefix_blocks = (prefix + 63) // 64 + # Keep 64 unprotected blocks so their retained budget stays at 16. + q, k, _ = make_attention_inputs(kv_len=prefix_blocks * 64 + 4096) + selector = SubBlockTopK({"target_sparsity": 0.75}) + selector.prepare((64, 64), 128, q.device) + selection = selector.select(q, k, 128**-0.5, prefix) + assert selection.indices.shape == (2, 4, 3, prefix_blocks + 16) + assert torch.all(selection.counts == prefix_blocks + 16) + assert torch.all(selection.indices[..., 1:] > selection.indices[..., :-1]) + assert torch.all(selection.indices[..., prefix_blocks:] >= prefix_blocks) + expected_prefix = torch.arange(prefix_blocks, dtype=torch.int32, device=q.device) + torch.testing.assert_close( + selection.indices[..., :prefix_blocks], expected_prefix.expand(2, 4, 3, prefix_blocks), atol=0, rtol=0 + ) + if prefix == 65: + compiled = torch.compile(selector.select, fullgraph=True) + actual = compiled(q, k, 128**-0.5, prefix) + torch.testing.assert_close(actual.indices, selection.indices, atol=0, rtol=0) + torch.testing.assert_close(actual.counts, selection.counts, atol=0, rtol=0) + + +@pytest.mark.cpu +@pytest.mark.parametrize("prefix,retained", [(0, 80), (514, 89), (19234, 301)]) +def test_cosmos3_prefix_budget(prefix, retained): + """The documented 301-block workload budgets unprotected candidates only.""" + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + prefix_blocks = (prefix + 63) // 64 + assert prefix_blocks + SubBlockTopK.block_budget(301 - prefix_blocks, 0.75) == retained + + +@pytest.mark.cuda +@pytest.mark.parametrize( + "kv_heads,block_size,prefix,sparsity", + [ + (1, (64, 64), 65, 0.75), + (2, (128, 64), 0, 0.875), + (4, (64, 128), 129, 0.5), + ], +) +@torch.inference_mode() +def test_pinned_routing_reference(cuda_selector, kv_heads, block_size, prefix, sparsity): + from tests.helpers.block_sparse import pinned_subblock_reference + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + q, k, _ = make_attention_inputs(q_len=129, kv_len=4226, kv_heads=kv_heads) + scale = q.shape[-1] ** -0.5 + scores, expected = pinned_subblock_reference(q, k, scale, sparsity, prefix, block_size) + selector = SubBlockTopK({"target_sparsity": sparsity}) + selector.prepare(block_size, q.shape[-1], q.device) + torch.testing.assert_close(selector._compute_scores(q, k, scale, block_size), scores, atol=1e-5, rtol=1e-5) + actual = selector.select(q, k, scale, prefix) + torch.testing.assert_close(actual.indices, expected.indices, atol=0, rtol=0) + torch.testing.assert_close(actual.counts, expected.counts, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "candidates,sparsity,expected", [(0, 0.75, 0), (3, 0.75, 3), (64, 0.75, 16), (292, 0.75, 80), (301, 0.75, 80)] +) +def test_candidate_budget_without_gpu(candidates, sparsity, expected): + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + assert SubBlockTopK.block_budget(candidates, sparsity) == expected + + +@pytest.mark.cuda +@pytest.mark.parametrize("sparsity", [0.75, 0.95, 0.98]) +@torch.inference_mode() +def test_pinned_routing_ties_and_high_sparsity(cuda_selector, sparsity): + from tests.helpers.block_sparse import pinned_subblock_reference + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + q, k, _ = make_attention_inputs(q_len=17, kv_len=16385, heads=2, kv_heads=1) + # Head zero has tied scores for all full blocks. Head one retains random scores. + q[:, :, 0] = 0 + scale = q.shape[-1] ** -0.5 + _, expected = pinned_subblock_reference(q, k, scale, sparsity, 65, (64, 64)) + selector = SubBlockTopK({"target_sparsity": sparsity}) + selector.prepare((64, 64), q.shape[-1], q.device) + actual = selector.select(q, k, scale, 65) + torch.testing.assert_close(actual.indices, expected.indices, atol=0, rtol=0) + torch.testing.assert_close(actual.counts, expected.counts, atol=0, rtol=0) + + +@pytest.mark.cuda +@pytest.mark.parametrize("kind", ["nan", "mixed", "positive_inf", "negative_inf", "extreme_finite"]) +@torch.inference_mode() +def test_nonfinite_score_rows_have_complete_unique_selection(cuda_selector, kind): + from vllm_omni.diffusion.attention.block_selection.subblock_topk import _fused_topk + + scores = torch.arange(32, device="cuda", dtype=torch.float32).repeat(2, 1) + if kind == "nan": + scores[0].fill_(float("nan")) + elif kind == "mixed": + scores[0, 17] = float("nan") + elif kind == "extreme_finite": + scores[0, ::2] = -torch.finfo(torch.float32).max + scores[0, 1::2] = torch.finfo(torch.float32).max + else: + scores[0, 17] = float("inf") if kind == "positive_inf" else -float("inf") + actual = _fused_topk(scores, 8) + assert torch.all((actual >= 0) & (actual < 32)) + assert torch.all(actual[:, 1:] > actual[:, :-1]) + if kind != "extreme_finite": + torch.testing.assert_close(actual[0], torch.arange(8, device="cuda", dtype=torch.int32)) + torch.testing.assert_close(actual[1:], _fused_topk(scores[1:].contiguous(), 8)) + + +@pytest.mark.cuda +@torch.inference_mode() +def test_finite_input_pooling_overflow_has_valid_counts_and_prefix(cuda_selector): + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + q = torch.zeros(1, 64, 1, 64, dtype=torch.float16, device="cuda") + q[:, :16] = 46400 + k = torch.ones(1, 2048, 1, 64, dtype=q.dtype, device=q.device) + assert torch.isfinite(q).all() and torch.isfinite(k).all() + selector = SubBlockTopK({"target_sparsity": 0.75}) + selector.prepare((64, 64), 64, q.device) + scores = selector._compute_scores(q, k, 1.0, (64, 64)) + assert not torch.isfinite(scores).all() + for prefix in (0, 65): + selection = selector.select(q, k, 1.0, prefix) + capacity = selection.indices.shape[-1] + assert torch.all(selection.counts == capacity) + expected = torch.arange(capacity, device=q.device, dtype=torch.int32).expand_as(selection.indices) + torch.testing.assert_close(selection.indices, expected) diff --git a/tests/diffusion/attention/test_block_sparse.py b/tests/diffusion/attention/test_block_sparse.py new file mode 100644 index 00000000000..62e2fd6d382 --- /dev/null +++ b/tests/diffusion/attention/test_block_sparse.py @@ -0,0 +1,879 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import math + +import pytest +import torch + +from tests.helpers.block_sparse import make_attention_inputs, selected_attention_reference +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, PackedPaddingMetadata +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend +from vllm_omni.diffusion.attention.capabilities import ( + CompilationMode, + ExecutionContext, + PackingMode, + ParallelStrategy, + SupportStatus, +) +from vllm_omni.diffusion.data import ( + AttentionConfig, + BlockSparseAttentionSpec, + build_attention_config, + parse_attention_config, +) + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model] + + +def test_inline_configuration(monkeypatch): + monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA") + config = build_attention_config( + { + "per_role": { + "self": { + "name": "block_sparse", + "config": { + "selection": { + "name": "block_topk", + "config": {"target_sparsity": 0.5}, + }, + "backend": {"require": "FLASH_ATTN", "implementation": "auto"}, + }, + }, + "cosmos3.gen": {"backend": "TORCH_SDPA"}, + } + } + ) + assert config.resolve_with_source("cosmos3.gen", "self")[0].backend == "TORCH_SDPA" + spec = config.resolve_with_source("another", "self")[0] + assert isinstance(spec, BlockSparseAttentionSpec) + assert spec.selection == { + "name": "block_topk", + "config": {"target_sparsity": 0.5}, + } + assert spec.block_size == (64, 64) + assert spec.implementation == "auto" + assert config.default.backend == "TORCH_SDPA" + with pytest.raises(ValueError, match="mutually exclusive"): + parse_attention_config({"default": {"name": "block_sparse"}}, attention_backend="FLASH_ATTN") + with pytest.raises(ValueError, match="Cannot mix"): + AttentionConfig(default={"name": "block_sparse", "backend": "FLASH_ATTN"}) + + +@pytest.mark.parametrize( + "config", + [ + {"block_size": [0, 64]}, + {"block_size": [64]}, + {"block_size": [True, 64]}, + {"unknown_option": 1}, + {"selection": {"name": "unknown"}}, + {"backend": {"require": "UNKNOWN"}}, + {"backend": {}}, + {"backend": {"require": ""}}, + {"backend": {"require": ["FLASH_ATTN"]}}, + {"backend": {"require": True}}, + {"backend": {"prefer": []}}, + {"backend": {"prefer": ["FLASH_ATTN"], "implementation": "unknown"}}, + {"backend": {"require": "FLASH_ATTN", "prefer": ["FLASH_ATTN"]}}, + {"selection": {"name": "block_topk", "config": {"target_sparsity": float("nan")}}}, + {"selection": {"name": "block_topk", "config": {"target_sparsity": True}}}, + {"selection": {"name": "block_topk", "config": {"target_sparsity": 1}}}, + ], +) +def test_unsupported_configuration(config): + with pytest.raises((TypeError, ValueError)): + BlockSparseAttentionSpec(name="block_sparse", config=config) + + +@pytest.fixture +def hopper(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0): + pytest.skip("Requires Hopper SM90") + pytest.importorskip("flash_attn.cute") + + +def make_impl(heads=4, kv_heads=2, sparsity=0.75, block_size=(64, 64), head_size=128): + spec = BlockSparseAttentionSpec( + name="block_sparse", + config={ + "block_size": block_size, + "selection": {"name": "block_topk", "config": {"target_sparsity": sparsity}}, + }, + ) + from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention + + return BlockSparseAttention( + heads, + kv_heads, + head_size, + head_size**-0.5, + False, + None, + spec, + adapter=FlashAttentionBackend.get_block_sparse_adapter()(), + ) + + +@pytest.mark.cuda +@torch.inference_mode() +def test_selected_key_numerics_and_compile(hopper): + kv_heads = 2 + impl = make_impl(kv_heads=kv_heads) + q, k, v = make_attention_inputs(kv_heads=kv_heads) + metadata = AttentionMetadata(extra={"protected_kv_prefix": 65}) + selection = impl.selector.select(q, k, impl.scale, 65) + indices = selection.indices + assert torch.all(indices[..., :2] == torch.tensor([0, 1], device=q.device)) + assert torch.all(indices[..., 1:] > indices[..., :-1]) + expected = selected_attention_reference(q, k, v, selection, impl.scale, impl.block_size) + impl.prepare_request(q, k, v, metadata) + eager = impl.forward_cuda(q, k, v, metadata) + torch.testing.assert_close(eager.float(), expected, atol=0.004, rtol=0.02) + compiled = torch.compile(impl.forward_cuda, fullgraph=True) + actual = compiled(q, k, v, metadata) + torch.testing.assert_close(actual, eager, atol=0, rtol=0) + # Same prepared layer, different values: intermediates must not retain request data. + torch.testing.assert_close(compiled(q, k, -v, metadata), -eager, atol=0, rtol=0) + result = impl.resolve_execution_path(ExecutionContext(platform="cuda", require_fullgraph=True), q, k, v, metadata) + assert result.path == "fa4_block_sparse" and result.support.status is SupportStatus.SUPPORTED + assert result.compilation_mode is CompilationMode.CUSTOM_OP + + +@pytest.mark.cuda +@torch.inference_mode() +def test_capability_rejections(hopper): + impl = make_impl() + q, k, v = make_attention_inputs() + for metadata in [ + AttentionMetadata(attn_mask=torch.ones(1, device=q.device, dtype=torch.bool)), + AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 129], device=q.device)}), + AttentionMetadata(extra={"protected_kv_prefix": 1090}), + AttentionMetadata(extra={"kv_cache_dtype": "fp8"}), + ]: + with pytest.raises(ValueError): + impl.forward_cuda(q, k, v, metadata) + result = impl.resolve_execution_path(ExecutionContext(platform="cuda"), q, k, v, metadata) + assert result.support.status is SupportStatus.UNSUPPORTED and result.support.reason + with pytest.raises(ValueError, match="matching Q/K/V dtypes"): + impl.forward_cuda(q.float(), k, v) + result = impl.resolve_execution_path( + ExecutionContext(platform="cuda", parallel_strategy=ParallelStrategy.RING), + q, + k, + v, + None, + ) + assert result.support.status is SupportStatus.UNSUPPORTED + with pytest.raises(ValueError, match="native head mapping"): + impl.forward_cuda(q, k.repeat_interleave(2, 2), v.repeat_interleave(2, 2)) + + +@pytest.mark.cuda +@torch.inference_mode() +def test_full_attention_layer_prepares_inline_spec(hopper, monkeypatch): + from types import SimpleNamespace + + from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + from vllm_omni.diffusion.config import set_current_diffusion_config + + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + cfg = SimpleNamespace( + diffusion_attention_config=AttentionConfig(per_role={"cosmos3.gen": {"name": "block_sparse"}}), + parallel_config=SimpleNamespace(ring_degree=1), + dtype=torch.bfloat16, + diffusion_kv_cache_dtype="float", + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + with set_current_diffusion_config(cfg): + layer = layer_mod.Attention(4, 128, False, 0.1, num_kv_heads=2, role="cosmos3.gen", role_category="cross") + from vllm_omni.diffusion.attention.block_sparse import BlockSparseBackend + + assert layer.backend_pref == "FLASH_ATTN" + assert layer.attn_backend is BlockSparseBackend + assert not layer.attn_backend.supports_piecewise_spans + assert not layer.attn_backend.supports_paged_kv + assert layer.attn_backend.supports_packed_mask_free() + assert not layer.attn_backend.supports_multi_doc_packed_varlen() + assert layer.attention.adapter.provider == "FLASH_ATTN" + assert layer.attention.selector.sparsity == 0.75 + q = torch.randn(1, 65, 4, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 130, 2, 128, device="cuda", dtype=q.dtype) + context = ExecutionContext(platform="cuda", require_fullgraph=True, kv_cache_dtype="float") + pending = layer.resolve_execution_path(context, q, k, k, None) + assert pending.support.status is SupportStatus.UNSUPPORTED + assert "not prepared" in pending.support.reason + actual = layer(q, k, k, None) + assert layer.resolve_execution_path(context, q, k, k, None).support.status is SupportStatus.SUPPORTED + assert actual.shape == q.shape + torch.testing.assert_close(torch.compile(layer, fullgraph=True)(q, k, k, None), actual, atol=0, rtol=0) + + +def test_checkpoint_calibration_preserves_block_sparse_spec(): + from types import SimpleNamespace + + from vllm_omni.diffusion.data import OmniDiffusionConfig + + cfg = OmniDiffusionConfig(diffusion_attention_config={"default": {"name": "block_sparse"}}) + cfg._propagate_skip_softmax_calibration(SimpleNamespace(get=lambda _: None)) + assert isinstance(cfg.diffusion_attention_config.default, BlockSparseAttentionSpec) + + +def test_provider_neutral_config_and_missing_adapter(): + from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role + + cfg = AttentionConfig( + default={ + "name": "block_sparse", + "config": { + "backend": {"require": "FLASHINFER_ATTN", "implementation": "future-kernel-id"}, + }, + } + ) + assert cfg.default.implementation == "future-kernel-id" + with pytest.raises(ValueError, match="FLASHINFER_ATTN: no adapter"): + get_attn_backend_for_role("self", 128, cfg) + # IDs are validated by an adapter, not by the common schema. + cfg = AttentionConfig( + default={ + "name": "block_sparse", + "config": { + "backend": {"require": "FLASH_ATTN", "implementation": "future-kernel-id"}, + }, + } + ) + with pytest.raises(ValueError, match="FA4 has no kernel-ID"): + get_attn_backend_for_role("self", 128, cfg) + + +def test_new_provider_adapter_receives_opaque_id(monkeypatch): + from vllm_omni.diffusion.attention import selector + from vllm_omni.diffusion.attention.backends.flashinfer_attn import FlashInferAttentionBackend + + accepted = [] + + class FutureAdapter: + @staticmethod + def validate_selection(implementation, head_size): + accepted.append((implementation, head_size)) + if implementation == "unsupported-id": + raise ValueError("unsupported kernel ID") + + monkeypatch.setattr(FlashInferAttentionBackend, "get_block_sparse_adapter", classmethod(lambda cls: FutureAdapter)) + + def reject_dense_resolution(*args): + pytest.fail("Sparse selection must not use dense backend support checks") + + monkeypatch.setattr(selector, "_cached_get_backend_cls", reject_dense_resolution) + cfg = AttentionConfig( + default={ + "name": "block_sparse", + "config": { + "backend": {"require": "FLASHINFER_ATTN", "implementation": "future-kernel-id"}, + }, + } + ) + backend, resolved = selector.get_attn_backend_for_role("self", 128, cfg) + assert backend is FlashInferAttentionBackend + assert accepted == [("future-kernel-id", 128)] + assert resolved is not None + assert resolved.implementation == "future-kernel-id" + cfg = AttentionConfig( + default={ + "name": "block_sparse", + "config": { + "backend": {"require": "FLASHINFER_ATTN", "implementation": "unsupported-id"}, + }, + } + ) + with pytest.raises(ValueError, match="unsupported kernel ID"): + selector.get_attn_backend_for_role("self", 128, cfg) + + +@pytest.mark.parametrize("providers", [[], ["FLASH_ATTN"], ["FLASH_ATTN", "FLASHINFER_ATTN"]]) +@pytest.mark.parametrize("location", ["default", "per_role"]) +def test_provider_preference_is_rejected_before_loading(monkeypatch, providers, location): + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + def unexpected_load(*args, **kwargs): + pytest.fail("Rejected preference lists must not load a provider") + + monkeypatch.setattr(DiffusionAttentionBackendEnum, "get_class", unexpected_load) + entry = {"name": "block_sparse", "config": {"backend": {"prefer": providers}}} + config = {"default": entry} if location == "default" else {"per_role": {"cosmos3.gen": entry}} + with pytest.raises(ValueError, match=r"backend.prefer is not supported; use backend.require"): + AttentionConfig(**config) + + +@pytest.mark.parametrize("error_type", [ImportError, ValueError, RuntimeError]) +def test_pinned_provider_errors_propagate(monkeypatch, error_type): + from vllm_omni.diffusion.attention.backends.flash_attn import FA4SparseAdapter + from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role + + error = error_type("provider failure") + + def fail(*args): + raise error + + monkeypatch.setattr(FA4SparseAdapter, "validate_selection", fail) + config = AttentionConfig(default={"name": "block_sparse", "config": {"backend": {"require": "FLASH_ATTN"}}}) + with pytest.raises(error_type) as caught: + get_attn_backend_for_role("self", 128, config) + assert caught.value is error + + +@pytest.mark.cuda +@pytest.mark.parametrize("block_size,kv_heads", [((64, 128), 2), ((128, 64), 4), ((128, 128), 2), ((256, 128), 4)]) +@torch.inference_mode() +def test_direct_larger_geometry(hopper, monkeypatch, block_size, kv_heads): + from vllm_omni.diffusion.attention.backends import flash_attn + from vllm_omni.diffusion.attention.block_selection.subblock_topk import SubBlockTopK + + impl = make_impl(kv_heads=kv_heads, block_size=block_size) + q, k, v = make_attention_inputs(kv_heads=kv_heads, q_len=257, kv_len=2177) + metadata = AttentionMetadata(extra={"protected_kv_prefix": 129}) + selection = impl.selector.select(q, k, impl.scale, 129) + indices = selection.indices + block_q, block_kv = block_size + prefix_blocks = math.ceil(129 / block_kv) + assert indices.shape == ( + 2, + 4, + math.ceil(257 / block_q), + prefix_blocks + SubBlockTopK.block_budget(math.ceil(2177 / block_kv) - prefix_blocks, 0.75), + ) + torch.testing.assert_close( + indices[..., :prefix_blocks], + torch.arange(prefix_blocks, device=q.device, dtype=torch.int32).expand_as(indices[..., :prefix_blocks]), + ) + kernel, sparse_type, dependency_version = flash_attn.FA4SparseAdapter._load_api() + observed = [] + + def checked_kernel(*args, **kwargs): + pattern = kwargs["block_sparse_tensors"] + observed.append(pattern.block_size) + # The kernel receives the configured geometry and unchanged logical pattern. + assert pattern.block_size == block_size + torch.testing.assert_close(pattern.mask_block_idx, indices, atol=0, rtol=0) + return kernel(*args, **kwargs) + + monkeypatch.setattr( + flash_attn.FA4SparseAdapter, + "_load_api", + staticmethod(lambda: (checked_kernel, sparse_type, dependency_version)), + ) + impl.prepare_request(q, k, v, metadata) + eager = impl.forward(q, k, v, metadata) + torch.testing.assert_close( + eager.float(), selected_attention_reference(q, k, v, selection, impl.scale, block_size), atol=0.004, rtol=0.02 + ) + if block_size == (64, 128): + compiled = torch.compile(impl.forward, fullgraph=True, dynamic=False) + torch.testing.assert_close(compiled(q, k, v, metadata), eager, atol=0, rtol=0) + assert observed + + +@pytest.mark.cuda +@torch.inference_mode() +def test_geometry_rejection_is_owned_by_strategy(hopper): + # Positive geometry is accepted by the schema; constraints belong to execution. + spec = BlockSparseAttentionSpec(name="block_sparse", config={"block_size": [96, 128]}) + assert spec.config["block_size"] == [96, 128] + with pytest.raises(ValueError, match="SubBlock scorer"): + make_impl(block_size=(96, 128)) + + +def test_adapter_preparation_does_not_probe_synthetic_shapes(monkeypatch): + from vllm_omni.diffusion.attention.backends.flash_attn import FA4SparseAdapter + + def unexpected_probe(*args, **kwargs): + pytest.fail("Preparation must defer hardware and shape support to actual execution") + + monkeypatch.setattr(FA4SparseAdapter, "_load_api", staticmethod(lambda: (None, None, "test"))) + monkeypatch.setattr(FA4SparseAdapter, "_run", unexpected_probe) + monkeypatch.setattr(torch.cuda, "get_device_capability", unexpected_probe) + adapter = FA4SparseAdapter() + adapter.prepare("auto", 192, 12, 3, torch.device("cuda:0"), (128, 128)) + assert adapter.dependency_version == "test" + + +@pytest.fixture +def isolated_compile_cache(): + # Independent model geometries should not exhaust one shared Dynamo frame cache. + torch.compiler.reset() + yield + torch.compiler.reset() + + +@pytest.mark.cuda +@pytest.mark.parametrize( + "head_size,value_size,dtype,kv_heads", + [ + (64, 32, torch.float16, 1), + (96, 64, torch.bfloat16, 2), + (192, 128, torch.float16, 4), + ], +) +@torch.inference_mode() +def test_actual_request_dimensions_and_dtype(hopper, isolated_compile_cache, head_size, value_size, dtype, kv_heads): + impl = make_impl(head_size=head_size, kv_heads=kv_heads, block_size=(128, 128)) + q, k, v = make_attention_inputs( + kv_heads=kv_heads, kv_len=2177, head_size=head_size, value_size=value_size, dtype=dtype + ) + # Exercise scoring with a noncontiguous feature dimension, including padded Triton head tiles. + q, k, v = (torch.stack((t, t), dim=-1)[..., 0] for t in (q, k, v)) + metadata = AttentionMetadata(extra={"protected_kv_prefix": 129}) + selection = impl.selector.select(q, k, impl.scale, 129) + expected = selected_attention_reference(q, k, v, selection, impl.scale, impl.block_size) + impl.prepare_request(q, k, v, metadata) + actual = impl.forward(q, k, v, metadata) + assert actual.shape == (*q.shape[:-1], value_size) + torch.testing.assert_close(actual.float(), expected, atol=0.004, rtol=0.02) + if head_size == 96: + compiled = torch.compile(impl.forward, fullgraph=True) + torch.testing.assert_close(compiled(q, k, v, metadata), actual, atol=0, rtol=0) + + +@pytest.mark.cuda +@pytest.mark.parametrize("compiled", [False, True]) +@torch.inference_mode() +def test_provider_rejection_propagates_during_preparation(hopper, isolated_compile_cache, monkeypatch, compiled): + impl = make_impl() + q, k, v = make_attention_inputs() + + def reject(*args, **kwargs): + raise ValueError("provider rejected actual request") + + monkeypatch.setattr(impl.adapter, "execute", reject) + context = ExecutionContext(platform="cuda") + pending = impl.resolve_execution_path(context, q, k, v, None) + assert pending.support.status is SupportStatus.UNSUPPORTED + assert "not prepared" in pending.support.reason + run = ( + torch.compile(impl.forward, backend="eager", fullgraph=True, dynamic=True) if compiled else impl.prepare_request + ) + with pytest.raises(ValueError, match="provider rejected actual request"): + run(q, k, v) + failed = impl.resolve_execution_path(context, q, k, v, None) + assert failed.support.status is SupportStatus.UNSUPPORTED + assert "provider rejected actual request" in failed.support.reason + + +@pytest.mark.cuda +@pytest.mark.parametrize("change", ["length", "strides", "dtype", "value_size", "prefix"]) +@torch.inference_mode() +def test_eager_preparation_is_specific_to_request(hopper, change): + impl = make_impl() + q, k, v = make_attention_inputs() + context = ExecutionContext(platform="cuda") + impl.forward(q, k, v) + assert impl.resolve_execution_path(context, q, k, v, None).support.status is SupportStatus.SUPPORTED + metadata = None + if change == "length": + q = q[:, :-1] + elif change == "strides": + q = torch.stack((q, q), dim=-1)[..., 0] + elif change == "dtype": + q, k, v = (t.to(torch.float16) for t in (q, k, v)) + elif change == "value_size": + v = v[..., :64] + else: + metadata = AttentionMetadata(extra={"protected_kv_prefix": 65}) + pending = impl.resolve_execution_path(context, q, k, v, metadata) + assert pending.support.status is SupportStatus.UNSUPPORTED + assert "not prepared" in pending.support.reason + actual = impl.forward(q, k, v, metadata) + assert actual.shape == (*q.shape[:-1], v.shape[-1]) + assert impl.resolve_execution_path(context, q, k, v, metadata).support.status is SupportStatus.SUPPORTED + + +@pytest.mark.cuda +@pytest.mark.parametrize("failure_stage", ["execute", "synchronize"]) +@torch.inference_mode() +def test_preparation_failure_is_recorded_and_retryable(hopper, monkeypatch, failure_stage): + impl = make_impl() + q, k, v = make_attention_inputs() + context = ExecutionContext(platform="cuda") + + def fail(*args, **kwargs): + raise RuntimeError("provider preparation failure") + + with monkeypatch.context() as patch: + target = impl.adapter if failure_stage == "execute" else torch.accelerator + patch.setattr(target, failure_stage, fail) + # Resolution must never trial-execute a kernel. + assert "not prepared" in impl.resolve_execution_path(context, q, k, v, None).support.reason + with pytest.raises(RuntimeError, match="provider preparation failure"): + impl.forward(q, k, v) + failed = impl.resolve_execution_path(context, q, k, v, None) + assert failed.support.status is SupportStatus.UNSUPPORTED + assert "provider preparation failure" in failed.support.reason + impl.prepare_request(q, k, v) + assert impl.resolve_execution_path(context, q, k, v, None).support.status is SupportStatus.SUPPORTED + + +@pytest.mark.cuda +@torch.inference_mode() +def test_dynamic_compilation_does_not_specialize_prepared_geometry(hopper, isolated_compile_cache): + from torch._dynamo.testing import CompileCounterWithBackend + + impl = make_impl() + # Exceed Dynamo's default recompilation limit without resetting its cache. + inputs = [make_attention_inputs(q_len=length) for length in range(65, 77)] + expected = [impl.forward(*tensors) for tensors in inputs] + counter = CompileCounterWithBackend("inductor") + compiled = torch.compile(impl.forward, backend=counter, fullgraph=True, dynamic=True) + for tensors, eager in zip(inputs, expected): + torch.testing.assert_close(compiled(*tensors), eager, atol=0, rtol=0) + assert counter.frame_count == 1 + # A new geometry is prepared by the runtime handler, without graph fallback + # or another graph. Capability queries remain read-only until that execution. + tensors = make_attention_inputs(q_len=77) + context = ExecutionContext(platform="cuda") + assert impl.resolve_execution_path(context, *tensors, None).support.status is SupportStatus.UNSUPPORTED + actual = compiled(*tensors) + assert impl.resolve_execution_path(context, *tensors, None).support.status is SupportStatus.SUPPORTED + torch.testing.assert_close(actual, impl.forward(*tensors), atol=0, rtol=0) + assert counter.frame_count == 1 + + +@pytest.mark.cuda +@pytest.mark.parametrize("fullgraph", [False, True]) +@torch.inference_mode() +def test_compiled_first_request_prepares_without_graph_break(hopper, isolated_compile_cache, fullgraph): + from torch._dynamo.testing import CompileCounterWithBackend + + impl = make_impl() + tensors = make_attention_inputs() + counter = CompileCounterWithBackend("inductor") + compiled = torch.compile(impl.forward, backend=counter, fullgraph=fullgraph, dynamic=True) + context = ExecutionContext(platform="cuda") + assert impl.resolve_execution_path(context, *tensors, None).support.status is SupportStatus.UNSUPPORTED + output = compiled(*tensors) + assert impl.resolve_execution_path(context, *tensors, None).support.status is SupportStatus.SUPPORTED + torch.testing.assert_close(output, impl.forward(*tensors), atol=0, rtol=0) + assert counter.frame_count == 1 + + +def test_sparse_preflight_does_not_inherit_dense_provider_capabilities(): + from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention, BlockSparseBackend + from vllm_omni.diffusion.attention.selector import get_attn_backend_for_capability + + config = AttentionConfig(per_role={"self": {"name": "block_sparse"}}) + backend = get_attn_backend_for_capability("self", config) + assert backend is BlockSparseBackend + assert backend.get_impl_cls() is BlockSparseAttention + assert not backend.supports_attention_mask() + assert backend.supports_packed_mask_free() + assert not backend.supports_multi_doc_packed_varlen() + assert not backend.supports_piecewise_spans + assert not backend.supports_prefix_kv_slicing + assert not backend.supports_paged_kv + assert not backend.accept_output_buffer + result = backend.resolve_capabilities(ExecutionContext(platform="cuda", require_fullgraph=True)) + assert result.support.status is SupportStatus.UNSUPPORTED + assert "preparation" in result.support.reason + # Dense selection still exposes the provider's own declarations. + dense = get_attn_backend_for_capability("self", AttentionConfig(default={"backend": "FLASH_ATTN"})) + assert dense is FlashAttentionBackend + assert dense.supports_piecewise_spans and dense.supports_paged_kv + + +@pytest.mark.cuda +@pytest.mark.parametrize("fullgraph", [False, True]) +@torch.inference_mode() +def test_regional_compilation_reuses_graph_across_sparse_owners(hopper, isolated_compile_cache, fullgraph): + from torch._dynamo.testing import CompileCounterWithBackend + + from vllm_omni.diffusion.compile import regionally_compile + + calls = [] + + class SparseTestBlock(torch.nn.Module): + def __init__(self, index): + super().__init__() + self.attention = make_impl() + execute = self.attention.adapter.execute + + def tracked_execute(*args, **kwargs): + calls.append(index) + # Give each owner a distinct result as well as a dispatch trace. + return execute(*args, **kwargs) * (index + 1) + + self.attention.adapter.execute = tracked_execute + + def forward(self, q, k, v): + return self.attention.forward(q, k, v) + + model = torch.nn.Module() + model._repeated_blocks = ["SparseTestBlock"] + model.blocks = torch.nn.ModuleList(SparseTestBlock(i) for i in range(12)) + inputs = [make_attention_inputs(q_len=length) for length in (65, 97)] + expected = [[block(*tensors) for block in model.blocks] for tensors in inputs] + counter = CompileCounterWithBackend("inductor") + # Exactly the production regional-compilation helper: separate torch.compile + # wrappers for the same forward code on distinct block instances. + regionally_compile(model, backend=counter, fullgraph=fullgraph, dynamic=True) + for tensors, outputs in zip(inputs, expected): + calls.clear() + for block, reference in zip(model.blocks, outputs): + torch.testing.assert_close(block(*tensors), reference, atol=0, rtol=0) + assert calls == list(range(12)) + # No cache resets or raised recompilation limit between owners or lengths. + assert counter.frame_count == 1 + + +def make_padded_request(q_length=65, kv_length=1089, prefix=65): + q, k, v = ( + tensor[:1] for tensor in make_attention_inputs(q_len=q_length + 31, kv_len=kv_length + 47, value_size=64) + ) + q[:, q_length:] = float("nan") + k[:, kv_length:] = float("nan") + v[:, kv_length:] = float("nan") + cu_q = torch.tensor([0, q_length, q.shape[1]], device=q.device, dtype=torch.int32) + cu_k = torch.tensor([0, kv_length, k.shape[1]], device=k.device, dtype=torch.int32) + metadata = AttentionMetadata( + packed_padding=PackedPaddingMetadata(q_length, kv_length, cu_q[:2], cu_k[:2]), + extra={ + "cu_seqlens_q": cu_q, + "cu_seqlens_k": cu_k, + "max_seqlen_q": q_length, + "max_seqlen_k": kv_length, + "valid_kv_length": kv_length, + "protected_kv_prefix": prefix, + }, + ) + return q, k, v, metadata + + +@pytest.mark.cuda +@pytest.mark.parametrize("q_length,kv_length,prefix", [(65, 1089, 65), (97, 2051, 0)]) +@torch.inference_mode() +def test_single_document_padding_preserves_routing_and_zeroes_tail(hopper, q_length, kv_length, prefix): + impl = make_impl() + q, k, v, metadata = make_padded_request(q_length, kv_length, prefix) + real = (q[:, :q_length], k[:, :kv_length], v[:, :kv_length]) + plain_metadata = AttentionMetadata(extra={"protected_kv_prefix": prefix}) + expected = impl.forward(*real, plain_metadata) + context = ExecutionContext(platform="cuda", packing_mode=PackingMode.PACKED_PADDING) + # Preparing the executed views also prepares the equivalent typed request. + assert impl.resolve_execution_path(context, q, k, v, metadata).support.status is SupportStatus.SUPPORTED + actual = impl.forward(q, k, v, metadata) + torch.testing.assert_close(actual[:, :q_length], expected, atol=0, rtol=0) + assert actual.shape == (1, q.shape[1], q.shape[2], v.shape[-1]) and actual.is_contiguous() + assert torch.count_nonzero(actual[:, q_length:]) == 0 + # Values in either padding region must have no influence, including on the + # last partially filled real query block's pooled routing scores. + q[:, q_length:] = 10000 + k[:, kv_length:] = -10000 + v[:, kv_length:] = 10000 + torch.testing.assert_close(impl.forward(q, k, v, metadata), actual, atol=0, rtol=0) + # The typed contract alone is sufficient; redundant extras are optional. + typed_only = AttentionMetadata(packed_padding=metadata.packed_padding, extra=plain_metadata.extra) + torch.testing.assert_close(impl.forward(q, k, v, typed_only), actual, atol=0, rtol=0) + larger = tuple(torch.cat((t, t.new_full((1, 17, *t.shape[2:]), float("nan"))), dim=1) for t in (q, k, v)) + extended = impl.forward(*larger, typed_only) + torch.testing.assert_close(extended[:, :q_length], expected, atol=0, rtol=0) + assert torch.count_nonzero(extended[:, q_length:]) == 0 + + +@pytest.mark.cuda +@torch.inference_mode() +def test_padding_preparation_uses_executed_views(hopper): + impl = make_impl() + q, k, v, metadata = make_padded_request() + context = ExecutionContext(platform="cuda", packing_mode=PackingMode.PACKED_PADDING) + assert impl.resolve_execution_path(context, q, k, v, metadata).support.status is SupportStatus.UNSUPPORTED + impl.prepare_request(q, k, v, metadata) + plain = AttentionMetadata(extra={"protected_kv_prefix": 65}) + assert ( + impl.resolve_execution_path( + ExecutionContext(platform="cuda"), q[:, :65], k[:, :1089], v[:, :1089], plain + ).support.status + is SupportStatus.SUPPORTED + ) + # The padded storage was never prepared as a real attention request. + assert ( + impl.resolve_execution_path(ExecutionContext(platform="cuda"), q, k, v, plain).support.status + is SupportStatus.UNSUPPORTED + ) + assert impl.resolve_execution_path(context, q, k, v, metadata).support.status is SupportStatus.SUPPORTED + assert impl.resolve_execution_path(context, q, k, v, plain).support.status is SupportStatus.UNSUPPORTED + assert ( + impl.resolve_execution_path( + ExecutionContext(platform="cuda", packing_mode=PackingMode.MULTI_DOCUMENT), q, k, v, metadata + ).support.status + is SupportStatus.UNSUPPORTED + ) + assert ( + impl.resolve_execution_path( + ExecutionContext( + platform="cuda", packing_mode=PackingMode.PACKED_PADDING, parallel_strategy=ParallelStrategy.RING + ), + q, + k, + v, + metadata, + ).support.status + is SupportStatus.UNSUPPORTED + ) + + +@pytest.mark.cuda +@pytest.mark.parametrize("field", ["q_length", "kv_length"]) +@pytest.mark.parametrize("bad_length", [True, 1.5, 0, -1, None, "cpu_scalar", "cuda_scalar", "overflow"]) +@torch.inference_mode() +def test_padding_rejects_invalid_host_lengths(hopper, field, bad_length, monkeypatch): + from dataclasses import replace + + impl = make_impl() + q, k, v, metadata = make_padded_request() + if bad_length == "cpu_scalar": + bad_length = torch.tensor(1) + elif bad_length == "cuda_scalar": + bad_length = torch.tensor(1, device=q.device) + elif bad_length == "overflow": + bad_length = max(q.shape[1], k.shape[1]) + 1 + metadata.packed_padding = replace(metadata.packed_padding, **{field: bad_length}) + monkeypatch.setattr(impl, "_execute", lambda *args: pytest.fail("Invalid padding must not execute")) + for run in (impl.forward, impl.prepare_request): + with pytest.raises(ValueError, match=field): + run(q, k, v, metadata) + result = impl.resolve_execution_path( + ExecutionContext(platform="cuda", packing_mode=PackingMode.PACKED_PADDING), q, k, v, metadata + ) + assert result.support.status is SupportStatus.UNSUPPORTED and field in result.support.reason + assert not impl._request_preparation + + +@pytest.mark.cuda +@pytest.mark.parametrize( + "case", + [ + "untyped", + "wrong_type", + "batch", + "prefix", + "cu_shape", + "cu_dtype", + "raw_multi_document", + "extra_length", + "extra_scalar", + "mask", + "piecewise", + ], +) +@torch.inference_mode() +def test_padding_rejects_inconsistent_or_unsupported_metadata(hopper, case, monkeypatch): + from dataclasses import replace + + impl = make_impl() + q, k, v, metadata = make_padded_request() + if case == "untyped": + metadata.packed_padding = None + elif case == "wrong_type": + metadata.packed_padding = {"q_length": 65, "kv_length": 1089} + elif case == "batch": + q, k, v = (t.expand(2, *t.shape[1:]) for t in (q, k, v)) + elif case == "prefix": + metadata.extra["protected_kv_prefix"] = 1090 # Within storage, beyond real K/V. + elif case == "cu_shape": + metadata.packed_padding = replace(metadata.packed_padding, cu_seqlens_q=metadata.extra["cu_seqlens_q"]) + elif case == "cu_dtype": + metadata.packed_padding = replace( + metadata.packed_padding, cu_seqlens_k=metadata.packed_padding.cu_seqlens_k.float() + ) + elif case == "raw_multi_document": + metadata.extra["cu_seqlens_q"] = torch.tensor([0, 20, 65, q.shape[1]], device=q.device, dtype=torch.int32) + elif case == "extra_length": + metadata.extra["max_seqlen_q"] = 64 + elif case == "extra_scalar": + metadata.extra["valid_kv_length"] = torch.tensor(1089, device=q.device) + elif case == "mask": + metadata.attn_mask = torch.ones(1, k.shape[1], device=q.device, dtype=torch.bool) + else: + metadata.full_attn_spans = [[[0, 65]]] + monkeypatch.setattr(impl, "_execute", lambda *args: pytest.fail("Invalid padding must not execute")) + for run in (impl.forward, impl.prepare_request): + with pytest.raises(ValueError): + run(q, k, v, metadata) + result = impl.resolve_execution_path( + ExecutionContext(platform="cuda", packing_mode=PackingMode.PACKED_PADDING), q, k, v, metadata + ) + assert result.support.status is SupportStatus.UNSUPPORTED and result.support.reason + assert not impl._request_preparation + + +@pytest.mark.cuda +@torch.inference_mode() +def test_padding_compile(hopper, isolated_compile_cache): + impl = make_impl() + q, k, v, metadata = make_padded_request() + impl.prepare_request(q, k, v, metadata) + expected = impl.forward(q, k, v, metadata) + compiled = torch.compile(impl.forward, fullgraph=True, dynamic=True) + torch.testing.assert_close(compiled(q, k, v, metadata), expected, atol=0, rtol=0) + # A different real length uses a different executed request despite having + # the same padded storage size. Compiled dispatch must prepare that request. + from dataclasses import replace + + changed = AttentionMetadata( + packed_padding=replace( + metadata.packed_padding, + q_length=64, + cu_seqlens_q=torch.tensor([0, 64], device=q.device, dtype=torch.int32), + ), + extra={"protected_kv_prefix": 65}, + ) + context = ExecutionContext(platform="cuda", packing_mode=PackingMode.PACKED_PADDING) + assert impl.resolve_execution_path(context, q, k, v, changed).support.status is SupportStatus.UNSUPPORTED + actual = compiled(q, k, v, changed) + assert impl.resolve_execution_path(context, q, k, v, changed).support.status is SupportStatus.SUPPORTED + plain = AttentionMetadata(extra={"protected_kv_prefix": 65}) + reference = impl.forward(q[:, :64], k[:, :1089], v[:, :1089], plain) + torch.testing.assert_close(actual[:, :64], reference, atol=0, rtol=0) + assert torch.count_nonzero(actual[:, 64:]) == 0 + + +@pytest.mark.cuda +@torch.inference_mode() +def test_cloned_sparse_layer_requires_own_preparation(hopper): + import copy + import gc + import weakref + + original = make_impl() + q, k, v = make_attention_inputs() + expected = original.forward(q, k, v) + cloned = copy.deepcopy(original) + assert cloned.adapter is not original.adapter and cloned.selector is not original.selector + context = ExecutionContext(platform="cuda") + assert original.resolve_execution_path(context, q, k, v, None).support.status is SupportStatus.SUPPORTED + assert cloned.resolve_execution_path(context, q, k, v, None).support.status is SupportStatus.UNSUPPORTED + original_ref = weakref.ref(original) + del original + gc.collect() + assert original_ref() is None + torch.testing.assert_close(cloned.forward(q, k, v), expected, atol=0, rtol=0) + assert cloned.resolve_execution_path(context, q, k, v, None).support.status is SupportStatus.SUPPORTED + + +@pytest.mark.parametrize("adapter_mode", [CompilationMode.EAGER_ONLY, CompilationMode.TRACEABLE]) +def test_shared_compilation_boundary_is_custom_op(adapter_mode): + from types import SimpleNamespace + + from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention + + impl = object.__new__(BlockSparseAttention) + impl.adapter = SimpleNamespace(compilation_mode=adapter_mode, provider="test", kernel_variant="test") + impl._normalize_request = lambda q, k, v, metadata: (q, k, v, 0, 1) + impl._request_signature = lambda *args: () + impl._request_preparation = {(): None} + result = impl.resolve_execution_path( + ExecutionContext(platform="cuda", require_fullgraph=True), None, None, None, None + ) + assert result.support.status is SupportStatus.SUPPORTED + assert result.compilation_mode is CompilationMode.CUSTOM_OP diff --git a/tests/diffusion/attention/test_block_sparse_adapters.py b/tests/diffusion/attention/test_block_sparse_adapters.py new file mode 100644 index 00000000000..9c027ce1957 --- /dev/null +++ b/tests/diffusion/attention/test_block_sparse_adapters.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Sparse adapter contract tests, independent of the block-selection strategy.""" + +import math + +import pytest +import torch + +from tests.helpers.block_sparse import make_attention_inputs, selected_attention_reference +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.diffusion.attention.block_selection.abstract import BlockSelection +from vllm_omni.diffusion.attention.capabilities import CompilationMode + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.fixture(params=["FLASH_ATTN"]) +def sparse_adapter(request): + # Add a provider here and its dependency/device setup below to run the same contract. + if not torch.cuda.is_available(): + pytest.skip("Requires CUDA") + if request.param == "FLASH_ATTN": + if torch.cuda.get_device_capability() != (9, 0): + pytest.skip("These FA4 contract cases are exercised on Hopper") + pytest.importorskip("flash_attn.cute") + adapter_cls = DiffusionAttentionBackendEnum[request.param].get_class().get_block_sparse_adapter() + assert adapter_cls is not None + torch.compiler.reset() + try: + yield adapter_cls() + finally: + torch.compiler.reset() + + +@pytest.mark.parametrize( + "kv_heads,dtype,pattern,compile_case", + [ + (1, torch.float16, "per_row", False), + (2, torch.bfloat16, "per_row", True), + (4, torch.float16, "per_row", False), + (4, torch.bfloat16, "full", False), + ], +) +@torch.inference_mode() +def test_selected_attention_contract(sparse_adapter, kv_heads, dtype, pattern, compile_case): + block_size = (64, 64) + q, k, v = make_attention_inputs(kv_heads=kv_heads, kv_len=193, value_size=64, dtype=dtype) + q, k, v = (torch.stack((t, t), dim=-1)[..., 0] for t in (q, k, v)) + scale = q.shape[-1] ** -0.5 + rows = (q.shape[0], q.shape[2], math.ceil(q.shape[1] / block_size[0])) + key_blocks = math.ceil(k.shape[1] / block_size[1]) + if pattern == "full": + indices = torch.arange(key_blocks, dtype=torch.int32, device=q.device).expand(*rows, key_blocks).contiguous() + counts = torch.full(rows, key_blocks, dtype=torch.int32, device=q.device) + else: + # Distinct query-head patterns, variable counts and a partially filled final KV block. + first = (torch.arange(math.prod(rows), device=q.device).reshape(rows) % 2).int() + indices = torch.stack((first, torch.full_like(first, key_blocks - 1)), dim=-1) + counts = first + 1 + selection = BlockSelection(indices, counts) + tensors = (q, k, v, indices, counts) + snapshots = tuple(t.clone() for t in tensors) + sparse_adapter.prepare("auto", q.shape[-1], q.shape[2], k.shape[2], q.device, block_size) + expected = selected_attention_reference(q, k, v, selection, scale, block_size) + + # Warm the actual inputs outside compilation, without a synthetic support probe. + result = sparse_adapter.execute(q, k, v, selection, scale, block_size) + assert result.shape == (*q.shape[:-1], v.shape[-1]) + assert result.dtype == q.dtype and result.device == q.device and result.is_contiguous() + torch.testing.assert_close(result.float(), expected, atol=0.004, rtol=0.02) + run = sparse_adapter.execute + if compile_case and sparse_adapter.compilation_mode is not CompilationMode.EAGER_ONLY: + run = torch.compile(run, fullgraph=True) + torch.testing.assert_close(run(q, k, v, selection, scale, block_size), result, atol=0, rtol=0) + + saved_result = result.clone() + torch.testing.assert_close(run(q, k, -v, selection, scale, block_size), -saved_result, atol=0, rtol=0) + # Change active counts without changing storage capacity: adapters must read the new pattern. + updated = BlockSelection(indices, torch.full_like(counts, indices.shape[-1])) + updated_expected = selected_attention_reference(q, k, v, updated, scale, block_size) + torch.testing.assert_close( + run(q, k, v, updated, scale, block_size).float(), updated_expected, atol=0.004, rtol=0.02 + ) + torch.testing.assert_close(result, saved_result, atol=0, rtol=0) + for original, snapshot in zip(tensors, snapshots): + torch.testing.assert_close(original, snapshot, atol=0, rtol=0) diff --git a/tests/diffusion/attention/test_block_sparse_owner.py b/tests/diffusion/attention/test_block_sparse_owner.py new file mode 100644 index 00000000000..5cd6854d47e --- /dev/null +++ b/tests/diffusion/attention/test_block_sparse_owner.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Owner identity/lifetime tests with a CPU executor; no FA4 dependency.""" + +import copy +import gc +import io +import weakref + +import pytest +import torch + +from vllm_omni.diffusion.attention.block_sparse import ( + BlockSparseAttention, + _block_sparse_request, + _request_owners, +) + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu] + + +class CPUOwner(BlockSparseAttention): + """Exercise the production handle/state protocol without initializing CUDA.""" + + def __init__(self, bias): + self.bias = bias + self._request_preparation = {} + self._register_request_owner() + + def _dispatch_request(self, query, key, value, prefix): + return query + self.bias + + +def dispatch(owner): + q = torch.zeros(1, 2, 1, 1) + return _block_sparse_request(q, q, q, 0, owner._request_owner_handle) + + +@pytest.mark.parametrize("clone_method", ["deepcopy", "torch_save"]) +def test_copied_owner_is_independent_and_survives_original_collection(clone_method): + original = CPUOwner(1) + original._request_preparation[("prepared",)] = None + original._request_preparation[("failed",)] = "previous failure" + original_id = original._request_owner_handle.item() + original_ref = weakref.ref(original) + stale_handle = original._request_owner_handle.clone() + if clone_method == "deepcopy": + cloned = copy.deepcopy(original) + else: + buffer = io.BytesIO() + torch.save(original, buffer) + buffer.seek(0) + # Only deserialize the object created in this test, never external data. + cloned = torch.load(buffer, weights_only=False) + clone_id = cloned._request_owner_handle.item() + assert clone_id != original_id and not cloned._request_preparation + assert _request_owners[original_id] is original and _request_owners[clone_id] is cloned + cloned.bias = 7 + cloned._request_preparation[("clone only",)] = None + assert ("clone only",) not in original._request_preparation + torch.testing.assert_close(dispatch(original), torch.ones(1, 2, 1, 1)) + torch.testing.assert_close(dispatch(cloned), torch.full((1, 2, 1, 1), 7.0)) + del original + gc.collect() + assert original_ref() is None and original_id not in _request_owners + torch.testing.assert_close(dispatch(cloned), torch.full((1, 2, 1, 1), 7.0)) + q = torch.zeros(1, 2, 1, 1) + with pytest.raises(RuntimeError, match="owner was released"): + _block_sparse_request(q, q, q, 0, stale_handle) + clone_ref = weakref.ref(cloned) + del cloned + gc.collect() + assert clone_ref() is None and clone_id not in _request_owners + + +def test_shallow_copy_rejected_without_changing_owner(): + owner = CPUOwner(3) + handle = owner._request_owner_handle + with pytest.raises(TypeError, match="copy.deepcopy"): + copy.copy(owner) + assert owner._request_owner_handle is handle + torch.testing.assert_close(dispatch(owner), torch.full((1, 2, 1, 1), 3.0)) + + +def test_deepcopy_preserves_aliases_within_copied_model(): + original = CPUOwner(1) + first, second = copy.deepcopy([original, original]) + assert first is second and first is not original + assert first._request_owner_handle.item() != original._request_owner_handle.item() diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index 082855b4392..d17582dacb5 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/tests/diffusion/diffusion_backend/test_diffusers_backend.py @@ -814,3 +814,19 @@ def test_t2i_random_weights( # Request config has incomplete width/height, so internal assertion in `send_diffusion_request` is incomplete. assert image.size == (512, 512) + + +@pytest.mark.parametrize("role", ["default", "self", "minimax_h3.dit"]) +@pytest.mark.cpu +@pytest.mark.core_model +def test_diffusers_rejects_sparse_method_during_construction(role, mocker): + sparse = {"name": "block_sparse", "config": {"backend": {"require": "FLASH_ATTN"}}} + config = ( + AttentionConfig(default=sparse) + if role == "default" + else AttentionConfig(default={"backend": "FLASH_ATTN"}, per_role={role: sparse}) + ) + load = mocker.patch.object(DiffusersAdapterPipeline, "_load_pipeline_from_pretrained") + with pytest.raises(ValueError, match="does not support block_sparse"): + DiffusersAdapterPipeline(od_config=_make_od_config(diffusion_attention_config=config)) + load.assert_not_called() diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 56c33b9e005..ab833983ee7 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -1050,7 +1050,8 @@ def test_multi_control_attention_weights_target_outputs() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention class ControlValueAttention: - def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, metadata) -> torch.Tensor: + assert metadata.extra["protected_kv_prefix"] == 1 del k # Each per-control K/V layout is [text, control_i, target]. return v[:, 1:2].expand(q.shape[0], q.shape[1], -1, -1) @@ -1394,3 +1395,159 @@ def __call__(self, x, position_ids): ) _, offset_gen_pos = rotary.position_ids assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107] + + +@pytest.mark.cuda +@pytest.mark.parametrize("multi_control_sparse", [False, True]) +@torch.inference_mode() +def test_cosmos3_sparse_recipe_constructs_model_and_dispatches(monkeypatch, multi_control_sparse): + """Exercise model-owned roles and prefix metadata, not a hand-built Attention.""" + import json + from pathlib import Path + + from tests.helpers.block_sparse import selected_attention_reference + from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0): + pytest.skip("Requires Hopper SM90 and FA4") + pytest.importorskip("flash_attn.cute") + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + monkeypatch.setattr(transformer_cosmos3, "_is_sp_active", lambda: False) + recipe = json.loads((Path(__file__).resolve().parents[4] / "recipes/attention/fa4-subblock.json").read_text()) + if multi_control_sparse: + recipe["per_role"]["cosmos3.gen_multi_control"] = recipe["per_role"]["cosmos3.gen"] + cfg = SimpleNamespace( + tf_model_config=_tiny_cosmos3_config( + hidden_size=256, + intermediate_size=512, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=128, + rope_scaling={"mrope_section": [24, 20, 20]}, + ), + diffusion_attention_config=AttentionConfig(**recipe), + parallel_config=SimpleNamespace(ring_degree=1), + dtype=torch.bfloat16, + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + with set_current_diffusion_config(cfg): + model = transformer_cosmos3.Cosmos3VFMTransformer(cfg).to(device="cuda", dtype=cfg.dtype) + # No checkpoint is needed to exercise construction and attention dispatch. + for parameter in model.parameters(): + parameter.fill_(1) if parameter.ndim == 1 else parameter.normal_(std=0.02) + understanding = model.language_model.layers[0].self_attn + generation = model.gen_layers[0].cross_attention + assert understanding.attn.role == "cosmos3.und" + assert understanding.attn.role_category == "self" + assert not isinstance(understanding.attn.attention, BlockSparseAttention) + assert generation.attn.role == "cosmos3.gen" + assert generation.attn.role_category == "self" + assert isinstance(generation.attn.attention, BlockSparseAttention) + assert generation.multi_control_attn.role == "cosmos3.gen_multi_control" + assert isinstance(generation.multi_control_attn.attention, BlockSparseAttention) == multi_control_sparse + + calls: list[str] = [] + + def check_dispatch(layer, args, output): + q, k, v, metadata = args + assert metadata.extra["protected_kv_prefix"] == prefix_length + torch.testing.assert_close(k[:, :prefix_length], k_und) + torch.testing.assert_close(v[:, :prefix_length], v_und) + calls.append(layer.role) + if isinstance(layer.attention, BlockSparseAttention): + impl = layer.attention + selection = impl.selector.select(q, k, impl.scale, prefix_length) + protected_blocks = (prefix_length + impl.block_size[1] - 1) // impl.block_size[1] + expected_prefix = torch.arange(protected_blocks, device=q.device, dtype=torch.int32) + torch.testing.assert_close( + selection.indices[..., :protected_blocks], + expected_prefix.expand(*selection.indices.shape[:-1], protected_blocks), + ) + expected = selected_attention_reference(q, k, v, selection, impl.scale, impl.block_size) + torch.testing.assert_close(output.float(), expected, atol=0.004, rtol=0.02) + + generation.attn.register_forward_hook(check_dispatch) + generation.multi_control_attn.register_forward_hook(check_dispatch) + hidden = torch.randn(1, 2049, 256, device="cuda", dtype=cfg.dtype) + cos = torch.ones(1, hidden.shape[1], 1, 128, device="cuda", dtype=cfg.dtype) + sin = torch.zeros_like(cos) + for prefix_length in (0, 65): + if prefix_length: + text = hidden[:, :prefix_length] + _, k_und, v_und = understanding(text, cos[:, :prefix_length], sin[:, :prefix_length]) + else: + k_und = v_und = hidden.new_empty(1, 0, 1, 128) + calls.clear() + output = generation(hidden, k_und, v_und, cos, sin) + assert output.shape == hidden.shape + assert calls == ["cosmos3.gen"] + calls.clear() + output = generation( + hidden, + k_und, + v_und, + cos, + sin, + control_token_sizes=(64, 64), + control_weights=(0.25, 0.75), + ) + assert output.shape == hidden.shape + assert calls == ["cosmos3.gen_multi_control", "cosmos3.gen_multi_control"] + + +@pytest.mark.parametrize( + "overrides,expected", + [ + ({"self": {"backend": "TORCH_SDPA"}}, ("SDPA", "SDPA", "SDPA")), + ({"cross": {"backend": "TORCH_SDPA"}}, ("FLASH_ATTN", "FLASH_ATTN", "FLASH_ATTN")), + ( + { + "self": {"backend": "TORCH_SDPA"}, + "cross": {"backend": "TORCH_SDPA"}, + "cosmos3.gen": {"backend": "FLASH_ATTN"}, + "cosmos3.gen_multi_control": {"backend": "FLASH_ATTN"}, + }, + ("SDPA", "FLASH_ATTN", "FLASH_ATTN"), + ), + ], +) +def test_cosmos3_exact_roles_preserve_legacy_self_category(monkeypatch, overrides, expected): + from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + cfg = SimpleNamespace( + tf_model_config=_tiny_cosmos3_config( + hidden_size=256, + intermediate_size=512, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=128, + rope_scaling={"mrope_section": [24, 20, 20]}, + ), + diffusion_attention_config=AttentionConfig(default={"backend": "FLASH_ATTN"}, per_role=overrides), + parallel_config=SimpleNamespace(ring_degree=1), + dtype=torch.bfloat16, + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + with set_current_diffusion_config(cfg): + model = Cosmos3VFMTransformer(cfg) + generation = model.gen_layers[0].cross_attention + layers = (model.language_model.layers[0].self_attn.attn, generation.attn, generation.multi_control_attn) + assert tuple(layer.role for layer in layers) == ("cosmos3.und", "cosmos3.gen", "cosmos3.gen_multi_control") + assert all(layer.role_category == "self" for layer in layers) + assert tuple(layer.attn_backend.get_name() for layer in layers) == expected diff --git a/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py index 9a530bf9006..c3a6ac71586 100644 --- a/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py +++ b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py @@ -677,3 +677,35 @@ def test_resolve_sigma_positions_uses_release_metadata_schedule(): assert steps == 8 == len(positions) - 1 with pytest.raises(OmniClientError, match="must be 8"): pipeline._resolve_sigma_positions("t2va", _sampling(num_inference_steps=9)) + + +@pytest.mark.parametrize( + "default,category,exact,topk,error", + [ + ("FASTVIDEO_VSA", "FASTVIDEO_VSA", "TORCH_SDPA", None, "FASTVIDEO_VSA"), + ("TORCH_SDPA", "TORCH_SDPA", "FASTVIDEO_VSA", None, None), + ("TORCH_SDPA", "TORCH_SDPA", "FASTVIDEO_VSA", 2, "fixed fastvideo_vsa_topk"), + ("TORCH_SDPA", "FASTVIDEO_VSA", None, None, None), + ("TORCH_SDPA", "FASTVIDEO_VSA", None, 2, "fixed fastvideo_vsa_topk"), + ], +) +def test_v2_dit_role_precedence(default, category, exact, topk, error): + roles = {"self": AttentionSpec(backend=category, fastvideo_vsa_topk=topk if exact is None else None)} + if exact is not None: + roles["minimax_h3.dit"] = AttentionSpec(backend=exact, fastvideo_vsa_topk=topk) + config = _od_config(backend=default, per_role=roles) + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + if error: + with pytest.raises(ValueError, match=error): + spec.check_serving_contract(partition="fl2va", od_config=config) + else: + spec.check_serving_contract(partition="fl2va", od_config=config) + + +def test_v2_legacy_backend_only_config(): + from types import SimpleNamespace + + config = SimpleNamespace(diffusion_attention_backend="FASTVIDEO_VSA") + FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).check_serving_contract( + partition="fl2va", od_config=config + ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 4405997f73e..23956a58863 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -3576,3 +3576,173 @@ def diffuse(**kwargs): pipeline.decode_to_mp4.assert_not_called() finally: registry.close() + + +@torch.inference_mode() +def test_unpadded_sage_token_refiner_preserves_legacy_dispatch(monkeypatch): + from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention.backends import sage_attn + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as transformer + + config = AttentionConfig( + default={"backend": "TORCH_SDPA"}, + per_role={"minimax_h3.token_refiner": {"backend": "SAGE_ATTN"}}, + ) + arch = transformer.MiniMaxH3DiTArchConfig( + num_layers=1, token_refiner_num_layers=1, hidden_size=64, num_attention_heads=2, attention_head_dim=32 + ) + # No distributed process group or optional Sage kernel is needed to check + # model-owned metadata and dispatch through the real Sage implementation. + monkeypatch.setattr(transformer, "QKVParallelLinear", lambda **kwargs: SimpleNamespace(num_heads=2, num_kv_heads=2)) + monkeypatch.setattr(transformer, "RowParallelLinear", lambda *args, **kwargs: nn.Identity()) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + calls = [] + + def sage_kernel(q, k, v, **kwargs): + calls.append(kwargs) + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), scale=kwargs["sm_scale"] + ).transpose(1, 2) + + monkeypatch.setattr(sage_attn, "sageattn", sage_kernel) + cfg = SimpleNamespace( + diffusion_attention_config=config, + parallel_config=SimpleNamespace(ring_degree=1), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + with set_current_diffusion_config(cfg): + attn = transformer.MiniMaxH3Attention( + arch, + None, + prefix="token_refiner.blocks.0.attn", + role="minimax_h3.token_refiner", + role_category="self", + skip_sequence_parallel=True, + ) + assert attn.attention.attn_backend is sage_attn.SageAttentionBackend + # Call CUDA dispatch on CPU tensors with only the external kernel replaced. + monkeypatch.setattr(attn.attention.attention, "forward", attn.attention.attention.forward_cuda) + q, k, v = (torch.randn(7, 2, 32) for _ in range(3)) + cu = torch.tensor([0, 7], dtype=torch.int32) + output = attn._run_packed_attention(q, k, v, cu_seqlens=cu, max_seqlen=7, packed_total=7) + assert output.shape == q.shape and len(calls) == 1 + # Padded and multi-request layouts still receive their existing runtime rejection. + with pytest.raises(ValueError, match="does not support attn_mask"): + attn._run_packed_attention(q, k, v, cu_seqlens=cu, max_seqlen=6, packed_total=7) + with pytest.raises(ValueError, match="does not isolate multi-document"): + attn._run_packed_attention(q, k, v, cu_seqlens=cu, max_seqlen=7, packed_total=7, num_requests=2) + assert len(calls) == 1 + + +@pytest.mark.cuda +@torch.inference_mode() +def test_minimax_sparse_recipe_constructs_model_and_dispatches(monkeypatch): + from pathlib import Path + + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + + from tests.helpers.block_sparse import selected_attention_reference + from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention, BlockSparseBackend + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as transformer + + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0): + pytest.skip("Requires Hopper SM90 and FA4") + pytest.importorskip("flash_attn.cute") + monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(transformer, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + recipe = json.loads( + (Path(__file__).resolve().parents[4] / "recipes/attention/minimax-h3-fa4-subblock.json").read_text() + ) + cfg = SimpleNamespace( + tf_model_config={ + "num_layers": 1, + "token_refiner_num_layers": 1, + "hidden_size": 256, + "num_attention_heads": 2, + "attention_head_dim": 128, + "ffn_hidden_size": 512, + "latents_dim": 2, + "audio_latents_dim": 2, + "patch_size": (1, 1, 1), + "text_dim": 256, + "timestep_input_dim": 32, + "time_embed_hidden_size": 64, + "time_embed_dim": 128, + "adaln_out_features": 18 * 256, + "final_adaln_out_features": 2 * 256, + "rope_inv_freq_len": 16, + }, + diffusion_attention_config=AttentionConfig(**recipe), + parallel_config=SimpleNamespace(ring_degree=1, ulysses_degree=1), + dtype=torch.bfloat16, + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + with set_current_diffusion_config(cfg): + model = transformer.MiniMaxH3DiTModel(cfg).to(device="cuda") + # Real model construction and projections, with initialized tiny weights; + # no checkpoint, kernel substitute or hand-assigned attention role. + for weight in model.parameters(): + weight.fill_(1) if weight.ndim == 1 else weight.normal_(std=0.02) + dit = model.blocks[0].attn + refiner = model.token_refiner.blocks[0].attn + assert dit.attention.role == "minimax_h3.dit" and dit.attention.role_category == "self" + assert refiner.attention.role == "minimax_h3.token_refiner" and refiner.attention.role_category == "self" + assert isinstance(dit.attention.attention, BlockSparseAttention) + assert dit.attention.attn_backend is BlockSparseBackend + assert not isinstance(refiner.attention.attention, BlockSparseAttention) + assert refiner.attention.attn_backend.get_name() == "FLASH_ATTN" + assert not transformer._attention_isolates_packed_requests(dit.attention) + calls = [] + real_length = 1089 + + def check_sparse_dispatch(layer, args, output): + q, k, v, metadata = args + calls.append(layer.role) + assert metadata.attn_mask is None + assert metadata.packed_padding.q_length == metadata.packed_padding.kv_length == real_length + assert metadata.extra["cu_seqlens_q"].shape == (3,) + assert metadata.extra["protected_kv_prefix"] == 0 + impl = layer.attention + real = (q[:, :real_length], k[:, :real_length], v[:, :real_length]) + selection = impl.selector.select(real[0], real[1], impl.scale, 0) + assert selection.indices.shape[-1] < (real_length + 63) // 64 + reference = selected_attention_reference(*real, selection, impl.scale, impl.block_size) + torch.testing.assert_close(output[:, :real_length].float(), reference, atol=0.004, rtol=0.02) + assert torch.count_nonzero(output[:, real_length:]) == 0 and output.is_contiguous() + + dit.attention.register_forward_hook(check_sparse_dispatch) + refiner.attention.register_forward_hook(lambda layer, args, output: calls.append(layer.role)) + text = torch.randn(97, 256, device="cuda", dtype=cfg.dtype) + refined = model.token_refiner( + text, + cu_seqlens=torch.tensor([0, 97], device=text.device, dtype=torch.int32), + max_seqlen=97, + ) + assert refined.shape == text.shape and calls == ["minimax_h3.token_refiner"] + hidden = torch.randn(1152, 256, device="cuda", dtype=cfg.dtype) + hidden[real_length:] = float("nan") + cu = torch.tensor([0, real_length, hidden.shape[0]], device=hidden.device, dtype=torch.int32) + calls.clear() + actual = dit(hidden, rope_table=None, cu_seqlens=cu, max_seqlen=real_length) + assert calls == ["minimax_h3.dit"] and actual.shape == hidden.shape + hidden[real_length:] = 10000 + repeated = dit(hidden, rope_table=None, cu_seqlens=cu, max_seqlen=real_length) + torch.testing.assert_close(repeated[:real_length], actual[:real_length], atol=0, rtol=0) + with pytest.raises(ValueError, match="does not isolate multi-document"): + dit(hidden, rope_table=None, cu_seqlens=cu, max_seqlen=real_length, num_requests=2) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py index c89c37db9a2..1aadf8065ef 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py @@ -10,6 +10,7 @@ import torch from safetensors.torch import save_file +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec from vllm_omni.diffusion.models.minimax_h3.fasth3 import ( FASTH3_BASE_MODEL, FASTH3_BASE_SCHEDULE, @@ -459,7 +460,7 @@ def test_a_vsa_variant_accepts_normalized_attention_config(tmp_path): sparse = tmp_path / "vsa" / "adapter_model.safetensors" _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))}) od_config = SimpleNamespace( - diffusion_attention_config=SimpleNamespace(default=SimpleNamespace(backend="FASTVIDEO_VSA")), + diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FASTVIDEO_VSA")), parallel_config=SimpleNamespace(sequence_parallel_size=1), ) _load(sparse.parent).check_serving_contract( @@ -467,8 +468,9 @@ def test_a_vsa_variant_accepts_normalized_attention_config(tmp_path): ) -def test_a_vsa_variant_reads_the_backend_the_dit_will_actually_resolve(tmp_path): - # The 50 DiT blocks carry attention role "self", so a per_role entry - not +@pytest.mark.parametrize("role", ["self", "minimax_h3.dit"]) +def test_a_vsa_variant_reads_the_backend_the_dit_will_actually_resolve(tmp_path, role): + # The exact DiT role or legacy self category override - not # the default - decides whether the compression gates ever reach a sparse # kernel. sparse = tmp_path / "vsa" / "adapter_model.safetensors" @@ -477,9 +479,9 @@ def test_a_vsa_variant_reads_the_backend_the_dit_will_actually_resolve(tmp_path) contract = {"partition": "fl2va", "video_shift": 12.0, "audio_shift": 3.0} dense_dit = SimpleNamespace( - diffusion_attention_config=SimpleNamespace( - default=SimpleNamespace(backend="FASTVIDEO_VSA"), - per_role={"self": SimpleNamespace(backend="TRTLLM_ATTN")}, + diffusion_attention_config=AttentionConfig( + default=AttentionSpec(backend="FASTVIDEO_VSA"), + per_role={role: AttentionSpec(backend="TRTLLM_ATTN")}, ), parallel_config=SimpleNamespace(sequence_parallel_size=1), ) @@ -487,9 +489,9 @@ def test_a_vsa_variant_reads_the_backend_the_dit_will_actually_resolve(tmp_path) fusion.check_serving_contract(od_config=dense_dit, **contract) sparse_dit = SimpleNamespace( - diffusion_attention_config=SimpleNamespace( + diffusion_attention_config=AttentionConfig( default=None, - per_role={"self": SimpleNamespace(backend="FASTVIDEO_VSA")}, + per_role={role: AttentionSpec(backend="FASTVIDEO_VSA")}, ), parallel_config=SimpleNamespace(sequence_parallel_size=1), ) diff --git a/tests/helpers/block_sparse.py b/tests/helpers/block_sparse.py new file mode 100644 index 00000000000..76e8a08b3af --- /dev/null +++ b/tests/helpers/block_sparse.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Provider-independent input generation and selected-attention reference.""" + +import math + +import torch + +from vllm_omni.diffusion.attention.block_selection.abstract import BlockSelection + + +def make_attention_inputs( + heads=4, kv_heads=2, q_len=129, kv_len=1089, head_size=128, value_size=None, dtype=torch.bfloat16 +): + generator = torch.Generator(device="cuda").manual_seed(17) + q = torch.randn(2, q_len, heads, head_size, device="cuda", dtype=dtype, generator=generator) + k = torch.randn(2, kv_len, kv_heads, head_size, device=q.device, dtype=q.dtype, generator=generator) + v = torch.randn(2, kv_len, kv_heads, value_size or head_size, device=q.device, dtype=q.dtype, generator=generator) + return q, k, v + + +def selected_attention_reference(q, k, v, selection: BlockSelection, scale, block_size): + """FP32 masked SDPA over exactly the active entries of each selection row.""" + block_q, block_kv = block_size + indices, counts = selection + active = torch.arange(indices.shape[-1], device=q.device) < counts[..., None] + # Padding entries are unspecified by the contract; never use them as indices. + safe_indices = torch.where(active, indices, 0).long() + mask = torch.zeros(*indices.shape[:-1], math.ceil(k.shape[1] / block_kv), device=q.device, dtype=torch.int32) + mask.scatter_add_(-1, safe_indices, active.int()) + mask = mask.bool().repeat_interleave(block_q, 2).repeat_interleave(block_kv, 3) + mask = mask[:, :, : q.shape[1], : k.shape[1]] + return torch.nn.functional.scaled_dot_product_attention( + q.float().transpose(1, 2), + k.float().transpose(1, 2), + v.float().transpose(1, 2), + attn_mask=mask, + scale=scale, + enable_gqa=True, + ).transpose(1, 2) + + +def pinned_subblock_reference(q, k, scale, sparsity, prefix, block_size): + """Independent Torch oracle for the pinned SubBlock recipe. + + Recipe: SGLang 704808ed27cef61e210100e7581c115db3ee8401; + threshold search: kernels.py at 367e3700cfb6a0b03b2fa41a4524febf18ec1f15. + Includes this integration's native GQA and additive protected-prefix policy. + Does not call production pooling, scoring, budget or top-k helpers. + """ + bq, bk = block_size + nq, nk = math.ceil(q.shape[1] / bq), math.ceil(k.shape[1] / bk) + + def pool(x, blocks, width, factor): + cells = blocks * (width // 16) + result = torch.zeros(x.shape[0], x.shape[2], cells, x.shape[3], device=x.device, dtype=x.dtype) + for cell in range(math.ceil(x.shape[1] / 16)): + result[:, :, cell] = (x[:, cell * 16 : (cell + 1) * 16].float().mean(1) * factor).to(x.dtype) + return result.float() + + qp = pool(q, nq, bq, scale * math.log2(math.e)) + kp = pool(k, nk, bk, 1).repeat_interleave(q.shape[2] // k.shape[2], dim=1) + dots = qp @ kp.transpose(-1, -2) + dots[..., math.ceil(q.shape[1] / 16) :, :] = -float("inf") + dots[..., math.ceil(k.shape[1] / 16) :] = -float("inf") + pairs = dots.reshape(q.shape[0], q.shape[2], nq, bq // 16, nk, bk // 16) + scores = torch.logsumexp(pairs.permute(0, 1, 2, 4, 3, 5).flatten(-2) * math.log(2), -1) + protected = math.ceil(prefix / bk) + candidates = nk - protected + retained = min(candidates, 8 * math.ceil(math.ceil((1 - sparsity) * candidates) / 8)) + rows = [] + for row in scores.reshape(-1, nk): + values = row[protected:] + if retained == candidates: + chosen = torch.arange(candidates, device=q.device) + else: + lo, hi = values.min(), values.max() + 1 + clo, chi = float(candidates), 0.0 + ratio = math.log2(candidates / retained) + iterations = 16 if ratio <= 2.5 else 24 if ratio <= 3.5 else 32 + for _ in range(iterations): + fraction = min(max((clo - retained) / (clo - chi if clo - chi > 0.5 else 1), 0.05), 0.95) + threshold = lo + (hi - lo) * fraction + count = float((values >= threshold).sum()) + if count >= retained: + lo, clo = threshold, count + else: + hi, chi = threshold, count + chosen = torch.where(values >= lo)[0][:retained] + rows.append(torch.cat((torch.arange(protected, device=q.device), chosen + protected))) + indices = torch.stack(rows).reshape(*scores.shape[:-1], protected + retained).int() + return scores, BlockSelection( + indices, torch.full(indices.shape[:-1], indices.shape[-1], device=q.device, dtype=torch.int32) + ) diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index ed6a9a452c6..9ebd45376b1 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -7,13 +7,96 @@ import torch +from vllm_omni.diffusion.attention.block_selection.abstract import BlockSelection from vllm_omni.diffusion.attention.capabilities import ( + CompilationMode, ExecutionContext, ExecutionPathResult, ) from vllm_omni.platforms import current_omni_platform +class BlockSparseAdapter(ABC): + """Provider contract for execution over a shared selected-block pattern. + + The shared method owns scoring, selection, budgets and protected prefixes. + Adapters preserve its geometry, selected keys and native Q/KV head mapping. + They may convert metadata representation, but must not subdivide blocks, + replicate/merge selections, expand K/V or fall back to dense attention. + + Instantiate without arguments, then prepare before executing. Prepared + adapter state describes the kernel/configuration; request tensors and mutable + workspaces must not be shared between invocations. Implementations live with + their provider backend and opt in through get_block_sparse_adapter(). + """ + + provider: str + kernel_variant: str + compilation_mode: CompilationMode + dependency_version: str # Populated by prepare(). + + @staticmethod + @abstractmethod + def validate_selection(implementation: str, head_size: int) -> None: + """Check the provider-scoped ID and known head size (-1 means unknown). + + Preserve opaque IDs unchanged. Raise ValueError for unsupported choices + or ImportError for missing dependencies; never substitute another ID. + This is a preliminary selection check, not complete operation support. + """ + raise NotImplementedError + + @abstractmethod + def prepare( + self, + implementation: str, + head_size: int, + num_heads: int, + num_kv_heads: int, + device: torch.device, + block_size: tuple[int, int], + ) -> None: + """Prepare provider configuration and dependencies without synthetic inputs. + + Runs outside compiled execution and records dependency_version. The + shared method warms actual request tensors through execute() and records + support only after successful synchronization. Loading dependencies here + does not establish support for any request geometry. + """ + raise NotImplementedError + + def validate_inputs(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> None: + """Reject provider-specific input constraints before kernel execution. + + Shared layout, Q/K dimensions, head mapping and device are checked by + the orchestration layer. Override for kernel dtype or V-size limits. + """ + + @abstractmethod + def execute( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + selection: BlockSelection, + scale: float, + block_size: tuple[int, int], + ) -> torch.Tensor: + """Execute the prepared operation without discovery or fallback. + + Q/K/V use BSHD layout. selection follows BlockSelection: each row has + its own count of sorted unique logical KV block IDs; padding is ignored. + block_size is the prepared (Bq, Bkv). Sequence lengths bound tails and + prefix protection is already encoded. Empty rows are not permitted. + + Normalize over the selected eligible keys only, preserving per-query-head + selections and native GQA. Return contiguous [B, Sq, Hq, Dv] output with + the input dtype and device. Do not mutate input tensors or retain request + buffers. Honor compilation_mode through the provider's execution boundary. + """ + raise NotImplementedError + + class AttentionBackend(ABC): """Abstract class for diffusion attention backends.""" @@ -29,6 +112,15 @@ class AttentionBackend(ABC): # avoid a slower masked-attention plan when tail padding is not semantic. supports_prefix_kv_slicing: bool = False + @classmethod + def get_block_sparse_adapter(cls) -> type[BlockSparseAdapter] | None: + """Return an adapter consuming the shared selected pattern, or None. + + Adapters validate kernel IDs, prepare provider state and execute Q/K/V + with a BlockSelection. Selection and budgets belong to the strategy. + """ + return None + @classmethod def supports_packed_mask_free(cls) -> bool: """Whether [real, pad] packed layouts can run without attn_mask. diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 2ba6b415a7d..2b4462dcb23 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -5,17 +5,24 @@ from collections.abc import Callable from dataclasses import replace from functools import cache, partial +from importlib.metadata import version from typing import NamedTuple import torch from vllm.logger import init_logger -from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + BlockSparseAdapter, +) from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( piecewise_attn, run_paged_piecewise_plan, ) +from vllm_omni.diffusion.attention.block_selection.abstract import BlockSelection from vllm_omni.diffusion.attention.capabilities import ( CapabilityResult, CompilationMode, @@ -81,6 +88,98 @@ def _fa4_dense_attention_fake( _fa4_dense_attention_op = torch.ops.vllm_omni.fa4_dense_attention +if not hasattr(torch.ops.vllm_omni, "fa4_selected_attention"): + + @torch.library.custom_op("vllm_omni::fa4_selected_attention", mutates_args=()) + def _fa4_selected_attention_op( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + indices: torch.Tensor, + counts: torch.Tensor, + block_q: int, + block_kv: int, + ) -> torch.Tensor: + query, key, value = (t.contiguous() if t.stride(-1) != 1 else t for t in (query, key, value)) + return FA4SparseAdapter._run(query, key, value, indices, scale, (block_q, block_kv), counts) + + @_fa4_selected_attention_op.register_fake + def _fa4_selected_attention_fake(query, key, value, scale, indices, counts, block_q, block_kv): + return query.new_empty((*query.shape[:-1], value.shape[-1])) + + +_fa4_selected_attention_op = torch.ops.vllm_omni.fa4_selected_attention + + +class FA4SparseAdapter(BlockSparseAdapter): + """Provider-specific preparation and sparse metadata conversion only.""" + + provider = "FLASH_ATTN" + kernel_variant = "fa4" + compilation_mode = CompilationMode.CUSTOM_OP + + @staticmethod + @cache + def _load_api(): + """Cache dependency callables only; request tensors remain invocation-local.""" + from flash_attn.cute import flash_attn_func + from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch + + return flash_attn_func, BlockSparseTensorsTorch, version("flash-attn-4") + + @classmethod + def _run(cls, query, key, value, indices, scale, block_size, counts): + """Pass the logical sparse geometry directly to FA4, without retile conversion.""" + flash_attn_func, sparse_type, _ = cls._load_api() + sparse = sparse_type( + mask_block_cnt=counts, + mask_block_idx=indices, + full_block_cnt=torch.zeros_like(counts), + full_block_idx=indices, + block_size=block_size, + ) + result = flash_attn_func( + query, + key, + value, + softmax_scale=scale, + causal=False, + num_splits=1, + pack_gqa=False, + block_sparse_tensors=sparse, + ) + return (result[0] if isinstance(result, tuple) else result).contiguous() + + @staticmethod + def validate_selection(implementation: str, head_size: int) -> None: + if implementation != "auto": + raise ValueError("FA4 has no kernel-ID selection interface; use implementation='auto'") + + def prepare( + self, + implementation: str, + head_size: int, + num_heads: int, + num_kv_heads: int, + device: torch.device, + block_size: tuple[int, int], + ) -> None: + self.validate_selection(implementation, head_size) + _, _, self.dependency_version = self._load_api() + + def execute( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + selection: BlockSelection, + scale: float, + block_size: tuple[int, int], + ) -> torch.Tensor: + return _fa4_selected_attention_op(query, key, value, scale, selection.indices, selection.counts, *block_size) + + _PACKED_KEYS = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k") @@ -173,6 +272,10 @@ class FlashAttentionBackend(AttentionBackend): supports_piecewise_spans: bool = True supports_paged_kv: bool = True + @classmethod + def get_block_sparse_adapter(cls) -> type[FA4SparseAdapter]: + return FA4SparseAdapter + @classmethod def supports_packed_mask_free(cls) -> bool: # CUDA: forward_cuda dispatches the packed cu_seqlens varlen path diff --git a/vllm_omni/diffusion/attention/block_selection/__init__.py b/vllm_omni/diffusion/attention/block_selection/__init__.py new file mode 100644 index 00000000000..4244ce4b4e1 --- /dev/null +++ b/vllm_omni/diffusion/attention/block_selection/__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/attention/block_selection/abstract.py b/vllm_omni/diffusion/attention/block_selection/abstract.py new file mode 100644 index 00000000000..06074123814 --- /dev/null +++ b/vllm_omni/diffusion/attention/block_selection/abstract.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Contract between block selection and sparse attention execution.""" + +from abc import ABC, abstractmethod +from typing import Any, NamedTuple + +import torch + + +class BlockSelection(NamedTuple): + """Selected logical KV blocks, without changing attention mathematics. + + indices: contiguous int32 [B, Hq, ceil(Sq / Bq), capacity]. + counts: contiguous int32 [B, Hq, ceil(Sq / Bq)]. + Both tensors use the input device. Each row's first counts entries are + sorted, unique, in-range KV block IDs; remaining entries are ignored. + Every row must contain at least one block. Selectors must include every + block intersecting the protected KV token prefix, or reject the request. + Sequence tails are bounded by input lengths. No K/V head expansion occurs. + """ + + indices: torch.Tensor + counts: torch.Tensor + + +def validate_protected_kv_prefix(prefix: int, key_length: int) -> None: + """Require a token prefix within the key sequence, including either endpoint.""" + if isinstance(prefix, bool) or not isinstance(prefix, int) or not 0 <= prefix <= key_length: + raise ValueError("protected_kv_prefix must be an integer within the key sequence") + + +class BlockSelector(ABC): + """Prepared selection policy; request tensors must remain invocation-local. + + Implementations own scoring/approximation, budgets, rounding and geometry + restrictions. A scorer configuration is optional and interpreted by the + implementation; fused selection need not materialize a score matrix. + Algorithms requiring corrections for omitted keys need a separate attention + method, rather than extending this selected-key normalization contract. + """ + + @abstractmethod + def __init__(self, options: dict[str, Any]) -> None: + """Bind strategy options; prepare() binds execution geometry separately.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def normalize_config(cls, options: dict[str, Any]) -> dict[str, Any]: + """Normalize strategy-owned serializable options without executing GPU work.""" + raise NotImplementedError + + @abstractmethod + def prepare(self, block_size: tuple[int, int], head_size: int, device: torch.device) -> None: + """Reject unsupported geometry/device choices before execution.""" + raise NotImplementedError + + @abstractmethod + def validate_request(self, query: torch.Tensor, key: torch.Tensor, protected_prefix: int) -> None: + """Check strategy input constraints and mandatory-block feasibility. + + Called by capability resolution before execution, and by select() for + direct callers. Shared layout, head mapping and device checks belong + to the orchestration layer; prepared geometry is checked in prepare(). + """ + raise NotImplementedError + + @abstractmethod + def select(self, query: torch.Tensor, key: torch.Tensor, scale: float, protected_prefix: int) -> BlockSelection: + """Return a pattern meeting BlockSelection, including nonempty rows.""" + raise NotImplementedError diff --git a/vllm_omni/diffusion/attention/block_selection/registry.py b/vllm_omni/diffusion/attention/block_selection/registry.py new file mode 100644 index 00000000000..10db1ec1d28 --- /dev/null +++ b/vllm_omni/diffusion/attention/block_selection/registry.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Available selection methods, imported lazily for configuration parsing.""" + +from collections.abc import Mapping +from importlib import import_module + +from .abstract import BlockSelector + +# New methods implement BlockSelector and add their import path here. +BLOCK_SELECTORS: dict[str, str] = { + "block_topk": "vllm_omni.diffusion.attention.block_selection.subblock_topk.SubBlockTopK", +} + + +def get_block_selector(name: str | None) -> type[BlockSelector]: + if not isinstance(name, str) or name not in BLOCK_SELECTORS: + raise ValueError(f"Unknown block selector: {name!r}; available: {sorted(BLOCK_SELECTORS)}") + module, class_name = BLOCK_SELECTORS[name].rsplit(".", 1) + return getattr(import_module(module), class_name) + + +def normalize_selection(selection=None): + selection = {"name": "block_topk"} if selection is None else selection + if not isinstance(selection, Mapping) or selection.keys() - {"name", "config"}: + raise ValueError("selection requires name and optional config") + selector = get_block_selector(selection.get("name")) + options = selection.get("config", {}) + if not isinstance(options, Mapping): + raise ValueError("selection.config must be a mapping") + return {"name": selection["name"], "config": selector.normalize_config(dict(options))} diff --git a/vllm_omni/diffusion/attention/block_selection/subblock_topk.py b/vllm_omni/diffusion/attention/block_selection/subblock_topk.py new file mode 100644 index 00000000000..4d76d864296 --- /dev/null +++ b/vllm_omni/diffusion/attention/block_selection/subblock_topk.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Routing kernels adapted from SGLang (Apache-2.0), commit +# 367e3700cfb6a0b03b2fa41a4524febf18ec1f15, +# python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/kernels.py. +# Recipe reference: 704808ed27cef61e210100e7581c115db3ee8401. +# Modified to integrate GQA and protected-prefix routing in vLLM-Omni. +"""SubBlock pooled scoring and approximate top-k selection.""" + +import math + +import torch +from vllm.triton_utils import tl, triton + +from .abstract import BlockSelection, BlockSelector, validate_protected_kv_prefix + + +@torch.library.custom_op("vllm_omni::select_subblocks", mutates_args=()) +def select_subblocks( + query: torch.Tensor, + key: torch.Tensor, + scale: float, + sparsity: float, + protected_prefix: int, + block_q: int, + block_kv: int, +) -> torch.Tensor: + return SubBlockTopK._route(query, key, scale, sparsity, protected_prefix, (block_q, block_kv)) + + +@select_subblocks.register_fake +def _select_subblocks_fake(query, key, scale, sparsity, protected_prefix, block_q, block_kv): + prefix_blocks = (protected_prefix + block_kv - 1) // block_kv + key_blocks = (key.shape[1] + block_kv - 1) // block_kv + return torch.empty( + ( + query.shape[0], + query.shape[2], + (query.shape[1] + block_q - 1) // block_q, + prefix_blocks + SubBlockTopK.block_budget(key_blocks - prefix_blocks, sparsity), + ), + dtype=torch.int32, + device=query.device, + ) + + +class SubBlockTopK(BlockSelector): + """Keep protected blocks plus a rounded budget of unprotected candidates. + + target_sparsity applies only to the unprotected blocks. Prefix protection + never consumes that candidate budget; rounding and protection both reduce + the overall achieved sparsity. A prefix protects every intersecting block, + including its non-prefix tokens, and an entirely protected sequence is dense. + """ + + @staticmethod + def block_budget(num_blocks, sparsity): + """Budget for unprotected blocks: ceil retained fraction, align to eight, cap at N.""" + if num_blocks < 0 or not 0 <= sparsity < 1: + raise ValueError("SubBlock requires a nonnegative candidate count and 0 <= sparsity < 1") + if num_blocks == 0: + return 0 + return min(num_blocks, math.ceil(max(1, math.ceil((1 - sparsity) * num_blocks)) / 8) * 8) + + @staticmethod + def _validate_geometry(block_size): + """The pooled-score Triton tiles must contain whole logical blocks.""" + if len(block_size) != 2 or any( + isinstance(size, bool) + or not isinstance(size, int) + or size < 16 + or size % 16 + or size & (size - 1) + or size // 16 > tile + for size, tile in zip(block_size, (block_m, block_n)) + ): + raise ValueError( + "The SubBlock scorer requires power-of-two block dimensions from 16 to 2048 " + "(16-token pooling cells dividing its 128-cell score tiles)" + ) + + @classmethod + def _route(cls, query, key, scale, sparsity, protected_prefix, block_size): + block_q, block_kv = block_size + rows = (query.shape[0], query.shape[2], math.ceil(query.shape[1] / block_q)) + key_blocks = math.ceil(key.shape[1] / block_kv) + prefix_blocks = math.ceil(protected_prefix / block_kv) + kept = prefix_blocks + cls.block_budget(key_blocks - prefix_blocks, sparsity) + # These patterns are known without pooling, scoring or top-k selection. + if kept == key_blocks: + indices = torch.arange(kept, dtype=torch.int32, device=query.device) + return indices.expand(*rows, kept).contiguous() + scores = cls._compute_scores(query, key, scale, block_size) + return cls._select_blocks(scores, kept, prefix_blocks) + + @staticmethod + def _compute_scores(query, key, scale, block_size): + block_q, block_kv = block_size + n_q, n_k = block_q // 16, block_kv // 16 + b, sq, hq, d = query.shape + _, sk, hk, _ = key.shape + gq, gk = math.ceil(sq / block_q), math.ceil(sk / block_kv) + qp = torch.empty(b * hq, gq * n_q, d, device=query.device, dtype=query.dtype) + kp = torch.empty(b * hk, gk * n_k, d, device=key.device, dtype=key.dtype) + _fused_pool(query, gq * n_q, 16, qp, scale=scale * math.log2(math.e)) + _fused_pool(key, gk * n_k, 16, kp) + scores = torch.empty(b * hq, gq, gk, device=query.device, dtype=torch.float32) + _fused_scores( + qp, + kp, + scores, + n_k=n_k, + n_valid=math.ceil(sk / 16), + n_q=n_q, + m_valid=math.ceil(sq / 16), + heads_per_kv=hq // hk, + ) + return scores.view(b, hq, gq, gk) + + @staticmethod + def _select_blocks(scores, kept, prefix_blocks): + # Reserve protected blocks structurally, outside approximate top-k. + rows, key_blocks = scores.shape[:-1], scores.shape[-1] + remaining = kept - prefix_blocks + candidates = scores[..., prefix_blocks:].contiguous().view(-1, key_blocks - prefix_blocks) + chosen = (_fused_topk(candidates, remaining) + prefix_blocks).view(*rows, remaining) + if not prefix_blocks: + return chosen + prefix = torch.arange(prefix_blocks, dtype=torch.int32, device=scores.device) + return torch.cat((prefix.expand(*rows, prefix_blocks), chosen), dim=-1) + + @classmethod + def normalize_config(cls, options): + if options.keys() - {"target_sparsity"}: + raise ValueError("block_topk config only supports target_sparsity") + sparsity = options.get("target_sparsity", 0.75) + if isinstance(sparsity, bool) or not isinstance(sparsity, (int, float)) or not 0 <= sparsity < 1: + raise ValueError("Require 0 <= target_sparsity < 1") + return {"target_sparsity": float(sparsity)} + + def __init__(self, options): + options = self.normalize_config(options) + self.sparsity = options["target_sparsity"] + + def prepare(self, block_size, head_size, device): + self._validate_geometry(block_size) + if device.type != "cuda" or head_size <= 0: + raise ValueError("SubBlock top-k requires CUDA and a positive head dimension") + self.block_size = block_size + + def validate_request(self, query, key, protected_prefix): + if key.shape[1] == 0: + raise ValueError("SubBlock requires nonempty keys") + if query.dtype not in (torch.float16, torch.bfloat16) or key.dtype != query.dtype: + raise ValueError("SubBlock top-k requires matching FP16 or BF16 Q/K") + + validate_protected_kv_prefix(protected_prefix, key.shape[1]) + + def select(self, query, key, scale, protected_prefix): + self.validate_request(query, key, protected_prefix) + indices = select_subblocks(query, key, scale, self.sparsity, protected_prefix, *self.block_size) + counts = torch.full(indices.shape[:-1], indices.shape[-1], dtype=torch.int32, device=indices.device) + return BlockSelection(indices, counts) + + +@triton.jit +def _pool_kernel( + input_ptr, + output_ptr, + stride_xb, + stride_xt, + stride_xh, + stride_xd, + stride_yl, + stride_yn, + sequence_length, + heads, + subblock_size: tl.constexpr, + head_dim: tl.constexpr, + block_d: tl.constexpr, + scale, +): + """Mean-pool token sub-blocks in the input dtype, masking token and head tails.""" + cell = tl.program_id(0) + batch_head = tl.program_id(1) + b = batch_head // heads + h = batch_head % heads + offs_t = cell * subblock_size + tl.arange(0, subblock_size) + offs_d = tl.arange(0, block_d) + mask = offs_t < sequence_length + x = tl.load( + input_ptr + b * stride_xb + offs_t[:, None] * stride_xt + h * stride_xh + offs_d[None, :] * stride_xd, + mask=mask[:, None] & (offs_d[None, :] < head_dim), + other=0.0, + ).to(tl.float32) + cnt = tl.sum(mask.to(tl.float32), axis=0) + acc = tl.sum(x, axis=0) / tl.maximum(cnt, 1.0) * scale + tl.store( + output_ptr + batch_head * stride_yl + cell * stride_yn + offs_d, + acc.to(output_ptr.dtype.element_ty), + mask=offs_d < head_dim, + ) + + +def _fused_pool(x, n_cells, sub, out, scale=1.0): + """Pool [B, sequence_length, heads, head_dim] into [B*heads, cells, head_dim].""" + B, sequence_length, heads, head_dim = x.shape + _pool_kernel[(n_cells, B * heads)]( + x, + out, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + out.stride(0), + out.stride(1), + sequence_length, + heads, + subblock_size=sub, + head_dim=head_dim, + block_d=max(16, triton.next_power_of_2(head_dim)), + scale=scale, + num_warps=1, + ) + return out + + +_NEG = tl.constexpr(-1.0e30) +_LN2 = tl.constexpr(0.6931471805599453) + + +@triton.jit +def _score_kernel( + query_ptr, + key_ptr, + output_ptr, + stride_qm, + stride_ql, + stride_kn, + stride_kl, + stride_om, + stride_on, + stride_ol, + pooled_q_len, + valid_q_cells, + valid_k_cells, + key_blocks, + query_blocks, + block_m: tl.constexpr, + block_n: tl.constexpr, + key_subblocks: tl.constexpr, + query_subblocks: tl.constexpr, + head_dim: tl.constexpr, + block_d: tl.constexpr, + heads_per_kv: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + pid_l = tl.program_id(2) + + offs_m = pid_m * block_m + tl.arange(0, block_m) + offs_n = pid_n * block_n + tl.arange(0, block_n) + offs_d = tl.arange(0, block_d) + + q = tl.load( + query_ptr + pid_l * stride_ql + offs_m[:, None] * stride_qm + offs_d[None, :], + mask=(offs_m[:, None] < pooled_q_len) & (offs_d[None, :] < head_dim), + other=0.0, + ) + k = tl.load( + key_ptr + (pid_l // heads_per_kv) * stride_kl + offs_n[:, None] * stride_kn + offs_d[None, :], + mask=(offs_n[:, None] < valid_k_cells) & (offs_d[None, :] < head_dim), + other=0.0, + ) + acc = tl.dot(q, tl.trans(k), out_dtype=tl.float32) # [block_m, block_n] + + # a sub-cell past the last real one must not contribute to its group's log-sum-exp + acc = tl.where(offs_n[None, :] < valid_k_cells, acc, _NEG) + # same on the query side: with NQ > 1 the last query block can own sub-cells that + # are entirely padding, and those pool to zero -- an exp2(0) = 1 term that would + # otherwise be folded into the block's score. + acc = tl.where(offs_m[:, None] < valid_q_cells, acc, _NEG) + + acc = tl.reshape(acc, (block_m, block_n // key_subblocks, key_subblocks)) + m = tl.max(acc, axis=2) + s = tl.sum(tl.exp2(acc - m[:, :, None]), axis=2) + lse = m + tl.log2(s) + lse = tl.where(m > _NEG / 2, lse, _NEG) # whole group was padding + + # Fold the query_subblocks query sub-cells of a query block together. Log-sum-exp is + # associative, so reducing key_subblocks then NQ is the same one log-sum-exp over all + # NQ*key_subblocks sub-block pairs -- and two stages keeps both reductions on an axis + # that is already contiguous in registers. + if query_subblocks > 1: + lse = tl.reshape(lse, (block_m // query_subblocks, query_subblocks, block_n // key_subblocks)) + m2 = tl.max(lse, axis=1) + s2 = tl.sum(tl.exp2(lse - m2[:, None, :]), axis=1) + lse = tl.where(m2 > _NEG / 2, m2 + tl.log2(s2), _NEG) + + # exp2/log2 internally (they map to the hardware instructions), then back to natural + # log units so the fused and reference backends return the same numbers, not just the + # same ranking. One multiply in registers. + out = lse * _LN2 + out = out.to(output_ptr.dtype.element_ty) + + offs_o = pid_n * (block_n // key_subblocks) + tl.arange(0, block_n // key_subblocks) + offs_q = pid_m * (block_m // query_subblocks) + tl.arange(0, block_m // query_subblocks) + tl.store( + output_ptr + pid_l * stride_ol + offs_q[:, None] * stride_om + offs_o[None, :] * stride_on, + out, + mask=(offs_q[:, None] < query_blocks) & (offs_o[None, :] < key_blocks), + ) + + +block_m = block_n = 128 # score tile; must hold whole blocks, so a multiple of n_q and n_k + + +def _fused_scores(qp, kp, out, *, n_k, n_valid, n_q, m_valid, heads_per_kv=1): + """Compute FP32 block scores from FP16 or BF16 pooled query/key sub-blocks.""" + L, pooled_q_len, head_dim = qp.shape + N = kp.shape[1] + Mout, Nout = out.shape[1], out.shape[2] + grid = (triton.cdiv(pooled_q_len, block_m), triton.cdiv(N, block_n), L) + _score_kernel[grid]( + qp, + kp, + out, + qp.stride(1), + qp.stride(0), + kp.stride(1), + kp.stride(0), + out.stride(1), + out.stride(2), + out.stride(0), + pooled_q_len, + m_valid, + n_valid, + Nout, + Mout, + block_m=block_m, + block_n=block_n, + key_subblocks=n_k, + query_subblocks=n_q, + head_dim=head_dim, + block_d=max(16, triton.next_power_of_2(head_dim)), + heads_per_kv=heads_per_kv, + num_warps=4, + num_stages=3, + ) + return out + + +@triton.jit +def _topk_kernel(scores_ptr, output_ptr, columns, topk, block_size: tl.constexpr, iterations: tl.constexpr): + """Select blocks using bounded threshold search and compact their indices.""" + row = tl.program_id(0) + offs = tl.arange(0, block_size) + m = offs < columns + s = tl.load(scores_ptr + row * columns + offs, mask=m, other=-float("inf")).to(tl.float32) + lo = tl.min(tl.where(m, s, float("inf"))) + hi = tl.max(tl.where(m, s, -float("inf"))) + 1.0 + clo = tl.sum(m.to(tl.int32), axis=0).to(tl.float32) + chi = 0.0 + for _ in tl.static_range(iterations): + den = clo - chi + t = (clo - topk) / tl.where(den > 0.5, den, 1.0) + t = tl.minimum(tl.maximum(t, 0.05), 0.95) # keep the step inside the bracket + mid = lo + (hi - lo) * t + cnt = tl.sum(((s >= mid) & m).to(tl.int32), axis=0).to(tl.float32) + take = cnt >= topk + lo = tl.where(take, mid, lo) + clo = tl.where(take, cnt, clo) + hi = tl.where(take, hi, mid) + chi = tl.where(take, chi, cnt) + sel = (s >= lo) & m + # Nonfinite scores have no reliable ranking (finite inputs can overflow + # during pooling). Use the first k candidates for the entire affected row. + # Also cover finite-score threshold arithmetic that cannot fill the budget. + nonfinite = tl.sum((m & ((s != s) | (tl.abs(s) == float("inf")))).to(tl.int32), axis=0) > 0 + insufficient = tl.sum(sel.to(tl.int32), axis=0) < topk + sel = tl.where(nonfinite | insufficient, m & (offs < topk), sel) + pos = tl.cumsum(sel.to(tl.int32), axis=0) - 1 + tl.store(output_ptr + row * topk + pos, offs.to(tl.int32), mask=sel & (pos < topk)) + + +def _topk_iters(columns, k): + """Use more threshold-search iterations for higher sparsity.""" + L = math.log2(max(columns, 1) / max(k, 1)) + return 16 if L <= 2.5 else 24 if L <= 3.5 else 32 + + +def _fused_topk(scores2d, k): + """Return ascending int32 indices; prefix-sum compaction preserves column order.""" + rows, columns = scores2d.shape + out = torch.empty(rows, k, dtype=torch.int32, device=scores2d.device) + _topk_kernel[(rows,)]( + scores2d, + out, + columns, + k, + block_size=triton.next_power_of_2(columns), + iterations=_topk_iters(columns, k), + num_warps=4 if columns >= 1024 else 2, # short rows do not fill four warps + ) + return out diff --git a/vllm_omni/diffusion/attention/block_sparse.py b/vllm_omni/diffusion/attention/block_sparse.py new file mode 100644 index 00000000000..10f7dd9a65f --- /dev/null +++ b/vllm_omni/diffusion/attention/block_sparse.py @@ -0,0 +1,410 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Compose block selection with a provider-specific sparse attention adapter.""" + +import copy +import itertools +import math +import weakref + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + BlockSparseAdapter, + PackedPaddingMetadata, +) +from vllm_omni.diffusion.attention.block_selection.abstract import validate_protected_kv_prefix +from vllm_omni.diffusion.attention.block_selection.registry import get_block_selector +from vllm_omni.diffusion.attention.capabilities import ( + CapabilityResult, + CompilationMode, + ExecutionPathResult, + MaskMode, + PackingMode, + ParallelStrategy, +) +from vllm_omni.diffusion.config import get_current_diffusion_config_or_none +from vllm_omni.diffusion.data import BlockSparseAttentionSpec + +# Compiled graphs receive a CPU scalar owner handle as a runtime tensor input, +# never a per-layer Python constant or a snapshot of readiness. +# Weak references do not extend model/request lifetimes; tokens are never reused. +_request_owners: weakref.WeakValueDictionary[int, "BlockSparseAttention"] = weakref.WeakValueDictionary() +_request_owner_ids = itertools.count() + + +@torch.library.custom_op("vllm_omni::block_sparse_request", mutates_args=()) +def _block_sparse_request( + query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, prefix: int, owner_handle: torch.Tensor +) -> torch.Tensor: + owner = _request_owners.get(owner_handle.item()) + if owner is None: + raise RuntimeError("Block-sparse request owner was released; keep the attention layer alive during execution") + return owner._dispatch_request(query, key, value, prefix) + + +@_block_sparse_request.register_fake +def _block_sparse_request_fake(query, key, value, prefix, owner_handle): + # No readiness lookup, allocation of real buffers, or trial execution while tracing. + return query.new_empty((*query.shape[:-1], value.shape[-1])) + + +class BlockSparseBackend(AttentionBackend): + """Method-level preflight capabilities, independent of the dense provider. + + Only producer-validated single-document suffix padding is supported. Masks, + multi-document packing, piecewise spans, prefix slicing and paged KV retain + AttentionBackend's conservative False defaults. The bound execution adapter + remains on BlockSparseAttention and establishes actual request support. + """ + + @staticmethod + def get_name() -> str: + return "BLOCK_SPARSE" + + @classmethod + def supports_packed_mask_free(cls) -> bool: + # Shared orchestration trims both Q and K/V using the typed contract. + return True + + @staticmethod + def get_impl_cls() -> type["BlockSparseAttention"]: + return BlockSparseAttention + + @staticmethod + def get_metadata_cls() -> type[AttentionMetadata]: + return AttentionMetadata + + @staticmethod + def get_builder_cls(): + return None + + @staticmethod + def get_supported_head_sizes() -> list[int]: + # No duplicate provider support table: actual tensors require preparation. + return [] + + @classmethod + def resolve_capabilities(cls, context): + return ExecutionPathResult( + backend=cls.get_name(), + path="block_sparse_unprepared", + support=CapabilityResult.unsupported( + "Block-sparse support requires actual request preparation; " + "query the bound Attention layer after eager warmup" + ), + compilation_mode=CompilationMode.EAGER_ONLY, + platform=context.platform, + kernel_variant=None, + parallel_strategy=context.parallel_strategy, + ) + + +class BlockSparseAttention(AttentionImpl[AttentionMetadata]): + def __init__( + self, + num_heads, + num_kv_heads, + head_size, + scale, + causal, + layout, + spec: BlockSparseAttentionSpec, + *, + adapter: BlockSparseAdapter, + ): + self.block_size = spec.block_size + selection = spec.selection + self.selector = get_block_selector(selection["name"])(selection["config"]) + if causal or layout not in (None, "BSHD", "BSND"): + raise ValueError("Block-sparse attention requires noncausal BSHD attention") + if head_size <= 0: + raise ValueError("Block-sparse attention requires a positive Q/K head dimension") + if min(num_heads, num_kv_heads) <= 0 or num_heads % num_kv_heads: + raise ValueError("Block-sparse attention requires positive Q heads divisible by KV heads") + if not math.isfinite(scale) or scale <= 0: + raise ValueError("Block-sparse attention requires a finite positive scale") + if not torch.cuda.is_available(): + raise ValueError("The initial Block-sparse attention router requires CUDA") + self.device = torch.device("cuda", torch.accelerator.current_device_index()) + self.num_heads, self.num_kv_heads, self.scale = num_heads, num_kv_heads, scale + self.head_size = head_size + cfg = get_current_diffusion_config_or_none() + if cfg is not None: + parallel = getattr(cfg, "parallel_config", None) + if any( + (getattr(parallel, name, 1) or 1) > 1 + for name in ( + "sequence_parallel_size", + "ring_degree", + "allgather_degree", + ) + ) or getattr(parallel, "use_hsdp", False): + raise ValueError("Block-sparse attention parallel wrappers require separate validation") + if getattr(cfg, "fa_deterministic", False): + raise ValueError("Block-sparse attention does not support fa_deterministic") + if getattr(cfg, "diffusion_kv_cache_dtype", None) not in (None, "auto", "float"): + raise ValueError("Block-sparse attention does not support KV quantization") + if getattr(cfg, "diffusion_kv_mode", "dense_legacy") != "dense_legacy": + raise ValueError("Block-sparse attention does not support paged KV") + self.selector.prepare(self.block_size, head_size, self.device) + self.adapter = adapter + # Geometry-only outcomes: never retain request tensors or exception tracebacks. + # None denotes successful execution and synchronization; strings record failures. + self._request_preparation: dict[tuple, str | None] = {} + adapter.prepare(spec.implementation, head_size, num_heads, num_kv_heads, self.device, self.block_size) + self._register_request_owner() + init_logger(__name__).info( + "Configured block-sparse attention (request preparation pending): provider=%s kernel=%s dependency=%s " + "native heads=%s/%s selector=%s block_size=%s", + adapter.provider, + adapter.kernel_variant, + adapter.dependency_version, + num_heads, + num_kv_heads, + type(self.selector).__name__, + self.block_size, + ) + + def _register_request_owner(self): + owner_id = next(_request_owner_ids) + # A CPU scalar is a runtime graph input, not a per-layer Python constant. + # Read it only inside the opaque op to avoid specializing owner identity. + self._request_owner_handle = torch.tensor(owner_id, dtype=torch.int64, device="cpu") + _request_owners[owner_id] = self + + def __getstate__(self): + # Handles are process-local identity, and readiness is evidence from this + # instance's actual execution. Neither belongs in copied/serialized state. + return { + name: value + for name, value in self.__dict__.items() + if name not in ("_request_owner_handle", "_request_preparation") + } + + def __setstate__(self, state): + self.__dict__.update(state) + self._request_preparation = {} + self._register_request_owner() + + def __deepcopy__(self, memo): + cloned = type(self).__new__(type(self)) + memo[id(self)] = cloned + cloned.__setstate__(copy.deepcopy(self.__getstate__(), memo)) + return cloned + + def __copy__(self): + raise TypeError( + "Shallow copying block-sparse attention is unsupported; use copy.deepcopy() or construct a new layer" + ) + + @staticmethod + def _validate_real_length(name, length, storage_length): + if isinstance(length, bool) or not isinstance(length, int): + raise ValueError(f"{name} must be a host integer") + if not 0 < length <= storage_length: + raise ValueError(f"{name} must be positive and within its tensor sequence length") + + @staticmethod + def _validate_cumulative_lengths(name, lengths, sizes): + # Values are producer-validated. Inspect only tensor metadata: reading + # device scalars would introduce an unnecessary host synchronization. + if ( + not isinstance(lengths, torch.Tensor) + or lengths.ndim != 1 + or lengths.shape[0] not in sizes + or lengths.dtype not in (torch.int32, torch.int64) + ): + raise ValueError(f"{name} must be a one-dimensional integer tensor with length in {sizes}") + + def _normalize_request(self, query, key, value, metadata): + """Validate and trim a producer-validated B=1 [real, padding] request. + + Returns executed Q/K/V, protected prefix and the original query length. + Cumulative-length values are never read; the typed producer contract is + authoritative. Raw packed metadata alone cannot establish this contract. + """ + if any(t.ndim != 4 for t in (query, key, value)): + raise ValueError("Block-sparse attention requires four-dimensional BSHD Q/K/V") + if ( + query.shape[0] != key.shape[0] + or key.shape[:3] != value.shape[:3] + or query.shape[-1] != self.head_size + or key.shape[-1] != self.head_size + or min(query.shape[0], query.shape[1], key.shape[1], value.shape[-1]) <= 0 + or query.shape[2] != self.num_heads + or key.shape[2] != self.num_kv_heads + ): + raise ValueError("Incompatible Block-sparse attention Q/K/V geometry; native head mapping is required") + if query.dtype != key.dtype or key.dtype != value.dtype: + raise ValueError("Block-sparse attention requires matching Q/K/V dtypes") + if any(t.device != self.device for t in (query, key, value)): + raise ValueError("Block-sparse attention inputs must use the prepared CUDA device") + original_q_length = query.shape[1] + prefix = 0 + if metadata is not None: + if any( + getattr(metadata, name) is not None + for name in ("attn_mask", "joint_attn_mask", "full_attn_spans", "query_ranges") + ): + raise ValueError("Block-sparse attention does not support masks or piecewise attention") + extra = metadata.extra + if not isinstance(extra, dict): + raise ValueError("Block-sparse attention metadata.extra must be a dictionary") + if extra.get("kv_cache_dtype") not in (None, "auto", "float"): + raise ValueError("Block-sparse attention does not support quantized KV") + padding = metadata.packed_padding + packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k", "valid_kv_length") + if padding is None: + if any(extra.get(name) is not None for name in packed_keys): + raise ValueError("Block-sparse packed extras require typed PackedPaddingMetadata") + else: + if not isinstance(padding, PackedPaddingMetadata): + raise ValueError("packed_padding must be PackedPaddingMetadata") + if query.shape[0] != 1: + raise ValueError("Block-sparse packed padding requires B=1 and a single real document") + self._validate_real_length("q_length", padding.q_length, query.shape[1]) + self._validate_real_length("kv_length", padding.kv_length, key.shape[1]) + self._validate_cumulative_lengths("packed_padding.cu_seqlens_q", padding.cu_seqlens_q, (2,)) + self._validate_cumulative_lengths("packed_padding.cu_seqlens_k", padding.cu_seqlens_k, (2,)) + # MiniMax also supplies the original [0, real, total] tensors. + # Accept these redundant extras only under the typed contract; + # four or more boundaries imply unsupported multi-document data. + for name in ("cu_seqlens_q", "cu_seqlens_k"): + if extra.get(name) is not None: + self._validate_cumulative_lengths(name, extra[name], (2, 3)) + for name, real_length in ( + ("max_seqlen_q", padding.q_length), + ("max_seqlen_k", padding.kv_length), + ("valid_kv_length", padding.kv_length), + ): + if extra.get(name) is not None: + self._validate_real_length(name, extra[name], real_length) + if extra[name] != real_length: + raise ValueError(f"{name} disagrees with PackedPaddingMetadata") + # Trim Q too: padding in its final real block would otherwise + # contaminate the pooled routing scores of valid query rows. + query = query[:, : padding.q_length] + key = key[:, : padding.kv_length] + value = value[:, : padding.kv_length] + prefix = extra.get("protected_kv_prefix", 0) + validate_protected_kv_prefix(prefix, key.shape[1]) + self.adapter.validate_inputs(query, key, value) + self.selector.validate_request(query, key, prefix) + return query, key, value, prefix, original_q_length + + def resolve_execution_path(self, context, query, key, value, metadata): + reason = None + if ( + context.platform != "cuda" + or context.parallel_strategy is not ParallelStrategy.NONE + or context.outer_boundaries + or context.paged_kv + or context.piecewise + or context.mask_mode is not MaskMode.NONE + or context.packing_mode not in (PackingMode.NONE, PackingMode.PACKED_PADDING) + or context.kv_cache_dtype not in (None, "auto", "float") + or context.causal is True + ): + reason = ( + "Block-sparse attention requires noncausal, mask-free CUDA attention with at most typed suffix padding " + "without parallel wrappers" + ) + else: + try: + if context.packing_mode is PackingMode.PACKED_PADDING and ( + metadata is None or not isinstance(metadata.packed_padding, PackedPaddingMetadata) + ): + raise ValueError("PACKED_PADDING requires typed PackedPaddingMetadata") + query, key, value, prefix, _ = self._normalize_request(query, key, value, metadata) + signature = self._request_signature(query, key, value, prefix) + reason = self._request_preparation.get(signature, self._unprepared_reason()) + except ValueError as exc: + reason = str(exc) + return ExecutionPathResult( + backend=self.adapter.provider, + path=f"{self.adapter.kernel_variant}_block_sparse", + support=CapabilityResult.unsupported(reason) if reason else CapabilityResult.supported(), + compilation_mode=CompilationMode.CUSTOM_OP, + platform=context.platform, + kernel_variant=self.adapter.kernel_variant, + parallel_strategy=context.parallel_strategy, + ) + + def _request_signature(self, query, key, value, prefix): + # Instance configuration is fixed. Include every request property that can + # affect routing capacity or kernel specialization, including Q/K/V strides. + return ( + self.block_size, + self.scale, + prefix, + *((tuple(t.shape), tuple(t.stride()), t.dtype, t.device) for t in (query, key, value)), + ) + + @staticmethod + def _unprepared_reason(): + return ( + "Block-sparse request is not prepared: run an eager forward or prepare_request() " + "with this geometry to establish execution support" + ) + + def _prepare_and_execute(self, query, key, value, prefix, signature): + if torch.compiler.is_compiling(): + raise RuntimeError(self._unprepared_reason()) + with torch.accelerator.device_index(self.device.index): + try: + output = self._execute(query, key, value, prefix) + # A launch alone does not establish support: surface asynchronous + # kernel errors before publishing a successful preparation outcome. + torch.accelerator.synchronize(self.device) + except Exception as exc: + self._request_preparation[signature] = ( + f"Block-sparse request preparation failed ({type(exc).__name__}): {exc}. " + "Correct the request or dependency issue and retry eager preparation" + ) + raise + self._request_preparation[signature] = None + return output + + def prepare_request(self, query, key, value, metadata=None): + """Explicit eager preparation to establish support for this signature. + + This is repeatable, including after a failed attempt. It establishes kernel + execution support, not approximation quality or generation-quality parity. + """ + query, key, value, prefix, _ = self._normalize_request(query, key, value, metadata) + signature = self._request_signature(query, key, value, prefix) + with torch.inference_mode(): + self._prepare_and_execute(query, key, value, prefix, signature) + + def forward(self, query, key, value, metadata=None): + query, key, value, prefix, original_q_length = self._normalize_request(query, key, value, metadata) + # Dispatch is opaque to Dynamo: exact readiness signatures are evaluated + # against real tensors at runtime and cannot specialize symbolic lengths. + output = _block_sparse_request(query, key, value, prefix, self._request_owner_handle) + if query.shape[1] != original_q_length: + tail = output.new_zeros((output.shape[0], original_q_length - query.shape[1], *output.shape[2:])) + output = torch.cat((output, tail), dim=1) + return output.contiguous() + + def _dispatch_request(self, query, key, value, prefix): + """Runtime owner for both eager and compiled request preparation. + + This executes outside Dynamo, including for fullgraph=True. A new actual + request is executed and synchronized before publishing support. + """ + signature = self._request_signature(query, key, value, prefix) + if signature not in self._request_preparation or self._request_preparation[signature] is not None: + return self._prepare_and_execute(query, key, value, prefix, signature) + return self._execute(query, key, value, prefix) + + def _execute(self, query, key, value, prefix): + selection = self.selector.select(query, key, self.scale, prefix) + return self.adapter.execute(query, key, value, selection, self.scale, self.block_size) + + def forward_cuda(self, query, key, value, attn_metadata=None): + return self.forward(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index d239525c245..a0fa23fd0f8 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -31,6 +31,7 @@ from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role from vllm_omni.diffusion.config import get_current_diffusion_config_or_none +from vllm_omni.diffusion.data import BlockSparseAttentionSpec from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.layout import assert_backend_layout_supported from vllm_omni.diffusion.distributed.parallel_state import get_sp_group @@ -114,6 +115,7 @@ def __init__( self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads self.head_size = head_size + self.attn_impl_cls: type[AttentionImpl] self._has_custom_attention = custom_attention is not None # Resolve backend via role-aware config. @@ -187,7 +189,8 @@ def __init__( ) self.attn_spec = spec if spec is not None: - backend_kwargs = spec.backend_kwargs() + if not isinstance(spec, BlockSparseAttentionSpec): + backend_kwargs = spec.backend_kwargs() self.backend_pref = spec.backend self.backend_explicit = True logger.debug("Attention(role=%s) → backend=%s", role, spec.backend) @@ -197,33 +200,56 @@ def __init__( self.backend_pref = attn_backend_cls.get_name() logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref) - 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, - softmax_scale=softmax_scale, - causal=causal, - num_kv_heads=num_kv_heads, - qkv_layout=qkv_layout, - prefix=prefix, - backend_kwargs=backend_kwargs, - role=role, - backend_explicit=self.backend_explicit, - ) + self.attn_backend = attn_backend_cls + if isinstance(spec, BlockSparseAttentionSpec): + from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention, BlockSparseBackend + + if impl_overrides and attn_backend_cls.get_name() in impl_overrides: + raise ValueError("Block-sparse attention does not support dense implementation overrides") + adapter_cls = attn_backend_cls.get_block_sparse_adapter() + if adapter_cls is None: + raise ValueError("Selected provider has no adapter for the shared block selection") + self.attention = BlockSparseAttention( + num_heads, + num_heads if num_kv_heads is None else num_kv_heads, + head_size, + softmax_scale, + causal, + qkv_layout, + spec, + adapter=adapter_cls(), + ) + self.attn_impl_cls = BlockSparseAttention + # Expose the selected method's capabilities, not the dense + # provider's packed/piecewise/paged declarations. + self.attn_backend = BlockSparseBackend + else: + 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, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + qkv_layout=qkv_layout, + prefix=prefix, + backend_kwargs=backend_kwargs, + role=role, + backend_explicit=self.backend_explicit, + ) # Compatibility kernels run inside shared dispatch, between the # parallel strategy's input preparation and output restoration. - self.sdpa_fallback: AttentionImpl | None = SDPABackend.get_impl_cls()( + self.sdpa_fallback = SDPABackend.get_impl_cls()( num_heads=num_heads, head_size=head_size, softmax_scale=softmax_scale, @@ -483,6 +509,9 @@ def _forward_impl( "Only the startup KV memory profile may execute before paged KV initialization." ) use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None + if isinstance(self.attn_spec, BlockSparseAttentionSpec): + if use_paged_attention or (strategy is not self._no_parallel_strategy and strategy.enabled): + raise ValueError("Block-sparse attention does not support paged KV or sequence parallel execution") if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False): backend_name = ( self.attn_backend.get_name() if self.attn_backend is not None else type(self.attention).__name__ diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index cafe4da4666..5df27b61d0a 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/vllm_omni/diffusion/attention/selector.py @@ -21,7 +21,7 @@ ) if TYPE_CHECKING: - from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec + from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseAttentionSpec logger = init_logger(__name__) @@ -98,13 +98,29 @@ def _log_backend_resolution( ) +def _resolve_sparse_backend(spec: BlockSparseAttentionSpec, head_size: int): + """Load the pinned adapter; actual request support is established by preparation. + + Do not translate dependency/initialization failures into compatibility + rejections. Ordered provider selection is not supported by this lifecycle. + """ + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + backend_cls = DiffusionAttentionBackendEnum[spec.backend].get_class() + adapter_cls = backend_cls.get_block_sparse_adapter() + if adapter_cls is None: + raise ValueError(f"{spec.backend}: no adapter for the shared block selection") + adapter_cls.validate_selection(spec.implementation, head_size) + return backend_cls, spec + + def get_attn_backend_for_role( role: str, head_size: int, attention_config: AttentionConfig | None = None, role_category: str | None = None, allow_trtllm_default: bool = True, -) -> tuple[type[AttentionBackend], AttentionSpec | None]: +) -> tuple[type[AttentionBackend], AttentionSpec | BlockSparseAttentionSpec | None]: """ Get attention backend for a specific attention role. @@ -137,7 +153,12 @@ def get_attn_backend_for_role( ) if spec is not None: - backend_cls = _cached_get_backend_cls(spec.backend, head_size) + from vllm_omni.diffusion.data import BlockSparseAttentionSpec + + if isinstance(spec, BlockSparseAttentionSpec): + backend_cls, spec = _resolve_sparse_backend(spec, head_size) + else: + backend_cls = _cached_get_backend_cls(spec.backend, head_size) _log_backend_resolution( role=role, role_category=role_category, @@ -162,7 +183,7 @@ def get_attn_backend_for_capability( attention_config: AttentionConfig | None = None, role_category: str | None = None, ) -> type[AttentionBackend]: - """Resolve the backend class for capability queries such as mask support. + """Resolve the selected method's capability surface, not just its provider. SequenceParallel auto-pad does not know ``head_size``. Explicit ``AttentionConfig`` selections are loaded from the registry without @@ -178,6 +199,12 @@ def get_attn_backend_for_capability( ) if spec is not None: from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + from vllm_omni.diffusion.data import BlockSparseAttentionSpec + + if isinstance(spec, BlockSparseAttentionSpec): + from vllm_omni.diffusion.attention.block_sparse import BlockSparseBackend + _resolve_sparse_backend(spec, HEAD_SIZE_UNKNOWN) + return BlockSparseBackend return DiffusionAttentionBackendEnum[spec.backend.upper()].get_class() return _cached_get_backend_cls(None, HEAD_SIZE_UNKNOWN) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0811d2634d1..0107fa7688d 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1505,7 +1505,7 @@ def _propagate_skip_softmax_calibration(self, tf_config: "TransformerConfig") -> cfg = getattr(self, "diffusion_attention_config", None) if not isinstance(cfg, AttentionConfig): return - specs = [s for s in (cfg.default, *cfg.per_role.values()) if s is not None] + specs = [s for s in (cfg.default, *cfg.per_role.values()) if isinstance(s, AttentionSpec)] from vllm_omni.diffusion.attention.backends.trtllm_calibration import ( propagate_skip_softmax_calibration, @@ -2128,6 +2128,80 @@ def backend_kwargs(self) -> dict[str, Any] | None: return kw or None +@dataclass +class BlockSparseAttentionSpec: + """Configuration for block selection and sparse attention execution.""" + + name: str + config: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + self._validate_fields() + from vllm_omni.diffusion.attention.block_selection.registry import normalize_selection + + self.config = { + "block_size": self._normalize_block_size(), + "selection": normalize_selection(copy.deepcopy(self.config.get("selection"))), + "backend": self._normalize_backend(), + } + + def _validate_fields(self) -> None: + if self.name != "block_sparse": + raise ValueError("BlockSparseAttentionSpec requires name='block_sparse'") + if not isinstance(self.config, Mapping): + raise TypeError("Block-sparse attention config must be a mapping") + unknown = self.config.keys() - {"block_size", "selection", "backend"} + if unknown: + raise ValueError(f"Unknown block_sparse config fields: {sorted(unknown)}") + + def _normalize_block_size(self) -> list[int]: + block_size = self.config.get("block_size", [64, 64]) + if ( + not isinstance(block_size, (list, tuple)) + or len(block_size) != 2 + or any(isinstance(size, bool) or not isinstance(size, int) or size <= 0 for size in block_size) + ): + raise ValueError("block_size must contain two positive integers [Bq, Bkv]") + return list(block_size) + + def _normalize_backend(self) -> dict[str, Any]: + backend = self.config.get("backend", {"require": "FLASH_ATTN"}) + if not isinstance(backend, Mapping) or backend.keys() - {"require", "prefer", "implementation"}: + raise ValueError("Invalid block_sparse backend fields") + if "prefer" in backend: + raise ValueError( + "block_sparse backend.prefer is not supported; use backend.require to pin one provider " + "until request-time provider selection is implemented" + ) + provider = backend.get("require") + if not isinstance(provider, str) or not provider: + raise ValueError("block_sparse backend.require must be a provider name") + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + if provider not in DiffusionAttentionBackendEnum.__members__: + raise ValueError(f"Unknown attention provider: {provider}") + implementation = backend.get("implementation", "auto") + if not isinstance(implementation, str) or not implementation: + raise ValueError("implementation must be a nonempty backend-provided ID or 'auto'") + return copy.deepcopy(dict(backend)) + + @property + def backend(self) -> str: + return self.config["backend"]["require"] + + @property + def implementation(self) -> str: + return self.config["backend"].get("implementation", "auto") + + @property + def block_size(self) -> tuple[int, int]: + return tuple(self.config["block_size"]) + + @property + def selection(self) -> dict[str, Any]: + return copy.deepcopy(self.config["selection"]) + + @dataclass class AttentionConfig: """Per-role attention backend configuration. @@ -2139,14 +2213,14 @@ class AttentionConfig: 4. platform default — unchanged platform logic """ - default: AttentionSpec | None = None - per_role: dict[str, AttentionSpec] = field(default_factory=dict) + default: AttentionSpec | BlockSparseAttentionSpec | None = None + per_role: dict[str, AttentionSpec | BlockSparseAttentionSpec] = field(default_factory=dict) def __post_init__(self) -> None: if self.default is not None: self.default = self._coerce_spec_or_none(self.default, "default") - normalized_per_role: dict[str, AttentionSpec] = {} + normalized_per_role: dict[str, AttentionSpec | BlockSparseAttentionSpec] = {} for role_key, spec_data in self._normalize_per_role_mapping(self.per_role).items(): spec = self._coerce_spec_or_none(spec_data, f"per_role[{role_key!r}]") if spec is not None: @@ -2154,17 +2228,21 @@ def __post_init__(self) -> None: self.per_role = normalized_per_role @staticmethod - def _coerce_spec(spec_data: Any, field_name: str) -> AttentionSpec: - if isinstance(spec_data, AttentionSpec): + def _coerce_spec(spec_data: Any, field_name: str) -> AttentionSpec | BlockSparseAttentionSpec: + if isinstance(spec_data, (AttentionSpec, BlockSparseAttentionSpec)): return spec_data if isinstance(spec_data, str): return AttentionSpec(backend=spec_data) if isinstance(spec_data, Mapping): + if "name" in spec_data or "config" in spec_data: + if spec_data.keys() - {"name", "config"}: + raise ValueError("Cannot mix backend-specific fields with name/config") + return BlockSparseAttentionSpec(**dict(spec_data)) return AttentionSpec(**dict(spec_data)) raise TypeError(f"Expected str, dict, or AttentionSpec for {field_name}, got {type(spec_data)!r}") @classmethod - def _coerce_spec_or_none(cls, spec_data: Any, field_name: str) -> AttentionSpec | None: + def _coerce_spec_or_none(cls, spec_data: Any, field_name: str) -> AttentionSpec | BlockSparseAttentionSpec | None: spec = cls._coerce_spec(spec_data, field_name) if spec.backend.lower() == "auto": return None @@ -2194,7 +2272,7 @@ def _flatten_per_role_entry( normalized[role] = node return - spec_keys = {"backend", "skip_softmax", "quant", "fastvideo_vsa_topk", "block_sparse"} + spec_keys = {"backend", "skip_softmax", "quant", "fastvideo_vsa_topk", "block_sparse", "name", "config"} node_dict = dict(node) node_keys = set(node_dict) if node_keys & spec_keys: @@ -2215,7 +2293,7 @@ def resolve_with_source( self, role: str = "self", role_category: str | None = None, - ) -> tuple[AttentionSpec | None, str | None]: + ) -> tuple[AttentionSpec | BlockSparseAttentionSpec | None, str | None]: """Resolve the AttentionSpec and report which config entry matched.""" spec = self.per_role.get(role) if spec is not None: @@ -2264,7 +2342,7 @@ def parse_attention_config( if fastvideo_vsa_topk is not None: if normalized.default is None: raise ValueError("--fastvideo-vsa-topk requires --diffusion-attention-backend FASTVIDEO_VSA.") - if normalized.default.backend.upper() != "FASTVIDEO_VSA": + if not isinstance(normalized.default, AttentionSpec) or normalized.default.backend.upper() != "FASTVIDEO_VSA": raise ValueError("--fastvideo-vsa-topk is only valid with the FASTVIDEO_VSA backend.") normalized.default.fastvideo_vsa_topk = fastvideo_vsa_topk normalized.default.__post_init__() diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index b060ea4202a..b8f2f8612dc 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -622,6 +622,8 @@ def __init__( num_heads=self.num_heads, head_size=self.head_dim, causal=True, + role="cosmos3.und", + role_category="self", softmax_scale=1.0 / (self.head_dim**0.5), num_kv_heads=self.num_kv_heads, skip_sequence_parallel=True, @@ -730,6 +732,9 @@ def __init__( num_heads=self.num_heads, head_size=self.head_dim, causal=False, + # Keep legacy per_role.self routing; exact roles opt into overrides. + role="cosmos3.gen", + role_category="self", softmax_scale=1.0 / (self.head_dim**0.5), num_kv_heads=self.num_kv_heads, ) @@ -741,6 +746,8 @@ def __init__( num_heads=self.num_heads, head_size=self.head_dim, causal=False, + role="cosmos3.gen_multi_control", + role_category="self", softmax_scale=1.0 / (self.head_dim**0.5), num_kv_heads=self.num_kv_heads, prefix=f"{prefix}.multi_control_attn", @@ -766,7 +773,8 @@ def _forward_local( k_all = torch.cat([k_und, k], dim=1) v_all = torch.cat([v_und, v], dim=1) - out = self.attn(q, k_all, v_all) + metadata = AttentionMetadata(extra={"protected_kv_prefix": k_und.shape[1]}) + out = self.attn(q, k_all, v_all, metadata) return out.reshape(B, S_gen, -1) # -- SP path: framework Attention with joint_key/value ------------------- @@ -806,6 +814,7 @@ def _forward_sp( joint_key=k_und, joint_value=v_und, joint_strategy="front", + extra={"protected_kv_prefix": k_und.shape[1]}, ) out = self.attn(q, k, v, attn_metadata) return out.reshape(B, S_gen, -1) @@ -843,6 +852,7 @@ def _forward_multi_control( v_target = v[:, control_tokens:] control_outputs: list[torch.Tensor] = [] target_output: torch.Tensor | None = None + metadata = AttentionMetadata(extra={"protected_kv_prefix": k_und.shape[1]}) start = 0 for size, weight in zip(control_token_sizes, control_weights, strict=True): if size <= 0: @@ -854,7 +864,7 @@ def _forward_multi_control( q_pair = torch.cat([q_control, q_target], dim=1) k_pair = torch.cat([k_und, k_control, k_target], dim=1) v_pair = torch.cat([v_und, v_control, v_target], dim=1) - pair_output = self.multi_control_attn(q_pair, k_pair, v_pair) + pair_output = self.multi_control_attn(q_pair, k_pair, v_pair, metadata) control_outputs.append(pair_output[:, :size]) weighted_target = pair_output[:, size:] * weight target_output = weighted_target if target_output is None else target_output + weighted_target diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py index fcc19123ba8..49614ef8255 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Cosmos3 Edge transformer variant with a Nemotron dense UND backbone.""" from __future__ import annotations @@ -150,6 +150,8 @@ def __init__( num_heads=self.num_heads, head_size=self.head_dim, causal=True, + role="cosmos3.und", + role_category="self", softmax_scale=1.0 / (self.head_dim**0.5), num_kv_heads=self.num_kv_heads, skip_sequence_parallel=True, diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index f3e1a7432f4..2fd71bdebaf 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -22,7 +22,7 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline from torch import nn -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import BlockSparseAttentionSpec, DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import ( BasePipelineUtils, get_pipeline_utils_for_config, @@ -242,6 +242,15 @@ def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: def _raise_unsupported_features(self) -> None: """Raise an error for incompatible feature switches.""" + attention_config = self.od_config.diffusion_attention_config + specs = [("default", attention_config.default), *attention_config.per_role.items()] + for role, spec in specs: + if isinstance(spec, BlockSparseAttentionSpec): + raise ValueError( + f"Diffusers adapter does not support block_sparse attention ({role}); " + "select a supported dense configuration or use a native model integration." + ) + if self.od_config.diffusion_offload_config is not None: raise NotImplementedError( "diffusion_offload_config is not supported with the diffusers backend. " diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3.py b/vllm_omni/diffusion/models/minimax_h3/fasth3.py index 9f04c204e4b..97a3681142c 100644 --- a/vllm_omni/diffusion/models/minimax_h3/fasth3.py +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3.py @@ -129,7 +129,7 @@ # The attention role MiniMaxH3Attention gives its 50 DiT blocks. The # compression gates live on exactly these layers, so this is the role whose # resolved backend decides whether the artifact runs sparse. -_H3_DIT_ATTENTION_ROLE = "self" +_H3_DIT_ATTENTION_ROLE = "minimax_h3.dit" # Adapter block prefix -> native block prefix. _BLOCK_PREFIXES = ( @@ -138,24 +138,21 @@ ) -def _resolve_dit_attention_backend(od_config: Any) -> str: - """The backend the 50-block H3 DiT will actually resolve to. - - The DiT's attention layers carry role ``"self"``, so a ``per_role`` entry - overrides the default for exactly the layers the compression gates live on. - Reading only the default would accept a config that runs the sparse student - dense, and reject a per-role-only config that is correct. - """ +def _resolve_dit_attention_spec(od_config: Any): + """Resolve the same exact role and category as the DiT attention layers.""" attention_config = getattr(od_config, "diffusion_attention_config", None) - per_role = getattr(attention_config, "per_role", None) or {} - spec = per_role.get(_H3_DIT_ATTENTION_ROLE) + if attention_config is None: + return None + spec, _ = attention_config.resolve_with_source(_H3_DIT_ATTENTION_ROLE, "self") + return spec + + +def _resolve_dit_attention_backend(od_config: Any) -> str: + """The DiT's resolved backend, with support for legacy backend-only configs.""" + spec = _resolve_dit_attention_spec(od_config) if spec is not None: return str(getattr(spec, "backend", "") or "").upper() - backend = str(getattr(od_config, "diffusion_attention_backend", "") or "").upper() - if backend: - return backend - default_spec = getattr(attention_config, "default", None) - return str(getattr(default_spec, "backend", "") or "").upper() + return str(getattr(od_config, "diffusion_attention_backend", "") or "").upper() class FastH3AdapterError(ValueError): diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py b/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py index 0382e92208b..47e12ddc2b1 100644 --- a/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py @@ -15,7 +15,7 @@ from vllm_omni.errors import OmniClientError from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific -from .fasth3 import FASTH3_BASE_MODEL, _resolve_dit_attention_backend +from .fasth3 import FASTH3_BASE_MODEL, _resolve_dit_attention_backend, _resolve_dit_attention_spec if TYPE_CHECKING: from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -96,9 +96,7 @@ def check_serving_contract(self, *, partition: str, od_config: OmniDiffusionConf raise ValueError("FastH3 V2 is a full checkpoint; additional LoRA adapters are unsupported") if _resolve_dit_attention_backend(od_config) != "FASTVIDEO_VSA": raise ValueError("FastH3 V2 requires --diffusion-attention-backend FASTVIDEO_VSA") - attention_config = getattr(od_config, "diffusion_attention_config", None) - per_role = getattr(attention_config, "per_role", None) or {} - spec = per_role.get("self") or getattr(attention_config, "default", None) + spec = _resolve_dit_attention_spec(od_config) if getattr(spec, "fastvideo_vsa_topk", None) is not None: raise ValueError("FastH3 V2 pins VSA sparsity=0.8; remove the fixed fastvideo_vsa_topk override") parallel = getattr(od_config, "parallel_config", None) 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 115264a4896..803b5636c8d 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -597,6 +597,10 @@ def _run_packed_attention( "max_seqlen_q": max_seqlen, "max_seqlen_k": max_seqlen, "valid_kv_length": used, + # Initial sparse experiment: conditioning may be sparsified too. + # Neither valid_kv_length nor the video prefix is a conditioning + # boundary (the latter also includes generated target audio). + "protected_kv_prefix": 0, # Opt the NPU flash backend into the packed varlen path so the # quadratic full_qk mask is never materialized. Ring attention # is excluded: it keeps the aligned padding rows for its @@ -900,6 +904,8 @@ def __init__( arch, quant_config, prefix=f"{prefix}.attn", + role="minimax_h3.dit", + role_category="self", ) self.mlp = MiniMaxH3MLP( arch, From 20d98fa1476aff1af60ea9924f8435f1d36b6690 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Thu, 1 Oct 2026 05:05:57 -0700 Subject: [PATCH 2/8] [Fix] Pin FA4 b33 and validate native sparse geometry Signed-off-by: Rahul Steiger --- docs/design/feature/fa4_subblock_poc.md | 44 ++++++++++++++----- .../attention_backends/dense_backends.md | 3 +- pyproject.toml | 2 +- .../attention/test_block_sparse_adapters.py | 20 ++++++--- .../attention/test_flash_attn_compile.py | 33 ++++++++++++++ 5 files changed, 82 insertions(+), 20 deletions(-) diff --git a/docs/design/feature/fa4_subblock_poc.md b/docs/design/feature/fa4_subblock_poc.md index ff8f1169d75..9bbf4c2e61e 100644 --- a/docs/design/feature/fa4_subblock_poc.md +++ b/docs/design/feature/fa4_subblock_poc.md @@ -13,6 +13,16 @@ interface lets other providers consume the same selection; separate drafts cover [TRTLLM (#8337)](https://github.com/vllm-project/vllm-omni/pull/8337), and [cuDNN (#8338)](https://github.com/vllm-project/vllm-omni/pull/8338). +## Installation + +Install the CUDA 13 extra from this checkout: + +```bash +pip install '.[fa4]' +``` + +This pins the published `flash-attn-4[cu13]==4.0.0b33` wheel. + ## Configuration Pass the JSON contents of a recipe to `--diffusion-attention-config`: @@ -22,7 +32,7 @@ Pass the JSON contents of a recipe to `--diffusion-attention-config`: | `recipes/attention/fa4-subblock.json` | `cosmos3.gen` | Understanding and multi-control | | `recipes/attention/minimax-h3-fa4-subblock.json` | `minimax_h3.dit` | Token refiner | -Both recipes use dense `FLASH_ATTN` by default, 64×64 routing blocks, and +Both recipes target Hopper and use dense `FLASH_ATTN` by default, 64×64 routing blocks, and `target_sparsity: 0.75`. Sparse execution pins the `FLASH_ATTN` adapter with `implementation: auto`, which uses FA4's sparse entry point. FA4 exposes no kernel-ID selector. Provider preference lists are not supported. @@ -77,13 +87,22 @@ output rows are restored as zeros. MiniMax supplies this metadata automatically. ## Current validation -The latest combined run used **vLLM-Omni 0.30.0 on GH200**, PyTorch -2.13.0+cu130, and FA4 revision -`e9cf2c1651d2303191eb40a739a3c135fda00999` -(`4.0.0b33.dev13+ge9cf2c1`) with CuTe DSL 4.7.1. +Validation used the vLLM-Omni 0.30.0 image, published FA4 `4.0.0b33`, +PyTorch `2.13.0+cu130` and CuTe DSL `4.7.1`. + +| GPU | Results | Sparse block size | +| --- | --- | --- | +| GH200 | 745 passed, 12 skipped | 64×64 | +| GB200 (SM100) | 84 passed, 6 skipped | 256×128 | + +Hardware-specific tests account for all skips except two unavailable MiniMax +APIs on GH200. Native dense and sparse compilation checks passed on both GPUs. +For the tested Blackwell path, explicitly set `block_size: [256, 128]`; +the recipes' 64×64 geometry targets Hopper. -**669 tests passed and two were skipped** in 96 seconds at `b076b99fd` -(before this documentation-only update): +These runs validate correctness and compilation. Earlier quality/performance +measurements used a development FA4 revision on Hopper and do not establish +results for this wheel or Blackwell's different selection granularity. ```bash python -m pytest -q \ @@ -91,6 +110,9 @@ python -m pytest -q \ tests/diffusion/attention/test_block_sparse.py \ tests/diffusion/attention/test_block_sparse_owner.py \ tests/diffusion/attention/test_block_sparse_adapters.py \ + tests/diffusion/attention/test_block_sparse_ops.py \ + tests/diffusion/attention/test_flash_attn.py \ + tests/diffusion/attention/test_flash_attn_compile.py \ tests/diffusion/attention/test_attention_config.py \ tests/diffusion/attention/test_selector.py \ tests/diffusion/models/cosmos3/test_cosmos3_transformer.py \ @@ -100,11 +122,9 @@ python -m pytest -q \ tests/diffusion/diffusion_backend/test_diffusers_backend.py ``` -Coverage includes independent routing and selected-key attention references, -head mapping, tails, prefix protection, padding, invalid inputs, owner lifetime, -fullgraph/regional compilation, small Cosmos3 and MiniMax model dispatch, and -dense Diffusers end-to-end execution. The two skips concern MiniMax APIs absent -from this branch. +Coverage includes selected-key numerical references, MHA/MQA/GQA, tails, +prefixes, padding, owner isolation, dynamic/fullgraph compilation, model-role +integration and dense regressions. For an attention-call benchmark, including routing and metadata conversion: diff --git a/docs/user_guide/diffusion/attention_backends/dense_backends.md b/docs/user_guide/diffusion/attention_backends/dense_backends.md index 2cb7019bedb..7e954a25fc0 100644 --- a/docs/user_guide/diffusion/attention_backends/dense_backends.md +++ b/docs/user_guide/diffusion/attention_backends/dense_backends.md @@ -39,8 +39,7 @@ Install the optional CUDA 13 extra: pip install 'vllm-omni[fa4]' ``` -Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on -Blackwell. On Blackwell, `FLASH_ATTN` is FA4 only: Hopper-only FA2/FA3 +The extra pins `4.0.0b33`. On Blackwell, `FLASH_ATTN` is FA4 only: Hopper-only FA2/FA3 wheels are not used, even if they import. If FA4 is missing, explicit `FLASH_ATTN` is rejected and automatic selection continues to other Blackwell backends. diff --git a/pyproject.toml b/pyproject.toml index c65255f9932..8ce7eb1eac1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,7 +120,7 @@ quack = [ # Optional Blackwell-native FlashAttention-4 backend. Keep this opt-in because # the CuTe/CUTLASS dependency is CUDA 13 and Blackwell specific. fa4 = [ - "flash-attn-4[cu13]==4.0.0b18", + "flash-attn-4[cu13]==4.0.0b33", ] # FastVideo VSA kernels for the supported Linux/Python 3.12 wheel environment. diff --git a/tests/diffusion/attention/test_block_sparse_adapters.py b/tests/diffusion/attention/test_block_sparse_adapters.py index 9c027ce1957..38aeeecf345 100644 --- a/tests/diffusion/attention/test_block_sparse_adapters.py +++ b/tests/diffusion/attention/test_block_sparse_adapters.py @@ -21,8 +21,8 @@ def sparse_adapter(request): if not torch.cuda.is_available(): pytest.skip("Requires CUDA") if request.param == "FLASH_ATTN": - if torch.cuda.get_device_capability() != (9, 0): - pytest.skip("These FA4 contract cases are exercised on Hopper") + if torch.cuda.get_device_capability() not in ((9, 0), (10, 0), (10, 3)): + pytest.skip("Requires Hopper or datacenter Blackwell") pytest.importorskip("flash_attn.cute") adapter_cls = DiffusionAttentionBackendEnum[request.param].get_class().get_block_sparse_adapter() assert adapter_cls is not None @@ -42,10 +42,20 @@ def sparse_adapter(request): (4, torch.bfloat16, "full", False), ], ) +@pytest.mark.parametrize("value_size", [64, 128]) @torch.inference_mode() -def test_selected_attention_contract(sparse_adapter, kv_heads, dtype, pattern, compile_case): - block_size = (64, 64) - q, k, v = make_attention_inputs(kv_heads=kv_heads, kv_len=193, value_size=64, dtype=dtype) +def test_selected_attention_contract(sparse_adapter, kv_heads, dtype, pattern, compile_case, value_size): + # FA4 b33 adapts tiles to 64x64 on Hopper. Blackwell defaults to + # tile_n=128 and two 128-row Q stages for these sequence lengths. + blackwell = torch.cuda.get_device_capability()[0] == 10 + block_size = (256, 128) if blackwell else (64, 64) + q, k, v = make_attention_inputs( + kv_heads=kv_heads, + q_len=513 if blackwell else 129, + kv_len=577 if blackwell else 193, + value_size=value_size, + dtype=dtype, + ) q, k, v = (torch.stack((t, t), dim=-1)[..., 0] for t in (q, k, v)) scale = q.shape[-1] ** -0.5 rows = (q.shape[0], q.shape[2], math.ceil(q.shape[1] / block_size[0])) diff --git a/tests/diffusion/attention/test_flash_attn_compile.py b/tests/diffusion/attention/test_flash_attn_compile.py index 5e485bd5b5d..d6c60883321 100644 --- a/tests/diffusion/attention/test_flash_attn_compile.py +++ b/tests/diffusion/attention/test_flash_attn_compile.py @@ -306,3 +306,36 @@ def test_real_fa4_custom_op_unequal_value_dimension_schema(noncontiguous): torch.ops.vllm_omni.fa4_dense_attention.default, (query, key, value, 80**-0.5, False, False), ) + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("kv_heads", [1, 2, 4]) +@pytest.mark.parametrize("head_size", [64, 128]) +@torch.inference_mode() +def test_real_fa4_dense_hopper_fullgraph(monkeypatch, kv_heads, head_size): + """Exercise the optional FA4 wheel on Hopper, where dense defaults to FA3.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Requires Hopper") + from flash_attn.cute import flash_attn_func + from torch.nn.attention import SDPBackend, sdpa_kernel + + # Select the real FA4 entry point without changing production Hopper defaults. + monkeypatch.setattr(fa, "flash_attn_func", flash_attn_func) + + def run(q, k, v): + return torch.ops.vllm_omni.fa4_dense_attention(q, k, v, head_size**-0.5, False, False) + + compiled = torch.compile(run, fullgraph=True, dynamic=True) + for q_len, kv_len in ((129, 257), (257, 193), (129, 257)): + q = torch.randn(1, q_len, 4, head_size, device="cuda", dtype=torch.bfloat16) + k, v = (torch.randn(1, kv_len, kv_heads, head_size, device=q.device, dtype=q.dtype) for _ in range(2)) + with sdpa_kernel(SDPBackend.MATH): + expected = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), + k.transpose(1, 2).float(), + v.transpose(1, 2).float(), + enable_gqa=True, + ).transpose(1, 2) + eager = run(q, k, v) + torch.testing.assert_close(eager.float(), expected, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(compiled(q, k, v), eager) From caa00523c590002a353c700c8b2ff7ba023ae6c2 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Thu, 1 Oct 2026 05:18:18 -0700 Subject: [PATCH 3/8] [Fix] Use rank-local Cosmos3 generation attention heads Signed-off-by: Rahul Steiger --- .../cosmos3/test_cosmos3_transformer.py | 88 +++++++++++++++++++ .../models/cosmos3/transformer_cosmos3.py | 9 +- 2 files changed, 93 insertions(+), 4 deletions(-) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index ab833983ee7..bf7aaf71973 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -1551,3 +1551,91 @@ def test_cosmos3_exact_roles_preserve_legacy_self_category(monkeypatch, override assert tuple(layer.role for layer in layers) == ("cosmos3.und", "cosmos3.gen", "cosmos3.gen_multi_control") assert all(layer.role_category == "self" for layer in layers) assert tuple(layer.attn_backend.get_name() for layer in layers) == expected + + +@pytest.mark.cuda +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("sparse_role", ["cosmos3.gen", "cosmos3.gen_multi_control"]) +@torch.inference_mode() +def test_cosmos3_sparse_tensor_parallel_local_heads(monkeypatch, tp_rank, sparse_role): + """Run real sharded projections and sparse attention; exclude the output collective.""" + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + + from tests.helpers.block_sparse import selected_attention_reference + from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0): + pytest.skip("Requires Hopper SM90 and FA4") + pytest.importorskip("flash_attn.cute") + for module in (linear, parameter): + monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: tp_rank) + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(transformer_cosmos3, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(transformer_cosmos3, "_is_sp_active", lambda: False) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + cfg = SimpleNamespace( + diffusion_attention_config=AttentionConfig( + default={"backend": "FLASH_ATTN"}, + per_role={ + sparse_role: { + "name": "block_sparse", + "config": { + "block_size": [64, 64], + "selection": {"name": "block_topk", "config": {"target_sparsity": 0.75}}, + "backend": {"require": "FLASH_ATTN", "implementation": "auto"}, + }, + } + }, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=2, ring_degree=1), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + with set_current_diffusion_config(cfg): + attention = transformer_cosmos3.Cosmos3CrossAttention( + hidden_size=512, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=128, + rms_norm_eps=1e-6, + ).to(device="cuda", dtype=torch.bfloat16) + # Validate the attention boundary on each TP rank without simulating all-reduce. + attention.to_out = nn.Identity() + for weight in attention.parameters(): + weight.fill_(1) if weight.ndim == 1 else weight.normal_(std=0.02) + for layer in (attention.attn, attention.multi_control_attn): + assert layer.num_heads == 2 and layer.num_kv_heads == 1 + calls = [] + + def check_local_attention(layer, args, output): + q, k, v, metadata = args + assert q.shape[2] == 2 and k.shape[2] == v.shape[2] == 1 + assert metadata.extra["protected_kv_prefix"] == 65 + impl = layer.attention + selection = impl.selector.select(q, k, impl.scale, 65) + expected = selected_attention_reference(q, k, v, selection, impl.scale, impl.block_size) + torch.testing.assert_close(output.float(), expected, atol=0.004, rtol=0.02) + calls.append(layer.role) + + target = attention.attn if sparse_role == "cosmos3.gen" else attention.multi_control_attn + target.register_forward_hook(check_local_attention) + hidden = torch.randn(1, 1153, 512, device="cuda", dtype=torch.bfloat16) + k_und, v_und = (torch.randn(1, 65, 1, 128, device=hidden.device, dtype=hidden.dtype) for _ in range(2)) + cos = torch.ones(1, hidden.shape[1], 1, 128, device=hidden.device, dtype=hidden.dtype) + kwargs = ( + {} + if sparse_role == "cosmos3.gen" + else { + "control_token_sizes": (64, 64), + "control_weights": (0.25, 0.75), + } + ) + output = attention(hidden, k_und, v_und, cos, torch.zeros_like(cos), **kwargs) + assert output.shape == (1, 1153, 256) + assert calls == [sparse_role] * (1 if sparse_role == "cosmos3.gen" else 2) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index b8f2f8612dc..7be296e8fb4 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -728,28 +728,29 @@ def __init__( self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps) self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps) + # Attention consumes the rank-local heads produced by the TP projections. self.attn = FrameworkAttention( - num_heads=self.num_heads, + num_heads=self.num_heads_local, head_size=self.head_dim, causal=False, # Keep legacy per_role.self routing; exact roles opt into overrides. role="cosmos3.gen", role_category="self", softmax_scale=1.0 / (self.head_dim**0.5), - num_kv_heads=self.num_kv_heads, + num_kv_heads=self.num_kv_heads_local, ) # Multi-control attention operates on one full [control_i, target] # sequence at a time. Keep those sequences replicated when Ulysses is # active; sharding the concatenated [control_1, ..., control_N, target] # layout would split the per-control ranges across ranks. self.multi_control_attn = FrameworkAttention( - num_heads=self.num_heads, + num_heads=self.num_heads_local, head_size=self.head_dim, causal=False, role="cosmos3.gen_multi_control", role_category="self", softmax_scale=1.0 / (self.head_dim**0.5), - num_kv_heads=self.num_kv_heads, + num_kv_heads=self.num_kv_heads_local, prefix=f"{prefix}.multi_control_attn", skip_sequence_parallel=True, ) From f5a8344057db0805cfe1a45057e2663f4f5ddb4c Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Thu, 1 Oct 2026 05:25:31 -0700 Subject: [PATCH 4/8] Route sparse selection through platform policy Signed-off-by: Rahul Steiger --- .../feature/attention_backend_selection.md | 13 +- docs/design/feature/fa4_subblock_poc.md | 2 + .../diffusion/attention/test_block_sparse.py | 57 ++++++- tests/diffusion/attention/test_selector.py | 144 +++++++++++++++++- vllm_omni/diffusion/attention/selector.py | 19 +-- vllm_omni/platforms/interface.py | 41 +++++ 6 files changed, 258 insertions(+), 18 deletions(-) diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md index 3e62e2cd924..f55c18d3fe1 100644 --- a/docs/design/feature/attention_backend_selection.md +++ b/docs/design/feature/attention_backend_selection.md @@ -45,7 +45,8 @@ qualified class paths. `register_diffusion_backend()` may replace a path at runtime without changing the public enum value. The active `OmniPlatform` owns compatibility policy through -`get_diffusion_attn_backend_cls()`. It must: +`resolve_diffusion_attn_backend()`. Dense requests delegate to each platform’s +existing `get_diffusion_attn_backend_cls()` implementation. It must: - validate explicit selections and fail with an actionable error when the requested kernel cannot run; @@ -56,6 +57,16 @@ The active `OmniPlatform` owns compatibility policy through The selector must not duplicate device capability or package-availability policy that belongs to the platform. +For `method="block_sparse"`, the platform requires an explicit provider and +checks the shared method's CUDA requirement before loading it. It resolves the +registered backend and calls its adapter's `validate_selection()` with the +unchanged implementation ID and head size (`-1` for capability queries). +Dense availability and head-size checks do not establish sparse support and +are not applied to this path. Missing adapters and provider errors propagate; +there is no fallback. Platforms may override the method-aware entry point. +Actual tensor compatibility is established during sparse request preparation, +not by selection or capability queries. + ## Typed backend options Backend-specific settings remain typed fields on `AttentionSpec` rather than diff --git a/docs/design/feature/fa4_subblock_poc.md b/docs/design/feature/fa4_subblock_poc.md index 9bbf4c2e61e..6ef8883920d 100644 --- a/docs/design/feature/fa4_subblock_poc.md +++ b/docs/design/feature/fa4_subblock_poc.md @@ -70,6 +70,8 @@ The selector and installed kernel must both accept the requested geometry; provider support is checked against the installed implementation rather than a copied hardware/dtype/head-size allowlist. +Sparse provider selection follows the [platform selection contract](attention_backend_selection.md#registry-and-platform-boundary). + `BlockSparseAttention` owns routing, validation, and request preparation. The first actual execution of each request signature runs and synchronizes the kernel before recording support. Capability queries are read-only and report diff --git a/tests/diffusion/attention/test_block_sparse.py b/tests/diffusion/attention/test_block_sparse.py index 62e2fd6d382..2954078f85b 100644 --- a/tests/diffusion/attention/test_block_sparse.py +++ b/tests/diffusion/attention/test_block_sparse.py @@ -522,7 +522,7 @@ def test_dynamic_compilation_does_not_specialize_prepared_geometry(hopper, isola impl = make_impl() # Exceed Dynamo's default recompilation limit without resetting its cache. - inputs = [make_attention_inputs(q_len=length) for length in range(65, 77)] + inputs = [make_attention_inputs(q_len=65 + 64 * index, kv_len=1089 + 64 * index) for index in range(12)] expected = [impl.forward(*tensors) for tensors in inputs] counter = CompileCounterWithBackend("inductor") compiled = torch.compile(impl.forward, backend=counter, fullgraph=True, dynamic=True) @@ -531,7 +531,7 @@ def test_dynamic_compilation_does_not_specialize_prepared_geometry(hopper, isola assert counter.frame_count == 1 # A new geometry is prepared by the runtime handler, without graph fallback # or another graph. Capability queries remain read-only until that execution. - tensors = make_attention_inputs(q_len=77) + tensors = make_attention_inputs(q_len=833, kv_len=1857) context = ExecutionContext(platform="cuda") assert impl.resolve_execution_path(context, *tensors, None).support.status is SupportStatus.UNSUPPORTED actual = compiled(*tensors) @@ -540,6 +540,55 @@ def test_dynamic_compilation_does_not_specialize_prepared_geometry(hopper, isola assert counter.frame_count == 1 +@pytest.mark.cuda +@pytest.mark.parametrize("kv_heads", [1, 2, 4]) +@torch.inference_mode() +def test_fa4_dynamic_adapter_preserves_changing_selections(hopper, isolated_compile_cache, kv_heads): + from torch._dynamo.testing import CompileCounterWithBackend + + from vllm_omni.diffusion.attention.block_selection.abstract import BlockSelection + + adapter = FlashAttentionBackend.get_block_sparse_adapter()() + block_size = (64, 64) + scale = 128**-0.5 + adapter.prepare("auto", 128, 4, kv_heads, torch.device("cuda"), block_size) + + # Scale and block geometry are fixed configuration, not dynamic inputs. + # Passing scale as a float argument makes Dynamo try symbolic scalar + # compilation before retrying with a specialized constant. + def execute(q, k, v, selection): + return adapter.execute(q, k, v, selection, 128**-0.5, (64, 64)) + + counter = CompileCounterWithBackend("inductor") + compiled = torch.compile(execute, backend=counter, fullgraph=True, dynamic=True) + for q_len, kv_len in ((129, 321), (257, 449), (385, 577), (129, 321)): + q, k, v = make_attention_inputs(kv_heads=kv_heads, q_len=q_len, kv_len=kv_len) + rows = (q.shape[0], q.shape[2], math.ceil(q_len / block_size[0])) + first = (torch.arange(math.prod(rows), device=q.device).reshape(rows) % 2).int() + counts = first + 1 + # Include the partial last KV block and poison inactive storage. + last = torch.where(counts == 2, math.ceil(kv_len / block_size[1]) - 1, -999).int() + indices = torch.stack((first, last), dim=-1) + initial_result = None + for active_first in (first, 1 - first, first): + # Same shapes/counts, different per-batch/head/query-row selections. + selection = BlockSelection(torch.stack((active_first, last), dim=-1), counts) + expected = selected_attention_reference(q, k, v, selection, scale, block_size) + eager = adapter.execute(q, k, v, selection, scale, block_size) + actual = compiled(q, k, v, selection) + torch.testing.assert_close(actual.float(), expected, atol=0.004, rtol=0.02) + torch.testing.assert_close(actual, eager, atol=0, rtol=0) + assert actual.is_contiguous() and actual.dtype == q.dtype and actual.device == q.device + if initial_result is None: + initial_result = actual.clone() + elif torch.equal(selection.indices, indices): + torch.testing.assert_close(actual, initial_result, atol=0, rtol=0) + else: + assert not torch.equal(actual, initial_result) + # No compiler retries or recompilation across lengths and selections. + assert counter.frame_count == 1 + + @pytest.mark.cuda @pytest.mark.parametrize("fullgraph", [False, True]) @torch.inference_mode() @@ -611,7 +660,9 @@ def forward(self, q, k, v): model = torch.nn.Module() model._repeated_blocks = ["SparseTestBlock"] model.blocks = torch.nn.ModuleList(SparseTestBlock(i) for i in range(12)) - inputs = [make_attention_inputs(q_len=length) for length in (65, 97)] + # Cross both Q and KV block boundaries, then revisit the first geometry. + inputs = [make_attention_inputs(q_len=q_len, kv_len=kv_len) for q_len, kv_len in ((65, 1089), (193, 1217))] + inputs.append(inputs[0]) expected = [[block(*tensors) for block in model.blocks] for tensors in inputs] counter = CompileCounterWithBackend("inductor") # Exactly the production regional-compilation helper: separate torch.compile diff --git a/tests/diffusion/attention/test_selector.py b/tests/diffusion/attention/test_selector.py index 7b16fb754fc..ae025e102eb 100644 --- a/tests/diffusion/attention/test_selector.py +++ b/tests/diffusion/attention/test_selector.py @@ -105,9 +105,7 @@ def fake_get_cls(**kwargs): platform_calls.append(kwargs) return "fake.module.Backend" - fake_platform = SimpleNamespace(get_diffusion_attn_backend_cls=fake_get_cls) - - monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform) + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.resolve_diffusion_attn_backend", fake_get_cls) def fake_load_backend(path): load_calls.append(path) @@ -136,8 +134,7 @@ def fake_get_cls(**kwargs): platform_calls.append(kwargs) return "fake.module.MustNotBeCalled" - fake_platform = SimpleNamespace(get_diffusion_attn_backend_cls=fake_get_cls) - monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform) + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.resolve_diffusion_attn_backend", fake_get_cls) config = AttentionConfig(default=AttentionSpec(backend="CUDNN_ATTN")) backend = selector.get_attn_backend_for_capability(role="self", attention_config=config) @@ -172,3 +169,140 @@ def test_load_backend_cls_reports_missing_class(monkeypatch): with pytest.raises(AttributeError, match="Class MissingBackend not found in module"): selector._load_backend_cls("fake.module.MissingBackend") + + +@pytest.fixture +def sparse_platform(monkeypatch): + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum + + class Platform(OmniPlatform): + _omni_enum = OmniPlatformEnum.CUDA + + @classmethod + def get_diffusion_attn_backend_cls(cls, *args, **kwargs): + pytest.fail("Sparse resolution must not run dense provider checks") + + calls = [] + + class Adapter: + @staticmethod + def validate_selection(implementation, head_size): + calls.append((implementation, head_size)) + + class Backend(_ConfiguredBackend): + @classmethod + def get_block_sparse_adapter(cls): + return Adapter + + monkeypatch.setattr(DiffusionAttentionBackendEnum, "get_class", lambda self: Backend) + monkeypatch.setattr(DiffusionAttentionBackendEnum, "get_path", lambda self: "fake.module.Backend") + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", Platform) + return Platform, Backend, Adapter, calls + + +@pytest.mark.parametrize("capability_query", [False, True]) +def test_sparse_resolution_uses_platform_and_preserves_method(monkeypatch, sparse_platform, capability_query): + from vllm_omni.diffusion.attention.block_sparse import BlockSparseBackend + + platform, backend, _, calls = sparse_platform + loaded = [] + + def load(path): + loaded.append(path) + return backend + + monkeypatch.setattr(selector, "_load_backend_cls", load) + config = AttentionConfig( + default={ + "name": "block_sparse", + "config": {"backend": {"require": "FLASH_ATTN", "implementation": "future-provider-id"}}, + } + ) + if capability_query: + result = selector.get_attn_backend_for_capability("self", config) + assert result is BlockSparseBackend + else: + result, spec = selector.get_attn_backend_for_role("self", 192, config) + assert result is backend + assert spec is config.default + assert calls == [("future-provider-id", -1 if capability_query else 192)] + assert loaded == ["fake.module.Backend"] + + # A platform-specific rejection must be honored on both selection paths. + def reject(cls, *args, **kwargs): + raise ValueError("platform policy rejected sparse execution") + + monkeypatch.setattr(platform, "resolve_diffusion_attn_backend", classmethod(reject)) + with pytest.raises(ValueError, match="platform policy rejected"): + if capability_query: + selector.get_attn_backend_for_capability("self", config) + else: + selector.get_attn_backend_for_role("self", 192, config) + + +@pytest.mark.parametrize("platform_name", ["ROCM", "NPU", "XPU", "MUSA", "UNSPECIFIED"]) +def test_sparse_platform_rejects_before_loading_provider(monkeypatch, sparse_platform, platform_name): + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + from vllm_omni.platforms.interface import OmniPlatformEnum + + platform, _, _, calls = sparse_platform + monkeypatch.setattr(platform, "_omni_enum", OmniPlatformEnum[platform_name]) + + def unexpected_load(self): + pytest.fail("Unsupported method/platform must fail before importing a provider") + + monkeypatch.setattr(DiffusionAttentionBackendEnum, "get_class", unexpected_load) + with pytest.raises(ValueError, match="block_sparse requires CUDA"): + platform.resolve_diffusion_attn_backend("FLASH_ATTN", 128, method="block_sparse") + assert calls == [] + + +def test_sparse_platform_requires_adapter(monkeypatch, sparse_platform): + platform, backend, _, calls = sparse_platform + monkeypatch.setattr(backend, "get_block_sparse_adapter", classmethod(lambda cls: None)) + with pytest.raises(ValueError, match="FLASH_ATTN: no adapter"): + platform.resolve_diffusion_attn_backend("FLASH_ATTN", 128, method="block_sparse") + assert calls == [] + + +@pytest.mark.parametrize("error_type", [ImportError, ValueError, RuntimeError]) +def test_sparse_platform_preserves_adapter_errors(monkeypatch, sparse_platform, error_type): + platform, _, adapter, _ = sparse_platform + error = error_type("provider failure") + + def fail(*args): + raise error + + monkeypatch.setattr(adapter, "validate_selection", staticmethod(fail)) + with pytest.raises(error_type) as caught: + platform.resolve_diffusion_attn_backend("FLASH_ATTN", 128, method="block_sparse") + assert caught.value is error + + +@pytest.mark.parametrize("backend", [None, "FLASH_ATTN"]) +@pytest.mark.parametrize("allow_default", [False, True]) +def test_method_resolution_preserves_dense_platform_policy(backend, allow_default): + from vllm_omni.platforms.interface import OmniPlatform + + calls = [] + + class Platform(OmniPlatform): + @classmethod + def get_diffusion_attn_backend_cls(cls, selected_backend, head_size, allow_trtllm_default): + calls.append((selected_backend, head_size, allow_trtllm_default)) + return "platform.override.Backend" + + assert Platform.resolve_diffusion_attn_backend(backend, 128, allow_default) == "platform.override.Backend" + assert calls == [(backend, 128, allow_default)] + + +@pytest.mark.parametrize( + ("backend", "method", "message"), + [(None, "block_sparse", "requires an explicit backend"), ("FLASH_ATTN", "unknown", "Unknown diffusion")], +) +def test_method_resolution_rejects_invalid_requests(sparse_platform, backend, method, message): + platform, _, _, calls = sparse_platform + with pytest.raises(ValueError, match=message): + platform.resolve_diffusion_attn_backend(backend, 128, method=method) + assert calls == [] diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index 5df27b61d0a..cf6d92c0644 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/vllm_omni/diffusion/attention/selector.py @@ -65,7 +65,7 @@ def _cached_get_backend_cls( """ from vllm_omni.platforms import current_omni_platform - backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls( + backend_cls_path = current_omni_platform.resolve_diffusion_attn_backend( selected_backend=backend_name, head_size=head_size, allow_trtllm_default=allow_trtllm_default, @@ -104,14 +104,15 @@ def _resolve_sparse_backend(spec: BlockSparseAttentionSpec, head_size: int): Do not translate dependency/initialization failures into compatibility rejections. Ordered provider selection is not supported by this lifecycle. """ - from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum - - backend_cls = DiffusionAttentionBackendEnum[spec.backend].get_class() - adapter_cls = backend_cls.get_block_sparse_adapter() - if adapter_cls is None: - raise ValueError(f"{spec.backend}: no adapter for the shared block selection") - adapter_cls.validate_selection(spec.implementation, head_size) - return backend_cls, spec + from vllm_omni.platforms import current_omni_platform + + backend_cls_path = current_omni_platform.resolve_diffusion_attn_backend( + selected_backend=spec.backend, + head_size=head_size, + method="block_sparse", + implementation=spec.implementation, + ) + return _load_backend_cls(backend_cls_path), spec def get_attn_backend_for_role( diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index d0de4290233..51765b31b4b 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -127,6 +127,47 @@ def get_diffusion_attn_backend_cls( """ raise NotImplementedError + @classmethod + def resolve_diffusion_attn_backend( + cls, + selected_backend: str | None, + head_size: int, + allow_trtllm_default: bool = True, + *, + method: str = "dense", + implementation: str = "auto", + ) -> str: + """Resolve a provider for the requested execution method. + + Dense requests retain each platform's existing selection policy. Sparse + requests use the adapter's selection checks, never dense-kernel support + checks. Actual sparse tensor support remains request-preparation work. + Platforms may override this entry point for method-specific policy. + """ + if method == "dense": + return cls.get_diffusion_attn_backend_cls( + selected_backend=selected_backend, + head_size=head_size, + allow_trtllm_default=allow_trtllm_default, + ) + if method != "block_sparse": + raise ValueError(f"Unknown diffusion attention method: {method}") + if selected_backend is None: + raise ValueError("block_sparse requires an explicit backend") + # The current shared selector/orchestration requires CUDA, independently + # of which provider executes its selected blocks. + if cls._omni_enum != OmniPlatformEnum.CUDA: + raise ValueError(f"block_sparse requires CUDA; current platform is {cls._omni_enum.value}") + + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + backend = DiffusionAttentionBackendEnum[selected_backend.upper()] + adapter_cls = backend.get_class().get_block_sparse_adapter() + if adapter_cls is None: + raise ValueError(f"{backend.name}: no adapter for the shared block selection") + adapter_cls.validate_selection(implementation, head_size) + return backend.get_path() + @classmethod def validate_diffusion_attn_backend(cls, selected_backend: str) -> None: """Reject an explicitly selected backend this platform cannot run. From be9ab0e84808c502bb790dd12cf63a8befc1c72c Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Fri, 2 Oct 2026 06:31:43 -0700 Subject: [PATCH 5/8] Fix mocked Sage dispatch test on unsupported GPUs Signed-off-by: Rahul Steiger --- .../models/minimax_h3/test_minimax_h3_contract.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 23956a58863..390e2ddd7b1 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -3581,6 +3581,7 @@ def diffuse(**kwargs): @torch.inference_mode() def test_unpadded_sage_token_refiner_preserves_legacy_dispatch(monkeypatch): from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention import selector as selector_mod from vllm_omni.diffusion.attention.backends import sage_attn from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention from vllm_omni.diffusion.config import set_current_diffusion_config @@ -3599,6 +3600,14 @@ def test_unpadded_sage_token_refiner_preserves_legacy_dispatch(monkeypatch): monkeypatch.setattr(transformer, "QKVParallelLinear", lambda **kwargs: SimpleNamespace(num_heads=2, num_kv_heads=2)) monkeypatch.setattr(transformer, "RowParallelLinear", lambda *args, **kwargs: nn.Identity()) monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + + # The kernel is mocked, so backend lookup must not depend on the host GPU + # or optional package availability. Keep role resolution and Sage dispatch real. + def resolve_backend(backend_name, head_size): + assert backend_name == "SAGE_ATTN" and head_size == 32 + return sage_attn.SageAttentionBackend + + monkeypatch.setattr(selector_mod, "_cached_get_backend_cls", resolve_backend) calls = [] def sage_kernel(q, k, v, **kwargs): From 73f0c8d68f952cd0acd87a0cf683ec26b64a8bc7 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Tue, 6 Oct 2026 08:48:26 -0700 Subject: [PATCH 6/8] [Core] Support strict Ulysses for static block-sparse attention Allow pure strict Ulysses through the shared Attention wrapper without requiring an attention strategy. Prepare sparse adapters with post-all-to-all head counts, retain full heads for replicated attention, and reject unsupported parallel combinations and inactive sharded execution. Let Cosmos3 remove synthetic suffix padding before sparse selection and restore zero rows before reverse communication. Preserve the protected understanding KV prefix and the ordinary dense mask path. Compose capability queries from actual local tensors without launching kernels or collectives. Validate static per-role sparse configuration against single-device outputs using real NCCL and FA4, padded and unpadded sequences, repeated requests, regional compilation, and model/layerwise CPU offloading. Sequence lengths exceed the selector's minimum retained budget to exercise sparse selection. Validation: - Two Hopper GPUs, FA4 4.0.0b33: 4 passed (156.11s). - Focused sparse attention, SP-hook, capability, adapter, parallel admission, and Cosmos3 regression tests passed after extraction fixes. - Changed-file pre-commit and git diff --check passed. Environment-specific validation launchers and logs remain local. Signed-off-by: Rahul Steiger --- docs/design/feature/fa4_subblock_poc.md | 40 +- .../diffusion/attention/test_block_sparse.py | 29 ++ .../attention/test_block_sparse_parallel.py | 37 ++ .../cosmos3/test_cosmos3_sparse_ulysses.py | 404 ++++++++++++++++++ vllm_omni/diffusion/attention/block_sparse.py | 32 +- vllm_omni/diffusion/attention/layer.py | 76 +++- vllm_omni/diffusion/distributed/sp_plan.py | 4 + .../diffusion/hooks/sequence_parallel.py | 40 +- .../models/cosmos3/transformer_cosmos3.py | 25 +- 9 files changed, 650 insertions(+), 37 deletions(-) create mode 100644 tests/diffusion/attention/test_block_sparse_parallel.py create mode 100644 tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py diff --git a/docs/design/feature/fa4_subblock_poc.md b/docs/design/feature/fa4_subblock_poc.md index 6ef8883920d..5269b800b0c 100644 --- a/docs/design/feature/fa4_subblock_poc.md +++ b/docs/design/feature/fa4_subblock_poc.md @@ -87,8 +87,45 @@ Typed `PackedPaddingMetadata` supports a single document with suffix padding at batch size one. Real Q/K/V lengths are trimmed before routing, and padded output rows are restored as zeros. MiniMax supplies this metadata automatically. +## Strict Ulysses + +The Cosmos3 recipe also supports pure strict Ulysses with ordinary static +`per_role` configuration. Set `ulysses_degree=2` and `ulysses_mode="strict"`. +No attention strategy or denoising schedule is required. Q and KV head counts +must both be divisible by the Ulysses degree. Replicated multi-control attention +keeps its full local heads and does not enter the Ulysses wrapper. + +The shared attention layer uses the existing all-to-all and prepares sparse +kernels with the resulting local head counts. Cosmos3 marks synthetic sequence +padding; it is removed after all-to-all, before block selection, and zero rows +are restored before the reverse communication. Real understanding tokens remain +a protected KV prefix. Other models must supply equivalent padding metadata or +use sequences that need no padding; this does not enable arbitrary sparse masks. + +Capability queries for the composed wrapper require `inputs_are_local=True` +and the actual post-communication tensors. They execute no collectives or +kernels. Local request preparation is still required. The wrapper reports +`EAGER_ONLY`; this does not promise fullgraph capture of communication. Regional +compilation may compile the local work and leave graph breaks around collectives. + +Run the two-rank Cosmos3 validation on two Hopper GPUs with FA4 `4.0.0b33`: + +```bash +CUDA_VISIBLE_DEVICES=0,1 python -m pytest -v \ + tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py +``` + +The tests spawn both ranks themselves. They compare single-device and two-rank +outputs, exercise padded and unpadded requests, verify sparse dispatch per layer, +and cover eager/regional execution with no, model, and layerwise offloading. + ## Current validation +Strict Ulysses: all four cases in `test_cosmos3_sparse_ulysses.py` passed on +two Hopper GPUs with FA4 `4.0.0b33`, including eager/regional execution and both +CPU offload modes. Test lengths exceed the selector's minimum retained-block +budget, so the parity check exercises actual sparse selection. + Validation used the vLLM-Omni 0.30.0 image, published FA4 `4.0.0b33`, PyTorch `2.13.0+cu130` and CuTe DSL `4.7.1`. @@ -143,7 +180,8 @@ because they have not been refreshed for the current implementation. ## Limitations and follow-up Arbitrary masks, causal sparse attention, multi-document packing, -paged or quantized KV, and sequence-parallel wrappers are unsupported. Diffusers +paged or quantized KV, Ring, AllGather-KV, advanced Ulysses, and hybrid +parallel modes are unsupported. Diffusers rejects sparse configurations before loading weights; sparse execution needs a native model integration. diff --git a/tests/diffusion/attention/test_block_sparse.py b/tests/diffusion/attention/test_block_sparse.py index 2954078f85b..078ac5f52f7 100644 --- a/tests/diffusion/attention/test_block_sparse.py +++ b/tests/diffusion/attention/test_block_sparse.py @@ -928,3 +928,32 @@ def test_shared_compilation_boundary_is_custom_op(adapter_mode): ) assert result.support.status is SupportStatus.SUPPORTED assert result.compilation_mode is CompilationMode.CUSTOM_OP + + +@pytest.mark.cuda +@torch.inference_mode() +def test_sparse_ulysses_head_geometry_and_replicated_execution(hopper, monkeypatch): + from types import SimpleNamespace + + from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + from vllm_omni.diffusion.config import set_current_diffusion_config + + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kw: NoParallelAttention()) + config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(default={"name": "block_sparse"}), + parallel_config=SimpleNamespace(ring_degree=1, ulysses_degree=2, sequence_parallel_size=2), + diffusion_kv_cache_dtype=None, + ) + with set_current_diffusion_config(config): + with pytest.raises(ValueError, match="heads divisible"): + layer_mod.Attention(4, 128, False, 0.1, num_kv_heads=1) + sharded = layer_mod.Attention(4, 128, False, 0.1, num_kv_heads=2) + replicated = layer_mod.Attention(4, 128, False, 0.1, num_kv_heads=2, skip_sequence_parallel=True) + assert (sharded.attention.num_heads, sharded.attention.num_kv_heads) == (2, 1) + assert (replicated.attention.num_heads, replicated.attention.num_kv_heads) == (4, 2) + q = torch.randn(1, 64, 4, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 64, 2, 128, device="cuda", dtype=q.dtype) + with pytest.raises(ValueError, match="active sharded region"): + sharded(q, k, k) + assert replicated(q, k, k).shape == q.shape diff --git a/tests/diffusion/attention/test_block_sparse_parallel.py b/tests/diffusion/attention/test_block_sparse_parallel.py new file mode 100644 index 00000000000..d093e01511c --- /dev/null +++ b/tests/diffusion/attention/test_block_sparse_parallel.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Sparse wrapper admission is independent of attention strategies.""" + +from types import SimpleNamespace + +import pytest + +from vllm_omni.diffusion.attention.block_sparse import validate_block_sparse_parallel + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("degree", [1, 2, 4]) +def test_local_or_strict_ulysses(degree): + config = SimpleNamespace(parallel_config=SimpleNamespace(ulysses_degree=degree, sequence_parallel_size=degree)) + assert validate_block_sparse_parallel(config) == degree + + +@pytest.mark.parametrize( + "options", + [ + {"ring_degree": 2}, + {"allgather_degree": 2}, + {"sequence_parallel_size": 4}, + {"ulysses_mode": "advanced_uaa"}, + {"use_hsdp": True}, + {"tensor_parallel_size": 2}, + {"pipeline_parallel_size": 2}, + {"data_parallel_size": 2}, + {"cfg_parallel_size": 2}, + ], +) +def test_unvalidated_parallel_combinations_are_rejected(options): + parallel = SimpleNamespace(**{"ulysses_degree": 2, "sequence_parallel_size": 2, **options}) + with pytest.raises(ValueError, match="Ulysses"): + validate_block_sparse_parallel(SimpleNamespace(parallel_config=parallel)) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py b/tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py new file mode 100644 index 00000000000..686f626459e --- /dev/null +++ b/tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py @@ -0,0 +1,404 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Static sparse Ulysses: native FA4 head shards and two-rank Cosmos3 parity.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.fixture(autouse=True) +def require_hopper(): + from vllm.platforms import current_platform + + capability = current_platform.get_device_capability() + if capability is None or capability.major != 9 or capability.minor != 0: + pytest.skip("Requires Hopper SM90 and the 64x64 FA4 sparse kernel") + pytest.importorskip("flash_attn.cute") + + +def sparse_spec(): + return { + "name": "block_sparse", + "config": { + "block_size": [64, 64], + "selection": {"name": "block_topk", "config": {"target_sparsity": 0.75}}, + "backend": {"require": "FLASH_ATTN", "implementation": "auto"}, + }, + } + + +def move_inputs(values, device): + return { + key: value.to(device=device, dtype=torch.bfloat16 if value.is_floating_point() else value.dtype) + if isinstance(value, torch.Tensor) + else value + for key, value in values.items() + } + + +@torch.inference_mode() +def test_sparse_head_shards_and_padding_on_one_gpu(monkeypatch): + """Validate the sparse math separately from the two-GPU communication gate.""" + from vllm_omni.diffusion.attention import layer as layer_mod + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + from vllm_omni.diffusion.attention.capabilities import ( + CompilationMode, + ExecutionContext, + ParallelStrategy, + SupportStatus, + ) + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig, DiffusionParallelConfig + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kw: NoParallelAttention()) + + def executor(degree): + config = SimpleNamespace( + parallel_config=DiffusionParallelConfig(ulysses_degree=degree), + diffusion_attention_config=AttentionConfig(default=sparse_spec()), + diffusion_kv_cache_dtype=None, + ) + with set_current_diffusion_config(config): + return layer_mod.Attention( + num_heads=4, + num_kv_heads=2, + head_size=128, + causal=False, + softmax_scale=128**-0.5, + ) + + q = torch.randn(1, 1026, 4, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 1029, 2, 128, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + baseline = executor(1) + expected = baseline( + q[:, :-1].contiguous(), + k[:, :-1].contiguous(), + v[:, :-1].contiguous(), + AttentionMetadata(extra={"protected_kv_prefix": 3}), + ) + parts = [] + for rank in range(2): + + class HeadShard(NoParallelAttention): + @property + def enabled(self): + return True + + @property + def name(self): + return "ulysses" + + def pre_attention(self, query, key, value, metadata): + return ( + query[:, :, 2 * rank : 2 * rank + 2].contiguous(), + key[:, :, rank : rank + 1].contiguous(), + value[:, :, rank : rank + 1].contiguous(), + metadata, + None, + ) + + attention = executor(2) + assert attention.attention.num_heads == 2 + assert attention.attention.num_kv_heads == 1 + attention.parallel_strategy = HeadShard() + with override_forward_context(ForwardContext(sp_plan_hooks_applied=True, _sp_shard_depth=1)): + local = ( + q[:, :, 2 * rank : 2 * rank + 2].contiguous()[:, :-1].contiguous(), + k[:, :, rank : rank + 1].contiguous()[:, :-1].contiguous(), + v[:, :, rank : rank + 1].contiguous()[:, :-1].contiguous(), + ) + metadata = AttentionMetadata(extra={"protected_kv_prefix": 3}) + context = ExecutionContext(platform="cuda", require_fullgraph=True) + pending = attention.resolve_execution_path(context, *local, metadata, inputs_are_local=True) + assert pending.support.status is SupportStatus.UNSUPPORTED + assert "not prepared" in pending.support.reason + parts.append( + attention(q, k, v, AttentionMetadata(extra={"protected_kv_prefix": 3, "ulysses_sp_padding": 1})) + ) + result = attention.resolve_execution_path(context, *local, metadata, inputs_are_local=True) + assert result.support.status is SupportStatus.SUPPORTED + assert result.parallel_strategy is ParallelStrategy.ULYSSES + assert result.compilation_mode is CompilationMode.EAGER_ONLY + assert result.requested_support(context).status is SupportStatus.UNSUPPORTED + raw = attention.resolve_execution_path(context, q, k, v, metadata) + assert raw.support.status is SupportStatus.UNSUPPORTED + assert "post-communication" in raw.support.reason + # The local kernel must not claim ownership of the wrapper. + kernel = attention.attention.resolve_execution_path( + ExecutionContext(platform="cuda", parallel_strategy=ParallelStrategy.ULYSSES), *local, metadata + ) + assert kernel.support.status is SupportStatus.UNSUPPORTED + + def unexpected_execution(*args, **kwargs): + raise AssertionError("Capability queries must not execute kernels or collectives") + + monkeypatch.setattr(attention.parallel_strategy, "pre_attention", unexpected_execution) + monkeypatch.setattr(attention.attention, "_execute", unexpected_execution) + assert ( + attention.resolve_execution_path(context, *local, metadata, inputs_are_local=True).support.status + is SupportStatus.SUPPORTED + ) + changed = tuple(t[:, :-1] for t in local) + assert ( + attention.resolve_execution_path(context, *changed, metadata, inputs_are_local=True).support.status + is SupportStatus.UNSUPPORTED + ) + attention.use_ring = True + assert ( + attention.resolve_execution_path(context, *local, metadata, inputs_are_local=True).support.status + is SupportStatus.UNSUPPORTED + ) + attention.use_ring = False + with override_forward_context( + ForwardContext( + omni_diffusion_config=SimpleNamespace(parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa")), + sp_plan_hooks_applied=True, + _sp_shard_depth=1, + ) + ): + rejected = attention.resolve_execution_path(context, *local, metadata, inputs_are_local=True) + assert rejected.support.status is SupportStatus.UNSUPPORTED + assert "strict Ulysses" in rejected.support.reason + with pytest.raises(ValueError, match="strict Ulysses"): + attention(q, k, v, metadata) + actual = torch.cat(parts, dim=2) + torch.testing.assert_close(actual[:, :-1], expected, atol=0.005, rtol=0.016) + assert torch.count_nonzero(actual[:, -1]) == 0 + + +@torch.inference_mode() +def _worker(rank, offload_mode, init_method): + from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules + + bootstrap_vllm_layer_custom_op_modules() + import torch.distributed as dist + from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + + from tests.diffusion.models.cosmos3.test_cosmos3_transformer import _tiny_cosmos3_config + from vllm_omni.diffusion.compile import regionally_compile + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig, DiffusionParallelConfig + from vllm_omni.diffusion.distributed import parallel_state + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.diffusion.hooks import apply_sequence_parallel + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 as cosmos + from vllm_omni.platforms import current_omni_platform + + torch.set_num_threads(1) + torch.manual_seed(0) + device = torch.device(f"cuda:{rank}") + offloader = None + with pytest.MonkeyPatch.context() as patch: + current_omni_platform.set_device(device) + parallel_state.init_distributed_environment( + world_size=2, rank=rank, local_rank=rank, distributed_init_method=init_method, backend="nccl" + ) + parallel_state.initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2) + try: + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + + for module in (linear, parameter): + patch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0) + patch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1) + common = dict( + diffusion_attention_config=AttentionConfig( + default="FLASH_ATTN", + per_role={"cosmos3.gen": sparse_spec()}, + ), + dtype=torch.bfloat16, + diffusion_kv_cache_dtype=None, + ) + steps = (0, 1) + patch.setattr(cosmos, "get_tensor_model_parallel_world_size", lambda: 1) + model_config = _tiny_cosmos3_config(num_hidden_layers=2, hidden_size=16, num_attention_heads=4) + model_config.update( + hidden_size=512, head_dim=128, intermediate_size=512, rope_scaling={"mrope_section": [24, 20, 20]} + ) + + def configuration(degree): + return SimpleNamespace( + **{**common, "parallel_config": DiffusionParallelConfig(ulysses_degree=degree)}, + num_gpus=degree, + tf_model_config=model_config, + enforce_eager=True, + diffusion_compile_granularity="regional", + ) + + def context(config, step): + return override_forward_context( + ForwardContext( + omni_diffusion_config=config, + denoise_step_idx=step, + sp_plan_hooks_applied=config.num_gpus > 1, + ) + ) + + with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda"))): + reference_config, sp_config = configuration(1), configuration(2) + with set_current_diffusion_config(reference_config), context(reference_config, 0): + reference = cosmos.Cosmos3VFMTransformer(reference_config).to(device=device, dtype=common["dtype"]) + reference.post_load_weights() + reference.eval() + for weight in reference.parameters(): + weight.fill_(1) if weight.ndim == 1 else weight.normal_(std=0.02) + with set_current_diffusion_config(sp_config), context(sp_config, 0): + model = cosmos.Cosmos3VFMTransformer(sp_config).to(device=device, dtype=common["dtype"]) + model.load_state_dict(reference.state_dict()) + model.post_load_weights() + model.eval() + apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=2), model._sp_plan) + assert model.gen_layers[0].cross_attention.attn.parallel_strategy.name == "ulysses" + # More than eight candidate blocks: exercise actual sparse selection. + shapes = (32, 33) + inputs = [ + move_inputs( + dict( + hidden_states=torch.randn(1, 2, 1, side, side), + timestep=torch.ones(1), + text_ids=torch.zeros(1, 3, dtype=torch.long), + text_mask=torch.ones(1, 3, dtype=torch.long), + video_shape=(1, side, side), + ), + device, + ) + for side in shapes + ] + references = {} + for request in range(2): + for index, values in enumerate(inputs): + reference.reset_cache() + for step in steps: + with context(reference_config, step): + references[request, index, step] = reference( + **{**values, "text_ids": values["text_ids"] + request} + ) + del reference + identities = {name: id(weight) for name, weight in model.named_parameters()} + if offload_mode == "model": + model.enable_model_cpu_offload(device=device, pin_memory=True) + elif offload_mode == "layerwise": + from vllm_omni.diffusion.offloader.base import OffloadConfig + from vllm_omni.diffusion.offloader.config import OffloadStrategy + from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend + + class Pipeline(torch.nn.Module): + _dit_modules = ["transformer.language_model", "transformer"] + _encoder_modules = [] + _vae_modules = [] + _resident_modules = [] + + pipeline = Pipeline() + pipeline.transformer = model + offloader = LayerWiseOffloadBackend(OffloadConfig(strategy=OffloadStrategy.LAYER_WISE), device) + offloader.enable(pipeline) + assert len(offloader._dit_hooks) == 4 + graphs, executions = [], [] + sparse_calls = [] + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention + from vllm_omni.diffusion.attention.capabilities import ( + ExecutionContext, + ParallelStrategy, + SupportStatus, + ) + + owners = { + id(executor.attention): (index, executor) + for index, block in enumerate(model.gen_layers) + for executor in (block.cross_attention.attn,) + if isinstance(executor.attention, BlockSparseAttention) + } + dispatch = BlockSparseAttention._dispatch_request + + def record_sparse(impl, query, key, value, prefix): + output = dispatch(impl, query, key, value, prefix) + index, executor = owners[id(impl)] + # Observe real execution inside the opaque request op, + # not graph capture or fake-tensor preparation. + result = executor.resolve_execution_path( + ExecutionContext(platform="cuda"), + query, + key, + value, + AttentionMetadata(extra={"protected_kv_prefix": prefix}), + inputs_are_local=True, + ) + assert result.parallel_strategy is ParallelStrategy.ULYSSES + assert result.support.status is SupportStatus.SUPPORTED + sparse_calls.append(index) + return output + + patch.setattr(BlockSparseAttention, "_dispatch_request", record_sparse) + + def backend(graph, example_inputs): + compiled = torch._inductor.compile(graph, example_inputs) + graphs.append(graph) + + def run(*args): + executions.append(graph) + return compiled(*args) + + return run + + for compiled in (False, True): + if compiled: + regionally_compile(model, backend=backend, dynamic=True, fullgraph=False) + for request in range(2): + for index, values in enumerate(inputs): + model.reset_cache() + for step in steps: + before = len(executions) + sparse_calls.clear() + with context(sp_config, step): + actual = model(**{**values, "text_ids": values["text_ids"] + request}) + torch.testing.assert_close( + actual, + references[request, index, step], + atol=0.005, + rtol=0.016, + ) + if compiled: + assert len(executions) > before + expected_calls = [0, 1] + assert sparse_calls == expected_calls, (rank, request, step, sparse_calls) + if request == 0: + graph_count = len(graphs) + else: + assert len(graphs) == graph_count, "Repeated request recompiled warmed regions" + assert identities == {name: id(weight) for name, weight in model.named_parameters()} + assert any( + "vllm_omni.block_sparse_request" in str(node.target) + for graph in graphs + for node in graph.graph.nodes + ) + if offloader is not None: + offloader.disable() + if offload_mode == "model": + model.disable_model_cpu_offload() + dist.barrier() + finally: + torch._dynamo.reset() + parallel_state.destroy_distributed_env() + + +@pytest.mark.parallel +@pytest.mark.parametrize("offload_mode", ["none", "model", "layerwise"]) +def test_static_sparse_ulysses_two_rank_parity(tmp_path, offload_mode): + from vllm_omni.platforms import current_omni_platform + + if current_omni_platform.get_device_count() < 2: + pytest.skip("Requires two Hopper GPUs for NCCL/FA4 validation") + torch.multiprocessing.spawn( + _worker, + args=(offload_mode, f"file://{tmp_path / 'rendezvous'}"), + nprocs=2, + ) diff --git a/vllm_omni/diffusion/attention/block_sparse.py b/vllm_omni/diffusion/attention/block_sparse.py index 10f7dd9a65f..c7ba39a0875 100644 --- a/vllm_omni/diffusion/attention/block_sparse.py +++ b/vllm_omni/diffusion/attention/block_sparse.py @@ -53,6 +53,27 @@ def _block_sparse_request_fake(query, key, value, prefix, owner_handle): return query.new_empty((*query.shape[:-1], value.shape[-1])) +def validate_block_sparse_parallel(config): + """Admit local execution or pure strict Ulysses; kernels consume local heads.""" + parallel = getattr(config, "parallel_config", None) + degree = getattr(parallel, "ulysses_degree", 1) or 1 + if ( + (getattr(parallel, "sequence_parallel_size", degree) or degree) != degree + or any((getattr(parallel, name, 1) or 1) != 1 for name in ("ring_degree", "allgather_degree")) + or getattr(parallel, "use_hsdp", False) + ): + raise ValueError("Block-sparse attention supports only local execution or pure strict Ulysses") + if degree > 1: + if getattr(parallel, "ulysses_mode", "strict") != "strict": + raise ValueError("Block-sparse attention requires strict Ulysses mode") + if any( + (getattr(parallel, name, 1) or 1) != 1 + for name in ("tensor_parallel_size", "pipeline_parallel_size", "data_parallel_size", "cfg_parallel_size") + ) or getattr(config, "enable_distributed_layerwise_offload", False): + raise ValueError("Block-sparse Ulysses does not support combined parallel modes") + return degree + + class BlockSparseBackend(AttentionBackend): """Method-level preflight capabilities, independent of the dense provider. @@ -135,16 +156,7 @@ def __init__( self.head_size = head_size cfg = get_current_diffusion_config_or_none() if cfg is not None: - parallel = getattr(cfg, "parallel_config", None) - if any( - (getattr(parallel, name, 1) or 1) > 1 - for name in ( - "sequence_parallel_size", - "ring_degree", - "allgather_degree", - ) - ) or getattr(parallel, "use_hsdp", False): - raise ValueError("Block-sparse attention parallel wrappers require separate validation") + validate_block_sparse_parallel(cfg) if getattr(cfg, "fa_deterministic", False): raise ValueError("Block-sparse attention does not support fa_deterministic") if getattr(cfg, "diffusion_kv_cache_dtype", None) not in (None, "auto", "float"): diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index a0fa23fd0f8..62312a30c13 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -21,6 +21,8 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend, SDPAImpl from vllm_omni.diffusion.attention.capabilities import ( + CapabilityResult, + CompilationMode, ExecutionContext, ExecutionPathResult, OuterBoundary, @@ -202,16 +204,26 @@ def __init__( self.attn_backend = attn_backend_cls if isinstance(spec, BlockSparseAttentionSpec): - from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention, BlockSparseBackend + from vllm_omni.diffusion.attention.block_sparse import ( + BlockSparseAttention, + BlockSparseBackend, + validate_block_sparse_parallel, + ) if impl_overrides and attn_backend_cls.get_name() in impl_overrides: raise ValueError("Block-sparse attention does not support dense implementation overrides") adapter_cls = attn_backend_cls.get_block_sparse_adapter() if adapter_cls is None: raise ValueError("Selected provider has no adapter for the shared block selection") + degree = validate_block_sparse_parallel(config) + self._sparse_ulysses_degree = 1 if skip_sequence_parallel else degree + if scatter_idx != 2 or gather_idx != 1: + raise ValueError("Block-sparse attention requires BSHD sequence/head axes") + if num_heads % self._sparse_ulysses_degree or self.num_kv_heads % self._sparse_ulysses_degree: + raise ValueError("Block-sparse Ulysses requires Q and KV heads divisible by ulysses_degree") self.attention = BlockSparseAttention( - num_heads, - num_heads if num_kv_heads is None else num_kv_heads, + num_heads // self._sparse_ulysses_degree, + self.num_kv_heads // self._sparse_ulysses_degree, head_size, softmax_scale, causal, @@ -439,6 +451,8 @@ def resolve_execution_path( key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata | None, + *, + inputs_are_local: bool = False, ) -> ExecutionPathResult: """Compose backend capabilities with outer Attention boundaries.""" boundaries = set(context.outer_boundaries) @@ -469,6 +483,27 @@ def resolve_execution_path( ) if not resolved_context.paged_kv: attn_metadata = self._with_kv_cache_dtype(attn_metadata) + if parallel_strategy is ParallelStrategy.ULYSSES and isinstance(self.attn_spec, BlockSparseAttentionSpec): + if get_ulysses_mode(default="strict") != "strict" or not inputs_are_local: + return ExecutionPathResult( + backend="BLOCK_SPARSE", + path="ulysses_local", + platform=resolved_context.platform, + kernel_variant=resolved_context.kernel_variant, + support=CapabilityResult.unsupported( + "Sparse capability queries require strict Ulysses and actual post-communication local inputs" + ), + compilation_mode=CompilationMode.EAGER_ONLY, + parallel_strategy=parallel_strategy, + ) + local = resolver( + replace(resolved_context, parallel_strategy=ParallelStrategy.NONE), + query, + key, + value, + attn_metadata, + ) + return replace(local, parallel_strategy=parallel_strategy, compilation_mode=CompilationMode.EAGER_ONLY) return resolver( resolved_context, query, @@ -509,9 +544,22 @@ def _forward_impl( "Only the startup KV memory profile may execute before paged KV initialization." ) use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None + sparse_ulysses = False if isinstance(self.attn_spec, BlockSparseAttentionSpec): - if use_paged_attention or (strategy is not self._no_parallel_strategy and strategy.enabled): - raise ValueError("Block-sparse attention does not support paged KV or sequence parallel execution") + sparse_ulysses = ( + strategy.name == "ulysses" + and strategy.enabled + and not self.use_ring + and self._sparse_ulysses_degree > 1 + ) + if sparse_ulysses and get_ulysses_mode(default="strict") != "strict": + raise ValueError("Block-sparse attention requires strict Ulysses mode") + if use_paged_attention or ( + strategy is not self._no_parallel_strategy and strategy.enabled and not sparse_ulysses + ): + raise ValueError("Block-sparse attention does not support paged KV or this sequence parallel mode") + if self._sparse_ulysses_degree > 1 and not sparse_ulysses: + raise ValueError("Block-sparse Ulysses requires an active sharded region and initialized Ulysses group") if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False): backend_name = ( self.attn_backend.get_name() if self.attn_backend is not None else type(self.attention).__name__ @@ -535,6 +583,20 @@ def _forward_impl( # For Ring: Concat joint_q query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata) + sparse_sp_padding = 0 + if sparse_ulysses and attn_metadata is not None: + sparse_sp_padding = attn_metadata.extra.get("ulysses_sp_padding", 0) + if sparse_sp_padding: + if ( + type(sparse_sp_padding) is not int + or not 0 < sparse_sp_padding < query.shape[1] + or attn_metadata.joint_strategy != "front" + ): + raise ValueError("Invalid sparse Ulysses suffix padding") + query = query[:, :-sparse_sp_padding].contiguous() + key = key[:, :-sparse_sp_padding].contiguous() + value = value[:, :-sparse_sp_padding].contiguous() + # Scheduler rows describe the logical sequence, while strict Ulysses # may append synthetic tokens solely to make the image shard divisible. # Remove those tokens after the all-to-all and put zero placeholders @@ -590,6 +652,10 @@ def _remove_paged_sp_padding(tensor: torch.Tensor) -> torch.Tensor: output_padding = out.new_zeros((out.shape[0], paged_sp_padding, *out.shape[2:])) out = torch.cat((out[:, :padding_start], output_padding, out[:, padding_start:]), dim=1) + if sparse_sp_padding: + padding = out.new_zeros((out.shape[0], sparse_sp_padding, *out.shape[2:])) + out = torch.cat((out, padding), dim=1) + # 3. Post-processing (Reverse Communication) # For Ulysses: AllToAll Output, and AllGather Joint Output out = strategy.post_attention(out, ctx) diff --git a/vllm_omni/diffusion/distributed/sp_plan.py b/vllm_omni/diffusion/distributed/sp_plan.py index 3e556870609..1eff557e16e 100644 --- a/vllm_omni/diffusion/distributed/sp_plan.py +++ b/vllm_omni/diffusion/distributed/sp_plan.py @@ -225,6 +225,9 @@ class SequenceParallelInput: ForwardContext so models can construct attention masks when needed. Note: Ring attention does not support attention mask, so auto_pad should only be used with Ulysses SP. + mask_free_padding: The consumer removes synthetic padding before attention + and restores it before reverse communication. Skip the backend mask + requirement only under this explicit model contract. shard_group: Optional key shared by tensors representing the same global sequence. Keyed groups track independent padding metadata; omitting it preserves the legacy single-sequence padding behavior. @@ -245,6 +248,7 @@ class SequenceParallelInput: split_output: bool = False auto_pad: bool = False shard_group: str | None = None + mask_free_padding: bool = False def __repr__(self) -> str: return ( diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 32f15de5f3d..c549cc0ed13 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/vllm_omni/diffusion/hooks/sequence_parallel.py @@ -396,6 +396,7 @@ def _maybe_validate_strict_divisibility(*, dim: int, seq_len: int) -> None: x, sp_input.split_dim, sp_input.shard_group, + mask_free_padding=sp_input.mask_free_padding, ) _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim)) return sp_shard(x, sp_input.split_dim, validate=False) @@ -422,6 +423,8 @@ def _shard_with_auto_pad( x: torch.Tensor, dim: int, shard_group: str | None, + *, + mask_free_padding: bool = False, ) -> torch.Tensor: """Pad and shard a tensor while recording its global sequence length. @@ -458,26 +461,27 @@ def _shard_with_auto_pad( # Keyed groups record their length even when no padding is needed. return sp_shard(x, dim, validate=False) - # Check backend compatibility - attention_config = None - if is_forward_context_available(): - od_config = get_forward_context().omni_diffusion_config - if od_config is not None: - attention_config = od_config.diffusion_attention_config + if not mask_free_padding: + # Check backend compatibility + attention_config = None + if is_forward_context_available(): + od_config = get_forward_context().omni_diffusion_config + if od_config is not None: + attention_config = od_config.diffusion_attention_config - attn_backend = get_attn_backend_for_capability( - role="self", - attention_config=attention_config, - ) - attention_spec = None - if attention_config is not None: - attention_spec, _ = attention_config.resolve_with_source(role="self") - if not attn_backend.supports_attention_mask(attention_spec): - raise ValueError( - f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). " - f"Cannot use {attn_backend.get_name()} which does not support attention_mask. " - f"Please switch to SDPA or Ascend attention backend." + attn_backend = get_attn_backend_for_capability( + role="self", + attention_config=attention_config, ) + attention_spec = None + if attention_config is not None: + attention_spec, _ = attention_config.resolve_with_source(role="self") + if not attn_backend.supports_attention_mask(attention_spec): + raise ValueError( + f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). " + f"Cannot use {attn_backend.get_name()} which does not support attention_mask. " + f"Please switch to SDPA or Ascend attention backend." + ) # Ring attention does not support attention_mask if get_ring_parallel_world_size() > 1: diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 7be296e8fb4..6964a8af92f 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -14,6 +14,7 @@ import math from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import replace from typing import TYPE_CHECKING, Any, NamedTuple import torch @@ -34,7 +35,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import BlockSparseAttentionSpec, OmniDiffusionConfig from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm @@ -798,9 +799,12 @@ def _forward_sp( gen_mask = None joint_mask = None + sparse_sp_padding = 0 if is_forward_context_available(): ctx = get_forward_context() - if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + if isinstance(getattr(self.attn, "attn_spec", None), BlockSparseAttentionSpec): + sparse_sp_padding = ctx.sp_padding_size + elif ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size gen_mask = torch.ones(B, padded_seq_len, dtype=torch.bool, device=q.device) gen_mask[:, ctx.sp_original_seq_len :] = False @@ -815,7 +819,7 @@ def _forward_sp( joint_key=k_und, joint_value=v_und, joint_strategy="front", - extra={"protected_kv_prefix": k_und.shape[1]}, + extra={"protected_kv_prefix": k_und.shape[1], "ulysses_sp_padding": sparse_sp_padding}, ) out = self.attn(q, k, v, attn_metadata) return out.reshape(B, S_gen, -1) @@ -1440,6 +1444,21 @@ def __init__( release_module_parameters_to_meta(layer) self.gen_layers.append(layer) + gen_attentions = [ + module + for module in self.gen_layers.modules() + if isinstance(module, FrameworkAttention) and getattr(module, "role", None) == "cosmos3.gen" + ] + if gen_attentions and all(isinstance(attn.attn_spec, BlockSparseAttentionSpec) for attn in gen_attentions): + # Sparse attention removes the synthetic suffix after all-to-all. + self._sp_plan = { + **self._sp_plan, + "gen_sp_prepare": { + index: replace(spec, mask_free_padding=True) + for index, spec in self._sp_plan["gen_sp_prepare"].items() + }, + } + self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None if mixed_precision_config is not None: self.mixed_precision_runtime = Cosmos3MixedPrecisionRuntime(mixed_precision_config) From b95ff49a739f6451558c0ae59473285ae46f45ed Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Tue, 6 Oct 2026 09:06:20 -0700 Subject: [PATCH 7/8] Test sparse Ulysses and offloading for Cosmos3 and MiniMax H3 Signed-off-by: Rahul Steiger --- docs/design/feature/fa4_subblock_poc.md | 33 ++- .../cosmos3/test_cosmos3_sparse_ulysses.py | 22 +- .../test_minimax_h3_sparse_ulysses.py | 262 ++++++++++++++++++ 3 files changed, 300 insertions(+), 17 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py diff --git a/docs/design/feature/fa4_subblock_poc.md b/docs/design/feature/fa4_subblock_poc.md index 5269b800b0c..3cdae0a2034 100644 --- a/docs/design/feature/fa4_subblock_poc.md +++ b/docs/design/feature/fa4_subblock_poc.md @@ -87,19 +87,24 @@ Typed `PackedPaddingMetadata` supports a single document with suffix padding at batch size one. Real Q/K/V lengths are trimmed before routing, and padded output rows are restored as zeros. MiniMax supplies this metadata automatically. -## Strict Ulysses +## Strict Ulysses and CPU offloading -The Cosmos3 recipe also supports pure strict Ulysses with ordinary static +The Cosmos3 and MiniMax H3 recipes support pure strict Ulysses with ordinary static `per_role` configuration. Set `ulysses_degree=2` and `ulysses_mode="strict"`. No attention strategy or denoising schedule is required. Q and KV head counts must both be divisible by the Ulysses degree. Replicated multi-control attention -keeps its full local heads and does not enter the Ulysses wrapper. +keeps its full local heads and does not enter the Ulysses wrapper. MiniMax H3's +token refiner likewise stays replicated and dense. The shared attention layer uses the existing all-to-all and prepares sparse kernels with the resulting local head counts. Cosmos3 marks synthetic sequence padding; it is removed after all-to-all, before block selection, and zero rows are restored before the reverse communication. Real understanding tokens remain -a protected KV prefix. Other models must supply equivalent padding metadata or +a protected KV prefix. MiniMax H3 instead uses its existing single-document +`PackedPaddingMetadata`: padding is trimmed after all-to-all by the sparse +executor and zero output rows are restored before the reverse communication. +Its packed total must divide evenly across ranks, and multi-request packed +batches remain unsupported. Other models must supply equivalent padding metadata or use sequences that need no padding; this does not enable arbitrary sparse masks. Capability queries for the composed wrapper require `inputs_are_local=True` @@ -108,11 +113,17 @@ kernels. Local request preparation is still required. The wrapper reports `EAGER_ONLY`; this does not promise fullgraph capture of communication. Regional compilation may compile the local work and leave graph breaks around collectives. -Run the two-rank Cosmos3 validation on two Hopper GPUs with FA4 `4.0.0b33`: +Both models reuse their existing model and layerwise CPU offloading. These modes +work with sparse attention on one GPU or with pure strict Ulysses, in eager or +regional compilation. Distributed layerwise offloading and combined parallel +modes are outside this validation. + +Run the validation on two Hopper GPUs with FA4 `4.0.0b33`: ```bash CUDA_VISIBLE_DEVICES=0,1 python -m pytest -v \ - tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py + tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py \ + tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py ``` The tests spawn both ranks themselves. They compare single-device and two-rank @@ -121,10 +132,12 @@ and cover eager/regional execution with no, model, and layerwise offloading. ## Current validation -Strict Ulysses: all four cases in `test_cosmos3_sparse_ulysses.py` passed on -two Hopper GPUs with FA4 `4.0.0b33`, including eager/regional execution and both -CPU offload modes. Test lengths exceed the selector's minimum retained-block -budget, so the parity check exercises actual sparse selection. +Ulysses and offloading: all six Cosmos3 cases and five MiniMax H3 cases passed +on Hopper GPUs with FA4 `4.0.0b33`. Coverage includes standalone model and +layerwise CPU offloading, two-rank Ulysses with each offload mode, and eager and +regional compilation. These tests use small randomly initialized models, not +full checkpoints. Test lengths exceed the selector's minimum retained-block +budget, so the parity checks exercise actual sparse selection. Validation used the vLLM-Omni 0.30.0 image, published FA4 `4.0.0b33`, PyTorch `2.13.0+cu130` and CuTe DSL `4.7.1`. diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py b/tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py index 686f626459e..cf20f2006bf 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_sparse_ulysses.py @@ -174,7 +174,7 @@ def unexpected_execution(*args, **kwargs): @torch.inference_mode() -def _worker(rank, offload_mode, init_method): +def _worker(rank, offload_mode, init_method, degree=2): from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules bootstrap_vllm_layer_custom_op_modules() @@ -199,9 +199,9 @@ def _worker(rank, offload_mode, init_method): with pytest.MonkeyPatch.context() as patch: current_omni_platform.set_device(device) parallel_state.init_distributed_environment( - world_size=2, rank=rank, local_rank=rank, distributed_init_method=init_method, backend="nccl" + world_size=degree, rank=rank, local_rank=rank, distributed_init_method=init_method, backend="nccl" ) - parallel_state.initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2) + parallel_state.initialize_model_parallel(sequence_parallel_size=degree, ulysses_degree=degree) try: from vllm.model_executor import parameter from vllm.model_executor.layers import linear @@ -243,7 +243,7 @@ def context(config, step): ) with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda"))): - reference_config, sp_config = configuration(1), configuration(2) + reference_config, sp_config = configuration(1), configuration(degree) with set_current_diffusion_config(reference_config), context(reference_config, 0): reference = cosmos.Cosmos3VFMTransformer(reference_config).to(device=device, dtype=common["dtype"]) reference.post_load_weights() @@ -255,8 +255,9 @@ def context(config, step): model.load_state_dict(reference.state_dict()) model.post_load_weights() model.eval() - apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=2), model._sp_plan) - assert model.gen_layers[0].cross_attention.attn.parallel_strategy.name == "ulysses" + if degree > 1: + apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=degree), model._sp_plan) + assert model.gen_layers[0].cross_attention.attn.parallel_strategy.name == "ulysses" # More than eight candidate blocks: exercise actual sparse selection. shapes = (32, 33) inputs = [ @@ -332,7 +333,9 @@ def record_sparse(impl, query, key, value, prefix): AttentionMetadata(extra={"protected_kv_prefix": prefix}), inputs_are_local=True, ) - assert result.parallel_strategy is ParallelStrategy.ULYSSES + assert result.parallel_strategy is ( + ParallelStrategy.ULYSSES if degree > 1 else ParallelStrategy.NONE + ) assert result.support.status is SupportStatus.SUPPORTED sparse_calls.append(index) return output @@ -402,3 +405,8 @@ def test_static_sparse_ulysses_two_rank_parity(tmp_path, offload_mode): args=(offload_mode, f"file://{tmp_path / 'rendezvous'}"), nprocs=2, ) + + +@pytest.mark.parametrize("offload_mode", ["model", "layerwise"]) +def test_static_sparse_offloading(tmp_path, offload_mode): + _worker(0, offload_mode, f"file://{tmp_path / 'rendezvous'}", degree=1) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py new file mode 100644 index 00000000000..af6c463b5c8 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real sparse MiniMax H3 execution with Ulysses and CPU offloading.""" + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.fixture(autouse=True) +def require_hopper(): + from vllm.platforms import current_platform + + capability = current_platform.get_device_capability() + if capability is None or capability.major != 9 or capability.minor != 0: + pytest.skip("Requires Hopper SM90 and FA4") + pytest.importorskip("flash_attn.cute") + + +def _configuration(degree): + from vllm_omni.diffusion.data import AttentionConfig, DiffusionParallelConfig + + return SimpleNamespace( + tf_model_config=dict( + num_layers=2, + token_refiner_num_layers=1, + hidden_size=256, + num_attention_heads=2, + attention_head_dim=128, + ffn_hidden_size=512, + latents_dim=2, + audio_latents_dim=2, + patch_size=(1, 1, 1), + text_dim=256, + timestep_input_dim=32, + time_embed_hidden_size=64, + time_embed_dim=128, + adaln_out_features=18 * 256, + final_adaln_out_features=2 * 256, + rope_inv_freq_len=16, + ), + diffusion_attention_config=AttentionConfig( + default="FLASH_ATTN", + per_role={ + "minimax_h3.dit": { + "name": "block_sparse", + "config": { + "block_size": [64, 64], + "selection": {"name": "block_topk", "config": {"target_sparsity": 0.75}}, + "backend": {"require": "FLASH_ATTN", "implementation": "auto"}, + }, + } + }, + ), + parallel_config=DiffusionParallelConfig(ulysses_degree=degree), + dtype=torch.bfloat16, + diffusion_kv_cache_dtype=None, + num_gpus=degree, + ) + + +def _inputs(real_length, device): + # Single document followed by suffix padding; both totals divide by U=2. + total = ((real_length + 63) // 64) * 64 + text = torch.arange(3, device=device) + audio = torch.arange(3, 7, device=device) + video = torch.arange(7, real_length, device=device) + tags = torch.full((total,), 2, device=device, dtype=torch.long) + tags[:3], tags[3:7], tags[real_length:] = 0, 1, -1 + return dict( + x=torch.randn(1, total, 2, device=device), + audio_x=torch.randn(1, total, 2, device=device), + img_position_ids=torch.zeros(1, total, 3, device=device, dtype=torch.long), + unique_timesteps=torch.tensor([0.5], device=device), + inverse_indices=torch.zeros(total, device=device, dtype=torch.long), + update_mask=torch.ones(video.numel(), device=device), + token_tags=tags, + prompt_embeds=torch.randn(3, 256, device=device, dtype=torch.bfloat16), + img_pos_info={"position_ids": video}, + audio_pos_info={"position_ids": audio}, + text_pos_info={"position_ids": text}, + img_pos_for_infer_output_info={"position_ids": video}, + packed_seq_params={ + "cu_seqlens_q": torch.tensor([0, real_length, total], device=device, dtype=torch.int32), + "max_seqlen_q": real_length, + }, + refiner_packed_seq_params={ + "cu_seqlens_q": torch.tensor([0, 3], device=device, dtype=torch.int32), + "max_seqlen_q": 3, + }, + ) + + +@torch.inference_mode() +def _run_case(rank, degree, offload_mode): + from vllm import distributed as vllm_distributed + from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + + from vllm_omni.diffusion.attention.block_sparse import BlockSparseAttention + from vllm_omni.diffusion.compile import regionally_compile + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.diffusion.hooks import apply_sequence_parallel + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.offloader.base import OffloadConfig + from vllm_omni.diffusion.offloader.config import OffloadStrategy + from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend + from vllm_omni.diffusion.offloader.sequential_backend import ( + ModelLevelOffloadBackend, + sequential_offload_component, + ) + + torch.set_num_threads(1) + torch.manual_seed(0) + device = torch.device(f"cuda:{rank}") + configs = (_configuration(1), _configuration(degree)) + + def context(config): + return override_forward_context( + ForwardContext( + omni_diffusion_config=config, + sp_plan_hooks_applied=config.num_gpus > 1, + ) + ) + + offloader = None + with ( + pytest.MonkeyPatch.context() as patch, + set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda"))), + ): + patch.setattr(vllm_distributed, "get_tensor_model_parallel_world_size", lambda: 1) + for module in (linear, parameter): + patch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0) + patch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1) + patch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1) + with set_current_diffusion_config(configs[0]), context(configs[0]): + reference = h3.MiniMaxH3DiTModel(configs[0]).to(device=device).eval() + for weight in reference.parameters(): + weight.fill_(1) if weight.ndim == 1 else weight.normal_(std=0.02) + reference.post_load_weights() + inputs = [_inputs(length, device) for length in (1024, 1089)] + with context(configs[0]): + expected = [reference(**values) for values in inputs] + with set_current_diffusion_config(configs[1]), context(configs[1]): + model = h3.MiniMaxH3DiTModel(configs[1]).to(device=device).eval() + model.load_state_dict(reference.state_dict()) + model.post_load_weights() + del reference + if degree > 1: + apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=degree), model._sp_plan) + assert model.token_refiner.blocks[0].attn.attention.skip_sequence_parallel + for block in model.blocks: + assert block.attn.attention.attention.num_heads == 2 // degree + # Use the production pipeline's discovery and offload lifecycle, + # without loading encoder/VAE checkpoints. + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.transformer = model + pipeline._dit_modules = ["transformer"] + pipeline._encoder_modules = [] + pipeline._vae_modules = [] + if offload_mode == "model": + offloader = ModelLevelOffloadBackend(OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL), device) + elif offload_mode == "layerwise": + offloader = LayerWiseOffloadBackend(OffloadConfig(strategy=OffloadStrategy.LAYER_WISE), device) + if offloader: + offloader.enable(pipeline) + assert offloader.enabled + if offload_mode == "layerwise": + assert len(offloader._dit_hooks) == 3 # Two DiT blocks plus token refiner. + dispatch = BlockSparseAttention._dispatch_request + sparse_calls = [] + + def record(impl, q, k, v, prefix): + assert q.shape[1] in (1024, 1089) # Padding removed before selection. + assert q.shape[2] == 2 // degree + sparse_calls.append(q.shape[1]) + return dispatch(impl, q, k, v, prefix) + + patch.setattr(BlockSparseAttention, "_dispatch_request", record) + graphs, executions = [], [] + + def backend(graph, args): + compiled = torch._inductor.compile(graph, args) + graphs.append(graph) + + def execute(*args): + executions.append(graph) + return compiled(*args) + + return execute + + try: + for compiled in (False, True): + if compiled: + regionally_compile(model, backend=backend, dynamic=True, fullgraph=False) + for repeat in range(2): + for values, output in zip(inputs, expected, strict=True): + sparse_calls.clear() + before = len(executions) + lifecycle = sequential_offload_component(model) if offload_mode == "model" else nullcontext() + with context(configs[1]), lifecycle: + actual = model(**values) + torch.testing.assert_close(actual, output, atol=0.005, rtol=0.02) + assert len(sparse_calls) == len(model.blocks) + if compiled: + assert len(executions) > before + if offload_mode == "model": + assert all(p.device.type == "cpu" for p in model.parameters()) + if repeat == 0: + warmed = len(graphs) + else: + assert len(graphs) == warmed + finally: + if offloader: + offloader.disable() + torch._dynamo.reset() + + +def _worker(rank, offload_mode, init_method): + from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules + + bootstrap_vllm_layer_custom_op_modules() + from vllm_omni.diffusion.distributed import parallel_state + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.set_device(torch.device(f"cuda:{rank}")) + parallel_state.init_distributed_environment( + world_size=2, + rank=rank, + local_rank=rank, + distributed_init_method=init_method, + backend="nccl", + ) + parallel_state.initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2) + try: + _run_case(rank, 2, offload_mode) + finally: + parallel_state.destroy_distributed_env() + + +@pytest.mark.parametrize("offload_mode", ["model", "layerwise"]) +def test_sparse_offloading(offload_mode): + _run_case(0, 1, offload_mode) + + +@pytest.mark.parallel +@pytest.mark.parametrize("offload_mode", ["none", "model", "layerwise"]) +def test_sparse_ulysses_offloading(tmp_path, offload_mode): + from vllm_omni.platforms import current_omni_platform + + if current_omni_platform.get_device_count() < 2: + pytest.skip("Requires two Hopper GPUs") + torch.multiprocessing.spawn(_worker, args=(offload_mode, f"file://{tmp_path / 'rendezvous'}"), nprocs=2) From 3afc47b12c6c230a647c95c290c68331e7d388be Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Thu, 8 Oct 2026 02:41:16 -0700 Subject: [PATCH 8/8] Clarify sparse attention composition and Ulysses offload scope Signed-off-by: Rahul Steiger --- .../feature/attention_backend_selection.md | 13 ++- docs/design/feature/fa4_subblock_poc.md | 89 +++++++++++++++++-- .../test_minimax_h3_sparse_ulysses.py | 5 +- 3 files changed, 98 insertions(+), 9 deletions(-) diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md index f55c18d3fe1..59b984186f7 100644 --- a/docs/design/feature/attention_backend_selection.md +++ b/docs/design/feature/attention_backend_selection.md @@ -14,7 +14,8 @@ separate from vLLM's autoregressive attention selector. The selection path has four responsibilities: -1. normalize user configuration into `AttentionConfig` and `AttentionSpec`; +1. normalize user configuration into `AttentionConfig` and a complete-backend + `AttentionSpec` or selected-block `BlockSparseAttentionSpec`; 2. resolve a spec for an attention role; 3. ask the active platform to validate an explicit backend or choose a hardware default; and @@ -29,8 +30,8 @@ The selection path has four responsibilities: 3. `default`; 4. platform default. -An explicit resolution returns both the backend class and its -`AttentionSpec`. A platform-default resolution returns the class and `None`. +An explicit resolution returns both the backend class and its typed spec. +A platform-default resolution returns the class and `None`. Layers must therefore treat the spec as optional and must not infer that a missing spec means SDPA. @@ -67,6 +68,12 @@ there is no fallback. Platforms may override the method-aware entry point. Actual tensor compatibility is established during sparse request preparation, not by selection or capability queries. +For model-facing capability queries, the selected-block method exposes +`BlockSparseBackend` rather than the provider's dense capability surface. The +execution adapter and any strict Ulysses wrapper establish the actual prepared +execution path. Scheduling must preserve that distinction; see the +[proposed composition with RFC #8382](fa4_subblock_poc.md#composition-with-scheduling-and-sparse-plugins). + ## Typed backend options Backend-specific settings remain typed fields on `AttentionSpec` rather than diff --git a/docs/design/feature/fa4_subblock_poc.md b/docs/design/feature/fa4_subblock_poc.md index 3cdae0a2034..b0008e425cd 100644 --- a/docs/design/feature/fa4_subblock_poc.md +++ b/docs/design/feature/fa4_subblock_poc.md @@ -40,6 +40,47 @@ kernel-ID selector. Provider preference lists are not supported. The selected role uses sparse attention on every invocation. There is no layer or denoising-step schedule. Exact roles take precedence over legacy `self` overrides; Cosmos3 multi-control requires its own explicit sparse override. +These are experimental execution examples; 75% is not a quality-qualified default. + +### Composition with scheduling and sparse plugins + +The current configuration selects either an `AttentionSpec` (a complete backend) +or a `BlockSparseAttentionSpec` (shared selection plus a provider-owned execution +adapter) for each role. `AttentionSpec.block_sparse` is an existing typed option +for complete backends such as RainFusion; it is not an alias for +`BlockSparseAttentionSpec`. Mixed `backend` and `name/config` fields are rejected. + +[RFC #8382](https://github.com/vllm-project/vllm-omni/issues/8382) proposes +`AttentionSpec.sparse` with `backend`, `start_step`, and `options`. That field is +not implemented here. The proposed composition, to agree before stabilizing the +public API, is: + +- Normalize the step policy into a dense method, a sparse method, and one shared + schedule. The sparse method may be a complete backend or this selected-block + method. Reuse `BlockSparseAttentionSpec` for the latter's geometry, selector, + and provider options; do not introduce a second set of equivalent knobs. +- Resolve role precedence once, then resolve both methods through platform-owned + dispatch. A provider plugin supplying a selected-block adapter uses the existing + backend registry and `get_block_sparse_adapter()` contract. A complete sparse + method retains its own backend contract and is not assumed to accept + `BlockSelection`. No second provider registry is needed. +- Keep scheduling outside the selector and kernel adapter. Exactly one policy + owns logical denoising progress; reject combinations with an independent backend + schedule unless that backend exposes an explicitly externally scheduled mode. + Sparse steps must not silently fall back to dense execution. +- Preserve model-facing metadata capability queries for every reachable method. + Require metadata valid for both phases (or model-owned phase-specific metadata) + rather than advertising only the dense backend's capabilities. Query execution + support for the chosen method using actual local tensors after communication; + selection-time capability queries launch no kernels or collectives, and sparse + preparation remains an explicit execution requirement. + +The [Cosmos3 strategy PoC (#8583)](https://github.com/vllm-project/vllm-omni/pull/8583) +demonstrates separate presets, layouts, and schedules on this foundation. It is +an example of composition, not an implementation of the proposed +`AttentionSpec.sparse` spelling or a finalized common schema. Parsing, +serialization, plugin dispatch, metadata queries, and missing/boundary step +handling need shared acceptance tests when that schema is implemented. ### What 75% sparsity means @@ -118,6 +159,23 @@ work with sparse attention on one GPU or with pure strict Ulysses, in eager or regional compilation. Distributed layerwise offloading and combined parallel modes are outside this validation. +| Execution combination | Scope and evidence | +| --- | --- | +| Single GPU, no offload | Local FA4 correctness and compilation tests, including dynamic/fullgraph execution | +| Single GPU, model or layerwise CPU offload | Small-model Cosmos3 and MiniMax H3 tests, eager and regional compilation | +| Pure strict Ulysses, no offload | Two-rank Hopper/FA4 small-model parity, padding, and sparse-dispatch tests for both models | +| Pure strict Ulysses, model or layerwise CPU offload | The same two-rank tests with each offload mode, eager and regional compilation | +| Fullgraph capture of the Ulysses wrapper, or full-model compilation with offload | Not established by this PoC; use eager or regional compilation | +| Distributed layerwise offload, Ring, AllGather-KV, advanced Ulysses, HSDP, or hybrid parallelism | Outside supported scope | +| Cosmos3 tensor parallelism | Rank-local projections and attention tested; full TP collectives and generation remain unvalidated | + +The two-rank evidence uses Hopper 64×64 geometry. Local Blackwell 256×128 +adapter tests do not qualify Blackwell Ulysses/offload combinations. Ordinary +layerwise CPU offload on each Ulysses rank is distinct from the distributed +layerwise offloader and its separate weight-transfer policies. With the current ordinary +layerwise offloader, MiniMax streams its DiT blocks and keeps the token refiner +resident on the device. + Run the validation on two Hopper GPUs with FA4 `4.0.0b33`: ```bash @@ -132,14 +190,24 @@ and cover eager/regional execution with no, model, and layerwise offloading. ## Current validation -Ulysses and offloading: all six Cosmos3 cases and five MiniMax H3 cases passed +After merging `main` at `a3d7a0444`, validation on one GH200 with the +vLLM-Omni `0.31.0rc1` ARM64 image and FA4 `4.0.0b33` passed 761 distinct tests +across the regression suite and focused reruns. Coverage includes sparse +selection/adapters, configuration, capability queries, sequence-parallel hooks, +Cosmos3/MiniMax integration, and single-GPU model/layerwise offloading. Eight +tests skipped: six require two GPUs and two concern unavailable MiniMax APIs. +Two-rank execution was not rerun after this merge; the following distributed +results describe the earlier foundation revision. + +Before integration with `main` at `a3d7a0444`, all six Cosmos3 cases and five +MiniMax H3 Ulysses/offloading cases passed on Hopper GPUs with FA4 `4.0.0b33`. Coverage includes standalone model and layerwise CPU offloading, two-rank Ulysses with each offload mode, and eager and regional compilation. These tests use small randomly initialized models, not full checkpoints. Test lengths exceed the selector's minimum retained-block budget, so the parity checks exercise actual sparse selection. -Validation used the vLLM-Omni 0.30.0 image, published FA4 `4.0.0b33`, +The earlier adapter validation used the vLLM-Omni 0.30.0 image, published FA4 `4.0.0b33`, PyTorch `2.13.0+cu130` and CuTe DSL `4.7.1`. | GPU | Results | Sparse block size | @@ -198,6 +266,17 @@ parallel modes are unsupported. Diffusers rejects sparse configurations before loading weights; sparse execution needs a native model integration. -Checkpoint generation quality and end-to-end latency remain unqualified. -A separate attention-strategies RFC will follow for mixing dense and sparse -attention across roles, layers, and denoising steps to improve quality. +Full-checkpoint generation quality remains unqualified for these static recipes. +The [Cosmos3 strategy PoC (#8583)](https://github.com/vllm-project/vllm-omni/pull/8583) +reports preliminary generation latency with the pinned wheel for one prompt and +seed; it does not establish general quality preservation or distributed throughput. + +The [upstream MiniMax H3 evaluation](https://www.lmsys.org/blog/2026-08-27-minimax-h3-h200) +motivates a mixed schedule: ten initial dense steps followed by SubBlock with +64×64 blocks. Its measured speed/similarity trade-off is upstream evidence, not +validation of this port's always-sparse recipes or of other block geometries. +Follow-up qualification should compare matched dense and mixed generations with +the pinned wheel, fixed prompts/inputs and seeds, viewable outputs, similarity +measurements, and repeated generation latency. Record conditioning-token pruning +and qualify Cosmos3 and each advertised geometry separately. Neither universal +quality preservation nor the necessity of a mixed schedule is claimed. diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py index 2f41ee3db40..8c317a68256 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_sparse_ulysses.py @@ -175,7 +175,10 @@ def context(config): offloader.enable(pipeline) assert offloader.enabled if offload_mode == "layerwise": - assert len(offloader._dit_hooks) == 3 # Two DiT blocks plus token refiner. + assert len(offloader._dit_hooks) == len(model.blocks) + # Ordinary layerwise offload keeps non-DiT-block state resident; + # only distributed layerwise offload streams nested submodules. + assert all(p.device == device for p in model.token_refiner.parameters()) dispatch = BlockSparseAttention._dispatch_request sparse_calls = []