1142 lines
43 KiB
Python
1142 lines
43 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""
|
|
Test script for FlashAttention backend with padding handling.
|
|
|
|
This script tests two main scenarios:
|
|
1. Case 1: Comparing padded vs unpadded inputs for batch_size=1
|
|
2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
|
|
"""
|
|
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
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]
|
|
|
|
is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
|
|
HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
|
|
flash_attn_func = fa.flash_attn_func # noqa: N813
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("device_major", "requested", "supported", "expected"),
|
|
[
|
|
(8, None, frozenset((2, 3, 4)), 2),
|
|
(9, None, frozenset((2, 3, 4)), 3),
|
|
(10, None, frozenset((2, 3, 4)), 4),
|
|
(12, None, frozenset((2, 3, 4)), 4),
|
|
(9, None, frozenset((2,)), 2),
|
|
(10, None, frozenset((2, 3)), 3),
|
|
(9, "2", frozenset((2, 3, 4)), 2),
|
|
(9, "4", frozenset((2, 3, 4)), 4),
|
|
],
|
|
)
|
|
@pytest.mark.cpu
|
|
def test_versioned_flash_attention_selection(device_major, requested, supported, expected):
|
|
assert fa._choose_vllm_flash_attn_version(device_major, requested, supported) == expected
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_versioned_flash_attention_selection_rejects_invalid_or_unavailable_version():
|
|
with pytest.raises(ValueError, match="must be 2, 3, or 4"):
|
|
fa._choose_vllm_flash_attn_version(9, "5", frozenset((2, 3, 4)))
|
|
with pytest.raises(RuntimeError, match="FlashAttention 4 is unavailable"):
|
|
fa._choose_vllm_flash_attn_version(9, "4", frozenset((2, 3)))
|
|
|
|
|
|
def create_attention_mask(batch_size: int, seq_len: int, valid_len: int, device: torch.device) -> torch.Tensor:
|
|
"""
|
|
Create attention mask where first valid_len tokens are valid (1) and rest are padding (0).
|
|
|
|
Args:
|
|
batch_size: Batch size
|
|
seq_len: Total sequence length (including padding)
|
|
valid_len: Number of valid (non-padded) tokens
|
|
|
|
Returns:
|
|
Attention mask of shape (batch_size, seq_len)
|
|
"""
|
|
mask = torch.zeros(batch_size, seq_len, dtype=torch.bool, device=device)
|
|
mask[:, :valid_len] = True
|
|
return mask
|
|
|
|
|
|
def pad_tensor(tensor: torch.Tensor, target_seq_len: int, pad_value: float = 0.0) -> torch.Tensor:
|
|
"""
|
|
Pad tensor along sequence dimension (dim=1).
|
|
|
|
Args:
|
|
tensor: Input tensor of shape (batch_size, seq_len, num_heads, head_dim)
|
|
target_seq_len: Target sequence length after padding
|
|
pad_value: Value to use for padding
|
|
|
|
Returns:
|
|
Padded tensor of shape (batch_size, target_seq_len, num_heads, head_dim)
|
|
"""
|
|
batch_size, seq_len, num_heads, head_dim = tensor.shape
|
|
if target_seq_len <= seq_len:
|
|
return tensor
|
|
|
|
padding = torch.full(
|
|
(batch_size, target_seq_len - seq_len, num_heads, head_dim), pad_value, dtype=tensor.dtype, device=tensor.device
|
|
)
|
|
return torch.cat([tensor, padding], dim=1)
|
|
|
|
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
|
|
@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires an accelerator")
|
|
def test_padding_equivalence():
|
|
"""
|
|
Case 1: Test that padded and unpadded inputs produce similar outputs.
|
|
|
|
- Input A: batch_size=1, hidden_states (1, 48), encoder_hidden_states (1, 16)
|
|
Concatenated length: 64, NO attention_mask
|
|
- Input B: Same data but padded: hidden_states (1, 58), encoder_hidden_states (1, 26)
|
|
Concatenated length: 84, WITH attention_mask
|
|
|
|
Expected: Output A and Output B should be very close.
|
|
"""
|
|
device = torch.device(current_omni_platform.device_type)
|
|
dtype = torch.bfloat16
|
|
|
|
# Configuration
|
|
batch_size = 1
|
|
hidden_seq_len = 48
|
|
encoder_seq_len = 16
|
|
pad_length = 10
|
|
num_heads = 8
|
|
head_dim = 64
|
|
|
|
# Initialize FlashAttention
|
|
fa_impl = FlashAttentionImpl(
|
|
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
|
|
)
|
|
|
|
# Create base tensors with random values (same for both A and B)
|
|
torch.manual_seed(42)
|
|
hidden_states_base = torch.randn(batch_size, hidden_seq_len, num_heads, head_dim, device=device, dtype=dtype)
|
|
encoder_hidden_states_base = torch.randn(
|
|
batch_size, encoder_seq_len, num_heads, head_dim, device=device, dtype=dtype
|
|
)
|
|
|
|
# ========== Input A: Unpadded, no attention mask ==========
|
|
query_a = torch.cat([hidden_states_base, encoder_hidden_states_base], dim=1)
|
|
key_a = query_a.clone()
|
|
value_a = query_a.clone()
|
|
|
|
attn_metadata_a = AttentionMetadata(attn_mask=None)
|
|
|
|
output_a = fa_impl.forward(query=query_a, key=key_a, value=value_a, attn_metadata=attn_metadata_a)
|
|
|
|
# ========== Input B: Padded with attention mask ==========
|
|
hidden_states_padded = pad_tensor(hidden_states_base, hidden_seq_len + pad_length)
|
|
encoder_hidden_states_padded = pad_tensor(encoder_hidden_states_base, encoder_seq_len + pad_length)
|
|
|
|
query_b = torch.cat([hidden_states_padded, encoder_hidden_states_padded], dim=1)
|
|
key_b = query_b.clone()
|
|
value_b = query_b.clone()
|
|
|
|
# Create attention mask
|
|
attn_mask_b = torch.cat(
|
|
[
|
|
create_attention_mask(batch_size, hidden_seq_len + pad_length, hidden_seq_len, device),
|
|
create_attention_mask(batch_size, encoder_seq_len + pad_length, encoder_seq_len, device),
|
|
],
|
|
dim=1,
|
|
)
|
|
|
|
attn_metadata_b = AttentionMetadata(attn_mask=attn_mask_b)
|
|
|
|
output_b = fa_impl.forward(query=query_b, key=key_b, value=value_b, attn_metadata=attn_metadata_b)
|
|
|
|
# Extract non-padded portion from output_b
|
|
output_b_unpadded = torch.cat(
|
|
[
|
|
output_b[:, :hidden_seq_len, :, :],
|
|
output_b[:, hidden_seq_len + pad_length : hidden_seq_len + pad_length + encoder_seq_len, :, :],
|
|
],
|
|
dim=1,
|
|
)
|
|
|
|
# Compare outputs
|
|
max_diff = torch.max(torch.abs(output_a - output_b_unpadded)).item()
|
|
mean_diff = torch.mean(torch.abs(output_a - output_b_unpadded)).item()
|
|
|
|
print("\n=== Case 1: Padding Equivalence Test ===")
|
|
print(f"Output A shape: {output_a.shape}")
|
|
print(f"Output B shape: {output_b.shape}")
|
|
print(f"Output B unpadded shape: {output_b_unpadded.shape}")
|
|
print(f"Max absolute difference: {max_diff:.6f}")
|
|
print(f"Mean absolute difference: {mean_diff:.6f}")
|
|
|
|
# Assert that outputs are close
|
|
# Using higher tolerance for bfloat16
|
|
assert max_diff < 0.1, f"Max difference {max_diff} exceeds threshold 0.1"
|
|
assert mean_diff < 0.01, f"Mean difference {mean_diff} exceeds threshold 0.01"
|
|
|
|
print("✓ Case 1 PASSED: Padded and unpadded outputs are very close!")
|
|
|
|
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
|
|
@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires an accelerator")
|
|
def test_fa_vs_sdpa():
|
|
"""
|
|
Case 2: Compare FlashAttention and SDPA backends with padding.
|
|
|
|
- batch_size=2
|
|
- hidden_states: (2, 48) padded to (2, 58)
|
|
- encoder_hidden_states: (2, 16) padded to (2, 26)
|
|
- Concatenated length: 84
|
|
- Compare FA and SDPA outputs
|
|
|
|
Expected: FA and SDPA outputs should be very close.
|
|
"""
|
|
device = torch.device(current_omni_platform.device_type)
|
|
dtype = torch.bfloat16
|
|
|
|
# Configuration
|
|
batch_size = 2
|
|
hidden_seq_len = 48
|
|
encoder_seq_len = 16
|
|
pad_length = 10
|
|
num_heads = 8
|
|
head_dim = 64
|
|
|
|
# Initialize both backends
|
|
fa_impl = FlashAttentionImpl(
|
|
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
|
|
)
|
|
|
|
sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
|
|
|
|
# Create base tensors
|
|
torch.manual_seed(123)
|
|
hidden_states_base = torch.randn(batch_size, hidden_seq_len, num_heads, head_dim, device=device, dtype=dtype)
|
|
encoder_hidden_states_base = torch.randn(
|
|
batch_size, encoder_seq_len, num_heads, head_dim, device=device, dtype=dtype
|
|
)
|
|
|
|
# Pad tensors
|
|
hidden_states_padded = pad_tensor(hidden_states_base, hidden_seq_len + pad_length)
|
|
encoder_hidden_states_padded = pad_tensor(encoder_hidden_states_base, encoder_seq_len + pad_length)
|
|
|
|
# Concatenate
|
|
query = torch.cat([hidden_states_padded, encoder_hidden_states_padded], dim=1)
|
|
key = query.clone()
|
|
value = query.clone()
|
|
|
|
# Create attention mask
|
|
attn_mask = torch.cat(
|
|
[
|
|
create_attention_mask(batch_size, hidden_seq_len + pad_length, hidden_seq_len, device),
|
|
create_attention_mask(batch_size, encoder_seq_len + pad_length, encoder_seq_len, device),
|
|
],
|
|
dim=1,
|
|
)
|
|
|
|
attn_metadata = AttentionMetadata(attn_mask=attn_mask)
|
|
|
|
# Run FlashAttention
|
|
output_fa = fa_impl.forward(query=query.clone(), key=key.clone(), value=value.clone(), attn_metadata=attn_metadata)
|
|
|
|
# Run SDPA
|
|
# SDPA expects 4D attention mask: (batch_size, 1, seq_len, seq_len) or (batch_size, seq_len)
|
|
# For causal=False, we need to convert 2D mask to 4D
|
|
if attn_mask is not None:
|
|
# Expand mask for SDPA: (batch_size, seq_len) -> (batch_size, 1, 1, seq_len)
|
|
attn_mask_4d = attn_mask.unsqueeze(1).unsqueeze(2)
|
|
# Convert bool to float: True -> 0.0, False -> -inf
|
|
attn_mask_float = torch.zeros_like(attn_mask_4d, dtype=dtype)
|
|
attn_mask_float.masked_fill_(~attn_mask_4d, float("-inf"))
|
|
attn_metadata_sdpa = AttentionMetadata(attn_mask=attn_mask_float)
|
|
else:
|
|
attn_metadata_sdpa = AttentionMetadata(attn_mask=None)
|
|
|
|
output_sdpa = sdpa_impl.forward(
|
|
query=query.clone(), key=key.clone(), value=value.clone(), attn_metadata=attn_metadata_sdpa
|
|
)
|
|
|
|
# Compare outputs (only compare valid regions)
|
|
output_fa_valid = torch.cat(
|
|
[
|
|
output_fa[:, :hidden_seq_len, :, :],
|
|
output_fa[:, hidden_seq_len + pad_length : hidden_seq_len + pad_length + encoder_seq_len, :, :],
|
|
],
|
|
dim=1,
|
|
)
|
|
output_sdpa_valid = torch.cat(
|
|
[
|
|
output_sdpa[:, :hidden_seq_len, :, :],
|
|
output_sdpa[:, hidden_seq_len + pad_length : hidden_seq_len + pad_length + encoder_seq_len, :, :],
|
|
],
|
|
dim=1,
|
|
)
|
|
|
|
max_diff = torch.max(torch.abs(output_fa_valid - output_sdpa_valid)).item()
|
|
mean_diff = torch.mean(torch.abs(output_fa_valid - output_sdpa_valid)).item()
|
|
|
|
print("\n=== Case 2: FA vs SDPA Comparison ===")
|
|
print(f"Batch size: {batch_size}")
|
|
print(f"FA output shape: {output_fa.shape}")
|
|
print(f"SDPA output shape: {output_sdpa.shape}")
|
|
print(f"Max absolute difference (valid region): {max_diff:.6f}")
|
|
print(f"Mean absolute difference (valid region): {mean_diff:.6f}")
|
|
|
|
# AITER's ROCm BF16 kernel is deterministic, but differs from ROCm SDPA by
|
|
# about 1.3% relative error. An absolute-only bound misclassifies that
|
|
# expected rounding difference for larger output values (for example,
|
|
# -2.65625 vs -2.6875). Keep the tighter historical bounds elsewhere.
|
|
if current_omni_platform.is_rocm():
|
|
torch.testing.assert_close(output_fa_valid, output_sdpa_valid, rtol=1e-2, atol=1.5e-2)
|
|
else:
|
|
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
|
|
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
|
|
|
|
print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")
|
|
|
|
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
|
|
@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires an accelerator")
|
|
def test_flash_attn_func_preferred_over_varlen():
|
|
"""Test flash_attn_func availability and basic forward call."""
|
|
if not HAS_FLASH_ATTN:
|
|
pytest.skip("No Flash Attention available")
|
|
|
|
if flash_attn_func is None:
|
|
pytest.skip("flash_attn_func not available, will use varlen fallback")
|
|
|
|
device = torch.device(current_omni_platform.device_type)
|
|
dtype = torch.bfloat16
|
|
|
|
num_heads, head_dim = 8, 64
|
|
fa_impl = FlashAttentionImpl(
|
|
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
|
|
)
|
|
|
|
torch.manual_seed(42)
|
|
q = torch.randn(1, 32, num_heads, head_dim, device=device, dtype=dtype)
|
|
k = q.clone()
|
|
v = q.clone()
|
|
|
|
attn_metadata = AttentionMetadata(attn_mask=None)
|
|
output = fa_impl.forward(q, k, v, attn_metadata)
|
|
|
|
assert output.shape == q.shape
|
|
assert not torch.isnan(output).any()
|
|
print("✓ flash_attn_func forward works correctly!")
|
|
|
|
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
|
|
@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires an accelerator")
|
|
@pytest.mark.parametrize("k_len", [40, 256])
|
|
def test_cross_attn_key_padding_vs_sdpa(k_len):
|
|
"""
|
|
Case 3: Cross-attention with a key-padding mask.
|
|
|
|
The mask covers only the key sequence; every query position must attend
|
|
to the valid keys. FA and SDPA outputs should be very close. k_len=256
|
|
covers equal Q/K lengths, which must not be mistaken for self-attention,
|
|
and valid key lengths differ per batch element.
|
|
"""
|
|
device = torch.device(current_omni_platform.device_type)
|
|
dtype = torch.bfloat16
|
|
|
|
batch_size = 2
|
|
q_len = 256
|
|
valid_lens = (17, 29)
|
|
num_heads = 8
|
|
head_dim = 64
|
|
|
|
fa_impl = FlashAttentionImpl(
|
|
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross"
|
|
)
|
|
sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
|
|
|
|
torch.manual_seed(7)
|
|
query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype)
|
|
key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
|
|
value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
|
|
|
|
attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device)
|
|
for i, valid_len in enumerate(valid_lens):
|
|
attn_mask[i, :valid_len] = True
|
|
attn_metadata = AttentionMetadata(attn_mask=attn_mask)
|
|
|
|
output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
|
|
output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
|
|
|
|
max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item()
|
|
mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item()
|
|
|
|
print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===")
|
|
print(f"Max absolute difference: {max_diff:.6f}")
|
|
print(f"Mean absolute difference: {mean_diff:.6f}")
|
|
|
|
# Match the platform-specific comparison used by test_fa_vs_sdpa above.
|
|
if current_omni_platform.is_rocm():
|
|
# AITER FlashAttention and PyTorch SDPA use different BF16 reduction
|
|
# implementations on ROCm, so individual values can differ by one BF16
|
|
# step even though the outputs agree within the expected relative error.
|
|
torch.testing.assert_close(output_fa, output_sdpa, rtol=1e-2, atol=1.5e-2)
|
|
else:
|
|
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
|
|
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
|
|
|
|
print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_varlen_masked_routing_by_role(monkeypatch):
|
|
"""The unpad route is picked by role, not Q/K length equality; runs without a GPU."""
|
|
calls = []
|
|
|
|
def fake_varlen_func(q, k, v, **kwargs):
|
|
calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"]))
|
|
return q
|
|
|
|
monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
|
|
monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
|
|
|
|
batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8
|
|
query = torch.randn(batch_size, seq_len, num_heads, head_dim)
|
|
metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]]))
|
|
|
|
cross_impl = FlashAttentionImpl(
|
|
num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross"
|
|
)
|
|
output = cross_impl.forward_cuda(query, query, query, metadata)
|
|
assert output.shape == query.shape
|
|
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
|
|
assert num_q_rows == batch_size * seq_len
|
|
assert cu_seqlens_q.tolist() == [0, 4, 8]
|
|
assert cu_seqlens_k.tolist() == [0, 2, 5]
|
|
|
|
self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False)
|
|
output = self_impl.forward_cuda(query, query, query, metadata)
|
|
assert output.shape == query.shape
|
|
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
|
|
assert num_q_rows == 5
|
|
assert cu_seqlens_q.tolist() == [0, 2, 5]
|
|
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 = []
|
|
|
|
def fake_varlen_func(q, k, v, **kwargs):
|
|
assert q.shape[1] == 4
|
|
assert k.shape[1] == 2
|
|
assert v.shape[1] == 2
|
|
calls.append(kwargs)
|
|
return q
|
|
|
|
monkeypatch.setattr(fa, "flash_attn_func", None)
|
|
monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
|
|
monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
|
|
|
|
impl = FlashAttentionImpl(num_heads=4, num_kv_heads=2, head_size=4, softmax_scale=0.5, causal=False)
|
|
query = torch.randn(1, 3, 4, 4)
|
|
key = torch.randn(1, 3, 2, 4)
|
|
value = torch.randn(1, 3, 2, 4)
|
|
metadata = AttentionMetadata(full_attn_spans=[[(0, 3)]])
|
|
|
|
output = impl.forward_cuda(query, key, value, metadata)
|
|
|
|
assert output.shape == query.shape
|
|
assert len(calls) == 1
|
|
assert calls[0]["max_seqlen_q"] == 3
|
|
assert calls[0]["max_seqlen_k"] == 3
|
|
assert calls[0]["causal"] is False
|
|
assert calls[0]["softmax_scale"] == 0.5
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch):
|
|
calls = []
|
|
|
|
def fake_varlen_func(q, k, v, **kwargs):
|
|
calls.append((q.shape, k.shape, v.shape, kwargs))
|
|
return q
|
|
|
|
monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
|
|
monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
|
|
|
|
impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
|
|
query = torch.randn(1, 8, 2, 4)
|
|
cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
|
|
metadata = AttentionMetadata(
|
|
attn_mask=torch.tensor([[True] * 6 + [False] * 2]),
|
|
extra={
|
|
"cu_seqlens_q": cu_seqlens,
|
|
"cu_seqlens_k": cu_seqlens,
|
|
"max_seqlen_q": 6,
|
|
"max_seqlen_k": 6,
|
|
},
|
|
)
|
|
|
|
output = impl.forward_cuda(query, query, query, metadata)
|
|
|
|
assert torch.equal(output, query)
|
|
assert len(calls) == 1
|
|
assert calls[0][0] == torch.Size([8, 2, 4])
|
|
assert calls[0][3]["cu_seqlens_q"] is cu_seqlens
|
|
assert calls[0][3]["max_seqlen_q"] == 6
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_packed_varlen_metadata_must_be_complete(monkeypatch):
|
|
monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
|
|
impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
|
|
query = torch.randn(1, 8, 2, 4)
|
|
metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 8], dtype=torch.int32)})
|
|
|
|
with pytest.raises(ValueError, match="Incomplete packed FlashAttention metadata"):
|
|
impl.forward_cuda(query, query, query, metadata)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# NPU mask-free packed attention (PR #5891).
|
|
#
|
|
# These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
|
|
# - 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 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:
|
|
return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
|
|
|
|
|
|
def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
|
|
"""Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
|
|
|
|
Restored automatically by monkeypatch. Defaults are no-op callables so the
|
|
top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
|
|
instance methods that would otherwise call them.
|
|
"""
|
|
fake = SimpleNamespace(
|
|
attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
|
|
attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
|
|
)
|
|
monkeypatch.setitem(sys.modules, "mindiesd", fake)
|
|
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
|
|
|
|
|
|
@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) -------------
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_resolve_packed_seq_accepts_real_plus_pad_contract():
|
|
impl = _npu_impl()
|
|
q = torch.randn(1, 8, 2, 4)
|
|
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
|
|
extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
|
|
|
|
assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_resolve_packed_seq_single_document_no_padding():
|
|
impl = _npu_impl()
|
|
q = torch.randn(1, 8, 2, 4)
|
|
cu = torch.tensor([0, 8], dtype=torch.int32)
|
|
extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
|
|
|
|
assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"case",
|
|
[
|
|
"causal",
|
|
"missing_key",
|
|
"batch_ne_1",
|
|
"cu_too_long",
|
|
"cu_qk_mismatch",
|
|
"max_seqlen_not_int",
|
|
"used_gt_total",
|
|
"used_zero",
|
|
"pad_longer_than_real",
|
|
],
|
|
)
|
|
@pytest.mark.cpu
|
|
def test_resolve_packed_seq_rejects_malformed_metadata(case):
|
|
"""Any metadata outside the exact [real, pad] contract returns None so the
|
|
caller falls back to the masked path."""
|
|
cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
|
|
cu2 = torch.tensor([0, 8], dtype=torch.int32)
|
|
cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
|
|
cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
|
|
base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
|
|
q = torch.randn(1, 8, 2, 4)
|
|
|
|
impl = _npu_impl(causal=(case == "causal"))
|
|
if case == "missing_key":
|
|
extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5} # 3 of 4 keys absent
|
|
elif case == "batch_ne_1":
|
|
q = torch.randn(2, 8, 2, 4)
|
|
extra = base
|
|
elif case == "cu_too_long":
|
|
extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
|
|
elif case == "cu_qk_mismatch":
|
|
extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
|
|
elif case == "max_seqlen_not_int":
|
|
extra = {**base, "max_seqlen_q": 5.0}
|
|
elif case == "used_gt_total":
|
|
extra = {**base, "max_seqlen_q": 9} # 9 > total length 8
|
|
elif case == "used_zero":
|
|
extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
|
|
elif case == "pad_longer_than_real":
|
|
# real=3, pad=5: the real document must be the longer one.
|
|
extra = {
|
|
"cu_seqlens_q": cu_short,
|
|
"cu_seqlens_k": cu_short,
|
|
"max_seqlen_q": 3,
|
|
"max_seqlen_k": 3,
|
|
}
|
|
else:
|
|
extra = base
|
|
|
|
assert impl._resolve_packed_seq_npu(q, q, extra) is None
|
|
|
|
|
|
# --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
|
|
|
|
|
|
def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
|
|
"""FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
|
|
|
|
Lets a test assert which branch ``forward_fa_npu`` dispatched to without
|
|
running any real kernel.
|
|
"""
|
|
impl = _npu_impl()
|
|
sentinel = torch.tensor([1.0])
|
|
impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
|
|
impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
|
|
return impl, sentinel
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_npu_env_unset_uses_varlen(monkeypatch):
|
|
monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
|
|
_fake_mindiesd(monkeypatch)
|
|
impl, _ = _npu_impl_with_mocked_paths()
|
|
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
|
|
|
|
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
|
|
|
|
impl._forward_varlen_packed_npu.assert_called_once()
|
|
impl._forward_prefix_kv_slice_npu.assert_not_called()
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_npu_env_laser_uses_slice(monkeypatch):
|
|
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
|
|
_fake_mindiesd(monkeypatch)
|
|
impl, _ = _npu_impl_with_mocked_paths()
|
|
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
|
|
|
|
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
|
|
|
|
impl._forward_prefix_kv_slice_npu.assert_called_once()
|
|
impl._forward_varlen_packed_npu.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
|
|
@pytest.mark.cpu
|
|
def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
|
|
"""Locks in the CURRENT dispatch behavior (PR #5891).
|
|
|
|
Any value other than ``ascend_laser_attention`` — including the valid
|
|
``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
|
|
the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
|
|
comment #1 on the PR and is intentionally captured here as a regression
|
|
net: update this test when the dispatch starts honoring/rejecting these.
|
|
"""
|
|
monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
|
|
_fake_mindiesd(monkeypatch)
|
|
impl, _ = _npu_impl_with_mocked_paths()
|
|
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
|
|
|
|
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
|
|
|
|
impl._forward_varlen_packed_npu.assert_called_once()
|
|
impl._forward_prefix_kv_slice_npu.assert_not_called()
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
|
|
"""Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
|
|
the existing masked attention_forward path; the new branches are skipped."""
|
|
fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
|
|
_fake_mindiesd(monkeypatch, attention_forward=fake_forward)
|
|
impl, _ = _npu_impl_with_mocked_paths()
|
|
metadata = AttentionMetadata(
|
|
attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
|
|
extra={}, # no npu_attn_varlen opt-in
|
|
)
|
|
q = torch.randn(1, 8, 2, 4)
|
|
|
|
out = impl.forward_fa_npu(q, q, q, metadata)
|
|
|
|
impl._forward_varlen_packed_npu.assert_not_called()
|
|
impl._forward_prefix_kv_slice_npu.assert_not_called()
|
|
fake_forward.assert_called_once()
|
|
assert out is fake_forward.return_value
|
|
|
|
|
|
# --- 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)],
|
|
)
|
|
@pytest.mark.cpu
|
|
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
|
|
|
|
|
|
@pytest.mark.cpu
|
|
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
|
|
|
|
|
|
@pytest.mark.cpu
|
|
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 --
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
|
|
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
|
|
captured: dict = {}
|
|
|
|
def fake_attention_forward(query, key, value, **kwargs):
|
|
captured["query"] = query
|
|
captured["key"] = key
|
|
captured["value"] = value
|
|
captured.update(kwargs)
|
|
return torch.ones_like(query) # deterministic stand-in kernel output
|
|
|
|
_fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
|
|
impl = _npu_impl() # softmax_scale == 0.5
|
|
q = torch.randn(1, 8, 2, 4)
|
|
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
|
|
extra = {
|
|
"cu_seqlens_q": cu,
|
|
"cu_seqlens_k": cu,
|
|
"max_seqlen_q": 5,
|
|
"max_seqlen_k": 5,
|
|
"laser_input_scale": 256.0,
|
|
}
|
|
|
|
out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
|
|
|
|
# K/V sliced to the valid prefix (used_k == 5); q keeps full length.
|
|
assert captured["key"].shape == (1, 5, 2, 4)
|
|
assert captured["value"].shape == (1, 5, 2, 4)
|
|
assert captured["query"].shape == (1, 8, 2, 4)
|
|
assert captured["attn_mask"] is None
|
|
# Pre-divide q/k/v by the factor and compensate the kernel scale by its
|
|
# square (exact for power-of-two); output scaled back by the factor.
|
|
torch.testing.assert_close(captured["query"], q / 256.0)
|
|
torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
|
|
assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
|
|
torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
|
|
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
|
|
captured: dict = {}
|
|
|
|
def fake_attention_forward(query, key, value, **kwargs):
|
|
captured["query"] = query
|
|
captured.update(kwargs)
|
|
return torch.ones_like(query)
|
|
|
|
_fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
|
|
impl = _npu_impl()
|
|
q = torch.randn(1, 8, 2, 4)
|
|
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
|
|
extra = {
|
|
"cu_seqlens_q": cu,
|
|
"cu_seqlens_k": cu,
|
|
"max_seqlen_q": 5,
|
|
"max_seqlen_k": 5,
|
|
# no laser_input_scale → no pre-scaling
|
|
}
|
|
|
|
out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
|
|
|
|
torch.testing.assert_close(captured["query"], q)
|
|
assert captured["scale"] == pytest.approx(0.5)
|
|
torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("Running FlashAttention Padding Tests...")
|
|
print("=" * 60)
|
|
|
|
# Try to run tests
|
|
if is_gpu:
|
|
try:
|
|
print("\n[Running Case 1: Padding Equivalence for FA]")
|
|
test_padding_equivalence()
|
|
except Exception as e:
|
|
print(f"✗ Case 1 failed: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
|
|
try:
|
|
print("\n[Running Case 2: FA vs SDPA]")
|
|
test_fa_vs_sdpa()
|
|
except Exception as e:
|
|
print(f"✗ Case 2 failed: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
else:
|
|
raise RuntimeError("CUDA is not available")
|
|
print("\n" + "=" * 60)
|
|
print("Test suite completed!")
|