wehub-resource-sync
453 lines
15 KiB
Python
453 lines
15 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.
|
|
|
|
from functools import partial
|
|
|
|
import pytest
|
|
from generate_startend_row_indices import (
|
|
generate_causal_blockwise_mask,
|
|
generate_causal_document_mask,
|
|
generate_document_mask,
|
|
generate_global_sliding_window_mask,
|
|
generate_none_mask,
|
|
generate_prefix_lm_causal_mask,
|
|
generate_prefix_lm_document_mask,
|
|
generate_qk_sparse_mask,
|
|
generate_random_eviction_mask,
|
|
generate_share_question_mask,
|
|
generate_sliding_window_mask,
|
|
startend_row_indices_to_attn_bias,
|
|
)
|
|
from test_util import attention_ref
|
|
|
|
import paddle
|
|
from paddle.nn.functional.flash_attention import flashmask_attention
|
|
|
|
# batch_size, seqlen_q, seqlen_k, nheads, nheads_kv
|
|
shape_cases = [
|
|
(2840, 32, 32, 16, 4),
|
|
(1, 300, 300, 16, 16),
|
|
# (2, 8192, 32768, 32, 4), # this will oom
|
|
# (2, 8192, 8192, 32, 4), # this will oom
|
|
(2, 8192, 8192, 14, 1),
|
|
(2, 16384, 16384, 4, 1),
|
|
(1, 1, 127, 1, 1),
|
|
(1, 128, 127, 1, 1),
|
|
(1, 127, 128, 1, 1),
|
|
(2, 16383, 16384, 4, 1),
|
|
(2, 16384, 16383, 4, 1),
|
|
(2, 1000, 1000, 4, 1),
|
|
(2, 2000, 2000, 4, 1),
|
|
(2, 3000, 3000, 4, 1),
|
|
(1, 4000, 4000, 1, 1),
|
|
(1, 8192, 32768 + 1024, 2, 1),
|
|
(1, 8192, 16384 + 1024, 2, 1),
|
|
# my case
|
|
]
|
|
|
|
|
|
# Generate all combinations for second param
|
|
def generate_shapes():
|
|
for batch_size, seqlen_q, seqlen_k, nheads, nheads_kv in shape_cases:
|
|
if nheads_kv == 1:
|
|
nheads_startend_row_indices_values = [1]
|
|
else:
|
|
nheads_startend_row_indices_values = [1, nheads_kv]
|
|
for nheads_startend_row_indices in nheads_startend_row_indices_values:
|
|
yield (
|
|
batch_size,
|
|
seqlen_q,
|
|
seqlen_k,
|
|
nheads,
|
|
nheads_kv,
|
|
nheads_startend_row_indices,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [paddle.bfloat16])
|
|
@pytest.mark.parametrize("fa_version", [3])
|
|
@pytest.mark.parametrize("d, dv", [(128, 128), (80, 80), (64, 64), (256, 256)])
|
|
@pytest.mark.parametrize(
|
|
"batch_size, seqlen_q, seqlen_k, nheads, nheads_kv, nheads_startend_row_indices",
|
|
list(generate_shapes()),
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"gen_startend_row_indices",
|
|
[
|
|
partial(generate_none_mask, causal=False), # full
|
|
partial(generate_none_mask, causal=True), # causal
|
|
partial(generate_sliding_window_mask), # sliding window
|
|
partial(generate_causal_document_mask), # causal document mask
|
|
partial(generate_document_mask), # document mask
|
|
partial(generate_share_question_mask), # share question mask
|
|
partial(generate_global_sliding_window_mask), # global sliding window
|
|
partial(generate_causal_blockwise_mask), # causal blockwise mask
|
|
partial(generate_prefix_lm_document_mask), # prefix lm document mask
|
|
partial(generate_prefix_lm_causal_mask), # prefix lm causal mask
|
|
partial(generate_qk_sparse_mask), # qk-sparse mask
|
|
partial(generate_random_eviction_mask), # random eviction mask
|
|
],
|
|
)
|
|
def test_flashmask(
|
|
batch_size,
|
|
seqlen_q,
|
|
seqlen_k,
|
|
nheads,
|
|
nheads_kv,
|
|
d,
|
|
dv,
|
|
nheads_startend_row_indices,
|
|
fa_version,
|
|
dtype,
|
|
gen_startend_row_indices,
|
|
softcap=0.0,
|
|
):
|
|
paddle.seed(2024)
|
|
assert nheads % nheads_kv == 0
|
|
q_ref = paddle.randn(shape=[batch_size, seqlen_q, nheads, d], dtype=dtype)
|
|
k_ref = paddle.randn(
|
|
shape=[batch_size, seqlen_k, nheads_kv, d], dtype=dtype
|
|
)
|
|
v_ref = paddle.randn(
|
|
shape=[batch_size, seqlen_k, nheads_kv, dv], dtype=dtype
|
|
)
|
|
|
|
q_ref.stop_gradient = False
|
|
k_ref.stop_gradient = False
|
|
v_ref.stop_gradient = False
|
|
|
|
q_bf16, k_bf16, v_bf16 = [x.detach().clone() for x in (q_ref, k_ref, v_ref)]
|
|
|
|
q_bf16.stop_gradient = False
|
|
k_bf16.stop_gradient = False
|
|
v_bf16.stop_gradient = False
|
|
|
|
q, k, v = [x.detach().clone() for x in (q_ref, k_ref, v_ref)]
|
|
|
|
q.stop_gradient = False
|
|
k.stop_gradient = False
|
|
v.stop_gradient = False
|
|
|
|
startend_row_indices, causal = gen_startend_row_indices(
|
|
batch_size, seqlen_q, seqlen_k, nheads_startend_row_indices
|
|
)
|
|
|
|
if startend_row_indices is None and causal and d == 80:
|
|
pytest.skip(
|
|
"Skipping because running headdim 80 with flash_attn in causal mask"
|
|
)
|
|
|
|
attn_bias = startend_row_indices_to_attn_bias(
|
|
startend_row_indices, seqlen_q, nheads, dtype, causal
|
|
)
|
|
|
|
out_ref, attn_ref = attention_ref(
|
|
q_ref, k_ref, v_ref, causal=causal, attn_bias=attn_bias
|
|
)
|
|
|
|
out_bf16, attn_bf16 = attention_ref(
|
|
q_bf16,
|
|
k_bf16,
|
|
v_bf16,
|
|
causal=causal,
|
|
attn_bias=attn_bias,
|
|
upcast=False,
|
|
reorder_ops=True,
|
|
)
|
|
|
|
# # Numerical error if we just do any arithmetic on out_ref
|
|
fwd_atol = 2 * (out_ref + 0.3 - 0.3 - out_ref).abs().max().item()
|
|
assert softcap == 0.0
|
|
rtol = 2 if softcap == 0.0 else 3
|
|
|
|
print(
|
|
f"Paddle naive bf16 Output max diff: {(out_bf16 - out_ref).abs().max().item()}"
|
|
)
|
|
print(
|
|
f"Paddle naive bf16 Output mean diff: {(out_bf16 - out_ref).abs().mean().item()}"
|
|
)
|
|
|
|
if fa_version == 2:
|
|
paddle.set_flags({'FLAGS_flash_attn_version': 2})
|
|
elif fa_version == 3:
|
|
paddle.set_flags({'FLAGS_flash_attn_version': 3})
|
|
else:
|
|
raise ValueError(f"Invalid flash attention version: {fa_version}")
|
|
|
|
out, lse = flashmask_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
startend_row_indices=startend_row_indices,
|
|
causal=causal,
|
|
return_softmax_lse=True,
|
|
)
|
|
print(f"flashmask Output max diff: {(out - out_ref).abs().max().item()}")
|
|
print(f"flashmask Output mean diff: {(out - out_ref).abs().mean().item()}")
|
|
# if not causal:
|
|
# print(f"LSE max diff: {(lse - lse_ref).abs().max().item()}")
|
|
# breakpoint()
|
|
|
|
# Check that FlashAttention's numerical error is at most twice the numerical error
|
|
# of a Pytorch implementation.
|
|
|
|
assert (out - out_ref).abs().max().item() <= rtol * (
|
|
out_bf16 - out_ref
|
|
).abs().max().item() + fwd_atol
|
|
|
|
g = paddle.randn(shape=out.shape, dtype=out.dtype)
|
|
out.backward(g)
|
|
out_ref.backward(g)
|
|
out_bf16.backward(g)
|
|
|
|
print(f"flashmask dQ max diff: {(q.grad - q_ref.grad).abs().max().item()}")
|
|
print(f"flashmask dK max diff: {(k.grad - k_ref.grad).abs().max().item()}")
|
|
print(f"flashmask dV max diff: {(v.grad - v_ref.grad).abs().max().item()}")
|
|
print(
|
|
f"flashmask dQ mean diff: {(q.grad - q_ref.grad).abs().mean().item()}"
|
|
)
|
|
print(
|
|
f"flashmask dK mean diff: {(k.grad - k_ref.grad).abs().mean().item()}"
|
|
)
|
|
print(
|
|
f"flashmask dV mean diff: {(v.grad - v_ref.grad).abs().mean().item()}"
|
|
)
|
|
|
|
print(
|
|
f"Paddle naive bf16 dQ max diff: {(q_bf16.grad - q_ref.grad).abs().max().item()}"
|
|
)
|
|
print(
|
|
f"Paddle naive bf16 dK max diff: {(k_bf16.grad - k_ref.grad).abs().max().item()}"
|
|
)
|
|
print(
|
|
f"Paddle naive bf16 dV max diff: {(v_bf16.grad - v_ref.grad).abs().max().item()}"
|
|
)
|
|
print(
|
|
f"Paddle naive bf16 dQ mean diff: {(q_bf16.grad - q_ref.grad).abs().mean().item()}"
|
|
)
|
|
print(
|
|
f"Paddle naive bf16 dK mean diff: {(k_bf16.grad - k_ref.grad).abs().mean().item()}"
|
|
)
|
|
print(
|
|
f"Paddle naive bf16 dV mean diff: {(v_bf16.grad - v_ref.grad).abs().mean().item()}"
|
|
)
|
|
|
|
dq_atol = 2 * (q_ref.grad + 0.3 - 0.3 - q_ref.grad).abs().max().item() + (
|
|
0 if softcap == 0 else 3e-4
|
|
)
|
|
assert (q.grad - q_ref.grad).abs().max().item() <= rtol * (
|
|
q_bf16.grad - q_ref.grad
|
|
).abs().max().item() + dq_atol
|
|
dk_atol = 2 * (k_ref.grad + 0.3 - 0.3 - k_ref.grad).abs().max().item() + (
|
|
0 if softcap == 0 else 3e-4
|
|
)
|
|
assert (k.grad - k_ref.grad).abs().max().item() <= rtol * (
|
|
k_bf16.grad - k_ref.grad
|
|
).abs().max().item() + dk_atol
|
|
dv_atol = 2 * (v_ref.grad + 0.3 - 0.3 - v_ref.grad).abs().max().item() + (
|
|
0 if softcap == 0 else 3e-4
|
|
)
|
|
assert (v.grad - v_ref.grad).abs().max().item() <= rtol * (
|
|
v_bf16.grad - v_ref.grad
|
|
).abs().max().item() + dv_atol
|
|
|
|
|
|
@pytest.mark.parametrize("d, dv", [(192, 128), (256, 128)])
|
|
@pytest.mark.parametrize("fa_version", [2, 3])
|
|
@pytest.mark.parametrize(
|
|
"gen_startend_row_indices",
|
|
[
|
|
partial(generate_none_mask, causal=False), # flash_attn path
|
|
partial(generate_sliding_window_mask), # flashmask path
|
|
],
|
|
)
|
|
def test_flashmask_forward_headdim_mismatch_raises(
|
|
d, dv, fa_version, gen_startend_row_indices
|
|
):
|
|
"""Test that headdim != headdim_v raises an error."""
|
|
|
|
if fa_version == 2:
|
|
paddle.set_flags({'FLAGS_flash_attn_version': 2})
|
|
elif fa_version == 3:
|
|
paddle.set_flags({'FLAGS_flash_attn_version': 3})
|
|
else:
|
|
raise ValueError(f"Invalid flash attention version: {fa_version}")
|
|
|
|
batch_size, seqlen_q, seqlen_k, nheads, nheads_startend_row_indices = (
|
|
1,
|
|
300,
|
|
300,
|
|
16,
|
|
16,
|
|
)
|
|
q = paddle.randn([batch_size, seqlen_q, nheads, d], dtype=paddle.bfloat16)
|
|
k = paddle.randn([batch_size, seqlen_k, nheads, d], dtype=paddle.bfloat16)
|
|
v = paddle.randn([batch_size, seqlen_k, nheads, dv], dtype=paddle.bfloat16)
|
|
startend_row_indices, causal = gen_startend_row_indices(
|
|
batch_size, seqlen_q, seqlen_k, nheads_startend_row_indices
|
|
)
|
|
if fa_version == 3:
|
|
if startend_row_indices is None:
|
|
# fallback to fa2
|
|
with pytest.raises(Exception, match="headdim != headdim_v"):
|
|
out, lse = flashmask_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
startend_row_indices=startend_row_indices,
|
|
causal=causal,
|
|
return_softmax_lse=True,
|
|
)
|
|
else:
|
|
# flashmask v3
|
|
if not (
|
|
(d > 128 and d <= 192 and dv > 96 and dv <= 128)
|
|
or (d <= 64 and dv <= 512)
|
|
):
|
|
with pytest.raises(
|
|
Exception,
|
|
match="headdim != headdim_v|V headdim is different from",
|
|
):
|
|
out, lse = flashmask_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
startend_row_indices=startend_row_indices,
|
|
causal=causal,
|
|
return_softmax_lse=True,
|
|
)
|
|
elif fa_version == 2:
|
|
with pytest.raises(Exception, match="headdim != headdim_v"):
|
|
out, lse = flashmask_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
startend_row_indices=startend_row_indices,
|
|
causal=causal,
|
|
return_softmax_lse=True,
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid flash attention version: {fa_version}")
|
|
|
|
|
|
@pytest.mark.parametrize("d, dv", [(192, 128), (256, 128)])
|
|
@pytest.mark.parametrize("fa_version", [2, 3])
|
|
def test_flashmask_backward_headdim_mismatch_raises(fa_version, d, dv):
|
|
"""Test that backward kernel raises when headdim != headdim_v."""
|
|
from paddle import _C_ops
|
|
|
|
paddle.set_flags({'FLAGS_flash_attn_version': fa_version})
|
|
batch_size, seqlen_q, seqlen_k, nheads = 1, 300, 300, 16
|
|
softmax_scale = d ** (-0.5)
|
|
|
|
if fa_version == 2:
|
|
q = paddle.randn(
|
|
[batch_size, seqlen_q, nheads, d], dtype=paddle.bfloat16
|
|
)
|
|
k = paddle.randn(
|
|
[batch_size, seqlen_k, nheads, d], dtype=paddle.bfloat16
|
|
)
|
|
v = paddle.randn(
|
|
[batch_size, seqlen_k, nheads, dv], dtype=paddle.bfloat16
|
|
)
|
|
out = paddle.randn(
|
|
[batch_size, seqlen_q, nheads, dv], dtype=paddle.bfloat16
|
|
)
|
|
softmax_lse = paddle.randn(
|
|
[batch_size, nheads, seqlen_q], dtype=paddle.float32
|
|
)
|
|
seed_offset = paddle.to_tensor([0, 0], dtype=paddle.int64)
|
|
dout = paddle.randn(
|
|
[batch_size, seqlen_q, nheads, dv], dtype=paddle.bfloat16
|
|
)
|
|
with pytest.raises(Exception, match="headdim != headdim_v"):
|
|
_C_ops.flash_attn_grad(
|
|
q, k, v, out, softmax_lse, seed_offset, None, dout, 0.0, False
|
|
)
|
|
elif fa_version == 3:
|
|
total_q = batch_size * seqlen_q
|
|
total_k = batch_size * seqlen_k
|
|
q = paddle.randn([total_q, nheads, d], dtype=paddle.bfloat16)
|
|
k = paddle.randn([total_k, nheads, d], dtype=paddle.bfloat16)
|
|
v = paddle.randn([total_k, nheads, dv], dtype=paddle.bfloat16)
|
|
out = paddle.randn([total_q, nheads, dv], dtype=paddle.bfloat16)
|
|
softmax_lse = paddle.randn(
|
|
[batch_size, nheads, seqlen_q], dtype=paddle.float32
|
|
)
|
|
cu_seqlens_q = paddle.to_tensor([0, seqlen_q], dtype=paddle.int32)
|
|
cu_seqlens_k = paddle.to_tensor([0, seqlen_k], dtype=paddle.int32)
|
|
dout = paddle.randn([total_q, nheads, dv], dtype=paddle.bfloat16)
|
|
with pytest.raises(Exception, match="headdim != headdim_v"):
|
|
_C_ops.flash_attn_v3_varlen_grad(
|
|
q,
|
|
k,
|
|
v,
|
|
out,
|
|
softmax_lse,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
None,
|
|
None, # seqused_q, seqused_k
|
|
dout,
|
|
softmax_scale,
|
|
seqlen_q,
|
|
seqlen_k,
|
|
False, # causal
|
|
-1,
|
|
-1, # window_size_left, window_size_right
|
|
0.0, # softcap
|
|
0, # sm_margin
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("d, dv", [(192, 128), (256, 128)])
|
|
def test_flash_attn_unpadded_grad_headdim_mismatch_raises(d, dv):
|
|
"""Test that flash_attn_unpadded_grad raises when headdim != headdim_v."""
|
|
from paddle import _C_ops
|
|
|
|
paddle.set_flags({'FLAGS_flash_attn_version': 2})
|
|
|
|
batch_size, seqlen_q, seqlen_k, nheads = 1, 300, 300, 16
|
|
total_q = batch_size * seqlen_q
|
|
total_k = batch_size * seqlen_k
|
|
q = paddle.randn([total_q, nheads, d], dtype=paddle.bfloat16)
|
|
k = paddle.randn([total_k, nheads, d], dtype=paddle.bfloat16)
|
|
v = paddle.randn([total_k, nheads, dv], dtype=paddle.bfloat16)
|
|
out = paddle.randn([total_q, nheads, dv], dtype=paddle.bfloat16)
|
|
softmax_lse = paddle.randn(
|
|
[batch_size, nheads, seqlen_q], dtype=paddle.float32
|
|
)
|
|
seed_offset = paddle.to_tensor([0, 0], dtype=paddle.int64)
|
|
cu_seqlens_q = paddle.to_tensor([0, seqlen_q], dtype=paddle.int32)
|
|
cu_seqlens_k = paddle.to_tensor([0, seqlen_k], dtype=paddle.int32)
|
|
dout = paddle.randn([total_q, nheads, dv], dtype=paddle.bfloat16)
|
|
with pytest.raises(Exception, match="headdim != headdim_v"):
|
|
_C_ops.flash_attn_unpadded_grad(
|
|
q,
|
|
k,
|
|
v,
|
|
cu_seqlens_q,
|
|
cu_seqlens_k,
|
|
out,
|
|
softmax_lse,
|
|
seed_offset,
|
|
None,
|
|
dout,
|
|
seqlen_q,
|
|
seqlen_k,
|
|
d ** (-0.5),
|
|
0.0,
|
|
False,
|
|
)
|