[Bugfix][NPU] Materialize causal mask in FlashAttention dense path (#7324)

Signed-off-by: Oliver7th <fangzhengyi@huawei.com>

tests/diffusion/attention/test_flash_attn.py

@@ -503,11 +503,12 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
 #   - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
 #     sync, no MindIE-SD import);
 #   - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
+#   - native torch_npu dispatch for right-down causal attention;
 #   - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
 #
-# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
-# into ``sys.modules`` (same pattern as ``test_paged_attention`` /
-# ``benchmarks/conftest``). No production code is exercised for real kernels.
+# They run on CPU without a real NPU by injecting fake ``mindiesd`` and
+# ``torch_npu`` modules into ``sys.modules`` (same pattern as
+# ``test_paged_attention`` / ``benchmarks/conftest``). No real kernel runs.
 
 
 def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
@@ -529,6 +530,15 @@ def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_var
     return fake
 
 
+def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None):
+    """Install the minimal torch_npu surface used by causal NPU attention."""
+    fake = SimpleNamespace(
+        npu_fusion_attention=npu_fusion_attention or Mock(return_value=(torch.zeros(1), None, None, None, 0, 0, 0)),
+    )
+    monkeypatch.setitem(sys.modules, "torch_npu", fake)
+    return fake
+
+
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 
 
@@ -685,7 +695,109 @@ def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
     assert out is fake_forward.return_value
 
 
-# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
+# --- Test group C: native NPU causal attention dispatch ----------------------
+
+
+@pytest.mark.parametrize(
+    ("query_length", "key_length", "expected_inner_precise"),
+    [(4, 4, 0), (2, 4, 0), (4, 2, 2)],
+)
+def test_npu_causal_uses_native_right_down_mode(
+    monkeypatch,
+    query_length,
+    key_length,
+    expected_inner_precise,
+):
+    expected_out = torch.randn(1, query_length, 2, 4)
+    fusion_attention = Mock(return_value=(expected_out, None, None, None, 0, 0, 0))
+    fake_torch_npu = _fake_torch_npu(monkeypatch, npu_fusion_attention=fusion_attention)
+    impl = FlashAttentionImpl(num_heads=8, head_size=4, softmax_scale=0.5, causal=True)
+    query = torch.randn(1, query_length, 2, 4)
+    key = torch.randn(1, key_length, 2, 4)
+
+    out = impl.forward_fa_npu(query, key, key)
+
+    fake_torch_npu.npu_fusion_attention.assert_called_once()
+    args, kwargs = fake_torch_npu.npu_fusion_attention.call_args
+    assert torch.equal(args[0], query)
+    assert torch.equal(args[1], key)
+    assert torch.equal(args[2], key)
+    assert all(tensor.is_contiguous() for tensor in args[:3])
+    assert kwargs["head_num"] == query.shape[2]
+    assert kwargs["input_layout"] == "BSND"
+    assert kwargs["scale"] == 0.5
+    assert kwargs["keep_prob"] == 1.0
+    assert kwargs["sparse_mode"] == 3
+    assert kwargs["inner_precise"] == expected_inner_precise
+
+    block_mask = kwargs["atten_mask"]
+    assert block_mask.shape == (2048, 2048)
+    assert block_mask.dtype == torch.bool
+    assert block_mask.is_contiguous()
+    expected_corner = torch.tensor(
+        [
+            [False, True, True, True],
+            [False, False, True, True],
+            [False, False, False, True],
+            [False, False, False, False],
+        ]
+    )
+    assert torch.equal(block_mask[:4, :4], expected_corner)
+    assert out is expected_out
+
+
+def test_npu_causal_composes_explicit_keep_mask(monkeypatch):
+    expected_out = torch.randn(1, 4, 2, 4)
+    fusion_attention = Mock(return_value=(expected_out, None, None, None, 0, 0, 0))
+    fake_torch_npu = _fake_torch_npu(monkeypatch, npu_fusion_attention=fusion_attention)
+    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=True)
+    query = torch.randn(1, 4, 2, 4)
+    key_keep_mask = torch.tensor([[True, False, True, True]])
+
+    out = impl.forward_fa_npu(
+        query,
+        query,
+        query,
+        AttentionMetadata(attn_mask=key_keep_mask),
+    )
+
+    fake_torch_npu.npu_fusion_attention.assert_called_once()
+    kwargs = fake_torch_npu.npu_fusion_attention.call_args.kwargs
+    assert kwargs["sparse_mode"] == 1
+    assert kwargs["inner_precise"] == 2
+    expected_block_mask = torch.tensor(
+        [
+            [
+                [
+                    [False, True, True, True],
+                    [False, True, True, True],
+                    [False, True, False, True],
+                    [False, True, False, False],
+                ]
+            ]
+        ]
+    )
+    assert torch.equal(kwargs["atten_mask"], expected_block_mask)
+    assert kwargs["atten_mask"].is_contiguous()
+    assert out is expected_out
+
+
+def test_npu_noncausal_without_explicit_mask_stays_unmasked(monkeypatch):
+    captured: dict = {}
+
+    def fake_attention_forward(query, key, value, **kwargs):
+        captured.update(kwargs)
+        return torch.zeros_like(query)
+
+    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
+    query = torch.randn(1, 4, 2, 4)
+
+    _npu_impl(causal=False).forward_fa_npu(query, query, query)
+
+    assert captured["attn_mask"] is None
+
+
+# --- Test group D: laser input pre-scaling in _forward_prefix_kv_slice_npu --
 
 
 def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):

