wehub-resource-sync
559 lines
18 KiB
Python
559 lines
18 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from op_test import get_cuda_version, is_custom_device
|
|
|
|
import paddle
|
|
import paddle.nn.functional as F
|
|
from paddle.nn.attention import (
|
|
SDPBackend,
|
|
_cur_sdpa_kernel_backends,
|
|
sdpa_kernel,
|
|
)
|
|
from paddle.nn.functional import scaled_dot_product_attention
|
|
|
|
|
|
def is_flashattn_supported():
|
|
if (
|
|
not paddle.base.core.is_compiled_with_cuda()
|
|
or get_cuda_version() < 11040
|
|
):
|
|
return False
|
|
|
|
if paddle.device.cuda.device_count() == 0:
|
|
return False
|
|
|
|
try:
|
|
capability = paddle.device.cuda.get_device_capability()
|
|
major, minor = capability[0], capability[1]
|
|
# Support sm8x or sm90
|
|
return (major == 8 and minor >= 0) or (major == 9 and minor == 0)
|
|
except:
|
|
return False
|
|
|
|
|
|
def attention_naive(q, k, v, causal=False):
|
|
"""Reference implementation for attention calculation."""
|
|
qt = paddle.transpose(q, [0, 2, 1, 3])
|
|
kt = paddle.transpose(k, [0, 2, 1, 3])
|
|
vt = paddle.transpose(v, [0, 2, 1, 3])
|
|
scale = 1.0 / np.sqrt(q.shape[-1])
|
|
s = paddle.matmul(qt * scale, paddle.transpose(kt, [0, 1, 3, 2]))
|
|
if causal:
|
|
mask = paddle.triu(paddle.ones_like(s) * -float('inf'), diagonal=1)
|
|
s = s + mask
|
|
p = F.softmax(s)
|
|
o = paddle.matmul(p, vt)
|
|
return paddle.transpose(o, [0, 2, 1, 3])
|
|
|
|
|
|
@unittest.skipIf(
|
|
paddle.is_compiled_with_xpu(),
|
|
"sdpa backend selection logic fails on XPU when testing CPU place",
|
|
)
|
|
class TestSDPAKernelCPU(unittest.TestCase):
|
|
"""Test sdpa_kernel on CPU specifically."""
|
|
|
|
def setUp(self):
|
|
self.place = paddle.CPUPlace()
|
|
self.shape = (2, 128, 8, 16)
|
|
self.dtype = 'float32'
|
|
|
|
def test_cpu_math_backend(self):
|
|
"""Test MATH backend on CPU."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
key = np.random.random(self.shape).astype(self.dtype)
|
|
value = np.random.random(self.shape).astype(self.dtype)
|
|
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
q_ = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
k_ = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
v_ = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
with sdpa_kernel(SDPBackend.MATH):
|
|
out = scaled_dot_product_attention(q, k, v)
|
|
|
|
ref_out = attention_naive(q_, k_, v_, causal=False)
|
|
np.testing.assert_allclose(
|
|
out.numpy(), ref_out.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
# Test backward
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.numpy(), q_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
np.testing.assert_allclose(
|
|
k.grad.numpy(), k_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
np.testing.assert_allclose(
|
|
v.grad.numpy(), v_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
def test_cpu_with_mask(self):
|
|
"""Test CPU with attention mask."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
# Create a mask
|
|
mask_shape = (self.shape[0], 1, self.shape[1], self.shape[1])
|
|
mask = np.random.random(mask_shape).astype(self.dtype)
|
|
m = paddle.to_tensor(mask, place=self.place, dtype=self.dtype)
|
|
|
|
with sdpa_kernel(SDPBackend.MATH):
|
|
out = scaled_dot_product_attention(q, q, q, attn_mask=m)
|
|
|
|
# Verify output shape and test backward
|
|
self.assertEqual(out.shape, q.shape)
|
|
out.backward()
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (paddle.is_compiled_with_cuda() or is_custom_device())
|
|
or paddle.is_compiled_with_rocm(),
|
|
"CUDA is not available, this test requires GPU support.",
|
|
)
|
|
class TestSDPAKernelBasic(unittest.TestCase):
|
|
"""Test basic functionality of sdpa_kernel context manager (defaults to available device)."""
|
|
|
|
def setUp(self):
|
|
self.shape = (2, 128, 8, 16)
|
|
self.dtype = 'float32'
|
|
|
|
def test_cur_sdpa_kernel_backends(self):
|
|
result = _cur_sdpa_kernel_backends()
|
|
self.assertIsInstance(result, list)
|
|
|
|
def test_single_backend(self):
|
|
"""Test with single backend."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
key = np.random.random(self.shape).astype(self.dtype)
|
|
value = np.random.random(self.shape).astype(self.dtype)
|
|
|
|
q = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
k = paddle.to_tensor(key, dtype=self.dtype, stop_gradient=False)
|
|
v = paddle.to_tensor(value, dtype=self.dtype, stop_gradient=False)
|
|
|
|
q_ = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
k_ = paddle.to_tensor(key, dtype=self.dtype, stop_gradient=False)
|
|
v_ = paddle.to_tensor(value, dtype=self.dtype, stop_gradient=False)
|
|
|
|
with sdpa_kernel(SDPBackend.MATH):
|
|
out = scaled_dot_product_attention(q, k, v)
|
|
|
|
ref_out = attention_naive(q_, k_, v_, causal=False)
|
|
np.testing.assert_allclose(
|
|
out.numpy(), ref_out.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
# Test backward
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.numpy(), q_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
np.testing.assert_allclose(
|
|
k.grad.numpy(), k_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
np.testing.assert_allclose(
|
|
v.grad.numpy(), v_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
def test_multiple_backends(self):
|
|
"""Test with multiple backends."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
key = np.random.random(self.shape).astype(self.dtype)
|
|
value = np.random.random(self.shape).astype(self.dtype)
|
|
|
|
q = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
k = paddle.to_tensor(key, dtype=self.dtype, stop_gradient=False)
|
|
v = paddle.to_tensor(value, dtype=self.dtype, stop_gradient=False)
|
|
|
|
q_ = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
k_ = paddle.to_tensor(key, dtype=self.dtype, stop_gradient=False)
|
|
v_ = paddle.to_tensor(value, dtype=self.dtype, stop_gradient=False)
|
|
|
|
# Test with multiple backends
|
|
backends = [SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]
|
|
with sdpa_kernel(backends):
|
|
out = scaled_dot_product_attention(q, k, v)
|
|
|
|
ref_out = attention_naive(q_, k_, v_, causal=False)
|
|
np.testing.assert_allclose(
|
|
out.numpy(), ref_out.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
# Test backward
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.numpy(), q_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
np.testing.assert_allclose(
|
|
k.grad.numpy(), k_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
np.testing.assert_allclose(
|
|
v.grad.numpy(), v_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
def test_multiple_backends_with_priority(self):
|
|
"""
|
|
Test set_priority=True with available backends (MATH, EFFICIENT).
|
|
"""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
key = np.random.random(self.shape).astype(self.dtype)
|
|
value = np.random.random(self.shape).astype(self.dtype)
|
|
|
|
q = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
k = paddle.to_tensor(key, dtype=self.dtype, stop_gradient=False)
|
|
v = paddle.to_tensor(value, dtype=self.dtype, stop_gradient=False)
|
|
|
|
q_ = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
k_ = paddle.to_tensor(key, dtype=self.dtype, stop_gradient=False)
|
|
v_ = paddle.to_tensor(value, dtype=self.dtype, stop_gradient=False)
|
|
|
|
backends = [SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]
|
|
|
|
with sdpa_kernel(backends, set_priority=True):
|
|
out = scaled_dot_product_attention(q, k, v)
|
|
|
|
ref_out = attention_naive(q_, k_, v_, causal=False)
|
|
np.testing.assert_allclose(
|
|
out.numpy(), ref_out.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.numpy(), q_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
np.testing.assert_allclose(
|
|
k.grad.numpy(), k_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
np.testing.assert_allclose(
|
|
v.grad.numpy(), v_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not is_flashattn_supported(),
|
|
"Priority test requires flash attention support (CUDA SM80+)",
|
|
)
|
|
class TestSDPAKernelPriority(unittest.TestCase):
|
|
"""Test priority settings for sdpa_kernel."""
|
|
|
|
def setUp(self):
|
|
self.shape = (2, 64, 4, 32)
|
|
self.dtype = 'float16'
|
|
|
|
def test_set_priority_true(self):
|
|
"""Test set_priority=True."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
q = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
q_ = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
|
|
backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]
|
|
with sdpa_kernel(backends, set_priority=True):
|
|
out = scaled_dot_product_attention(q, q, q)
|
|
|
|
# Verify output correctness
|
|
ref_out = attention_naive(q_, q_, q_, causal=False)
|
|
np.testing.assert_allclose(
|
|
out.numpy(), ref_out.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
# Test backward
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.numpy(), q_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
def test_set_priority_false(self):
|
|
"""Test set_priority=False (default)."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
q = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
q_ = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
|
|
|
|
backends = [SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]
|
|
with sdpa_kernel(backends, set_priority=False):
|
|
out = scaled_dot_product_attention(q, q, q)
|
|
|
|
ref_out = attention_naive(q_, q_, q_, causal=False)
|
|
np.testing.assert_allclose(
|
|
out.numpy(), ref_out.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
# Test backward
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.numpy(), q_.grad.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
|
|
class TestSDPAKernelExceptions(unittest.TestCase):
|
|
"""Test exception handling in sdpa_kernel."""
|
|
|
|
def test_invalid_backend_type(self):
|
|
"""Test with invalid backend type."""
|
|
with self.assertRaises(AssertionError), sdpa_kernel("invalid_backend"):
|
|
pass
|
|
|
|
def test_invalid_backend_in_list(self):
|
|
"""Test with invalid backend in list."""
|
|
with (
|
|
self.assertRaises(TypeError),
|
|
sdpa_kernel([SDPBackend.MATH, "invalid"]),
|
|
):
|
|
pass
|
|
|
|
def test_empty_backend_list(self):
|
|
"""Test with empty backend list."""
|
|
with self.assertRaises(ValueError), sdpa_kernel([]):
|
|
pass
|
|
|
|
|
|
@unittest.skipIf(
|
|
not is_flashattn_supported(),
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.4"
|
|
"and device's compute capability must be 8.x or 90",
|
|
)
|
|
class TestSDPAKernelGPU(unittest.TestCase):
|
|
"""Test sdpa_kernel on GPU with different backends."""
|
|
|
|
def setUp(self):
|
|
self.place = paddle.CUDAPlace(0)
|
|
self.shape = (2, 128, 8, 32)
|
|
self.dtype = 'float16'
|
|
|
|
def test_gpu_math_backend(self):
|
|
"""Test MATH backend on GPU."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
key = np.random.random(self.shape).astype(self.dtype)
|
|
value = np.random.random(self.shape).astype(self.dtype)
|
|
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
q_ = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
k_ = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
v_ = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
with sdpa_kernel(SDPBackend.MATH):
|
|
out = scaled_dot_product_attention(q, k, v)
|
|
|
|
# Convert to float32 for comparison
|
|
q_fp32 = q_.astype('float32')
|
|
k_fp32 = k_.astype('float32')
|
|
v_fp32 = v_.astype('float32')
|
|
ref_out = attention_naive(q_fp32, k_fp32, v_fp32, causal=False)
|
|
|
|
np.testing.assert_allclose(
|
|
out.astype('float32').numpy(), ref_out.numpy(), rtol=5e-3, atol=1e-3
|
|
)
|
|
|
|
# Test backward
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.astype('float32').numpy(),
|
|
q_.grad.numpy(),
|
|
rtol=5e-3,
|
|
atol=1e-3,
|
|
)
|
|
|
|
def test_flash_attention_backend(self):
|
|
"""Test FLASH_ATTENTION backend on GPU."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
key = np.random.random(self.shape).astype(self.dtype)
|
|
value = np.random.random(self.shape).astype(self.dtype)
|
|
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
q_ = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
k_ = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
v_ = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
try:
|
|
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
|
|
out = scaled_dot_product_attention(q, k, v)
|
|
|
|
# Convert to float32 for comparison
|
|
q_fp32 = q_.astype('float32')
|
|
k_fp32 = k_.astype('float32')
|
|
v_fp32 = v_.astype('float32')
|
|
ref_out = attention_naive(q_fp32, k_fp32, v_fp32, causal=False)
|
|
|
|
np.testing.assert_allclose(
|
|
out.astype('float32').numpy(),
|
|
ref_out.numpy(),
|
|
rtol=5e-3,
|
|
atol=1e-3,
|
|
)
|
|
|
|
# Test backward
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.astype('float32').numpy(),
|
|
q_.grad.numpy(),
|
|
rtol=5e-3,
|
|
atol=1e-3,
|
|
)
|
|
except RuntimeError:
|
|
# Flash attention might not be available
|
|
self.skipTest("Flash attention not available on this GPU")
|
|
|
|
def test_efficient_attention_backend(self):
|
|
"""Test EFFICIENT_ATTENTION backend on GPU."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
q_ = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
try:
|
|
with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
|
|
out = scaled_dot_product_attention(q, q, q)
|
|
|
|
# Convert to float32 for comparison
|
|
q_fp32 = q_.astype('float32')
|
|
ref_out = attention_naive(q_fp32, q_fp32, q_fp32, causal=False)
|
|
|
|
np.testing.assert_allclose(
|
|
out.astype('float32').numpy(),
|
|
ref_out.numpy(),
|
|
rtol=5e-3,
|
|
atol=1e-3,
|
|
)
|
|
|
|
# Test backward
|
|
out.backward()
|
|
ref_out.backward()
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.astype('float32').numpy(),
|
|
q_.grad.numpy(),
|
|
rtol=5e-3,
|
|
atol=1e-3,
|
|
)
|
|
except RuntimeError:
|
|
# Efficient attention might not be available
|
|
self.skipTest("Efficient attention not available on this GPU")
|
|
|
|
def test_all_backends_gpu(self):
|
|
"""Test all backends on GPU."""
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape).astype(self.dtype)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
backends = [
|
|
SDPBackend.FLASH_ATTENTION,
|
|
SDPBackend.EFFICIENT_ATTENTION,
|
|
SDPBackend.MATH,
|
|
]
|
|
|
|
with sdpa_kernel(backends):
|
|
out = scaled_dot_product_attention(q, q, q)
|
|
|
|
# Verify output shape and test backward
|
|
self.assertEqual(out.shape, q.shape)
|
|
out.backward()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|