diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 4ad1fa786d..461d4e318a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -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; diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu index 448b2476e4..5d6182b613 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu @@ -138,6 +138,7 @@ Status QkvToContext( auto q = reinterpret_cast(data.unpacked_qkv_buffer); auto k = reinterpret_cast(data.unpacked_qkv_buffer + q_size); auto v = reinterpret_cast(data.unpacked_qkv_buffer + q_size + k_size); + Status status = LaunchUnpackQKV(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(query), batch_size, num_heads, sequence_length, head_size); - #if DUMP_TENSOR_LEVEL > 0 + DUMP_TENSOR("query (BNSH)", reinterpret_cast(query), batch_size, num_heads, sequence_length, head_size); + if (LAYOUT_BNSH == kv_layout) { - DUMP_TENSOR("key", reinterpret_cast(key), batch_size, kv_num_heads, sequence_length, head_size); - DUMP_TENSOR("value", reinterpret_cast(value), batch_size, kv_num_heads, sequence_length, head_size); + DUMP_TENSOR("key (BNSH)", reinterpret_cast(key), batch_size, kv_num_heads, sequence_length, head_size); + DUMP_TENSOR("value (BNSH)", reinterpret_cast(value), batch_size, kv_num_heads, sequence_length, head_size); } else { - DUMP_TENSOR("key", reinterpret_cast(key), batch_size, sequence_length, kv_num_heads, head_size); - DUMP_TENSOR("value", reinterpret_cast(value), batch_size, sequence_length, kv_num_heads, head_size); + DUMP_TENSOR("key (BSNH)", reinterpret_cast(key), batch_size, sequence_length, kv_num_heads, head_size); + DUMP_TENSOR("value (BSNH)", reinterpret_cast(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(data.output), batch_size, num_heads, sequence_length, head_size); + return Status::OK(); } diff --git a/onnxruntime/test/python/transformers/test_sparse_attention.py b/onnxruntime/test/python/transformers/test_sparse_attention.py index c5e5d232bd..58c48aa133 100644 --- a/onnxruntime/test/python/transformers/test_sparse_attention.py +++ b/onnxruntime/test/python/transformers/test_sparse_attention.py @@ -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()