saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor flash attn test (#21028)

Browse source 
Code improvement and allow change to use other ep
This commit is contained in:
cloudhan 2024-06-15 01:23:28 +08:00 committed by GitHub
parent c66e920154
commit f4b22f89bc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 62 additions and 208 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

270
onnxruntime/test/python/transformers/test_flash_attn_cuda.py
View file

@ -24,10 +24,6 @@ from parameterized import parameterized
from onnxruntime import InferenceSession, OrtValue, SessionOptions
RED = "\033[31m"
GREEN = "\033[32m"
RESET = "\033[0m"
torch.manual_seed(0)
pipeline_mode = True # Reduces number of tests so pipeline doesn't time out
@ -46,6 +42,7 @@ class Config:
num_heads = 0
kv_num_heads = 0
head_size = 0
ep = "CUDAExecutionProvider"
def __init__(
self, batch_size, sequence_length, kv_sequence_length, past_sequence_length, num_heads, kv_num_heads, head_size
@ -59,10 +56,11 @@ class Config:
self.head_size = head_size
def __repr__(self):
short_ep = self.ep[: -len("ExecutionProvider")].lower()
return (
f"Config(batch_size={self.batch_size}, sequence_length={self.sequence_length}, "
f"kv_sequence_length={self.kv_sequence_length}, past_sequence_length={self.past_sequence_length}, "
f"num_heads={self.num_heads}, kv_num_heads={self.kv_num_heads}, head_size={self.head_size})"
f"num_heads={self.num_heads}, kv_num_heads={self.kv_num_heads}, head_size={self.head_size}, ep={short_ep})"
)
@ -74,6 +72,7 @@ class PromptConfig:
num_heads = 0
kv_num_heads = 0
head_size = 0
ep = "CUDAExecutionProvider"
def __init__(
self,
@ -94,10 +93,11 @@ class PromptConfig:
self.head_size = head_size
def __repr__(self):
short_ep = self.ep[: -len("ExecutionProvider")].lower()
return (
f"PromptConfig(batch_size={self.batch_size}, q_sequence_length={self.q_sequence_length}, "
f"kv_sequence_length={self.kv_sequence_length}, buffer_sequence_length={self.buffer_sequence_length}, "
f"num_heads={self.num_heads}, kv_num_heads={self.kv_num_heads}, head_size={self.head_size})"
f"num_heads={self.num_heads}, kv_num_heads={self.kv_num_heads}, head_size={self.head_size}, ep={short_ep})"
)
@ -738,7 +738,7 @@ def flash_attn_varlen_qkvpacked_func(qkv_unpad, cu_seqlens, token_offset, config
"cumulative_sequence_length": cu_seqlens.cpu().numpy(),
}
sess_options = SessionOptions()
ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"])
ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep])
ort_output = ort_session.run(None, ort_inputs)
output = torch.tensor(ort_output)
return output
@ -755,7 +755,7 @@ def mha_func(q, k, v, config):
"value": v.detach().cpu().numpy(),
}
sess_options = SessionOptions()
ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"])
ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep])
ort_output = ort_session.run(None, ort_inputs)
ort_output = numpy.array(ort_output)
output = torch.tensor(ort_output)
@ -807,7 +807,7 @@ def gqa_prompt_func(
"total_sequence_length": torch.tensor([config.q_sequence_length], dtype=torch.int32).detach().cpu().numpy(),
}
sess_options = SessionOptions()
ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"])
ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep])
io_binding = ort_session.io_binding()
if new_k is not None:
ort_inputs["key"] = new_k.detach().cpu().numpy()
@ -848,7 +848,7 @@ def gqa_prompt_func(
"total_sequence_length": torch.tensor([config.q_sequence_length], dtype=torch.int32).detach().cpu().numpy(),
}
sess_options = SessionOptions()
ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"])
ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep])
io_binding = ort_session.io_binding()
if new_k is not None:
ort_inputs["key"] = new_k.detach().cpu().numpy()
@ -915,7 +915,7 @@ def gqa_past_func(
.numpy(),
}
sess_options = SessionOptions()
ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"])
ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep])
io_binding = ort_session.io_binding()
if new_k is not None:
ort_inputs["key"] = new_k.detach().cpu().numpy()
@ -963,7 +963,7 @@ def gqa_past_func(
.numpy(),
}
sess_options = SessionOptions()
ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"])
ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep])
io_binding = ort_session.io_binding()
if new_k is not None:
ort_inputs["key"] = new_k.detach().cpu().numpy()
@ -1171,25 +1171,9 @@ def parity_check_mha(
out_ref, _ = attention_ref(q, k, v, None, None, 0.0, None, causal=False)
out_ref = out_ref.detach().cpu().numpy()
# Compare results
all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True)
correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET
print(
" B:",
config.batch_size,
" S:",
config.sequence_length,
" N:",
config.num_heads,
" kvN:",
config.kv_num_heads,
" h:",
config.head_size,
" Mean Error:",
numpy.mean(numpy.abs(out - out_ref)),
correct,
numpy.testing.assert_allclose(
out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=f" with {config} packed={packed}"
)
return all_close
def rotary_embedding(*args, **kwargs):
@ -1366,60 +1350,19 @@ def parity_check_gqa_prompt(
out = torch.reshape(out, (config.batch_size, config.q_sequence_length, config.num_heads, config.head_size))
out = out.detach().cpu().numpy()
err_msg = (
f" with {config}, causal={causal}, local={local}, past_format={past_format},"
f" rotary={rotary}, rotary_interleaved={rotary_interleaved}, packed={packed}"
)
# Make sure past-present buffer updating correctly
assert numpy.allclose(present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True)
assert numpy.allclose(present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True)
# Compare results
all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True)
correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET
print(
"KV-buffer",
" packed:",
packed,
" causal:",
causal,
" local:",
local,
" rotary:",
rotary,
" rotary_interleaved:",
rotary_interleaved,
"past kv format:",
"BSNH" if past_format == Formats.BSNH else "BNSH",
" B:",
config.batch_size,
" S:",
config.q_sequence_length,
" kv S:",
config.kv_sequence_length,
" N:",
config.num_heads,
" kv N:",
config.kv_num_heads,
" h:",
config.head_size,
" Mean Error:",
numpy.mean(numpy.abs(out - out_ref)),
correct,
numpy.testing.assert_allclose(
present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg
)
numpy.testing.assert_allclose(
present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg
)
if not all_close:
close_mask = numpy.isclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True)
not_close_mask = numpy.logical_not(close_mask)
# Values that are not close
out_not_close = out[not_close_mask]
out_ref_not_close = out_ref[not_close_mask]
# Indices that are not close
not_close_indices = numpy.argwhere(not_close_mask)
print("Values in 'out' that are not close:", out_not_close)
print("Corresponding values in 'out_ref':", out_ref_not_close)
print("Indices of values that are not close:", not_close_indices)
return all_close
numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg)
def parity_check_gqa_prompt_no_buff(
@ -1566,44 +1509,19 @@ def parity_check_gqa_prompt_no_buff(
out = torch.reshape(out, (config.batch_size, config.q_sequence_length, config.num_heads, config.head_size))
out = out.detach().cpu().numpy()
# Make sure past-present buffer updating correctly
assert numpy.allclose(present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True)
assert numpy.allclose(present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True)
# Compare results
all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True)
correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET
print(
"No buff",
" packed:",
packed,
" causal:",
causal,
" local:",
local,
" rotary:",
rotary,
" rotary_interleaved:",
rotary_interleaved,
"past kv format:",
"BSNH" if past_format == Formats.BSNH else "BNSH",
" B:",
config.batch_size,
" S:",
config.q_sequence_length,
" kv S:",
config.kv_sequence_length,
" N:",
config.num_heads,
" kv N:",
config.kv_num_heads,
" h:",
config.head_size,
" Mean Error:",
numpy.mean(numpy.abs(out - out_ref)),
correct,
err_msg = (
f" with {config}, causal={causal}, local={local}, past_format={past_format},"
f" rotary={rotary}, rotary_interleaved={rotary_interleaved}, packed={packed}"
)
return all_close
# Make sure past-present buffer updating correctly
numpy.testing.assert_allclose(
present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg
)
numpy.testing.assert_allclose(
present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg
)
numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg)
def parity_check_gqa_past(
@ -1769,44 +1687,19 @@ def parity_check_gqa_past(
out = torch.reshape(out, (config.batch_size, config.sequence_length, config.num_heads, config.head_size))
out = out.detach().cpu().numpy()
# Make sure past-present buffer updating correctly
assert numpy.allclose(present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True)
assert numpy.allclose(present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True)
# Compare results
all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True)
correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET
print(
"KV-buffer",
"past kv format:",
"BSNH" if past_format == Formats.BSNH else "BNSH",
" packed:",
packed,
" causal:",
causal,
" local:",
local,
" rotary:",
rotary,
" rotary_interleaved:",
rotary_interleaved,
" B:",
config.batch_size,
" S:",
config.sequence_length,
" kv S:",
config.kv_sequence_length,
" N:",
config.num_heads,
" kv N:",
config.kv_num_heads,
" h:",
config.head_size,
" Mean Error:",
numpy.mean(numpy.abs(out - out_ref)),
correct,
err_msg = (
f" with {config}, causal={causal}, local={local}, past_format={past_format},"
f" rotary={rotary}, rotary_interleaved={rotary_interleaved}, packed={packed}"
)
return all_close
# Make sure past-present buffer updating correctly
numpy.testing.assert_allclose(
present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg
)
numpy.testing.assert_allclose(
present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg
)
numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg)
def parity_check_gqa_past_no_buff(
@ -1978,40 +1871,11 @@ def parity_check_gqa_past_no_buff(
out = torch.reshape(out, (config.batch_size, config.sequence_length, config.num_heads, config.head_size))
out = out.detach().cpu().numpy()
# Compare results
all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True)
correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET
print(
"NO buff",
" packed:",
packed,
" causal:",
causal,
" local:",
local,
" rotary:",
rotary,
" rotary_interleaved:",
rotary_interleaved,
"past kv format:",
"BSNH" if past_format == Formats.BSNH else "BNSH",
" B:",
config.batch_size,
" S:",
config.sequence_length,
" kv S:",
config.kv_sequence_length,
" N:",
config.num_heads,
" kv N:",
config.kv_num_heads,
" h:",
config.head_size,
" Mean Error:",
numpy.mean(numpy.abs(out - out_ref)),
correct,
err_msg = (
f" with {config}, causal={causal}, local={local}, past_format={past_format},"
f" rotary={rotary}, rotary_interleaved={rotary_interleaved}, packed={packed}"
)
return all_close
numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg)
def packed_mha_test_cases():
@ -2064,8 +1928,7 @@ class TestMHA(unittest.TestCase):
if major < 8:
return
print("-------- TEST PACKED MHA ---------")
all_close = parity_check_mha(config, True)
self.assertTrue(all_close)
parity_check_mha(config, True)
@parameterized.expand(mha_test_cases())
def test_mha(self, _, config):
@ -2075,8 +1938,7 @@ class TestMHA(unittest.TestCase):
if major < 8:
return
print("-------- TEST MHA ---------")
all_close = parity_check_mha(config, False)
self.assertTrue(all_close)
parity_check_mha(config, False)
def gqa_no_past_memory_efficient_test_cases():
@ -2252,7 +2114,7 @@ class TestGQA(unittest.TestCase):
os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "1"
print("------- MEMORY EFFICIENT ATTENTION (PROMPT CASE) ---------")
all_close = parity_check_gqa_prompt(
parity_check_gqa_prompt(
config,
rtol=5e-3,
atol=5e-3,
@ -2261,8 +2123,7 @@ class TestGQA(unittest.TestCase):
rotary_interleaved=rotary_interleaved,
packed=packed,
)
self.assertTrue(all_close)
all_close = parity_check_gqa_prompt_no_buff(
parity_check_gqa_prompt_no_buff(
config,
rtol=5e-3,
atol=5e-3,
@ -2271,7 +2132,6 @@ class TestGQA(unittest.TestCase):
rotary_interleaved=rotary_interleaved,
packed=packed,
)
self.assertTrue(all_close)
@parameterized.expand(gqa_no_past_flash_attention_test_cases())
def test_gqa_no_past_flash_attention(self, _, config, local, rotary, rotary_interleaved, packed):
@ -2283,7 +2143,7 @@ class TestGQA(unittest.TestCase):
print("------- FLASH ATTENTION (PROMPT CASE) --------")
os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0"
all_close = parity_check_gqa_prompt(
parity_check_gqa_prompt(
config,
local=local,
past_format=Formats.BNSH,
@ -2291,8 +2151,7 @@ class TestGQA(unittest.TestCase):
rotary_interleaved=rotary_interleaved,
packed=packed,
)
self.assertTrue(all_close)
all_close = parity_check_gqa_prompt_no_buff(
parity_check_gqa_prompt_no_buff(
config,
local=local,
past_format=Formats.BNSH,
@ -2300,7 +2159,6 @@ class TestGQA(unittest.TestCase):
rotary_interleaved=rotary_interleaved,
packed=packed,
)
self.assertTrue(all_close)
@parameterized.expand(gqa_past_memory_efficient_test_cases())
def test_gqa_past_memory_efficient(self, _, config, rotary, rotary_interleaved, packed):
@ -2312,7 +2170,7 @@ class TestGQA(unittest.TestCase):
os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "1"
print("-------- MEMORY EFFICIENT (TOKEN GEN) --------")
all_close = parity_check_gqa_past(
parity_check_gqa_past(
config,
past_format=Formats.BNSH,
rtol=1e-3,
@ -2321,8 +2179,7 @@ class TestGQA(unittest.TestCase):
rotary_interleaved=rotary_interleaved,
packed=packed,
)
self.assertTrue(all_close)
all_close = parity_check_gqa_past_no_buff(
parity_check_gqa_past_no_buff(
config,
past_format=Formats.BNSH,
rtol=1e-3,
@ -2331,7 +2188,6 @@ class TestGQA(unittest.TestCase):
rotary_interleaved=rotary_interleaved,
packed=packed,
)
self.assertTrue(all_close)
@parameterized.expand(gqa_past_flash_attention_test_cases())
def test_gqa_past_flash_attention(self, _, config, local, rotary, rotary_interleaved, packed):
@ -2343,7 +2199,7 @@ class TestGQA(unittest.TestCase):
print("------- FLASH ATTENTION (TOKEN GEN) -------")
os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0"
all_close = parity_check_gqa_past(
parity_check_gqa_past(
config,
local=local,
past_format=Formats.BNSH,
@ -2353,8 +2209,7 @@ class TestGQA(unittest.TestCase):
rotary_interleaved=rotary_interleaved,
packed=packed,
)
self.assertTrue(all_close)
all_close = parity_check_gqa_past_no_buff(
parity_check_gqa_past_no_buff(
config,
local=local,
past_format=Formats.BNSH,
@ -2364,7 +2219,6 @@ class TestGQA(unittest.TestCase):
rotary_interleaved=rotary_interleaved,
packed=packed,
)
self.assertTrue(all_close)
if __name__ == "__main__":