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

docs/design/feature/attention_backend_selection.md

@@ -4,6 +4,9 @@ This document defines the selection and extension contract for diffusion
 attention backends. User-facing backend choices, installation, and tuning are
 in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).
 
+The [execution contract PoC](attention_execution_contract_poc.md) describes the
+experimental post-selection capability and lifecycle contract from RFC #7226.
+
 ## Scope
 
 The contract applies to DiT and other diffusion attention layers. It is

docs/design/feature/attention_execution_contract_poc.md

@@ -0,0 +1,133 @@
+# Diffusion attention execution contract PoC
+
+This prototype accompanies [RFC #7226](https://github.com/vllm-project/vllm-omni/issues/7226)
+and extends the [attention selection design](attention_backend_selection.md).
+It demonstrates path-specific capabilities with dense BF16 FA4, NPU/ROCm
+routing examples, and tensor-state lifetime with a test-only attention module.
+
+## Execution contract
+
+`ExecutionContext` describes the requested execution path. `ExecutionPathResult`
+reports its identity, support status and reason, and compilation mode.
+`AttentionBackend.resolve_capabilities()` provides conservative pre-construction
+results. After initialization, `Attention.resolve_execution_path()` supplies the
+active parallel, paged-KV, and HSDP context; the backend combines it with the
+selected kernel and normalized metadata.
+
+Resolution runs outside compiled execution. FA4 reports `SUPPORTED` and
+`CUSTOM_OP` for dense, noncausal BF16 without parallel or HSDP boundaries when
+its kernel accepts the head dimensions. Dimension validation delegates to FA4's
+architecture-specific rules through `backends/utils/fa.py`. Kernel rejections
+report `UNSUPPORTED` with an actionable reason. Missing private validators and
+other unmigrated paths report `UNMIGRATED` with advisory `EAGER_ONLY` defaults.
+`requested_support()` checks a fullgraph request; it does not enforce selection
+or change existing execution.
+
+FA4 execution uses an opaque custom op. Its fake output preserves Q's batch,
+sequence, and head count and uses V's head dimension. Output is contiguous even
+when Q is noncontiguous. Metadata normalization is
+shared with dispatch. Producers must update published mask semantics when masks
+change; unpublished masks remain runtime-dependent to avoid synchronization.
+Callers resolve again when execution metadata changes.
+
+## NPU and ROCm worked examples
+
+The same `FLASH_ATTN` backend resolves MindIE paths on NPU and AITER paths on
+ROCm. Both examples report `UNMIGRATED`: their route is inspected, while hardware
+correctness and compilation remain unvalidated. `EAGER_ONLY` is an advisory
+default, not a measured compiler limitation. Existing execution is unchanged.
+
+| Input | NPU / MindIE | ROCm / AITER |
+| --- | --- | --- |
+| Dense | `npu_dense` | `rocm_dense` |
+| Padding mask | `npu_masked`: expand the key mask to `[B, 1, Q, K]` | `rocm_masked_varlen`: unpad, call varlen attention, then restore rows |
+| Packed inputs | Opt-in `[real, pad]` contract; choose varlen or prefix-KV slicing from the environment | Forward complete packed-document boundaries to varlen attention |
+| Incomplete packing | Rebuild a mask when possible; otherwise reject before the kernel | Shared CUDA-like normalization rejects incomplete metadata |
+| Unpublished mask semantics | Route to the masked wrapper | Keep resolution runtime-dependent without reading mask values |
+
+NPU resolution reuses `_resolve_packed_seq_npu`; ROCm reuses the metadata
+normalization used by dispatch. AITER identity comes from initialization rather
+than caller-provided context. Quantized, paged, piecewise, and parallel paths
+remain outside these worked examples.
+
+The new L1 tests use CPU tensors and replace only vendor kernel calls. They
+check path identity, forwarded masks/boundaries, fallback rejection, and
+conservative fullgraph-request handling. They do not test vendor numerics or
+compiler support. Existing L1 CI already collects this test file.
+
+```bash
+python -m pytest tests/diffusion/attention/test_flash_attn.py \
+  -k 'npu_contract or rocm_contract' -m 'core_model and cpu' -q
+```
+
+### Hardware validation required before migration
+
+On an Ascend/MindIE or ROCm/AITER runner, record device, driver, PyTorch, vendor
+library, and compiler versions, then validate the corresponding rows above:
+
+1. Compare eager outputs with FP32 SDPA for dense and padding-mask inputs,
+   including unequal Q/K lengths. Compare valid query rows; do not assume both
+   platforms define padded-query outputs identically.
+2. Compare packed outputs with separate per-document references. On NPU, cover
+   both varlen and laser prefix-slicing modes, plus mask reconstruction and the
+   missing-mask error. A padding fallback does not establish multi-document
+   isolation. On ROCm, cover multiple real documents.
+3. Check output shape, dtype, device, finite values, and unchanged inputs.
+   Start BF16 comparisons at `atol=rtol=1e-2` and have the platform maintainer
+   confirm the tolerance against its reference tests.
+4. Run `torch.compile(fullgraph=True, dynamic=True)` with the platform's supported
+   compiler over multiple sequence lengths. Record graph breaks and recompiles;
+   choose `TRACEABLE`, `CUSTOM_OP`, or verified `EAGER_ONLY` from that evidence.
+   A custom-op path additionally needs schema/fake checks and repeated replay.
+
+Promote only the paths validated on that platform. These hardware checks have
+not been run locally; ROCm/NPU fullgraph support is not claimed.
+
+## State lifetime example
+
+`test_attention_state_lifetime.py` prepares an owned query-scale tensor eagerly
+and passes it explicitly to an opaque attention op. The compiled module retains
+its state after the caller drops its reference and releases ownership when the
+compiled callable is deleted. Equivalent instances reuse a graph with their own
+state values.
+
+This example tests ownership and release, dynamic replay across instances,
+custom-op schema/fake behavior, and Inductor execution. It adds no production
+backend or state registry. Planning keys, caching, fallback-policy declarations,
+opaque handles, CUDA graph capture, and TRTLLM integration are deferred.
+
+## Validation
+
+```bash
+python -m pytest tests/diffusion/attention/test_flash_attn_compile.py \
+  tests/diffusion/attention/test_flash_attn.py \
+  tests/diffusion/attention/test_attention_capabilities.py \
+  tests/diffusion/attention/test_attention_state_lifetime.py -q -rs
+```
+
+The focused suite covers capability decisions and metadata changes, existing
+padding/mask regressions, tensor-state lifetime, and real FA4 fullgraph execution
+against FP32 SDPA and eager FA4. Representative Q/K and V dimensions are
+(32, 32), (64, 64), (80, 48), (192, 128), and (256, 256), with batches 1 and 2
+and sequence lengths up to 1,024. Invalid dimensions are checked against actual
+kernel errors. These samples exercise the contract rather than define support.
+BF16 comparisons use `atol=rtol=1e-2`; `torch.library.opcheck` checks schema and
+fake-output correctness.
+
+Single-graph reuse is asserted in the tensor-state test. The FA4 numerical tests
+retain compiler caches across shapes within each case but do not require one
+graph across different batch sizes, Q/K length equality, or head dimensions.
+
+Tested environment: GB300, PyTorch `2.13.0+cu130`, FA4 `4.0.0b18`, CUTLASS DSL
+`4.6.2` with CUDA 13 libraries, Quack `0.6.4`, and TVM FFI `0.1.11`.
+Real-kernel tests require CUDA and CuTe FA4; their availability is checked at runtime.
+
+Before the rebase, the focused suite passed all 97 tests without skips, including
+11 CPU NPU/ROCm routing tests. After rebasing onto upstream `e3be42e05`, the
+expanded suite reports 99 passed and 11 failed in the local vLLM 0.28 environment.
+All 11 failures are import errors: upstream now requires `compute_layout_strides`
+from vLLM, which this environment lacks. Both noncontiguous FA4 regression cases
+and real-kernel schema/fake checks pass. The full suite must be rerun with the
+vLLM version required by upstream. NPU/ROCm device numerics and compilation remain
+unvalidated. Pre-commit passes with the CI hook skips. Upstream PyTorch/CUTLASS
+warnings remain.

docs/design/index.md

@@ -68,6 +68,7 @@ guides list selectable backends, platform defaults, installation, and tuning.
 The design contracts separate selection mechanics from backend algorithms:
 
 - [Attention Backend Selection](feature/attention_backend_selection.md)
+- [Attention Execution Contract PoC](feature/attention_execution_contract_poc.md)
 - [Skip-Softmax](feature/skip_softmax.md)
 
 #### CPU offloading

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

vllm_omni/diffusion/attention/backends/abstract.py

@@ -7,6 +7,10 @@ from typing import Any, Generic, Literal, TypeVar
 
 import torch
 
+from vllm_omni.diffusion.attention.capabilities import (
+    ExecutionContext,
+    ExecutionPathResult,
+)
 from vllm_omni.platforms import current_omni_platform
 
 
@@ -132,6 +136,15 @@ class AttentionBackend(ABC):
 
         return False
 
+    @classmethod
+    def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult:
+        """Return conservative pre-construction capabilities.
+
+        Unmigrated backends remain advisory so introducing the contract cannot
+        reject an execution path that works today.
+        """
+        return ExecutionPathResult.unmigrated(cls.get_name(), context)
+
 
 @dataclass(frozen=True, slots=True)
 class QueryRange:
@@ -218,6 +231,9 @@ class AttentionMetadata:
     #   "valid_kv_length": int — contiguous valid K/V prefix length. Local
     #     backends may slice it directly; ring backends trim each circulated
     #     K/V block from the same global prefix.
+    #   "attention_mask_mode": "none" | "padding" | "arbitrary" — producer
+    #     assertion used by capability planning. "none" means a present mask
+    #     is semantically a no-op; absent means its contents are runtime-dependent.
     #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
     #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
     #     Requires the [real, pad] two-document packing contract; see
@@ -277,6 +293,21 @@ class AttentionImpl(ABC, Generic[T]):
             return True
         return kv_cache_dtype in cls._supported_kv_cache_dtypes.get(platform_key, set())
 
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None,
+    ) -> ExecutionPathResult:
+        """Resolve the complete path after backend initialization.
+
+        Subclasses report verified results; this inspection API does not enforce them.
+        """
+        del query, key, value, attn_metadata
+        return ExecutionPathResult.unmigrated(type(self).__name__, context)
+
     def forward(
         self,
         query: torch.Tensor,

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -3,7 +3,9 @@
 
 import os
 from collections.abc import Callable
+from dataclasses import replace
 from functools import cache, partial
+from typing import NamedTuple
 
 import torch
 from vllm.logger import init_logger
@@ -14,6 +16,16 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
     run_paged_piecewise_plan,
 )
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    ExecutionPathResult,
+    MaskMode,
+    PackingMode,
+    ParallelStrategy,
+    SupportStatus,
+)
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.platforms import current_omni_platform
 
