From 767cc7977e04dc1f1ae7307e630429e9169af622 Mon Sep 17 00:00:00 2001 From: Oliver7th <97155135+Oliver7th@users.noreply.github.com> Date: Mon, 14 Sep 2026 09:46:29 +0800 Subject: [PATCH] [Bugfix][NPU] Materialize causal mask in FlashAttention dense path (#7324) Signed-off-by: Oliver7th --- tests/diffusion/attention/test_flash_attn.py | 120 +++++++++++++++++- .../attention/test_flash_attn_npu.py | 118 +++++++++++++++++ .../attention/backends/flash_attn.py | 59 ++++++++- 3 files changed, 291 insertions(+), 6 deletions(-) create mode 100644 tests/diffusion/attention/test_flash_attn_npu.py diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index be78ad690..183bbd94c 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/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): diff --git a/tests/diffusion/attention/test_flash_attn_npu.py b/tests/diffusion/attention/test_flash_attn_npu.py new file mode 100644 index 000000000..c971e251b --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index cef6a97f5..d1f09b51a 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/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