tests/diffusion/attention/test_flash_attn_npu.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
+
+
+def _bottom_right_causal_reference(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    scale: float,
+    key_keep_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * scale
+    query_positions = torch.arange(query.shape[1]).unsqueeze(1)
+    key_positions = torch.arange(key.shape[1]).unsqueeze(0)
+    keep_mask = key_positions <= query_positions + (key.shape[1] - query.shape[1])
+    if key_keep_mask is not None:
+        assert query.shape[0] == 1
+        keep_mask = keep_mask & key_keep_mask[0].bool().unsqueeze(0)
+    valid_rows = keep_mask.any(dim=-1)
+
+    output = torch.zeros(
+        query.shape[0],
+        query.shape[1],
+        query.shape[2],
+        value.shape[-1],
+        dtype=torch.float32,
+    )
+    valid_scores = scores[:, :, valid_rows].masked_fill(~keep_mask[valid_rows], float("-inf"))
+    probabilities = torch.softmax(valid_scores, dim=-1)
+    output[:, valid_rows] = torch.einsum("bhqk,bkhd->bqhd", probabilities, value.float())
+    return output
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native causal attention requires Ascend NPU")
+def test_npu_bottom_right_causal_fully_masked_rows_are_zero():
+    torch.manual_seed(0)
+    batch_size, query_length, key_length = 1, 4, 2
+    num_heads, head_size = 2, 64
+    scale = head_size**-0.5
+
+    query_cpu = torch.randn(batch_size, query_length, num_heads, head_size).to(torch.bfloat16)
+    key_cpu = torch.randn(batch_size, key_length, num_heads, head_size).to(torch.bfloat16)
+    value_cpu = torch.randn(batch_size, key_length, num_heads, head_size).to(torch.bfloat16)
+    expected = _bottom_right_causal_reference(query_cpu, key_cpu, value_cpu, scale)
+
+    device = torch.device(current_omni_platform.device_type)
+    query = query_cpu.to(device)
+    key = key_cpu.to(device)
+    value = value_cpu.to(device)
+    impl = FlashAttentionImpl(
+        num_heads=num_heads,
+        head_size=head_size,
+        softmax_scale=scale,
+        causal=True,
+    )
+
+    output = impl.forward_fa_npu(query, key, value)
+    output_cpu = output.float().cpu()
+
+    assert torch.isfinite(output_cpu).all()
+    assert torch.equal(output_cpu[:, : query_length - key_length], torch.zeros_like(output_cpu[:, :2]))
+    torch.testing.assert_close(
+        output_cpu[:, query_length - key_length :],
+        expected[:, query_length - key_length :],
+        rtol=2e-2,
+        atol=2e-2,
+    )
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native causal attention requires Ascend NPU")
+def test_npu_bottom_right_causal_composes_explicit_keep_mask():
+    torch.manual_seed(1)
+    batch_size, sequence_length = 1, 4
+    num_heads, head_size = 2, 64
+    scale = head_size**-0.5
+    key_keep_mask_cpu = torch.tensor([[True, False, True, True]])
+
+    query_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16)
+    key_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16)
+    value_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16)
+    expected = _bottom_right_causal_reference(
+        query_cpu,
+        key_cpu,
+        value_cpu,
+        scale,
+        key_keep_mask_cpu,
+    )
+
+    device = torch.device(current_omni_platform.device_type)
+    query = query_cpu.to(device)
+    key = key_cpu.to(device)
+    value = value_cpu.to(device)
+    impl = FlashAttentionImpl(
+        num_heads=num_heads,
+        head_size=head_size,
+        softmax_scale=scale,
+        causal=True,
+    )
+
+    output = impl.forward_fa_npu(
+        query,
+        key,
+        value,
+        AttentionMetadata(attn_mask=key_keep_mask_cpu.to(device)),
+    )
+    output_cpu = output.float().cpu()
+
+    assert torch.isfinite(output_cpu).all()
+    torch.testing.assert_close(output_cpu, expected, rtol=2e-2, atol=2e-2)

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import os
-from functools import partial
+from functools import cache, partial
 
 import torch
 from vllm.logger import init_logger