@@ -29,6 +41,133 @@ def _get_npu_compressed_causal_mask(device: torch.device) -> torch.Tensor:
     ).contiguous()
 
 
+if not hasattr(torch.ops.vllm_omni, "fa4_dense_attention"):
+
+    @torch.library.custom_op("vllm_omni::fa4_dense_attention", mutates_args=())
+    def _fa4_dense_attention_op(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        softmax_scale: float,
+        causal: bool,
+        deterministic: bool,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.utils.fa import flash_attn_func
+
+        if flash_attn_func is None:
+            raise RuntimeError("CuTe FlashAttention-4 is unavailable")
+        kwargs = {
+            "causal": causal,
+            "softmax_scale": softmax_scale,
+        }
+        if deterministic:
+            kwargs["deterministic"] = True
+        out = flash_attn_func(query, key, value, **kwargs)
+        return out[0] if isinstance(out, tuple) else out
+
+    @_fa4_dense_attention_op.register_fake
+    def _fa4_dense_attention_fake(
+        query,
+        key,
+        value,
+        softmax_scale,
+        causal,
+        deterministic,
+    ):
+        # FA4 returns contiguous output even when Q is noncontiguous; V sets its head dimension.
+        return query.new_empty((*query.shape[:-1], value.shape[-1]))
+
+
+_fa4_dense_attention_op = torch.ops.vllm_omni.fa4_dense_attention
+
+
+_PACKED_KEYS = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
+
+
+class _FlashAttentionMetadataPlan(NamedTuple):
+    attention_mask: torch.Tensor | None
+    mask_mode: MaskMode
+    full_attn_spans: list[list[tuple[int, int]]] | None
+    extra: dict
+    packing_mode: PackingMode
+
+
+def _normalize_flash_attention_metadata(
+    attn_metadata: AttentionMetadata | None,
+) -> _FlashAttentionMetadataPlan:
+    attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+    full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
+    extra = attn_metadata.extra if attn_metadata is not None else {}
+    published_mask_mode = extra.get("attention_mask_mode")
+    if attention_mask is None:
+        mask_mode = MaskMode.NONE
+    elif published_mask_mode is None:
+        mask_mode = MaskMode.UNKNOWN
+    else:
+        try:
+            mask_mode = MaskMode(published_mask_mode)
+        except ValueError as error:
+            raise ValueError(f"Unknown attention_mask_mode {published_mask_mode!r}") from error
+        if mask_mode is MaskMode.UNKNOWN:
+            raise ValueError("attention_mask_mode='unknown' is reserved for unpublished semantics")
+    # Piecewise dispatch takes precedence and does not consume packed metadata.
+    present_packed_keys = [key for key in _PACKED_KEYS if key in extra] if full_attn_spans is None else []
+    if present_packed_keys and len(present_packed_keys) != len(_PACKED_KEYS):
+        missing = sorted(set(_PACKED_KEYS) - set(present_packed_keys))
+        raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
+
+    packing_mode = PackingMode.NONE
+    if present_packed_keys:
+        cu_seqlens_q = extra["cu_seqlens_q"]
+        packing_mode = PackingMode.MULTI_DOCUMENT if cu_seqlens_q.shape[0] > 3 else PackingMode.PACKED_PADDING
+    return _FlashAttentionMetadataPlan(
+        attention_mask=attention_mask,
+        mask_mode=mask_mode,
+        full_attn_spans=full_attn_spans,
+        extra=extra,
+        packing_mode=packing_mode,
+    )
+
+
+def _flash_attention_execution_path(
+    context: ExecutionContext,
+) -> ExecutionPathResult:
+    if context.mask_mode is MaskMode.UNKNOWN:
+        return ExecutionPathResult.unmigrated(
+            "FLASH_ATTN",
+            context,
+            path="runtime_mask_dependent",
+        )
+    if not (
+        context.platform == "cuda"
+        and context.kernel_variant == "fa4"
+        and context.dtype in {"bfloat16", "torch.bfloat16"}
+        and context.causal is False
+        and context.mask_mode is MaskMode.NONE
+        and context.packing_mode is PackingMode.NONE
+        and not context.piecewise
+        and not context.paged_kv
+        and context.kv_cache_dtype is None
+        and context.parallel_strategy is ParallelStrategy.NONE
+        and not context.outer_boundaries
+    ):
+        return ExecutionPathResult.unmigrated(
+            "FLASH_ATTN",
+            context,
+            path="unverified",
+        )
+
+    return ExecutionPathResult(
+        backend="FLASH_ATTN",
+        path="fa4_dense",
+        support=CapabilityResult.supported(),
+        compilation_mode=CompilationMode.CUSTOM_OP,
+        platform=context.platform,
+        kernel_variant=context.kernel_variant,
+        parallel_strategy=context.parallel_strategy,
+    )
+
+
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supports_piecewise_spans: bool = True
@@ -74,6 +213,10 @@ class FlashAttentionBackend(AttentionBackend):
     def get_impl_cls() -> type["FlashAttentionImpl"]:
         return FlashAttentionImpl
 
+    @classmethod
+    def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult:
+        return _flash_attention_execution_path(replace(context, kernel_variant=None))
+
 
 class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
     # Per-platform attention quantization support.
@@ -108,11 +251,135 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         self.softmax_scale = softmax_scale
         self.qkv_layout = qkv_layout
         self.is_cross_attn = role == "cross"
+        from vllm_omni.diffusion.attention.backends.utils.fa import IS_AITER, IS_FLASH_ATTN_4
+
+        self._kernel_variant = "fa4" if IS_FLASH_ATTN_4 else "aiter" if IS_AITER else None
         cfg = get_current_diffusion_config_or_none()
         self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False
         if backend_kwargs:
             logger.warning("FlashAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        if context.platform == "npu":
+            return self._resolve_npu_execution_path(context, query, key, attn_metadata)
+        metadata = _normalize_flash_attention_metadata(attn_metadata)
+        context = replace(
+            context,
+            kernel_variant=self._kernel_variant,
+            dtype=str(query.dtype).removeprefix("torch."),
+            causal=self.causal,
+            mask_mode=metadata.mask_mode,
+            packing_mode=metadata.packing_mode,
+            piecewise=metadata.full_attn_spans is not None,
+            kv_cache_dtype=metadata.extra.get("kv_cache_dtype"),
+        )
+        if context.platform == "rocm":
+            return self._resolve_rocm_execution_path(context)
+        result = _flash_attention_execution_path(context)
+        if result.support.status is not SupportStatus.SUPPORTED:
+            return result
+
+        from vllm_omni.diffusion.attention.backends.utils.fa import validate_fa4_head_dims
+
+        try:
+            if query.dtype != key.dtype or query.dtype != value.dtype:
+                raise ValueError("Q, K, and V dtypes must match")
+            if query.device != key.device or query.device != value.device:
+                raise ValueError("Q, K, and V devices must match")
+            if query.shape[-1] != key.shape[-1]:
+                raise ValueError("Q and K head dimensions must match")
+            verified = validate_fa4_head_dims(query.shape[-1], value.shape[-1], 16 // value.element_size())
+        except (AssertionError, ValueError) as error:
+            return replace(
+                result,
+                support=CapabilityResult.unsupported(f"FA4: {error} Select compatible inputs or another backend."),
+            )
+        if not verified:
+            return replace(
+                result,
+                support=CapabilityResult.unmigrated("Selected FA4 kernel has no available head-dimension validator"),
+                compilation_mode=CompilationMode.EAGER_ONLY,
+            )
+        return result
+
+    @staticmethod
+    def _resolve_rocm_execution_path(context: ExecutionContext) -> ExecutionPathResult:
+        """Describe AITER dispatch; device and compiler validation is pending."""
+        result = ExecutionPathResult.unmigrated("FLASH_ATTN", context, path="rocm_unverified")
+        if (
+            context.kernel_variant != "aiter"
+            or context.parallel_strategy is not ParallelStrategy.NONE
+            or context.outer_boundaries
+            or context.paged_kv
+            or context.piecewise
+            or context.kv_cache_dtype is not None
+        ):
+            return result
+        if context.packing_mode is not PackingMode.NONE:
+            path = "rocm_packed_varlen"
+        elif context.mask_mode is MaskMode.UNKNOWN:
+            path = "rocm_runtime_mask_dependent"
+        elif context.mask_mode is not MaskMode.NONE:
+            path = "rocm_masked_varlen"
+        else:
+            path = "rocm_dense"
+        return replace(
+            result,
+            path=path,
+            support=CapabilityResult.unmigrated("AITER routing only; kernel and compilation validation require ROCm"),
+        )
+
+    def _resolve_npu_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        """Describe existing MindIE routing without claiming hardware validation.
+
+        NPU packing has its own opt-in and [real, pad] contract. In particular,
+        incomplete packed metadata can fall back to a mask, so CUDA metadata
+        normalization must not run here. No tensor values are inspected.
+        """
+        context = replace(context, kernel_variant="mindiesd")
+        result = ExecutionPathResult.unmigrated("FLASH_ATTN", context, path="npu_unverified")
+        extra = attn_metadata.extra if attn_metadata else {}
+        if (
+            context.parallel_strategy is not ParallelStrategy.NONE
+            or context.outer_boundaries
+            or context.paged_kv
+            or self.causal
+            or extra.get("kv_cache_dtype") not in (None, "float", "auto")
+            or (attn_metadata is not None and attn_metadata.full_attn_spans is not None)
+        ):
+            return result
+
+        path = "npu_masked" if attn_metadata is not None and attn_metadata.attn_mask is not None else "npu_dense"
+        if extra.get("npu_attn_varlen", False):
+            if self._resolve_packed_seq_npu(query, key, extra) is not None:
+                path = (
+                    "npu_prefix_kv_slice"
+                    if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention"
+                    else "npu_packed_varlen"
+                )
+            else:
+                # This route may reject the request if neither an explicit mask
+                # nor usable valid_kv_length is supplied; it is not support.
+                path = "npu_masked_fallback"
+        return replace(
+            result,
+            path=path,
+            support=CapabilityResult.unmigrated("NPU routing only; kernel and compilation validation require Ascend"),
+        )
+
     def _warn_fa_deterministic_non_dense(self, path: str) -> None:
         if not self.fa_deterministic:
             return
@@ -401,6 +668,7 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         """CUDA/ROCm/MUSA flash attention implementation."""
         from vllm_omni.diffusion.attention.backends.utils.fa import (
             HAS_FLASH_ATTN,
+            IS_FLASH_ATTN_4,
             flash_attn_func,
             flash_attn_varlen_func,
         )
@@ -412,9 +680,10 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
             )
 
-        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
-        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
-        extra = attn_metadata.extra if attn_metadata is not None else {}
+        metadata_plan = _normalize_flash_attention_metadata(attn_metadata)
+        attention_mask = metadata_plan.attention_mask
+        full_attn_spans = metadata_plan.full_attn_spans
+        extra = metadata_plan.extra
 
         # Try piecewise attention
         if full_attn_spans is not None:
@@ -443,12 +712,7 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 query_ranges=None if attn_metadata is None else attn_metadata.query_ranges,
             )
 
-        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
-        present_packed_keys = [key for key in packed_keys if key in extra]
-        if present_packed_keys:
-            if len(present_packed_keys) != len(packed_keys):
-                missing = sorted(set(packed_keys) - set(present_packed_keys))
-                raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
+        if metadata_plan.packing_mode is not PackingMode.NONE:
             self._warn_fa_deterministic_non_dense("packed-varlen")
             return self._forward_varlen_packed(
                 query,
@@ -460,7 +724,14 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 max_seqlen_k=extra["max_seqlen_k"],
             )
 
-        if attention_mask is not None and torch.any(~attention_mask):
+        use_masked_path = False
+        if attention_mask is not None:
+            use_masked_path = (
+                torch.any(~attention_mask)
+                if metadata_plan.mask_mode is MaskMode.UNKNOWN
+                else metadata_plan.mask_mode is not MaskMode.NONE
+            )
+        if use_masked_path:
             self._warn_fa_deterministic_non_dense("masked-varlen")
             return self._forward_varlen_masked(
                 query,
@@ -470,6 +741,15 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
             )
 
         if flash_attn_func is not None:
+            if IS_FLASH_ATTN_4:
+                return _fa4_dense_attention_op(
+                    query,
+                    key,
+                    value,
+                    self.softmax_scale,
+                    self.causal,
+                    self.fa_deterministic,
+                )
             fa_kwargs = {
                 "causal": self.causal,
                 "softmax_scale": self.softmax_scale,

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -59,6 +59,8 @@ def is_flash_attn_4_available() -> bool:
 # Bind via aliases so mypy does not treat each candidate import as a redefinition.
 flash_attn_func: FlashAttnFn | None = None
 flash_attn_varlen_func: FlashAttnFn | None = None
+IS_FLASH_ATTN_4 = False
+IS_AITER = False
 
 if current_omni_platform.is_rocm():
     # ROCm: try Aiter first
@@ -71,6 +73,7 @@ if current_omni_platform.is_rocm():
 
             flash_attn_func = _fa_func
             flash_attn_varlen_func = _fa_varlen
+            IS_AITER = True
     except (ImportError, ModuleNotFoundError):
         pass
 elif current_omni_platform.is_xpu():
@@ -102,6 +105,7 @@ else:
 
                 flash_attn_func = _fa_func
                 flash_attn_varlen_func = _fa_varlen
+                IS_FLASH_ATTN_4 = True
                 logger.info("Using CuTe FlashAttention-4 on Blackwell")
             except Exception as exc:
                 # Optional FA4 dependencies may be present but ABI-incompatible
@@ -166,6 +170,27 @@ else:
 HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None
 
 
+def validate_fa4_head_dims(head_dim: int, head_dim_v: int, alignment: int) -> bool:
+    """Use the selected FA4 implementation's constraints during path resolution.
+
+    Return False when its validator is unavailable; propagate kernel validation
+    errors for the backend to report. Keep the private FA4 API dependency here,
+    rather than duplicating architecture/version-specific dimension rules.
+    This is a planning-time call, never part of compiled tensor execution.
+    """
+    try:
+        from flash_attn.cute.interface import _get_device_arch, _validate_head_dims
+    except ImportError:
+        return False
+
+    arch = _get_device_arch() // 10
+    # FA4 routes SM80/SM120 through separate kernels and bypasses this validator.
+    if arch not in (9, 10, 11):
+        return False
+    _validate_head_dims(head_dim, head_dim_v, arch, alignment)
+    return True
+
+
 def _choose_vllm_flash_attn_version(
     device_major: int,
     requested: str | int | None,

vllm_omni/diffusion/attention/capabilities.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+
+class SupportStatus(str, Enum):
+    SUPPORTED = "supported"
+    UNSUPPORTED = "unsupported"
+    UNMIGRATED = "unmigrated"
+
+
+class CompilationMode(str, Enum):
+    TRACEABLE = "traceable"
+    CUSTOM_OP = "custom_op"
+    EAGER_ONLY = "eager_only"
+
+
+class PackingMode(str, Enum):
+    NONE = "none"
+    PACKED_PADDING = "packed_padding"
+    MULTI_DOCUMENT = "multi_document"
+
+
+class MaskMode(str, Enum):
+    NONE = "none"
+    PADDING = "padding"
+    ARBITRARY = "arbitrary"
+    UNKNOWN = "unknown"
+
+
+class ParallelStrategy(str, Enum):
+    NONE = "none"
+    ULYSSES = "ulysses"
+    RING = "ring"
+    HYBRID_ULYSSES_RING = "hybrid_ulysses_ring"
+    ALLGATHER_KV = "allgather_kv"
+
+
+class OuterBoundary(str, Enum):
+    HSDP = "hsdp"
+
+
+@dataclass(frozen=True, slots=True)
+class CapabilityResult:
+    status: SupportStatus
+    reason: str | None = None
+
+    def __post_init__(self) -> None:
+        if self.status is SupportStatus.UNSUPPORTED and not self.reason:
+            raise ValueError("UNSUPPORTED capability results require an actionable reason")
+
+    @classmethod
+    def supported(cls) -> CapabilityResult:
+        return cls(SupportStatus.SUPPORTED)
+
+    @classmethod
+    def unsupported(cls, reason: str) -> CapabilityResult:
+        return cls(SupportStatus.UNSUPPORTED, reason)
+
+    @classmethod
+    def unmigrated(cls, reason: str | None = None) -> CapabilityResult:
+        return cls(SupportStatus.UNMIGRATED, reason)
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionContext:
+    platform: str
+    kernel_variant: str | None = None
+    dtype: str | None = None
+    causal: bool | None = None
+    mask_mode: MaskMode = MaskMode.NONE
+    packing_mode: PackingMode = PackingMode.NONE
+    piecewise: bool = False
+    paged_kv: bool = False
+    kv_cache_dtype: str | None = None
+    parallel_strategy: ParallelStrategy = ParallelStrategy.NONE
+    require_fullgraph: bool = False
+    outer_boundaries: frozenset[OuterBoundary] = frozenset()
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionPathResult:
+    backend: str
+    path: str
+    support: CapabilityResult
+    compilation_mode: CompilationMode
+    platform: str
+    kernel_variant: str | None
+    parallel_strategy: ParallelStrategy
+
+    @classmethod
+    def unmigrated(
+        cls,
+        backend: str,
+        context: ExecutionContext,
+        *,
+        path: str = "unmigrated",
+    ) -> ExecutionPathResult:
+        return cls(
+            backend=backend,
+            path=path,
+            support=CapabilityResult.unmigrated(),
+            compilation_mode=CompilationMode.EAGER_ONLY,
+            platform=context.platform,
+            kernel_variant=context.kernel_variant,
+            parallel_strategy=context.parallel_strategy,
+        )
+
+    def requested_support(self, context: ExecutionContext) -> CapabilityResult:
+        """Validate request-level guarantees without changing path migration state."""
+        if context.require_fullgraph and self.support.status is SupportStatus.UNMIGRATED:
+            return CapabilityResult.unsupported(
+                f"{self.backend} path {self.path!r} has no verified fullgraph declaration; "
+                "use a verified path or disable the guaranteed-fullgraph request"
+            )
+        if self.support.status is not SupportStatus.SUPPORTED:
+            return self.support
+        if context.require_fullgraph and self.compilation_mode is CompilationMode.EAGER_ONLY:
+            return CapabilityResult.unsupported(
+                f"{self.backend} path {self.path!r} does not support guaranteed fullgraph compilation"
+            )
+        return CapabilityResult.supported()

vllm_omni/diffusion/attention/layer.py

@@ -19,7 +19,13 @@ from vllm.model_executor.models.utils import extract_layer_index
 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
-from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend, SDPAImpl
+from vllm_omni.diffusion.attention.capabilities import (
+    ExecutionContext,
+    ExecutionPathResult,
+    OuterBoundary,
+    ParallelStrategy,
+)
 from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
@@ -119,6 +125,8 @@ class Attention(nn.Module):
 
         config = get_current_diffusion_config_or_none()
         attention_config = config.diffusion_attention_config if config is not None else None
+        parallel_config = getattr(config, "parallel_config", None)
+        self._hsdp_compile_boundary_enabled = bool(getattr(parallel_config, "use_hsdp", False))
 
         from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 
@@ -131,6 +139,9 @@ class Attention(nn.Module):
             is DiffusionKVCacheMode.PAGED_SCHEDULER
         )
         self._scheduler_paged_kv = scheduler_paged_kv
+        self.attn_backend: type[AttentionBackend] | None
+        self.attention: AttentionImpl | nn.Module
+        self.sdpa_fallback: SDPAImpl | None
         if custom_attention is None:
             attn_backend_cls, spec = get_attn_backend_for_role(
                 role=role,
@@ -166,7 +177,6 @@ class Attention(nn.Module):
                     attn_backend_cls.get_name(),
                     dense_backend_name,
                 )
-            parallel_config = getattr(config, "parallel_config", None)
             allgather_degree = getattr(parallel_config, "allgather_degree", 1)
             # TODO: Move AllGather-KV compatibility into an AttentionBackend capability
             # so validation does not depend on backend names.
@@ -279,6 +289,7 @@ class Attention(nn.Module):
 
         if self.paged_kv_cache_role is None:
             return None
+        assert self.attn_backend is not None  # Custom attention cannot opt into paged KV.
         dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
         # Keep backend layout discovery under the same config context used by
         # upstream vLLM's attention-spec collector.  vLLM 0.29 moved the
@@ -378,17 +389,68 @@ class Attention(nn.Module):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
-        if torch.compiler.is_compiling() and is_forward_context_available():
-            od_config = get_forward_context().omni_diffusion_config
-            parallel_config = getattr(od_config, "parallel_config", None)
-            if getattr(parallel_config, "use_hsdp", False):
-                # Keep HSDP/FSDP2 parameter all-gather outside Inductor's
-                # attention graph; otherwise scheduler dependency analysis can
-                # fail on the fused attention region.
-                return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata)
+        if torch.compiler.is_compiling() and self._uses_hsdp_compile_boundary():
+            # Keep HSDP/FSDP2 parameter all-gather outside Inductor's
+            # attention graph; otherwise scheduler dependency analysis can
+            # fail on the fused attention region.
+            return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata)
 
         return self._forward_impl(query, key, value, attn_metadata)
 
