@@ -36,6 +36,9 @@ class FlashInferMetadata(AttentionMetadata):
 
 class FlashInferMetadataBuilder(AttentionMetadataBuilder):
 
+    def __init__(self) -> None:
+        pass
+
     def prepare(self):
         pass
 
@@ -81,6 +84,9 @@ def _mask_for_sample(attn_mask: torch.Tensor, sample: int, query_len: int, key_l
     if mask.dim() == 3:
         mask = mask[sample]
     elif mask.dim() == 4:
+        if mask.shape[1] != 1:
+            raise ValueError("FLASHINFER does not support per-head attention masks; "
+                             f"expected [B, 1, Q, K], got {attn_mask.shape}")
         mask = mask[sample, 0]
     else:
         raise ValueError(f"Unsupported FLASHINFER attention mask shape: {attn_mask.shape}")
 
@@ -17,6 +17,7 @@
 
 import fastvideo.envs as envs
 from fastvideo.attention import (DistributedAttention, LocalAttention)
+from fastvideo.attention.selector import effective_attention_backend
 from fastvideo.models.wan.config import WanVideoConfig
 from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather)
 from fastvideo.distributed.parallel_state import (get_sp_parallel_rank, get_sp_world_size)
@@ -378,12 +379,18 @@ def forward(
 class CausalWanTransformer3DModel(BaseDiT):
     _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
     _compile_conditions = WanVideoConfig()._compile_conditions
-    _supported_attention_backends = WanVideoConfig()._supported_attention_backends
+    _supported_attention_backends = tuple(
+        backend for backend in WanVideoConfig()._supported_attention_backends
+        if backend != AttentionBackendEnum.FLASHINFER)
     param_names_mapping = WanVideoConfig().param_names_mapping
     reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
     lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping
 
     def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None:
+        if effective_attention_backend(config) == AttentionBackendEnum.FLASHINFER:
+            raise NotImplementedError(
+                "FLASHINFER is supported only for dense Wan attention; causal Wan is not supported. "
+                "Use FLASH_ATTN or TORCH_SDPA for causal Wan.")
         super().__init__(config=config, hf_config=hf_config)
 
         inner_dim = config.num_attention_heads * config.attention_head_dim
 
@@ -5,6 +5,9 @@
 from typing import Literal
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+_WAN_ATTENTION_BACKENDS = DiTArchConfig._supported_attention_backends + (AttentionBackendEnum.FLASHINFER, )
 
 
 def is_blocks(n: str, m) -> bool:
@@ -13,6 +16,7 @@ def is_blocks(n: str, m) -> bool:
 
 @dataclass
 class WanVideoArchConfig(DiTArchConfig):
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = _WAN_ATTENTION_BACKENDS
     _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
 
     param_names_mapping: dict = field(
 
@@ -296,7 +296,8 @@ def __init__(self,
                                    causal=False,
                                    supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
                                                                  AttentionBackendEnum.TORCH_SDPA,
-                                                                 AttentionBackendEnum.ATTN_QAT_INFER))
+                                                                 AttentionBackendEnum.ATTN_QAT_INFER,
+                                                                 AttentionBackendEnum.FLASHINFER))
 
     def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int):
         r"""
 
@@ -78,7 +78,6 @@ def _sdpa_reference(query: torch.Tensor,
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("head_size", [64, 128, 256])
 def test_flashinfer_real_cuda_kernel_matches_sdpa(dtype: torch.dtype, head_size: int) -> None:
-    """Launch the real single-GPU FlashInfer kernel for every supported head size."""
     _require_flashinfer_cuda()
     torch.manual_seed(0)
     device = torch.device("cuda", torch.cuda.current_device())
@@ -98,7 +97,6 @@ def test_flashinfer_real_cuda_kernel_matches_sdpa(dtype: torch.dtype, head_size:
 
 @pytest.mark.parametrize("mode", ["causal", "cross_gqa", "causal_padding"])
 def test_flashinfer_real_cuda_kernel_attention_modes(mode: str) -> None:
-    """Exercise native causal, GQA/cross-attention, and combined custom masks."""
     _require_flashinfer_cuda()
     torch.manual_seed(1)
     device = torch.device("cuda", torch.cuda.current_device())
@@ -138,3 +136,117 @@ def test_flashinfer_real_cuda_kernel_attention_modes(mode: str) -> None:
 
     torch.cuda.synchronize(device)
     torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2)
+
+
+def test_metadata_builder_instantiates_through_backend_api() -> None:
+    from fastvideo.attention.backends.abstract import AttentionMetadataBuilder
+    from fastvideo.attention.backends.flashinfer import FlashInferBackend
+    builder = FlashInferBackend.get_builder_cls()()
+    assert isinstance(builder, AttentionMetadataBuilder)
+    builder.prepare()
+    mask = torch.ones(2, 3, dtype=torch.bool)
+    metadata = builder.build(current_timestep=7, attn_mask=mask)
+    assert metadata.current_timestep == 7
+    assert metadata.attn_mask is mask
+
+
+@pytest.mark.parametrize("shared_head", [False, True])
+@pytest.mark.parametrize("additive", [False, True])
+def test_shared_masks_preserve_each_sample(shared_head, additive) -> None:
+    mask = torch.tensor([[[True, False, True], [False, True, True]],
+                         [[False, True, False], [True, True, False]]])
+    source = torch.where(mask, 0.0, float("-inf")) if additive else mask
+    if shared_head:
+        source = source[:, None]
+    for sample in range(2):
+        torch.testing.assert_close(_mask_for_sample(source, sample, 2, 3), mask[sample])
+
+
+def test_per_head_mask_rejected_before_kernel(monkeypatch) -> None:
+    prefill = types.ModuleType("flashinfer.prefill")
+    def unexpected_kernel(*args, **kwargs):
+        pytest.fail("per-head mask reached FlashInfer kernel")
+    prefill.single_prefill_with_kv_cache = unexpected_kernel
+    monkeypatch.setitem(sys.modules, "flashinfer.prefill", prefill)
+    query = torch.randn(2, 3, 4, 64, dtype=torch.bfloat16)
+    metadata = FlashInferMetadata(current_timestep=0, attn_mask=torch.ones(2, 4, 3, 3, dtype=torch.bool))
+    impl = FlashInferImpl(num_heads=4, head_size=64, causal=False, softmax_scale=0.125)
+    with pytest.raises(ValueError, match="per-head"):
+        impl.forward(query, query, query, metadata)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("shared_head", [False, True])
+def test_shared_mask_cuda_matches_flash_attention(dtype, shared_head) -> None:
+    _require_flashinfer_cuda()
+    flash_attn = pytest.importorskip("flash_attn")
+    torch.manual_seed(19)
+    q = torch.randn(2, 16, 2, 64, device="cuda", dtype=dtype)
+    k, v = torch.randn_like(q), torch.randn_like(q)
+    mask = torch.zeros(2, 16, 16, device="cuda", dtype=torch.bool)
+    mask[0, :, :11] = True
+    mask[1, :, :7] = True
+    metadata = FlashInferMetadata(current_timestep=0, attn_mask=mask[:, None] if shared_head else mask)
+    impl = FlashInferImpl(num_heads=2, head_size=64, causal=False, softmax_scale=0.2)
+    actual = impl.forward(q, k, v, metadata)
+    expected = torch.cat([
+        flash_attn.flash_attn_func(q[i:i + 1], k[i:i + 1, :length], v[i:i + 1, :length], softmax_scale=0.2)
+        for i, length in enumerate([11, 7])
+    ])
+    tolerance = 3e-2 if dtype == torch.bfloat16 else 1e-2
+    torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("i2v", [False, True])
+def test_wan_attention_executes_flashinfer(monkeypatch, request, dtype, i2v) -> None:
+    _require_flashinfer_cuda()
+    request.getfixturevalue("distributed_setup")
+    import flashinfer.prefill as prefill
+    import fastvideo.envs as envs
+    from flash_attn import flash_attn_func
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.models.wan.config import WanVideoConfig
+    from fastvideo.models.wan.transformer import WanTransformerBlock
+    from fastvideo.platforms import AttentionBackendEnum
+
+    real_kernel = prefill.single_prefill_with_kv_cache
+    calls = []
+    tolerance = 3e-2 if dtype == torch.bfloat16 else 1e-2
+    def checked_kernel(q, k, v, **kwargs):
+        actual = real_kernel(q, k, v, **kwargs)
+        expected = flash_attn_func(q[None], k[None], v[None],
+                                  softmax_scale=kwargs["sm_scale"], causal=kwargs["causal"])[0]
+        torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance)
+        calls.append((q.shape, k.shape))
+        return actual
+    monkeypatch.setattr(prefill, "single_prefill_with_kv_cache", checked_kernel)
+    with envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASHINFER"), torch.device("cuda"):
+        block = WanTransformerBlock(
+            dim=256, ffn_dim=512, num_heads=2, cross_attn_norm=True,
+            added_kv_proj_dim=256 if i2v else None,
+            supported_attention_backends=WanVideoConfig()._supported_attention_backends,
+        ).to(dtype=dtype).eval()
+    assert block.attn1.backend == AttentionBackendEnum.FLASHINFER
+    assert block.attn2.attn.backend == AttentionBackendEnum.FLASHINFER
+    for param in block.parameters():
+        torch.nn.init.normal_(param, std=0.02)
+    q = torch.randn(1, 16, 2, 128, device="cuda", dtype=dtype)
+    k, v = torch.randn_like(q), torch.randn_like(q)
+    x = torch.randn(1, 16, 256, device="cuda", dtype=dtype)
+    context = torch.randn(1, 264 if i2v else 7, 256, device="cuda", dtype=dtype)
+    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
+        self_output, _ = block.attn1(q, k, v)
+        cross_output = block.attn2(x, context, None)
+    assert self_output.shape == q.shape
+    assert cross_output.shape == x.shape
+    assert len(calls) == (3 if i2v else 2)
+
+
+def test_wan_cuda_skip_precedes_distributed_setup(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    class Request:
+        def getfixturevalue(self, name):
+            pytest.fail("distributed setup ran before the CUDA availability check")
+    with pytest.raises(pytest.skip.Exception, match="NVIDIA CUDA GPU"):
+        test_wan_attention_executes_flashinfer(monkeypatch, Request(), torch.float16, False)
@@ -73,3 +73,50 @@ def test_flashinfer_warns_on_dtype_cast(monkeypatch, caplog) -> None:
     with caplog.at_level("WARNING"):
         CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.FLASHINFER, 128, torch.float32)
     assert any("cast" in record.message for record in caplog.records)
+
+
+def test_real_wan_config_reaches_flashinfer_dispatch(monkeypatch, env_overrides):
+    import fastvideo.envs as envs
+    import fastvideo.platforms as platforms
+    from fastvideo.attention.selector import get_attn_backend
+    from fastvideo.attention.backends.flashinfer import FlashInferBackend
+    from fastvideo.models.wan.config import WanVideoConfig
+    _patch_capability(monkeypatch, supported=True)
+    _install_fake_flashinfer(monkeypatch)
+    monkeypatch.setattr(platforms, "_current_platform", CudaPlatformBase())
+    env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASHINFER"))
+    backend = get_attn_backend(128, torch.bfloat16,
+                               supported_attention_backends=WanVideoConfig()._supported_attention_backends)
+    assert backend is FlashInferBackend
+
+
+def test_wan_capabilities_preserve_defaults_and_other_models():
+    from fastvideo.configs.models.dits.base import DiTArchConfig
+    from fastvideo.models.wan.config import WanVideoArchConfig
+    before = DiTArchConfig()._supported_attention_backends
+    after = WanVideoArchConfig()._supported_attention_backends
+    assert after == (*before, AttentionBackendEnum.FLASHINFER)
+    assert AttentionBackendEnum.FLASHINFER not in before
+
+
+@pytest.mark.parametrize("backend", [AttentionBackendEnum.FLASHINFER,
+                                    AttentionBackendEnum.FLASH_ATTN,
+                                    AttentionBackendEnum.TORCH_SDPA])
+def test_causal_wan_rejects_flashinfer_before_weight_allocation(monkeypatch, backend):
+    from fastvideo.attention.selector import _component_attention_backend_scope
+    from fastvideo.models.wan import causal_transformer
+    from fastvideo.models.wan.config import WanVideoConfig
+    class WeightAllocationReached(Exception):
+        pass
+    def stop_at_first_layer(*args, **kwargs):
+        raise WeightAllocationReached
+    monkeypatch.setattr(causal_transformer, "PatchEmbed", stop_at_first_layer)
+    model_cls = causal_transformer.CausalWanTransformer3DModel
+    assert AttentionBackendEnum.FLASHINFER not in model_cls._supported_attention_backends
+    with _component_attention_backend_scope(backend):
+        if backend == AttentionBackendEnum.FLASHINFER:
+            with pytest.raises(NotImplementedError, match="causal Wan is not supported"):
+                model_cls(WanVideoConfig(), {})
+        else:
+            with pytest.raises(WeightAllocationReached):
+                model_cls(WanVideoConfig(), {})