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

Fix Sparse Attention with Packed QKV inputs (#20591)

Browse source 
### Description
(1) Fix UnpackQKV kernel
(2) Update test_sparse_attention.py with packed QKV option
This commit is contained in:
Tianlei Wu 2024-05-07 10:50:01 -07:00 committed by GitHub
parent 478d3e0c62
commit d693aef39e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 218 additions and 378 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu
View file

@ -531,19 +531,22 @@ __global__ void UnpackQKV(const T* packed_qkv, T* unpacked_q, T* unpacked_k, T*
int offset = tid % d;
if (output_bnsh) { // output BNSH
int head_count = kv_num_heads;
T* unpacked;
if (offset < q_hidden) {
unpacked = unpacked_q;
head_count = num_heads;
} else if (offset < q_hidden + k_hidden) {
unpacked = unpacked_k;
offset -= q_hidden;
} else {
unpacked = unpacked_v;
offset -= (q_hidden + k_hidden);
}
int n = offset / head_size;
int h = offset % head_size;
int unpacked_i = INDEX_4D(head_count, sequence_length, head_size, b, n, s, h);
unpacked_q[unpacked_i] = packed_qkv[tid];
unpacked[unpacked_i] = packed_qkv[tid];
} else { // output BSNH
if (offset < q_hidden) {
int unpacked_i = b * sequence_length * num_heads * head_size + s * num_heads * head_size + offset;

16
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu
View file

@ -138,6 +138,7 @@ Status QkvToContext(
auto q = reinterpret_cast<T*>(data.unpacked_qkv_buffer);
auto k = reinterpret_cast<T*>(data.unpacked_qkv_buffer + q_size);
auto v = reinterpret_cast<T*>(data.unpacked_qkv_buffer + q_size + k_size);
Status status = LaunchUnpackQKV<T, LAYOUT_BNSH>(data.query, q, k, v, num_heads, kv_num_heads, head_size,
sequence_length, batch_size, stream, max_threads_per_block);
if (status != Status::OK()) {
@ -152,15 +153,15 @@ Status QkvToContext(
constexpr bool q_layout = LAYOUT_BNSH;
bool kv_layout = parameters.is_packed_qkv ? LAYOUT_BNSH : LAYOUT_BSNH;
DUMP_TENSOR("query", reinterpret_cast<const T*>(query), batch_size, num_heads, sequence_length, head_size);
#if DUMP_TENSOR_LEVEL > 0
DUMP_TENSOR("query (BNSH)", reinterpret_cast<const T*>(query), batch_size, num_heads, sequence_length, head_size);
if (LAYOUT_BNSH == kv_layout) {
DUMP_TENSOR("key", reinterpret_cast<const T*>(key), batch_size, kv_num_heads, sequence_length, head_size);
DUMP_TENSOR("value", reinterpret_cast<const T*>(value), batch_size, kv_num_heads, sequence_length, head_size);
DUMP_TENSOR("key (BNSH)", reinterpret_cast<const T*>(key), batch_size, kv_num_heads, sequence_length, head_size);
DUMP_TENSOR("value (BNSH)", reinterpret_cast<const T*>(value), batch_size, kv_num_heads, sequence_length, head_size);
} else {
DUMP_TENSOR("key", reinterpret_cast<const T*>(key), batch_size, sequence_length, kv_num_heads, head_size);
DUMP_TENSOR("value", reinterpret_cast<const T*>(value), batch_size, sequence_length, kv_num_heads, head_size);
DUMP_TENSOR("key (BSNH)", reinterpret_cast<const T*>(key), batch_size, sequence_length, kv_num_heads, head_size);
DUMP_TENSOR("value (BSNH)", reinterpret_cast<const T*>(value), batch_size, sequence_length, kv_num_heads, head_size);
}
#endif
@ -317,6 +318,9 @@ Status QkvToContext(
ORT_RETURN_IF_ERROR(sparse_attention_v1::run_sparse_attention_fp16(params));
}
}
DUMP_TENSOR("output", reinterpret_cast<const T*>(data.output), batch_size, num_heads, sequence_length, head_size);
return Status::OK();
}

573
onnxruntime/test/python/transformers/test_sparse_attention.py
View file

@ -8,8 +8,8 @@ Parity test and benchmark performance of SparseAttention. Requires Nvidia GPU of
Install required packages before running this script:
pip install matplotlib pandas onnx torch onnxruntime-gpu
"""
import math
import unittest
from typing import Optional
import torch
@ -65,10 +65,12 @@ class AttentionConfig:
self.dtype = dtype
def shape_dict(self):
return {
"query": (self.batch_size, self.sequence_length, self.num_heads * self.head_size),
"key": (self.batch_size, self.sequence_length, self.kv_num_heads * self.head_size),
"value": (self.batch_size, self.sequence_length, self.kv_num_heads * self.head_size),
shapes = {
"query": (
self.batch_size,
self.sequence_length,
(self.num_heads + 2 * self.kv_num_heads) * self.head_size,
),
"past_key": (self.batch_size, self.kv_num_heads, self.past_buffer_length, self.head_size),
"past_value": (self.batch_size, self.kv_num_heads, self.past_buffer_length, self.head_size),
"total_sequence_length": (1,),
@ -79,6 +81,16 @@ class AttentionConfig:
"sin_cache": (self.max_sequence_length, (math.floor(self.head_size / 16) * 16) // 2),
}
if not self.is_packed_qkv:
shapes.update(
{
"query": (self.batch_size, self.sequence_length, self.num_heads * self.head_size),
"key": (self.batch_size, self.sequence_length, self.kv_num_heads * self.head_size),
"value": (self.batch_size, self.sequence_length, self.kv_num_heads * self.head_size),
}
)
return shapes
def get_cos_sin_cache(self, dtype):
rotary_fraction = 1.0
rotary_dim = math.floor(int(rotary_fraction * self.head_size) / 16) * 16
@ -91,9 +103,14 @@ class AttentionConfig:
device = self.device
# Since bfloat16 is not supported in ORT python I/O binding API, we always use float16 as model inputs.
dtype = torch.float16
shape_dict = self.shape_dict()
# Always use non-packed qkv to generate same inputs for Torch and ORT.
packed = self.is_packed_qkv # Save the original value.
self.is_packed_qkv = False
shape_dict = self.shape_dict()
self.is_packed_qkv = packed # Restore the original value.
torch.manual_seed(123)
feeds = {
"query": torch.empty(shape_dict["query"], device=device, dtype=dtype).normal_(mean=0, std=0.1),
"key": torch.empty(shape_dict["key"], device=device, dtype=dtype).normal_(mean=0, std=0.1),
@ -103,6 +120,14 @@ class AttentionConfig:
"total_sequence_length": torch.tensor([self.total_sequence_length], dtype=torch.int32),
}
if packed:
query = feeds["query"].view(self.batch_size, self.sequence_length, self.num_heads, self.head_size)
key = feeds["key"].view(self.batch_size, self.sequence_length, self.kv_num_heads, self.head_size)
value = feeds["value"].view(self.batch_size, self.sequence_length, self.kv_num_heads, self.head_size)
feeds["query"] = torch.dstack((query, key, value)).reshape(self.batch_size, self.sequence_length, -1)
del feeds["key"]
del feeds["value"]
if self.do_rotary:
cos_cache, sin_cache = self.get_cos_sin_cache(dtype)
feeds["cos_cache"] = cos_cache
@ -127,6 +152,7 @@ class GroupQueryAttentionConfig(AttentionConfig):
device="cuda",
local_window_size: int = -1,
attention_mask=None,
is_packed_qkv=False,
):
super().__init__(
"GroupQueryAttention",
@ -141,6 +167,7 @@ class GroupQueryAttentionConfig(AttentionConfig):
do_rotary,
rotary_interleaved,
device,
is_packed_qkv=is_packed_qkv,
)
# local_window_size is for ORT only, not for Torch implementation.
self.local_window_size = local_window_size
@ -186,6 +213,7 @@ class SparseAttentionConfig(AttentionConfig):
do_rotary: bool = False,
rotary_interleaved: bool = False,
device="cuda",
is_packed_qkv=False,
):
super().__init__(
"SparseAttention",
@ -200,6 +228,7 @@ class SparseAttentionConfig(AttentionConfig):
do_rotary,
rotary_interleaved,
device,
is_packed_qkv=is_packed_qkv,
)
self.sparse_block_size = sparse_block_size
self.num_layout = num_layout
@ -239,10 +268,27 @@ class SparseAttentionConfig(AttentionConfig):
)
return feeds
def get_comparable_gqa_config(self, use_local=False, torch_use_sparse=False) -> GroupQueryAttentionConfig:
def get_comparable_ort_gqa_config(self, use_local=False) -> GroupQueryAttentionConfig:
return GroupQueryAttentionConfig(
self.batch_size,
self.sequence_length,
self.max_sequence_length,
self.past_sequence_length,
self.num_heads,
self.kv_num_heads,
self.head_size,
self.softmax_scale,
self.do_rotary,
self.rotary_interleaved,
self.device,
local_window_size=self.local_blocks * self.sparse_block_size if use_local else -1,
is_packed_qkv=self.is_packed_qkv,
)
def get_comparable_torch_gqa_config(self, use_sparse=False) -> GroupQueryAttentionConfig:
attention_mask = None
if torch_use_sparse:
if use_sparse is True:
attention_mask = self.dense_mask()[:, :, : self.total_sequence_length, : self.total_sequence_length]
if self.past_sequence_length > 0:
attention_mask = attention_mask[:, :, -self.sequence_length :, :]
@ -259,8 +305,8 @@ class SparseAttentionConfig(AttentionConfig):
self.do_rotary,
self.rotary_interleaved,
self.device,
local_window_size=self.local_blocks * self.sparse_block_size if use_local else -1,
attention_mask=attention_mask,
is_packed_qkv=False, # torch reference implementation does not support packed qkv.
)
@ -328,7 +374,11 @@ def create_sparse_attention_onnx_model(config: SparseAttentionConfig):
nodes.extend(
[
helper.make_node("Cast", [input], [input + suffix], f"Cast_{input}", to=TensorProto.BFLOAT16)
for input in ["query", "key", "value", "past_key", "past_value"]
for input in (
["query", "key", "value", "past_key", "past_value"]
if not config.is_packed_qkv
else ["query", "past_key", "past_value"]
)
]
)
if config.do_rotary:
@ -348,19 +398,30 @@ def create_sparse_attention_onnx_model(config: SparseAttentionConfig):
shape_dict = config.shape_dict()
graph_input = [
helper.make_tensor_value_info("query", io_float_type, list(shape_dict["query"])),
helper.make_tensor_value_info("key", io_float_type, list(shape_dict["key"])),
helper.make_tensor_value_info("value", io_float_type, list(shape_dict["value"])),
helper.make_tensor_value_info("past_key", io_float_type, list(shape_dict["past_key"])),
helper.make_tensor_value_info("past_value", io_float_type, list(shape_dict["past_value"])),
helper.make_tensor_value_info("block_mask", TensorProto.INT32, list(shape_dict["block_mask"])),
helper.make_tensor_value_info(
"total_sequence_length", TensorProto.INT32, list(shape_dict["total_sequence_length"])
),
helper.make_tensor_value_info(
"key_total_sequence_lengths", TensorProto.INT32, list(shape_dict["key_total_sequence_lengths"])
),
]
if not config.is_packed_qkv:
graph_input.extend(
[
helper.make_tensor_value_info("key", io_float_type, list(shape_dict["key"])),
helper.make_tensor_value_info("value", io_float_type, list(shape_dict["value"])),
]
)
graph_input.extend(
[
helper.make_tensor_value_info("past_key", io_float_type, list(shape_dict["past_key"])),
helper.make_tensor_value_info("past_value", io_float_type, list(shape_dict["past_value"])),
helper.make_tensor_value_info("block_mask", TensorProto.INT32, list(shape_dict["block_mask"])),
helper.make_tensor_value_info(
"total_sequence_length", TensorProto.INT32, list(shape_dict["total_sequence_length"])
),
helper.make_tensor_value_info(
"key_total_sequence_lengths", TensorProto.INT32, list(shape_dict["key_total_sequence_lengths"])
),
]
)
if config.do_rotary:
graph_input += [
helper.make_tensor_value_info("cos_cache", io_float_type, list(shape_dict["cos_cache"])),
@ -417,16 +478,27 @@ def create_group_query_attention_onnx_model(config: GroupQueryAttentionConfig):
shape_dict = config.shape_dict()
graph_input = [
helper.make_tensor_value_info("query", float_type, list(shape_dict["query"])),
helper.make_tensor_value_info("key", float_type, list(shape_dict["key"])),
helper.make_tensor_value_info("value", float_type, list(shape_dict["value"])),
helper.make_tensor_value_info("past_key", float_type, list(shape_dict["past_key"])),
helper.make_tensor_value_info("past_value", float_type, list(shape_dict["past_value"])),
helper.make_tensor_value_info("seqlens_k", TensorProto.INT32, list(shape_dict["seqlens_k"])),
helper.make_tensor_value_info(
"total_sequence_length", TensorProto.INT32, list(shape_dict["total_sequence_length"])
),
]
if not config.is_packed_qkv:
graph_input.extend(
[
helper.make_tensor_value_info("key", float_type, list(shape_dict["key"])),
helper.make_tensor_value_info("value", float_type, list(shape_dict["value"])),
]
)
graph_input.extend(
[
helper.make_tensor_value_info("past_key", float_type, list(shape_dict["past_key"])),
helper.make_tensor_value_info("past_value", float_type, list(shape_dict["past_value"])),
helper.make_tensor_value_info("seqlens_k", TensorProto.INT32, list(shape_dict["seqlens_k"])),
helper.make_tensor_value_info(
"total_sequence_length", TensorProto.INT32, list(shape_dict["total_sequence_length"])
),
]
)
if config.do_rotary:
graph_input += [
helper.make_tensor_value_info("cos_cache", float_type, list(shape_dict["cos_cache"])),
@ -562,7 +634,6 @@ class OrtGroupQueryAttention:
"""A wrapper of ORT GroupQueryAttention to test relevance and performance."""
def __init__(self, config: GroupQueryAttentionConfig):
device = config.device
cuda_provider_options = CudaSession.get_cuda_provider_options(
torch.cuda.current_device(), enable_cuda_graph=False, stream=torch.cuda.current_stream().cuda_stream
)
@ -570,7 +641,7 @@ class OrtGroupQueryAttention:
self.ort_session = create_session(onnx_model_str, cuda_provider_options=cuda_provider_options)
self.gpu_binding_manager = GpuBindingManager(
ort_session=self.ort_session,
device=device,
device=config.device,
stream=torch.cuda.current_stream().cuda_stream,
max_cuda_graphs=2,
)
@ -580,7 +651,7 @@ class OrtGroupQueryAttention:
)
self.feed_dict = config.random_inputs()
if ENABLE_DEBUG:
if ENABLE_DEBUG and not config.is_packed_qkv:
query = self.feed_dict["query"].view(
config.batch_size, config.sequence_length, config.num_heads, config.head_size
)
@ -604,7 +675,6 @@ class OrtSparseAttention:
"""A wrapper of ORT SparseAttention to test relevance and performance."""
def __init__(self, config: SparseAttentionConfig):
device = config.device
cuda_provider_options = CudaSession.get_cuda_provider_options(
torch.cuda.current_device(), enable_cuda_graph=False, stream=torch.cuda.current_stream().cuda_stream
)
@ -612,7 +682,7 @@ class OrtSparseAttention:
self.ort_session = create_session(onnx_model_str, cuda_provider_options=cuda_provider_options)
self.gpu_binding_manager = GpuBindingManager(
ort_session=self.ort_session,
device=device,
device=config.device,
stream=torch.cuda.current_stream().cuda_stream,
max_cuda_graphs=2,
)
@ -622,7 +692,7 @@ class OrtSparseAttention:
)
self.feed_dict = config.random_inputs()
if ENABLE_DEBUG:
if ENABLE_DEBUG and not config.is_packed_qkv:
query = self.feed_dict["query"].view(
config.batch_size, config.sequence_length, config.num_heads, config.head_size
)
@ -644,357 +714,120 @@ class OrtSparseAttention:
return self.gpu_binding.infer(self.feed_dict)
def run_one_relevance_test(config: SparseAttentionConfig):
# Run QGA
if not config.do_rotary: # config.past_sequence_length == 0:
gqa_config: GroupQueryAttentionConfig = config.get_comparable_gqa_config(torch_use_sparse=True)
obj = TorchGroupQueryAttention(gqa_config)
expected_out = obj.infer()
else:
gqa_config: GroupQueryAttentionConfig = config.get_comparable_gqa_config(use_local=False)
obj = OrtGroupQueryAttention(gqa_config)
ort_qga_outputs = obj.infer()
expected_out = ort_qga_outputs["output"].view(
config.batch_size, config.sequence_length, config.num_heads, config.head_size
)
class TestSparseAttention(unittest.TestCase):
@unittest.skipUnless(torch.cuda.is_available(), "cuda not available")
def test_sparse_attention(self):
major, minor = torch.cuda.get_device_capability()
sm = major * 10 + minor
# Run SparseAttention by ORT
obj = OrtSparseAttention(config)
ort_outputs = obj.infer()
ort_output = ort_outputs["output"]
actual_out = ort_output.view(config.batch_size, config.sequence_length, config.num_heads, config.head_size)
if sm not in [75, 80, 86, 89, 90]:
self.skipTest("SparseAttention is not supported on this GPU")
red_color = "\033[31m"
green_color = "\033[32m"
reset_color = "\033[0m"
if torch.allclose(expected_out, actual_out, atol=1e-2, rtol=0):
print(f"Relevance test {green_color}passed{reset_color}: {vars(config)}")
else:
print(f"Relevance test {red_color}failed{reset_color}: {vars(config)}")
print("ort_output", actual_out)
print("expected_out", expected_out)
print("diff", expected_out - actual_out)
exit(1)
self.run_relevance_test(sm)
def run_relevance_no_past(sm: int, device):
"""Test prompt prefilling without past kv cache."""
for seq_len in [1, 64, 127, 128, 192, 256]:
config = SparseAttentionConfig(
batch_size=3,
sequence_length=seq_len,
max_sequence_length=256,
past_sequence_length=0,
num_heads=8,
kv_num_heads=4,
head_size=128,
sparse_block_size=64,
num_layout=2,
local_blocks=2,
vert_stride=2,
softmax_scale=1.8 / (128**0.5),
device=device,
)
run_one_relevance_test(config)
if sm >= 80:
config.dtype = torch.bfloat16
run_one_relevance_test(config)
def run_relevance_past(sm: int, device, do_rotary: bool):
"""Test token generation with past kv cache."""
for past_seq_len in [1, 63, 64, 127, 128, 511]:
config = SparseAttentionConfig(
batch_size=3,
sequence_length=1,
max_sequence_length=512,
past_sequence_length=past_seq_len,
num_heads=8,
kv_num_heads=4,
head_size=128,
sparse_block_size=64,
num_layout=4,
local_blocks=2,
vert_stride=4,
softmax_scale=None,
do_rotary=do_rotary,
rotary_interleaved=(past_seq_len % 2 == 1),
device=device,
)
if do_rotary:
# When there is rotary, we use ORT GQA as reference: ORT GQA does not support mask so here we use dense.
config.local_blocks = config.max_blocks
run_one_relevance_test(config)
if sm >= 80:
config.dtype = torch.bfloat16
run_one_relevance_test(config)
def run_relevance_test(sm: int):
device_id = torch.cuda.current_device()
device = torch.device("cuda", device_id)
with torch.no_grad():
run_relevance_no_past(sm, device)
run_relevance_past(sm, device, do_rotary=False)
run_relevance_past(sm, device, do_rotary=True)
# ------------------------------------------------------------------
# Below are performance tests
def get_plot_algos(sm: int):
# GQA with local windows only works in sm=8x
if sm >= 80:
return {
"line_vals": ["torch_gqa", "ort_gqa", "ort_gqa_local", "ort_sparse_att"],
"line_names": ["TORCH-GQA", "ORT-GQA-Dense", "ORT-GQA-Local", "ORT-SparseAtt"],
"styles": [("red", "-"), ("blue", "-"), ("yellow", "-"), ("green", "-")],
}
else:
return {
"line_vals": ["torch_gqa", "ort_gqa", "ort_sparse_att"],
"line_names": ["TORCH-GQA", "ORT-GQA-Dense", "ORT-SparseAtt"],
"styles": [("red", "-"), ("blue", "-"), ("green", "-")],
}
def plot_prompt_performance(
sm: int,
batch_size=4,
num_heads=32,
max_seq_len=8192,
head_size=128,
sparse_block_size=64,
local_blocks=16,
vert_stride=8,
num_layout=8,
dtype=torch.float16,
):
import triton
algos = get_plot_algos(sm)
configs = [
triton.testing.Benchmark(
x_names=["sequence_length"],
x_vals=[2**i for i in range(4, 14)],
line_arg="provider",
ylabel="ms",
**algos,
plot_name=f"prompt-sm{sm}-batch{batch_size}-head{num_heads}-d{head_size}-local{local_blocks}-vert{vert_stride}-{dtype}",
args={"num_heads": num_heads, "batch_size": batch_size, "head_size": head_size, "dtype": dtype},
)
]
@triton.testing.perf_report(configs)
def benchmark(batch_size, num_heads, sequence_length, head_size, provider, dtype=torch.float16, device="cuda"):
warmup = 15
repeat = 100
config: SparseAttentionConfig = SparseAttentionConfig(
batch_size=batch_size,
sequence_length=sequence_length,
max_sequence_length=max_seq_len,
past_sequence_length=0,
num_heads=num_heads,
kv_num_heads=8,
head_size=head_size,
sparse_block_size=sparse_block_size,
num_layout=num_layout,
local_blocks=local_blocks,
vert_stride=vert_stride,
)
if provider in ["ort_gqa", "ort_gqa_local"]:
gqa_config = config.get_comparable_gqa_config(use_local=(provider == "ort_gqa_local"))
obj = OrtGroupQueryAttention(gqa_config)
elif provider == "ort_sparse_att":
obj = OrtSparseAttention(config)
else: # Torch GQA
assert provider == "torch_gqa"
if sequence_length > 2048: # out of memory
return 0
gqa_config = config.get_comparable_gqa_config(torch_use_sparse=True)
def run_one_relevance_test(self, config: SparseAttentionConfig):
if not config.do_rotary:
# Run QGA by Torch
gqa_config: GroupQueryAttentionConfig = config.get_comparable_torch_gqa_config(use_sparse=True)
obj = TorchGroupQueryAttention(gqa_config)
ms = triton.testing.do_bench(obj.infer, warmup=warmup, rep=repeat)
return ms
benchmark.run(save_path=".", print_data=True)
def plot_token_performance(
sm: int,
batch_size=4,
num_heads=32,
max_seq_len=8192,
head_size=128,
sparse_block_size=64,
local_blocks=16,
vert_stride=8,
num_layout=8,
dtype=torch.float16,
):
import triton
algos = get_plot_algos(sm)
configs = [
triton.testing.Benchmark(
x_names=["past_sequence_length"],
x_vals=[2**i for i in range(4, 13)] + [max_seq_len - 1],
line_arg="provider",
ylabel="ms",
**algos,
plot_name=f"token-sm{sm}-batch{batch_size}-head{num_heads}-d{head_size}-local{local_blocks}-vert{vert_stride}-{dtype}",
args={"num_heads": num_heads, "batch_size": batch_size, "head_size": head_size, "dtype": dtype},
)
]
@triton.testing.perf_report(configs)
def benchmark(batch_size, num_heads, past_sequence_length, head_size, provider, dtype=torch.float16, device="cuda"):
warmup = 15
repeat = 100
config: SparseAttentionConfig = SparseAttentionConfig(
batch_size=batch_size,
sequence_length=1,
max_sequence_length=max_seq_len,
past_sequence_length=past_sequence_length,
num_heads=num_heads,
kv_num_heads=8,
head_size=head_size,
sparse_block_size=sparse_block_size,
num_layout=num_layout,
local_blocks=local_blocks,
vert_stride=vert_stride,
)
if provider in ["ort_gqa", "ort_gqa_local"]:
gqa_config = config.get_comparable_gqa_config(use_local=(provider == "ort_gqa_local"))
obj = OrtGroupQueryAttention(gqa_config)
elif provider == "ort_sparse_att":
obj = OrtSparseAttention(config)
expected_out = obj.infer()
else:
assert provider == "torch_gqa"
if past_sequence_length > 2048: # out of memory
return 0
gqa_config = config.get_comparable_gqa_config(torch_use_sparse=True)
obj = TorchGroupQueryAttention(gqa_config)
# Run QGA by ORT
gqa_config: GroupQueryAttentionConfig = config.get_comparable_ort_gqa_config(use_local=False)
obj = OrtGroupQueryAttention(gqa_config)
ort_qga_outputs = obj.infer()
expected_out = ort_qga_outputs["output"].view(
config.batch_size, config.sequence_length, config.num_heads, config.head_size
)
ms = triton.testing.do_bench(obj.infer, warmup=warmup, rep=repeat)
return ms
# Run SparseAttention by ORT
obj = OrtSparseAttention(config)
ort_outputs = obj.infer()
ort_output = ort_outputs["output"]
actual_out = ort_output.view(config.batch_size, config.sequence_length, config.num_heads, config.head_size)
benchmark.run(save_path=".", print_data=True)
red_color = "\033[31m"
green_color = "\033[32m"
reset_color = "\033[0m"
passed = torch.allclose(expected_out, actual_out, atol=1e-2, rtol=0)
if passed:
print(f"Relevance test {green_color}passed{reset_color}: {vars(config)}")
else:
print(f"Relevance test {red_color}failed{reset_color}: {vars(config)}")
print("ort_output", actual_out)
print("expected_out", expected_out)
print("diff", expected_out - actual_out)
def run_performance_test(sm: int):
"""
Run performance tests for prompt and token generation.
self.assertTrue(passed)
Example results in Azure Standard_ND96isr_H100_v5 VM with NVIDIA H100-80GB-HBM3 GPU (sm=90):
def run_relevance_no_past(self, sm: int, device):
"""Test prompt prefilling without past kv cache."""
for seq_len in [1, 64, 127, 128, 192, 256]:
for packed_qkv in [False, True]:
config = SparseAttentionConfig(
batch_size=3,
sequence_length=seq_len,
max_sequence_length=256,
past_sequence_length=0,
num_heads=8,
kv_num_heads=4,
head_size=128,
sparse_block_size=64,
num_layout=2,
local_blocks=2,
vert_stride=2,
softmax_scale=1.8 / (128**0.5),
device=device,
is_packed_qkv=packed_qkv,
)
self.run_one_relevance_test(config)
prompt-sm90-batch4-head32-d128-local16-vert8-torch.float16:
sequence_length TORCH-GQA ORT-GQA-Dense ORT-GQA-Local ORT-SparseAtt
0 16.0 0.079877 0.006362 0.006403 0.042758
1 32.0 0.086920 0.016404 0.016686 0.044183
2 64.0 0.090727 0.020429 0.020409 0.045343
3 128.0 0.128148 0.032009 0.031984 0.051516
4 256.0 0.323933 0.074110 0.073920 0.068308
5 512.0 1.021856 0.162167 0.161951 0.109226
6 1024.0 3.596002 0.452629 0.452780 0.231653
7 2048.0 13.865088 1.499534 1.195749 0.515488
8 4096.0 0.000000 5.454785 2.669682 1.163233
9 8192.0 0.000000 22.068159 6.018604 2.772873
if sm >= 80 and not packed_qkv:
config.dtype = torch.bfloat16
self.run_one_relevance_test(config)
token-sm90-batch4-head32-d128-local16-vert8-torch.float16:
past_sequence_length TORCH-GQA ORT-GQA-Dense ORT-GQA-Local ORT-SparseAtt
0 16.0 0.104460 0.012652 0.012661 0.069549
1 32.0 0.113866 0.012776 0.012765 0.069024
2 64.0 0.124600 0.016791 0.012672 0.069397
3 128.0 0.108658 0.017900 0.018294 0.074844
4 256.0 0.115463 0.029409 0.029608 0.078911
5 512.0 0.149824 0.033968 0.033701 0.092998
6 1024.0 0.234050 0.042930 0.042951 0.116920
7 2048.0 0.390695 0.061462 0.043008 0.121555
8 4096.0 0.000000 0.097505 0.042948 0.134757
9 8191.0 0.000000 0.165861 0.043542 0.158796
def run_relevance_past(self, sm: int, device, do_rotary: bool):
"""Test token generation with past kv cache."""
for past_seq_len in [1, 63, 64, 127, 128, 511]:
for packed_qkv in [False, True]:
config = SparseAttentionConfig(
batch_size=3,
sequence_length=1,
max_sequence_length=512,
past_sequence_length=past_seq_len,
num_heads=8,
kv_num_heads=4,
head_size=128,
sparse_block_size=64,
num_layout=4,
local_blocks=2,
vert_stride=4,
softmax_scale=None,
do_rotary=do_rotary,
rotary_interleaved=(past_seq_len % 2 == 1),
device=device,
is_packed_qkv=packed_qkv,
)
if do_rotary:
# When there is rotary, we use ORT GQA as reference: ORT GQA does not support mask so here we use dense.
config.local_blocks = config.max_blocks
Example results in A100-SXM4-80GB (sm=80):
self.run_one_relevance_test(config)
prompt-sm80-batch4-head32-d128-local16-vert8-torch.float16:
sequence_length TORCH-GQA ORT-GQA-Dense ORT-GQA-Local ORT-SparseAtt
0 16.0 0.274839 0.008849 0.015198 0.054403
1 32.0 0.272238 0.022875 0.048804 0.055898
2 64.0 0.272420 0.027722 0.028318 0.073052
3 128.0 0.273514 0.085971 0.062785 0.068287
4 256.0 0.545428 0.108228 0.135093 0.095949
5 512.0 1.678597 0.278193 0.248580 0.167271
6 1024.0 6.021056 0.702882 0.701022 0.379936
7 2048.0 23.512320 2.331175 1.863045 0.895726
8 4096.0 0.000000 8.789178 4.526275 2.105048
9 8192.0 0.000000 39.664131 10.046236 5.219436
if sm >= 80 and not packed_qkv:
config.dtype = torch.bfloat16
self.run_one_relevance_test(config)
token-sm80-batch4-head32-d128-local16-vert8-torch.float16:
past_sequence_length TORCH-GQA ORT-GQA-Dense ORT-GQA-Local ORT-SparseAtt
0 16.0 0.299303 0.020081 0.018587 0.082479
1 32.0 0.301700 0.018655 0.041943 0.084583
2 64.0 0.305700 0.017825 0.018420 0.085265
3 128.0 0.303379 0.023213 0.023152 0.090508
4 256.0 0.304119 0.034438 0.035257 0.100197
5 512.0 0.306051 0.063312 0.045373 0.114726
6 1024.0 0.359197 0.092181 0.088628 0.145165
7 2048.0 0.599463 0.101573 0.062101 0.159452
8 4096.0 0.000000 0.196258 0.091019 0.180342
9 8191.0 0.000000 0.334519 0.065158 0.213508
Example results in Standard_NC4as_T4_v3 Azure VM with T4 GPU (sm=75):
prompt-sm75-batch4-head32-d128-local16-vert8-torch.float16:
sequence_length TORCH-GQA ORT-GQA-Dense ORT-SparseAtt
0 16.0 0.165154 3.003173 0.081945
1 32.0 0.184173 2.994347 0.089064
2 64.0 0.303300 3.023986 0.107418
3 128.0 0.887795 3.073728 0.174213
4 256.0 2.797654 3.246899 0.357869
5 512.0 10.055048 3.814039 0.893903
6 1024.0 37.849937 5.818439 2.658720
7 2048.0 148.641785 13.638480 7.202690
8 4096.0 0.000000 43.556847 17.680954
9 8192.0 0.000000 161.628540 44.336670
token-sm75-batch4-head32-d128-local16-vert8-torch.float16:
past_sequence_length TORCH-GQA ORT-GQA-Dense ORT-SparseAtt
0 16.0 0.144368 4.179228 0.137407
1 32.0 0.110353 2.996305 0.137509
2 64.0 0.145088 3.006860 0.165424
3 128.0 0.219500 3.036448 0.192001
4 256.0 0.347496 3.071341 0.249125
5 512.0 0.595842 3.135225 0.398726
6 1024.0 1.081216 3.261110 0.612744
7 2048.0 2.060307 3.515578 0.685670
8 4096.0 0.000000 4.022986 0.819707
9 8191.0 0.000000 5.024528 1.072912
"""
with torch.no_grad():
plot_prompt_performance(sm=sm)
plot_token_performance(sm=sm)
def run_relevance_test(self, sm: int):
device_id = torch.cuda.current_device()
device = torch.device("cuda", device_id)
with torch.no_grad():
self.run_relevance_no_past(sm, device)
self.run_relevance_past(sm, device, do_rotary=False)
self.run_relevance_past(sm, device, do_rotary=True)
if __name__ == "__main__":
torch.set_printoptions(precision=6, edgeitems=3, linewidth=150, profile="default", sci_mode=False)
major, minor = torch.cuda.get_device_capability()
sm = major * 10 + minor
s = torch.cuda.Stream()
with torch.cuda.stream(s):
run_relevance_test(sm)
run_performance_test(sm)
unittest.main()