[Core] Prototype diffusion attention execution contracts with FA4 (#7379)

tests/diffusion/attention/test_attention_capabilities.py

@@ -0,0 +1,359 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import sys
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_validate_fa4_head_dims = fa.validate_fa4_head_dims
+
+
+@pytest.fixture(autouse=True)
+def kernel_validator(monkeypatch):
+    # Capability unit tests do not require FA4 or a CUDA device. Real-kernel
+    # validation is exercised by test_flash_attn_compile.py.
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+
+def _impl(*, kernel_variant: str | None = "fa4", causal: bool = False):
+    impl = FlashAttentionImpl.__new__(FlashAttentionImpl)
+    impl._kernel_variant = kernel_variant
+    impl.causal = causal
+    return impl
+
+
+def _resolve(
+    impl,
+    *,
+    context: ExecutionContext | None = None,
+    attn_metadata: AttentionMetadata | None = None,
+    dtype=torch.bfloat16,
+):
+    tensor = torch.empty((1, 16, 8, 64), dtype=dtype)
+    return impl.resolve_execution_path(
+        context or ExecutionContext(platform="cuda"),
+        tensor,
+        tensor,
+        tensor,
+        attn_metadata,
+    )
+
+
+class _UnmigratedBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "UNMIGRATED_TEST"
+
+    @staticmethod
+    def get_impl_cls():
+        return None
+
+    @staticmethod
+    def get_metadata_cls():
+        return None
+
+    @staticmethod
+    def get_builder_cls():
+        return None
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return []
+
+
+def test_capability_result_is_tri_state():
+    assert CapabilityResult.supported().status is SupportStatus.SUPPORTED
+    assert CapabilityResult.unsupported("not available").status is SupportStatus.UNSUPPORTED
+    assert CapabilityResult.unmigrated().status is SupportStatus.UNMIGRATED
+    with pytest.raises(ValueError, match="actionable reason"):
+        CapabilityResult(SupportStatus.UNSUPPORTED)
+
+
+def test_unmigrated_backend_default_is_conservative():
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    result = _UnmigratedBackend.resolve_capabilities(context)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    requested = result.requested_support(context)
+    assert requested.status is SupportStatus.UNSUPPORTED
+    assert "no verified fullgraph declaration" in requested.reason
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize("mode", list(CompilationMode))
+def test_fullgraph_request_checks_compilation_mode(mode):
+    result = replace(_resolve(_impl()), compilation_mode=mode)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    expected = SupportStatus.UNSUPPORTED if mode is CompilationMode.EAGER_ONLY else SupportStatus.SUPPORTED
+    assert result.requested_support(context).status is expected
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.SUPPORTED
+
+
+def test_request_preserves_verified_rejection_reason():
+    result = replace(_resolve(_impl()), support=CapabilityResult.unsupported("unsupported head size"))
+    assert result.requested_support(ExecutionContext(platform="cuda", require_fullgraph=True)) == result.support
+
+
+def test_exact_cuda_fa4_dense_path_is_supported():
+    context = ExecutionContext(
+        platform="cuda",
+        require_fullgraph=True,
+    )
+    result = _resolve(_impl(), context=context)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+    assert result.compilation_mode is CompilationMode.CUSTOM_OP
+    assert result.requested_support(context).status is SupportStatus.SUPPORTED
+
+
+def test_causal_fa4_path_remains_unmigrated():
+    result = _resolve(_impl(causal=True))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_fa4_delegates_dimensions_to_kernel(monkeypatch):
+    head_dims = (80, 48)
+    validator = Mock(return_value=True)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", validator)
+    query = torch.empty((1, 16, 8, head_dims[0]), dtype=torch.bfloat16)
+    value = torch.empty((1, 16, 8, head_dims[1]), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, query, value, None)
+    assert result.support.status is SupportStatus.SUPPORTED
+    validator.assert_called_once_with(*head_dims, 8)
+
+
+def test_fa4_reports_kernel_rejection(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", Mock(side_effect=AssertionError("kernel dimension constraint")))
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "kernel dimension constraint" in result.support.reason
+    assert "another backend" in result.support.reason
+
+
+def test_fa4_missing_validator_remains_unmigrated(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: False)
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_fa4_rejects_mismatched_query_key_dimensions():
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    key = torch.empty((1, 16, 8, 32), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, key, key, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "Q and K head dimensions must match" in result.support.reason
+
+
+@pytest.mark.parametrize("arch", [90, 100, 120])
+def test_fa4_validator_adapter_uses_selected_architecture(monkeypatch, arch):
+    validator = Mock()
+    monkeypatch.setitem(
+        sys.modules,
+        "flash_attn.cute.interface",
+        SimpleNamespace(_get_device_arch=lambda: arch, _validate_head_dims=validator),
+    )
+    # Call the real adapter, not the fixture's stand-in.
+    verified = _validate_fa4_head_dims(80, 48, 8)
+    if arch == 120:
+        assert not verified
+        validator.assert_not_called()
+    else:
+        assert verified
+        validator.assert_called_once_with(80, 48, arch // 10, 8)
+
+
+def test_fa4_validator_adapter_handles_missing_private_api(monkeypatch):
+    monkeypatch.setitem(sys.modules, "flash_attn.cute.interface", SimpleNamespace())
+    assert not _validate_fa4_head_dims(64, 64, 8)
+
+
+def test_initialized_kernel_identity_overrides_caller_claim():
+    result = _resolve(
+        _impl(kernel_variant=None),
+        context=ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+        ),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_preconstruction_without_resolved_cuda_kernel_is_unmigrated():
+    result = FlashAttentionBackend.resolve_capabilities(ExecutionContext(platform="cuda"))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.path == "unverified"
+
+
+@pytest.mark.parametrize(
+    "changes",
+    [
+        {"platform": "rocm"},
+        {"paged_kv": True},
+        {"parallel_strategy": ParallelStrategy.ULYSSES},
+        {"parallel_strategy": ParallelStrategy.RING},
+        {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING},
+        {"outer_boundaries": frozenset({OuterBoundary.HSDP})},
+    ],
+)
+def test_unverified_flash_attention_variants_remain_unmigrated(changes):
+    result = _resolve(
+        _impl(),
+        context=replace(ExecutionContext(platform="cuda"), **changes),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_unverified_dtype_remains_unmigrated():
+    result = _resolve(_impl(), dtype=torch.float32)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize(
+    "attn_metadata",
+    [
+        AttentionMetadata(full_attn_spans=[[(0, 8)]]),
+        AttentionMetadata(
+            extra={
+                "cu_seqlens_q": torch.tensor([0, 16], dtype=torch.int32),
+                "cu_seqlens_k": torch.tensor([0, 16], dtype=torch.int32),
+                "max_seqlen_q": 16,
+                "max_seqlen_k": 16,
+            }
+        ),
+        AttentionMetadata(extra={"kv_cache_dtype": "fp8"}),
+    ],
+)
+def test_metadata_normalization_prevents_dense_support(attn_metadata):
+    result = _resolve(_impl(), attn_metadata=attn_metadata)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_unpublished_all_true_mask_remains_runtime_dependent():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "runtime_mask_dependent"
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_producer_published_noop_mask_resolves_dense_fa4():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+
+
+def test_fixed_shape_mask_semantic_change_invalidates_dense_resolution():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    impl = _impl()
+    dense = _resolve(impl, attn_metadata=metadata)
+    assert dense.support.status is SupportStatus.SUPPORTED
+
+    metadata.attn_mask[:, 8:] = False
+    metadata.extra["attention_mask_mode"] = "padding"
+    masked = _resolve(impl, attn_metadata=metadata)
+    assert masked.support.status is SupportStatus.UNMIGRATED
+    request = ExecutionContext(platform="cuda", require_fullgraph=True)
+    assert masked.requested_support(request).status is SupportStatus.UNSUPPORTED
+
+
+def test_piecewise_dispatch_ignores_unused_incomplete_packed_metadata(monkeypatch):
+    from vllm_omni.diffusion.attention.backends import flash_attn
+    from vllm_omni.diffusion.attention.backends.utils import fa
+
+    metadata = AttentionMetadata(full_attn_spans=[[(0, 16)]], extra={"max_seqlen_q": 16})
+    impl = _impl()
+    impl.fa_deterministic = False
+    impl.softmax_scale = 0.125
+    assert _resolve(impl, attn_metadata=metadata).support.status is SupportStatus.UNMIGRATED
+
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", lambda *args, **kwargs: query)
+    monkeypatch.setattr(flash_attn, "piecewise_attn", lambda *args, **kwargs: query)
+    assert impl.forward_cuda(query, query, query, metadata) is query
+
+
+@pytest.mark.parametrize("input_index", [1, 2])
+@pytest.mark.parametrize("mismatch", ["dtype", "device"])
+def test_fa4_rejects_incompatible_key_value(input_index, mismatch):
+    tensors = [torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) for _ in range(3)]
+    tensors[input_index] = (
+        tensors[input_index].to(dtype=torch.float16) if mismatch == "dtype" else tensors[input_index].to("meta")
+    )
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *tensors, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert f"{mismatch}s must match" in result.support.reason
+
+
+@pytest.mark.parametrize("skip_quant", [False, True])
+def test_layer_resolution_uses_effective_kv_quantization(skip_quant):
+    from vllm_omni.diffusion.attention.layer import Attention
+    from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+
+    layer = Attention.__new__(Attention)
+    torch.nn.Module.__init__(layer)
+    layer.attention = _impl(kernel_variant=None)
+    layer._hsdp_compile_boundary_enabled = False
+    layer.skip_sequence_parallel = False
+    layer._no_parallel_strategy = NoParallelAttention()
+    layer.parallel_strategy = layer._no_parallel_strategy
+    layer.use_ring = False
+    layer.paged_kv_cache_role = None
+    layer._kv_cache_dtype = "fp8"
+    layer._disable_kv_quant = False
+    layer._kv_cache_skip_steps = None
+    layer._kv_cache_skip_layers = {0} if skip_quant else None
+    layer.layer_idx = 0
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": "fp8"})
+    result = layer.resolve_execution_path(ExecutionContext(platform="npu"), query, query, query, metadata)
+    assert result.path == ("npu_dense" if skip_quant else "npu_unverified")
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert metadata.extra == {"kv_cache_dtype": "fp8"}

tests/diffusion/attention/test_attention_state_lifetime.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Stateful contract example, independent of an optional attention kernel.
+
+The prepared module owns tensor state and passes it explicitly to an opaque op.
+This deliberately does not model opaque Python/C++ planning handles.
+"""
+
+import gc
+import weakref
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@torch.library.custom_op("vllm_omni_test::planned_attention", mutates_args=())
+def _planned_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_scale: torch.Tensor,
+) -> torch.Tensor:
+    return F.scaled_dot_product_attention(query * query_scale, key, value).contiguous()
+
+
+@_planned_attention.register_fake
+def _planned_attention_fake(query, key, value, query_scale):
+    return query.new_empty((*query.shape[:-1], value.shape[-1]))
+
+
+class _PreparedAttention(torch.nn.Module):
+    """Initialization/planning runs eagerly; forward only executes tensors."""
+
+    def __init__(self, query_scale: torch.Tensor):
+        super().__init__()
+        self.register_buffer("query_scale", query_scale.detach().clone())
+
+    def forward(self, query, key, value):
+        return _planned_attention(query, key, value, self.query_scale)
+
+
+def _inputs(length):
+    generator = torch.Generator().manual_seed(length)
+    return tuple(torch.randn(2, 3, length, 8, generator=generator) for _ in range(3))
+
+
+def _reference(plan, inputs):
+    query, key, value = inputs
+    return F.scaled_dot_product_attention(query * plan.query_scale, key, value)
+
+
+def test_dynamic_replay_uses_each_instances_state_without_recompilation():
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    # Same tensor metadata, different values: state must be a graph input,
+    # never a Python instance identifier or a captured first-instance constant.
+    plans = [_PreparedAttention(torch.full((1, 3, 1, 1), scale)) for scale in (0.5, 2.0)]
+    compiled = torch.compile(
+        lambda plan, q, k, v: plan(q, k, v),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    for length in (5, 9, 13):
+        inputs = _inputs(length)
+        for plan in plans:
+            torch.testing.assert_close(compiled(plan, *inputs), _reference(plan, inputs))
+        assert not torch.allclose(_reference(plans[0], inputs), _reference(plans[1], inputs))
+    assert compile_count == 1
+
+
+def test_compiled_callable_retains_state_until_released():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    plan_ref = weakref.ref(plan)
+    state_ref = weakref.ref(plan.query_scale)
+    inputs = _inputs(5)
+    expected = _reference(plan, inputs)
+    compiled = torch.compile(plan, backend="eager", fullgraph=True, dynamic=True)
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    # Deleting the caller's reference must not invalidate subsequent replay.
+    del plan
+    gc.collect()
+    assert plan_ref() is not None
+    assert state_ref() is not None
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    del compiled
+    gc.collect()
+    assert plan_ref() is None
+    assert state_ref() is None
+
+
+def test_stateful_custom_op_schema_and_fake_implementation():
+    inputs = _inputs(5)
+    # Unequal Q/K lengths and V head dimensions exercise the fake output shape.
+    query, _, _ = inputs
+    key = torch.randn(2, 3, 7, 8)
+    value = torch.randn(2, 3, 7, 4)
+    scale = torch.full((1, 3, 1, 1), 0.5)
+    torch.library.opcheck(_planned_attention, (query, key, value, scale))
+
+
+def test_stateful_attention_compiles_with_inductor():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    compiled = torch.compile(plan, fullgraph=True, dynamic=True)
+    for length in (5, 9):
+        inputs = _inputs(length)
+        torch.testing.assert_close(compiled(*inputs), _reference(plan, inputs))

tests/diffusion/attention/test_flash_attn.py

@@ -21,6 +21,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
+from vllm_omni.diffusion.attention.capabilities import CompilationMode, ExecutionContext, SupportStatus
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -432,6 +433,35 @@ def test_varlen_masked_routing_by_role(monkeypatch):
     assert cu_seqlens_k.tolist() == [0, 2, 5]
 
 
+@pytest.mark.cpu
+def test_published_noop_mask_uses_dense_path(monkeypatch):
+    calls = []
+
+    def fake_dense_func(q, k, v, **_kwargs):
+        calls.append("dense")
+        return q
+
+    def fail_varlen(*_args, **_kwargs):
+        raise AssertionError("published no-op mask must not use varlen attention")
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_dense_func)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fail_varlen)
+
+    impl = FlashAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 4, 2, 8)
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 4), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert torch.equal(output, query)
+    assert calls == ["dense"]
+
+
 @pytest.mark.cpu
 def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     calls = []
@@ -550,6 +580,199 @@ def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None):
     return fake
 
 
+@pytest.fixture
+def npu_contract(monkeypatch, mocker):
+    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("NPU must not validate FA4"))
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(return_value=torch.zeros_like(query.squeeze(0)))
+    _fake_mindiesd(monkeypatch, attention_forward=dense, attention_forward_varlen=packed)
+    return _npu_impl(), query, dense, packed
+
+
+def _assert_platform_path(impl, platform, query, metadata, path):
+    context = ExecutionContext(platform=platform, require_fullgraph=True)
+    result = impl.resolve_execution_path(context, query, query, query, metadata)
+    assert result.path == path
+    assert result.kernel_variant == {"npu": "mindiesd", "rocm": "aiter"}[platform]
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+def _real_pad_metadata():
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    return AttentionMetadata(
+        extra={
+            "npu_attn_varlen": True,
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 5,
+            "max_seqlen_k": 5,
+        }
+    )
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("kv_cache_dtype", [None, "float", "auto"])
+def test_npu_contract_dense_dispatch(npu_contract, kv_cache_dtype):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": kv_cache_dtype})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_dense")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_mask_layout(npu_contract):
+    impl, query, dense, packed = npu_contract
+    mask = torch.arange(8)[None] < 5
+    metadata = AttentionMetadata(attn_mask=mask)
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask[:, None, None, :].expand(1, 1, 8, 8))
+
+
+@pytest.mark.cpu
+def test_npu_contract_packed_dispatch(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[3:] == ([0, 5, 8], [0, 5, 8])
+
+
+@pytest.mark.cpu
+def test_npu_contract_prefix_dispatch(npu_contract, monkeypatch):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    # The same metadata resolves differently after the kernel choice changes.
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    _assert_platform_path(impl, "npu", query, metadata, "npu_prefix_kv_slice")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.args[1], query[:, :5])
+    torch.testing.assert_close(dense.call_args.args[2], query[:, :5])
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_incomplete_packing_rebuilds_mask(npu_contract):
+    impl, query, dense, packed = npu_contract
+    # CUDA-like normalization rejects this metadata, but NPU can rebuild a mask.
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True, "max_seqlen_q": 5, "valid_kv_length": 5})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    mask = (torch.arange(8) < 5)[None, None, None, :].expand(1, 1, 8, 8)
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask)
+
+
+@pytest.mark.cpu
+def test_npu_contract_unusable_packing_does_not_run_unmasked(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    with pytest.raises(ValueError, match="refusing to run unmasked attention"):
+        impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_not_called()
+
+
+@pytest.fixture
+def rocm_contract(monkeypatch, mocker):
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(side_effect=lambda *args, **kwargs: torch.zeros_like(args[0] if args else kwargs["q"]))
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", True)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", dense)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", packed)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("ROCm must not validate FA4"))
+    return _npu_impl(), query, dense, packed
+
+
+@pytest.mark.cpu
+def test_rocm_contract_dense_dispatch(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    _assert_platform_path(impl, "rocm", query, None, "rocm_dense")
+    impl.forward_hip(query, query, query)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_masked_unpadding(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.arange(8)[None] < 5, extra={"attention_mask_mode": "padding"})
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_masked_varlen")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[0].shape == (5, 2, 4)
+    torch.testing.assert_close(packed.call_args.kwargs["cu_seqlens_k"], torch.tensor([0, 5], dtype=torch.int32))
+
+
+@pytest.mark.cpu
+def test_rocm_contract_multiple_packed_documents(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    cu = torch.tensor([0, 2, 4, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 4,
+            "max_seqlen_k": 4,
+        }
+    )
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_packed_varlen")
+    # NPU's [real, pad] path cannot consume this packing contract.
+    assert impl._resolve_packed_seq_npu(query, query, metadata.extra) is None
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.kwargs["cu_seqlens_q"] is cu
+    assert packed.call_args.kwargs["cu_seqlens_k"] is cu
+
+
+@pytest.mark.cpu
+def test_rocm_contract_unknown_mask_stays_runtime_dependent(rocm_contract, monkeypatch):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.ones(1, 8, dtype=torch.bool))
+    # Resolution may inspect shape/metadata, but must not read mask values.
+    with monkeypatch.context() as patch:
+        patch.setattr(torch, "any", lambda *_args, **_kwargs: pytest.fail("resolution read mask values"))
+        _assert_platform_path(impl, "rocm", query, metadata, "rocm_runtime_mask_dependent")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_requires_initialized_aiter(monkeypatch):
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", False)
+    impl = _npu_impl()
+    query = torch.zeros(1, 8, 2, 4)
+    result = impl.resolve_execution_path(
+        ExecutionContext(platform="rocm", kernel_variant="aiter"), query, query, query, None
+    )
+    assert result.path == "rocm_unverified"
+    assert result.kernel_variant is None
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 
 

tests/diffusion/attention/test_flash_attn_compile.py

@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import ForwardContext
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+
+@pytest.fixture(autouse=True)
+def isolated_compiler_cache():
+    # Each case changes kernel/configuration on the same Attention.forward code
+    # object. Keep unrelated cases from exhausting its Dynamo recompile limit;
+    # retain the cache across all shapes and replays within an individual test.
+    torch.compiler.reset()
+    yield
+    torch.compiler.reset()
+
+
+@dataclass
+class _ParallelStrategyStub:
+    name: str
+    enabled: bool = True
+
+
+def _make_attention_layer(monkeypatch, *, head_size=64):
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    monkeypatch.setattr(
+        attention_layer,
+        "get_attn_backend_for_role",
+        lambda **_kwargs: (FlashAttentionBackend, None),
+    )
+    monkeypatch.setattr(
+        attention_layer,
+        "build_parallel_attention_strategy",
+        lambda **_kwargs: NoParallelAttention(),
+    )
+    return attention_layer.Attention(num_heads=8, head_size=head_size, softmax_scale=head_size**-0.5, causal=False)
+
+
+@pytest.fixture
+def fake_fa4(monkeypatch, tmp_path):
+    marker = tmp_path / "kernel"
+    marker.write_text("loaded", encoding="utf-8")
+
+    @functools.cache
+    def cached_kernel_loader():
+        with open(marker, encoding="utf-8") as handle:
+            handle.read()
+
+    def fake_attention(query, key, value, **kwargs):
+        cached_kernel_loader()
+        return (
+            torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2),
+                key.transpose(1, 2),
+                value.transpose(1, 2),
+                scale=kwargs["softmax_scale"],
+                is_causal=kwargs["causal"],
+            )
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", True)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_attention)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+    return FlashAttentionImpl(
+        num_heads=8,
+        head_size=64,
+        softmax_scale=0.125,
+        causal=False,
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_fa4_dense_dispatch_is_opaque_to_dynamic_torch_compile(fake_fa4, noncontiguous):
+    """Fast Dynamo/custom-op regression; Inductor coverage is separate."""
+    impl = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q = q.transpose(1, 2).contiguous().transpose(1, 2)
+        assert not q.is_contiguous()
+    path = impl.resolve_execution_path(
+        ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+            dtype="bfloat16",
+            require_fullgraph=True,
+        ),
+        q,
+        q,
+        q,
+        None,
+    )
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    compiled = torch.compile(
+        lambda query, key, value: impl.forward_cuda(query, key, value),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    out = compiled(q, q, q)
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q2 = q2.transpose(1, 2).contiguous().transpose(1, 2)
+    out2 = compiled(q2, q2, q2)
+    assert out.is_contiguous()
+    assert out2.is_contiguous()
+
+    assert out.shape == q.shape
+    assert out2.shape == q2.shape
+    torch.testing.assert_close(out, impl.forward_cuda(q, q, q))
+    torch.testing.assert_close(out2, impl.forward_cuda(q2, q2, q2))
+    assert compile_count == 1
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_fa4_production_attention_entry_compiles_with_inductor(
+    monkeypatch,
+    fake_fa4,
+):
+    """Compile through Attention.forward with the default Inductor backend."""
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    layer = _make_attention_layer(monkeypatch)
+    layer.attention = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+
+    resolved_contexts = []
+    resolve_execution_path = fake_fa4.resolve_execution_path
+
+    def record_execution_path(context, query, key, value, attn_metadata):
+        resolved_contexts.append(context)
+        return resolve_execution_path(
+            context,
+            query,
+            key,
+            value,
+            attn_metadata,
+        )
+
+    monkeypatch.setattr(fake_fa4, "resolve_execution_path", record_execution_path)
+    context = ExecutionContext(
+        platform="cuda",
+        kernel_variant="fa4",
+        dtype="bfloat16",
+        paged_kv=True,
+        parallel_strategy=ParallelStrategy.ULYSSES,
+        require_fullgraph=True,
+    )
+    path = layer.resolve_execution_path(context, q, q, q, None)
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+    assert resolved_contexts[-1].paged_kv is False
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    out = compiled(q, q, q)
+
+    assert out.shape == q.shape
+    torch.testing.assert_close(out, layer(q, q, q))
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    torch.testing.assert_close(compiled(q2, q2, q2), layer(q2, q2, q2))
+
+    layer.paged_kv_cache_role = "self"
+    monkeypatch.setattr(layer, "_active_paged_kv_adapter", lambda: object())
+    paged_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].paged_kv is True
+    assert paged_path.support.status is SupportStatus.UNMIGRATED
+    layer.paged_kv_cache_role = None
+
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ulysses")
+    hybrid_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.HYBRID_ULYSSES_RING
+    assert hybrid_path.support.status is SupportStatus.UNMIGRATED
+    layer.use_ring = False
+    layer.parallel_strategy = NoParallelAttention()
+
+    layer._hsdp_compile_boundary_enabled = True
+    hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert hsdp_path.support.status is SupportStatus.UNMIGRATED
+    assert hsdp_path.compilation_mode is CompilationMode.EAGER_ONLY
+
+    layer._hsdp_compile_boundary_enabled = False
+    runtime_context = ForwardContext(
+        omni_diffusion_config=OmniDiffusionConfig(
+            parallel_config=DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=1),
+        ),
+    )
+    monkeypatch.setattr(attention_layer, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(attention_layer, "get_forward_context", lambda: runtime_context)
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ring")
+    runtime_hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+    assert runtime_hsdp_path.support.status is SupportStatus.UNMIGRATED
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(32, 32), (64, 64), (80, 48), (192, 128), (256, 256)])
+def test_real_fa4_fullgraph_matches_sdpa_across_shapes(monkeypatch, head_dims):
+    """Validate real FA4 numerics through the complete compiled attention layer."""
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major < 10:
+        pytest.skip("Requires Blackwell")
+    if not fa.IS_FLASH_ATTN_4 or fa.flash_attn_func is None:
+        pytest.skip("Requires CuTe FlashAttention-4")
+
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    # Tile boundaries, cross-attention lengths, batching, and a longer sequence.
+    for batch, q_length, kv_length in ((1, 16, 16), (1, 129, 257), (2, 257, 129), (1, 1024, 1024)):
+        query, key, value = (
+            torch.randn(batch, length, 8, dim, device="cuda", dtype=torch.bfloat16, generator=generator)
+            for length, dim in ((q_length, head_size), (kv_length, head_size), (kv_length, value_head_size))
+        )
+        path = layer.resolve_execution_path(context, query, key, value, None)
+        assert path.requested_support(context).status is SupportStatus.SUPPORTED
+        with sdpa_kernel(SDPBackend.MATH):
+            reference = torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2).float(),
+                key.transpose(1, 2).float(),
+                value.transpose(1, 2).float(),
+                scale=head_size**-0.5,
+            ).transpose(1, 2)
+        actual = compiled(query, key, value)
+        assert actual.dtype == query.dtype
+        assert actual.device == query.device
+        torch.testing.assert_close(actual.float(), reference, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(actual, layer(query, key, value))
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(192, 192), (65, 64), (64, 65)])
+def test_real_fa4_dimension_rejection_matches_kernel(monkeypatch, head_dims):
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major not in (10, 11) or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires SM100/SM110 CuTe FlashAttention-4")
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    query = torch.randn(1, 16, 8, head_size, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 16, 8, value_head_size, device="cuda", dtype=torch.bfloat16)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    path = layer.resolve_execution_path(context, query, query, value, None)
+    assert path.requested_support(context).status is SupportStatus.UNSUPPORTED
+    with pytest.raises(AssertionError) as error:
+        layer(query, query, value)
+    assert str(error.value) in path.support.reason
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_real_fa4_custom_op_unequal_value_dimension_schema(noncontiguous):
+    if not current_omni_platform.is_cuda() or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires CuTe FlashAttention-4")
+    query = torch.randn(1, 17, 8, 80, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(1, 25, 8, 80, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 25, 8, 48, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        query, key, value = (t.transpose(1, 2).contiguous().transpose(1, 2) for t in (query, key, value))
+        assert not query.is_contiguous()
+    torch.library.opcheck(
+        torch.ops.vllm_omni.fa4_dense_attention.default,
+        (query, key, value, 80**-0.5, False, False),
+    )

tests/diffusion/test_attention_compile.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
 
 from types import SimpleNamespace
 
@@ -14,6 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = False
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -38,6 +39,7 @@ def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
 
 def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = False
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -60,8 +62,10 @@ def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
     assert calls == ["impl"]
 
 
-def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
+@pytest.mark.parametrize("initialized_with_hsdp", [False, True])
+def test_attention_compile_boundary_without_diffusion_config(monkeypatch, initialized_with_hsdp):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = initialized_with_hsdp
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -80,4 +84,4 @@ def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
     with set_forward_context(vllm_config=None):
         assert Attention.forward(attention, query, query, query) is query
 
-    assert calls == ["impl"]
+    assert calls == (["boundary"] if initialized_with_hsdp else ["impl"])