wehub-resource-sync
344 lines
14 KiB
Python
344 lines
14 KiB
Python
# Copyright (c) 2024 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.
|
|
|
|
# [AUTO-GENERATED]
|
|
# Target file: python/paddle/nn/functional/sparse_attention.py
|
|
# Coverage target: sparse_attention function (both dynamic and static paths)
|
|
# 未覆盖行: static graph path (LayerHelper branch)
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.nn.functional.sparse_attention import sparse_attention
|
|
|
|
|
|
class TestSparseAttentionBasic(unittest.TestCase):
|
|
"""Test sparse_attention basic functionality.
|
|
测试 sparse_attention 基本功能。"""
|
|
|
|
def setUp(self):
|
|
paddle.disable_static()
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_basic(self):
|
|
"""Test sparse_attention with dense sparse pattern (all entries).
|
|
测试全密集稀疏模式的 sparse_attention。"""
|
|
try:
|
|
# batch=1, num_heads=1, seq_len=4, head_dim=2
|
|
query = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
key = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
value = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
# Dense CSR: every position attends to all positions
|
|
offset = paddle.to_tensor([[[0, 4, 8, 12, 16]]], dtype="int32")
|
|
columns = paddle.to_tensor(
|
|
[[[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]],
|
|
dtype="int32",
|
|
)
|
|
out = sparse_attention(query, key, value, offset, columns)
|
|
self.assertEqual(out.shape, [1, 1, 4, 2])
|
|
self.assertEqual(out.dtype, paddle.float32)
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_output_shape(self):
|
|
"""Test sparse_attention output shape matches input.
|
|
测试 sparse_attention 输出形状与输入匹配。"""
|
|
try:
|
|
batch_size, num_heads, seq_len, head_dim = 2, 4, 8, 16
|
|
query = paddle.randn(
|
|
[batch_size, num_heads, seq_len, head_dim], dtype="float32"
|
|
)
|
|
key = paddle.randn(
|
|
[batch_size, num_heads, seq_len, head_dim], dtype="float32"
|
|
)
|
|
value = paddle.randn(
|
|
[batch_size, num_heads, seq_len, head_dim], dtype="float32"
|
|
)
|
|
# Each position attends to 2 positions
|
|
nnz_per_row = 2
|
|
nnz_per_head = seq_len * nnz_per_row
|
|
offset = paddle.zeros(
|
|
[batch_size, num_heads, seq_len + 1], dtype="int32"
|
|
)
|
|
for b in range(batch_size):
|
|
for h in range(num_heads):
|
|
for s in range(seq_len):
|
|
offset[b, h, s + 1] = offset[b, h, s] + nnz_per_row
|
|
columns = paddle.zeros(
|
|
[batch_size, num_heads, nnz_per_head], dtype="int32"
|
|
)
|
|
# Each position attends to itself and the next position
|
|
for b in range(batch_size):
|
|
for h in range(num_heads):
|
|
for s in range(seq_len):
|
|
base = offset[b, h, s].item()
|
|
columns[b, h, base] = s
|
|
columns[b, h, base + 1] = min(s + 1, seq_len - 1)
|
|
out = sparse_attention(query, key, value, offset, columns)
|
|
self.assertEqual(
|
|
out.shape, [batch_size, num_heads, seq_len, head_dim]
|
|
)
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_with_key_padding_mask(self):
|
|
"""Test sparse_attention with key_padding_mask.
|
|
测试带 key_padding_mask 的 sparse_attention。"""
|
|
try:
|
|
query = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
key = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
value = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
offset = paddle.to_tensor([[[0, 2, 4, 6, 8]]], dtype="int32")
|
|
columns = paddle.to_tensor(
|
|
[[[0, 1, 0, 1, 2, 3, 2, 3]]], dtype="int32"
|
|
)
|
|
key_padding_mask = paddle.to_tensor([[1, 1, 1, 0]], dtype="float32")
|
|
out = sparse_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
offset,
|
|
columns,
|
|
key_padding_mask=key_padding_mask,
|
|
)
|
|
self.assertEqual(out.shape, [1, 1, 4, 2])
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_with_attn_mask(self):
|
|
"""Test sparse_attention with attention mask.
|
|
测试带 attention mask 的 sparse_attention。"""
|
|
try:
|
|
query = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
key = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
value = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
offset = paddle.to_tensor([[[0, 2, 4, 6, 8]]], dtype="int32")
|
|
columns = paddle.to_tensor(
|
|
[[[0, 1, 0, 1, 2, 3, 2, 3]]], dtype="int32"
|
|
)
|
|
attn_mask = paddle.to_tensor(
|
|
[
|
|
[1, 0, 1, 1],
|
|
[1, 1, 1, 1],
|
|
[1, 1, 1, 1],
|
|
[1, 1, 1, 1],
|
|
],
|
|
dtype="float32",
|
|
)
|
|
out = sparse_attention(
|
|
query, key, value, offset, columns, attn_mask=attn_mask
|
|
)
|
|
self.assertEqual(out.shape, [1, 1, 4, 2])
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_with_both_masks(self):
|
|
"""Test sparse_attention with both key_padding_mask and attn_mask.
|
|
测试同时带有 key_padding_mask 和 attn_mask 的 sparse_attention。"""
|
|
try:
|
|
query = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
key = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
value = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32"
|
|
)
|
|
offset = paddle.to_tensor([[[0, 2, 4, 6, 8]]], dtype="int32")
|
|
columns = paddle.to_tensor(
|
|
[[[0, 1, 0, 1, 2, 3, 2, 3]]], dtype="int32"
|
|
)
|
|
key_padding_mask = paddle.to_tensor([[1, 1, 1, 0]], dtype="float32")
|
|
attn_mask = paddle.to_tensor(
|
|
[
|
|
[1, 0, 1, 1],
|
|
[1, 1, 1, 1],
|
|
[1, 1, 1, 1],
|
|
[1, 1, 1, 1],
|
|
],
|
|
dtype="float32",
|
|
)
|
|
out = sparse_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
offset,
|
|
columns,
|
|
key_padding_mask=key_padding_mask,
|
|
attn_mask=attn_mask,
|
|
)
|
|
self.assertEqual(out.shape, [1, 1, 4, 2])
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_float64(self):
|
|
"""Test sparse_attention with float64 dtype.
|
|
测试 float64 数据类型的 sparse_attention。"""
|
|
try:
|
|
query = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float64"
|
|
)
|
|
key = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float64"
|
|
)
|
|
value = paddle.to_tensor(
|
|
[[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float64"
|
|
)
|
|
offset = paddle.to_tensor([[[0, 2, 4, 6, 8]]], dtype="int32")
|
|
columns = paddle.to_tensor(
|
|
[[[0, 1, 0, 1, 2, 3, 2, 3]]], dtype="int32"
|
|
)
|
|
out = sparse_attention(query, key, value, offset, columns)
|
|
self.assertEqual(out.dtype, paddle.float64)
|
|
self.assertEqual(out.shape, [1, 1, 4, 2])
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_multi_head(self):
|
|
"""Test sparse_attention with multiple heads.
|
|
测试多头 sparse_attention。"""
|
|
try:
|
|
batch, heads, seq, dim = 1, 4, 4, 8
|
|
query = paddle.randn([batch, heads, seq, dim], dtype="float32")
|
|
key = paddle.randn([batch, heads, seq, dim], dtype="float32")
|
|
value = paddle.randn([batch, heads, seq, dim], dtype="float32")
|
|
# Dense pattern: each row attends to all 4 positions
|
|
nnz_per_row = 4
|
|
nnz_per_head = seq * nnz_per_row
|
|
offset = paddle.zeros([batch, heads, seq + 1], dtype="int32")
|
|
for s in range(seq):
|
|
offset[0, :, s + 1] = offset[0, :, s] + nnz_per_row
|
|
columns = paddle.zeros([batch, heads, nnz_per_head], dtype="int32")
|
|
for h in range(heads):
|
|
for s in range(seq):
|
|
base = offset[0, h, s].item()
|
|
for k in range(nnz_per_row):
|
|
columns[0, h, base + k] = k
|
|
out = sparse_attention(query, key, value, offset, columns)
|
|
self.assertEqual(out.shape, [batch, heads, seq, dim])
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_diagonal_pattern(self):
|
|
"""Test sparse_attention with diagonal (causal-like) sparse pattern.
|
|
测试对角线（类似因果）稀疏模式的 sparse_attention。"""
|
|
try:
|
|
batch, heads, seq, dim = 1, 1, 4, 4
|
|
query = paddle.randn([batch, heads, seq, dim], dtype="float32")
|
|
key = paddle.randn([batch, heads, seq, dim], dtype="float32")
|
|
value = paddle.randn([batch, heads, seq, dim], dtype="float32")
|
|
# Causal pattern: position i can attend to positions 0..i
|
|
offset_np = np.zeros((batch, heads, seq + 1), dtype=np.int32)
|
|
nnz_per_pos = []
|
|
for s in range(seq):
|
|
nnz_per_pos.append(s + 1)
|
|
columns_list = []
|
|
for s in range(seq):
|
|
for k in range(s + 1):
|
|
columns_list.append(k)
|
|
offset_np[0, 0, 1:] = np.cumsum(nnz_per_pos)
|
|
columns_np = np.array(columns_list, dtype=np.int32).reshape(
|
|
1, 1, -1
|
|
)
|
|
offset = paddle.to_tensor(offset_np)
|
|
columns = paddle.to_tensor(columns_np)
|
|
out = sparse_attention(query, key, value, offset, columns)
|
|
self.assertEqual(out.shape, [batch, heads, seq, dim])
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"CUDA 11.3+ required for sparse_attention",
|
|
)
|
|
def test_sparse_attention_different_head_dim(self):
|
|
"""Test sparse_attention with different head dimensions.
|
|
测试不同头维度的 sparse_attention。"""
|
|
try:
|
|
for head_dim in [8, 32, 64, 128]:
|
|
query = paddle.randn([1, 2, 4, head_dim], dtype="float32")
|
|
key = paddle.randn([1, 2, 4, head_dim], dtype="float32")
|
|
value = paddle.randn([1, 2, 4, head_dim], dtype="float32")
|
|
nnz_per_row = 4
|
|
nnz_per_head = 4 * nnz_per_row
|
|
offset = paddle.zeros([1, 2, 5], dtype="int32")
|
|
for s in range(4):
|
|
offset[0, :, s + 1] = offset[0, :, s] + nnz_per_row
|
|
columns = paddle.zeros([1, 2, nnz_per_head], dtype="int32")
|
|
for h in range(2):
|
|
for s in range(4):
|
|
base = offset[0, h, s].item()
|
|
for k in range(nnz_per_row):
|
|
columns[0, h, base + k] = k
|
|
out = sparse_attention(query, key, value, offset, columns)
|
|
self.assertEqual(out.shape, [1, 2, 4, head_dim])
|
|
except RuntimeError as e:
|
|
self.skipTest(f"CUDA kernel not available: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|