@@ -19,6 +19,15 @@ from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
 
 
+@cache
+def _get_npu_compressed_causal_mask(device: torch.device) -> torch.Tensor:
+    """Return the shared block mask used by NPU right-down causal attention."""
+    return torch.triu(
+        torch.ones((2048, 2048), dtype=torch.bool, device=device),
+        diagonal=1,
+    ).contiguous()
+
+
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supports_piecewise_spans: bool = True
@@ -549,6 +558,53 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
+        attention_mask = attn_metadata.attn_mask if attn_metadata else None
+        if self.causal:
+            import torch_npu
+
+            if attention_mask is None:
+                # The compressed upper-triangular mask with sparse_mode=3
+                # applies FlashAttention's bottom-right causal alignment when
+                # Sq != Skv without materializing the full attention matrix.
+                npu_attention_mask = _get_npu_compressed_causal_mask(query.device)
+                sparse_mode = 3
+                # Explicitly enable invalid-row handling for Sq > Skv. Equal
+                # and shorter query sequences have no fully causal-masked rows.
+                inner_precise = 2 if query.shape[1] > key.shape[1] else 0
+            else:
+                # A custom keep-mask cannot be combined with the compressed
+                # sparse_mode=3 mask. Materialize the composed block-mask and
+                # use allMask mode instead. The explicit mask follows the
+                # framework convention (True=keep), while npu_fusion_attention
+                # uses True=block.
+                explicit_keep_mask = _maybe_reshape_attn_mask(
+                    query,
+                    key,
+                    attention_mask,
+                    mask_mode="full_qk",
+                ).to(device=query.device, dtype=torch.bool)
+                query_positions = torch.arange(query.shape[1], device=query.device).unsqueeze(1)
+                key_positions = torch.arange(key.shape[1], device=query.device).unsqueeze(0)
+                causal_block_mask = key_positions > query_positions + (key.shape[1] - query.shape[1])
+                npu_attention_mask = (causal_block_mask | ~explicit_keep_mask).contiguous()
+                sparse_mode = 1
+                # An explicit mask can create fully masked rows regardless of
+                # the Sq/Skv relationship.
+                inner_precise = 2
+
+            return torch_npu.npu_fusion_attention(
+                query.contiguous(),
+                key.contiguous(),
+                value.contiguous(),
+                head_num=query.shape[2],
+                input_layout="BSND",
+                atten_mask=npu_attention_mask,
+                scale=float(self.softmax_scale),
+                keep_prob=1.0,
+                inner_precise=inner_precise,
+                sparse_mode=sparse_mode,
+            )[0]
+
         from mindiesd import attention_forward
 
         # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
@@ -567,7 +623,6 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 out = self._forward_varlen_packed_npu(query, key, value, extra)
             if out is not None:
                 return out
-        attention_mask = attn_metadata.attn_mask if attn_metadata else None
         if attention_mask is None and extra.get("npu_attn_varlen", False):
             # Models skip mask construction when this opt-in is set
             # (FlashAttentionBackend.supports_packed_mask_free). The packed