+    def _uses_hsdp_compile_boundary(self) -> bool:
+        if self._hsdp_compile_boundary_enabled:
+            return True
+        if not is_forward_context_available():
+            return False
+        od_config = get_forward_context().omni_diffusion_config
+        parallel_config = getattr(od_config, "parallel_config", None)
+        return bool(getattr(parallel_config, "use_hsdp", False))
+
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        """Compose backend capabilities with outer Attention boundaries."""
+        boundaries = set(context.outer_boundaries)
+        if self._uses_hsdp_compile_boundary():
+            boundaries.add(OuterBoundary.HSDP)
+        active_strategy = self._get_active_parallel_strategy()
+        strategy_name = active_strategy.name
+        if self.use_ring and active_strategy.enabled and strategy_name == "ulysses":
+            parallel_strategy = ParallelStrategy.HYBRID_ULYSSES_RING
+        elif self.use_ring and active_strategy.enabled:
+            parallel_strategy = ParallelStrategy.RING
+        else:
+            parallel_strategy = {
+                "allgather_kv": ParallelStrategy.ALLGATHER_KV,
+                "ulysses": ParallelStrategy.ULYSSES,
+            }.get(strategy_name, ParallelStrategy.NONE)
+        resolved_context = replace(
+            context,
+            outer_boundaries=frozenset(boundaries),
+            paged_kv=self.is_paged_kv_active(),
+            parallel_strategy=parallel_strategy,
+        )
+        resolver = getattr(self.attention, "resolve_execution_path", None)
+        if not callable(resolver):
+            return ExecutionPathResult.unmigrated(
+                type(self.attention).__name__,
+                resolved_context,
+            )
+        if not resolved_context.paged_kv:
+            attn_metadata = self._with_kv_cache_dtype(attn_metadata)
+        return resolver(
+            resolved_context,
+            query,
+            key,
+            value,
+            attn_metadata,
+        )
+
     @torch.compiler.disable
     def _forward_hsdp_compile_boundary(
         self,
@@ -422,10 +484,11 @@ class Attention(nn.Module):
             )
         use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
         if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
-            assert self.attn_backend is not None
+            backend_name = (
+                self.attn_backend.get_name() if self.attn_backend is not None else type(self.attention).__name__
+            )
             raise NotImplementedError(
-                f"Diffusion paged KV requires an Omni backend with paged support; "
-                f"selected {self.attn_backend.get_name()}"
+                f"Diffusion paged KV requires an Omni backend with paged support; selected {backend_name}"
             )
         if use_paged_attention and strategy is not self._no_parallel_strategy:
             strategy_name = strategy.name
@@ -554,9 +617,11 @@ class Attention(nn.Module):
             self._scheduler_paged_kv
             and self.paged_kv_cache_role is not None
             and in_kv_memory_profile
-            and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN"
+            and self.attn_backend is not None
+            and self.attn_backend.get_name() == "FLASH_ATTN"
             and not current_omni_platform.supports_diffusion_dense_flash_attention()
         ):
+            assert self.sdpa_fallback is not None
             logger.warning_once(
                 "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. "
                 "Formal paged requests still use the platform-native paged attention backend."