[LLVMGPU] Add dense and paged attention runtime correctness tests (#24829) ## Summary This PR adds a CUDA end-to-end runtime correctness test for paged-KV attention as Stage 2 of [#24766](https://github.com/iree-org/iree/issues/24766). This change also fixes and enables the shared host-side attention checker used by the existing dense-attention e2e tests. Before this change, the reference softmax was numerically incorrect and the checker generated expected results without comparing them against device output. ## Changes - Adds CUDA paged-KV attention coverage with non-contiguous, distinct key and value page tables, `iree_linalg_ext.gather`, and `iree_linalg_ext.online_attention`. - Makes the page-table dimension dynamic (`tensor<4x?xi64>`), matching the runtime `NUM_PAGES` use case. - Fixes the shared stable-softmax reference and makes it compare device output with expected results. - Refactors the paged checker to gather paged K/V into dense tensors and reuse the shared attention reference and checker. - Passes the f16-rounded `1 / sqrt(head_dim)` scale to both the operation and host reference. ## Testing Validated on NVIDIA RTX PRO 6000 Blackwell Server Edition - All existing CPU dense-attention configurations. - CUDA paged attention with 4 pages. - The same compiled paged-attention VMFB with 2, 4, and 6 runtime pages. Attention e2e tests use: ```text --require_exact_results=false --acceptable_fp_delta=0.01 ``` --------- Signed-off-by: weimin023 <tnwilly@gmail.com>
diff --git a/tests/e2e/attention/CMakeLists.txt b/tests/e2e/attention/CMakeLists.txt index ed58cdd..356158f 100644 --- a/tests/e2e/attention/CMakeLists.txt +++ b/tests/e2e/attention/CMakeLists.txt
@@ -1,3 +1,8 @@ +# Attention e2e tests compare a device result against a host reference that +# computes attention in f32. The operation order differs and the f16 inputs are +# rounded, so bit-exact results are not achievable and every test has to state +# an explicit tolerance. 0.01 was measured to hold across all shapes here. + iree_generated_e2e_runner_test( NAME e2e_attention_cpu_f16_f16_f16_small @@ -16,6 +21,9 @@ "llvm-cpu" DRIVERS "local-task" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "hostonly" "local" @@ -39,6 +47,9 @@ "llvm-cpu" DRIVERS "local-task" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "hostonly" "local" @@ -62,6 +73,9 @@ "llvm-cpu" DRIVERS "local-task" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "hostonly" "local" @@ -86,6 +100,9 @@ "llvm-cpu" DRIVERS "local-task" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "hostonly" "local" @@ -109,6 +126,9 @@ "llvm-cpu" DRIVERS "local-task" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "hostonly" "local" @@ -116,6 +136,30 @@ iree_generated_e2e_runner_test( NAME + e2e_paged_attention_gpu_cuda_f16 + TEST_TYPE + attention + GENERATOR + "generate_e2e_paged_attention_tests.py" + TEST_RUNNER + iree_tools_testing_e2e_iree-e2e-attention-test + TARGET_BACKENDS + "cuda" + DRIVERS + "cuda" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" + LABELS + "noasan" + "nomsan" + "notsan" + "noubsan" + "requires-gpu-nvidia" +) + +iree_generated_e2e_runner_test( + NAME e2e_attention_cpu_f16_f16_f16_decode_large TEST_TYPE attention @@ -132,6 +176,9 @@ "llvm-cpu" DRIVERS "local-task" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "hostonly" "local" @@ -157,6 +204,9 @@ "llvm-cpu" DRIVERS "local-task" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "hostonly" "local" @@ -181,6 +231,9 @@ "llvm-cpu" DRIVERS "local-task" + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "hostonly" "local" @@ -216,6 +269,9 @@ "hip" COMPILER_FLAGS ${IREE_HIP_TEST_COMPILER_FLAGS} + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "noasan" "nomsan" @@ -244,6 +300,9 @@ "hip" COMPILER_FLAGS ${IREE_HIP_TEST_COMPILER_FLAGS} + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "noasan" "nomsan" @@ -272,6 +331,9 @@ "hip" COMPILER_FLAGS ${IREE_HIP_TEST_COMPILER_FLAGS} + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "noasan" "nomsan" @@ -301,6 +363,9 @@ "hip" COMPILER_FLAGS ${IREE_HIP_TEST_COMPILER_FLAGS} + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "noasan" "nomsan" @@ -329,6 +394,9 @@ "hip" COMPILER_FLAGS ${IREE_HIP_TEST_COMPILER_FLAGS} + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "noasan" "nomsan" @@ -357,6 +425,9 @@ "hip" COMPILER_FLAGS ${IREE_HIP_TEST_COMPILER_FLAGS} + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "noasan" "nomsan" @@ -387,6 +458,9 @@ "hip" COMPILER_FLAGS ${IREE_HIP_TEST_COMPILER_FLAGS} + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "noasan" "nomsan" @@ -416,6 +490,9 @@ "hip" COMPILER_FLAGS ${IREE_HIP_TEST_COMPILER_FLAGS} + RUNNER_ARGS + "--require_exact_results=false" + "--acceptable_fp_delta=0.01" LABELS "noasan" "nomsan"
diff --git a/tests/e2e/attention/generate_e2e_attention_tests.py b/tests/e2e/attention/generate_e2e_attention_tests.py index e32ecdb..69d3cc8 100644 --- a/tests/e2e/attention/generate_e2e_attention_tests.py +++ b/tests/e2e/attention/generate_e2e_attention_tests.py
@@ -9,6 +9,7 @@ import argparse import enum import dataclasses +import struct import typing import math @@ -80,47 +81,77 @@ scale: float +# Attention is only numerically well conditioned when the Q*K^T scores stay in a +# reasonable range. With a scale of 1.0 the scores grow with k1, softmax +# saturates towards one-hot, and near-ties let the device and the host reference +# select different positions. Use the standard 1/sqrt(k1) scale instead. +# +# The value is pre-rounded to f16 because the kernel truncates the scale to f16. +# Doing the rounding here keeps the host reference on the exact same value +# without emitting f16 arithmetic into the host-side calls module. +def default_scale(k1: int) -> float: + return struct.unpack("<e", struct.pack("<e", 1.0 / math.sqrt(k1)))[0] + + # Returns the list of TestShape's to use for the collection of shapes # identified by shapes_id. def get_test_shapes(shapes_id: ShapesId): if shapes_id == ShapesId.SMALL: return [ - TestShapeAndScale(batch=2, m=256, k1=64, k2=32, n=16, scale=1.0), + TestShapeAndScale( + batch=2, m=256, k1=64, k2=32, n=16, scale=default_scale(64) + ), ] if shapes_id == ShapesId.MEDIUM: return [ - TestShapeAndScale(batch=2, m=512, k1=128, k2=64, n=32, scale=1.0), + TestShapeAndScale( + batch=2, m=512, k1=128, k2=64, n=32, scale=default_scale(128) + ), ] if shapes_id == ShapesId.LARGE: return [ - TestShapeAndScale(batch=2, m=1024, k1=128, k2=128, n=64, scale=1.0), + TestShapeAndScale( + batch=2, m=1024, k1=128, k2=128, n=64, scale=default_scale(128) + ), ] # Decode: m = 1 (single token attending to cached KV) if shapes_id == ShapesId.DECODE_SMALL: return [ - TestShapeAndScale(batch=2, m=1, k1=128, k2=128, n=128, scale=1.0), + TestShapeAndScale( + batch=2, m=1, k1=128, k2=128, n=128, scale=default_scale(128) + ), ] if shapes_id == ShapesId.DECODE_MEDIUM: return [ - TestShapeAndScale(batch=2, m=1, k1=128, k2=2048, n=128, scale=1.0), + TestShapeAndScale( + batch=2, m=1, k1=128, k2=2048, n=128, scale=default_scale(128) + ), ] if shapes_id == ShapesId.DECODE_LARGE: return [ - TestShapeAndScale(batch=2, m=1, k1=128, k2=16384, n=128, scale=1.0), + TestShapeAndScale( + batch=2, m=1, k1=128, k2=16384, n=128, scale=default_scale(128) + ), ] # Prefill: m = k2 (self-attention on full sequence) if shapes_id == ShapesId.PREFILL_SMALL: return [ - TestShapeAndScale(batch=2, m=128, k1=128, k2=128, n=128, scale=1.0), + TestShapeAndScale( + batch=2, m=128, k1=128, k2=128, n=128, scale=default_scale(128) + ), ] if shapes_id == ShapesId.PREFILL_MEDIUM: return [ - TestShapeAndScale(batch=2, m=2048, k1=128, k2=2048, n=128, scale=1.0), + TestShapeAndScale( + batch=2, m=2048, k1=128, k2=2048, n=128, scale=default_scale(128) + ), ] if shapes_id == ShapesId.PREFILL_LARGE: # Currently not used due to time-out. return [ - TestShapeAndScale(batch=2, m=16384, k1=128, k2=16384, n=128, scale=1.0), + TestShapeAndScale( + batch=2, m=16384, k1=128, k2=16384, n=128, scale=default_scale(128) + ), ] raise ValueError(shapes_id) @@ -393,7 +424,7 @@ mask_type: MaskType, ): global call_id - func_name = f"{function.name}_{shapes_scale.batch}_{shapes_scale.m}_{shapes_scale.k1}_{shapes_scale.k2}_{shapes_scale.n}_{shapes_scale.k1}_{shapes_scale.scale}" + func_name = f"{function.name}_{shapes_scale.batch}_{shapes_scale.m}_{shapes_scale.k1}_{shapes_scale.k2}_{shapes_scale.n}_{shapes_scale.k1}_{shapes_scale.scale:g}" func_name = f"{func_name}_{call_id}" call_id = call_id + 1 @@ -476,11 +507,11 @@ # Export mask as i8 for the reference implementation. op = op + ( f" %mask_i8_export = hal.tensor.export %mask_i8_reshaped : tensor<{batch}x{m}x{k2}xi8> -> !hal.buffer_view\n" - f" call @attention_test.check_attention_results_with_mask(%device, %batch, %m, %k1, %k2, %n, %queryExtBufferView, %keyExtBufferView, %valueExtBufferView, %mask_i8_export, %resultExtBufferView) : (!hal.device, i64, i64, i64, i64, i64, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> ()\n" + f" call @attention_test.check_attention_results_with_mask(%device, %batch, %m, %k1, %k2, %n, %scale, %queryExtBufferView, %keyExtBufferView, %valueExtBufferView, %mask_i8_export, %resultExtBufferView) : (!hal.device, i64, i64, i64, i64, i64, f32, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> ()\n" ) else: op = op + ( - f" call @attention_test.check_attention_results(%device, %batch, %m, %k1, %k2, %n, %queryExtBufferView, %keyExtBufferView, %valueExtBufferView, %resultExtBufferView) : (!hal.device, i64, i64, i64, i64, i64, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> ()\n" + f" call @attention_test.check_attention_results(%device, %batch, %m, %k1, %k2, %n, %scale, %queryExtBufferView, %keyExtBufferView, %valueExtBufferView, %resultExtBufferView) : (!hal.device, i64, i64, i64, i64, i64, f32, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> ()\n" ) op = op + " return\n" @@ -606,8 +637,8 @@ # Declare the custom module that generates arguments. module_definition = module_definition + ( "func.func private @attention_test.generate_random_tensor(%device: !hal.device, %dim0: i64, %dim1: i64, %dim2: i64, %element_type: i32, %seed: i32) -> !hal.buffer_view\n" - "func.func private @attention_test.check_attention_results(%device: !hal.device, %batch: i64, %m: i64, %k1: i64, %k2: i64, %n: i64, %query: !hal.buffer_view, %key: !hal.buffer_view, %value: !hal.buffer_view, %result: !hal.buffer_view)\n" - "func.func private @attention_test.check_attention_results_with_mask(%device: !hal.device, %batch: i64, %m: i64, %k1: i64, %k2: i64, %n: i64, %query: !hal.buffer_view, %key: !hal.buffer_view, %value: !hal.buffer_view, %mask: !hal.buffer_view, %result: !hal.buffer_view)\n" + "func.func private @attention_test.check_attention_results(%device: !hal.device, %batch: i64, %m: i64, %k1: i64, %k2: i64, %n: i64, %scale: f32, %query: !hal.buffer_view, %key: !hal.buffer_view, %value: !hal.buffer_view, %result: !hal.buffer_view)\n" + "func.func private @attention_test.check_attention_results_with_mask(%device: !hal.device, %batch: i64, %m: i64, %k1: i64, %k2: i64, %n: i64, %scale: f32, %query: !hal.buffer_view, %key: !hal.buffer_view, %value: !hal.buffer_view, %mask: !hal.buffer_view, %result: !hal.buffer_view)\n" "\n" )
diff --git a/tests/e2e/attention/generate_e2e_paged_attention_tests.py b/tests/e2e/attention/generate_e2e_paged_attention_tests.py new file mode 100644 index 0000000..0fac0dc --- /dev/null +++ b/tests/e2e/attention/generate_e2e_paged_attention_tests.py
@@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +# Copyright 2026 The IREE Authors +# +# Licensed under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +"""Generator for the paged-KV CUDA e2e attention test. + +This follows the structure of `generate_e2e_attention_tests.py`. The batch, +page pool and head dimensions are fixed, while the number of pages per request +is dynamic so that the test exercises the same shape behavior a real model has. +""" + +import argparse +import math +import struct + +BATCH = 4 +NUM_PAGES = 4 +POOL_PAGES = 8 +PAGE_SIZE = 32 +HEAD_DIM = 128 + +# The attention op consumes an f16 scale, so round 1/sqrt(HEAD_DIM) to f16 here +# and reuse the exact same value for the host reference. Keeping both sides on +# one Python value avoids a silently diverging magic constant. +SCALE = struct.unpack("<e", struct.pack("<e", 1.0 / math.sqrt(HEAD_DIM)))[0] + +QUERY_TYPE = f"tensor<{BATCH}x1x1x{HEAD_DIM}xf16>" +KV_STORAGE_TYPE = f"tensor<{POOL_PAGES}x1x{PAGE_SIZE}x{HEAD_DIM}xf16>" +PAGE_TABLE_TYPE = f"tensor<{BATCH}x?xi64>" +GATHERED_KV_TYPE = f"tensor<{BATCH}x?x1x{PAGE_SIZE}x{HEAD_DIM}xf16>" +MASK_TYPE = f"tensor<{BATCH}x1x1x?x{PAGE_SIZE}xf32>" +OUTPUT_TYPE = f"tensor<{BATCH}x1x1x{HEAD_DIM}xf32>" +STATISTICS_TYPE = f"tensor<{BATCH}x1x1xf32>" + + +def gather_kv_pages(result, page_table, output): + """Emits a gather of one logical KV tensor out of the physical page pool.""" + return f"""\ + %{result} = iree_linalg_ext.gather dimension_map = [0] + ins(%kv_storage, %{page_table} : {KV_STORAGE_TYPE}, {PAGE_TABLE_TYPE}) + outs(%{output} : {GATHERED_KV_TYPE}) + -> {GATHERED_KV_TYPE} +""" + + +def generate_attention_mlir(output_path): + text = f"""\ +func.func @paged_attention( + %query: {QUERY_TYPE}, + %kv_storage: {KV_STORAGE_TYPE}, + %key_page_table: {PAGE_TABLE_TYPE}, + %value_page_table: {PAGE_TABLE_TYPE}) + -> {OUTPUT_TYPE} {{ + %zero_f32 = arith.constant 0.0 : f32 + %neg_inf_f32 = arith.constant 0xFF800000 : f32 + %scale = arith.constant {SCALE!r} : f16 + %c1 = arith.constant 1 : index + %num_pages = tensor.dim %key_page_table, %c1 : {PAGE_TABLE_TYPE} + %key_empty = tensor.empty(%num_pages) : {GATHERED_KV_TYPE} + %value_empty = tensor.empty(%num_pages) : {GATHERED_KV_TYPE} +{gather_kv_pages("key", "key_page_table", "key_empty")}\ +{gather_kv_pages("value", "value_page_table", "value_empty")}\ + // All-zero additive mask. It leaves real scores untouched and only exists so + // that the compiler is able to pad the dynamic key/value dimension. + %mask_empty = tensor.empty(%num_pages) : {MASK_TYPE} + %mask = linalg.fill ins(%zero_f32 : f32) + outs(%mask_empty : {MASK_TYPE}) -> {MASK_TYPE} + %output_empty = tensor.empty() : {OUTPUT_TYPE} + %max_empty = tensor.empty() : {STATISTICS_TYPE} + %sum_empty = tensor.empty() : {STATISTICS_TYPE} + %output_init = linalg.fill ins(%zero_f32 : f32) + outs(%output_empty : {OUTPUT_TYPE}) -> {OUTPUT_TYPE} + %max_init = linalg.fill ins(%neg_inf_f32 : f32) + outs(%max_empty : {STATISTICS_TYPE}) -> {STATISTICS_TYPE} + %sum_init = linalg.fill ins(%zero_f32 : f32) + outs(%sum_empty : {STATISTICS_TYPE}) -> {STATISTICS_TYPE} + %attention:3 = iree_linalg_ext.online_attention {{ + indexing_maps = [ + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d4)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d5, d1, d6, d4)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d5, d1, d6, d3)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> ()>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d5, d6)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2)>] + }} + ins(%query, %key, %value, %scale, %mask : + {QUERY_TYPE}, {GATHERED_KV_TYPE}, {GATHERED_KV_TYPE}, f16, {MASK_TYPE}) + outs(%output_init, %max_init, %sum_init : + {OUTPUT_TYPE}, {STATISTICS_TYPE}, {STATISTICS_TYPE}) {{ + ^bb0(%score: f32): + iree_linalg_ext.yield %score : f32 + }} -> {OUTPUT_TYPE}, {STATISTICS_TYPE}, {STATISTICS_TYPE} + %normalized_empty = tensor.empty() : {OUTPUT_TYPE} + %normalized = linalg.generic {{ + indexing_maps = [ + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>, + affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], + iterator_types = ["parallel", "parallel", "parallel", "parallel"] + }} ins(%attention#0, %attention#2 : {OUTPUT_TYPE}, {STATISTICS_TYPE}) + outs(%normalized_empty : {OUTPUT_TYPE}) {{ + ^bb0(%attn_value: f32, %sum: f32, %out: f32): + %result = arith.divf %attn_value, %sum : f32 + linalg.yield %result : f32 + }} -> {OUTPUT_TYPE} + return %normalized : {OUTPUT_TYPE} +}} +""" + with open(output_path, "w") as file: + file.write(text) + + +def page_table_entries(offset, stride): + """Builds a deterministic, non-contiguous logical-to-physical page mapping. + + The strides are coprime with POOL_PAGES so the mapping walks the whole page + pool, both strides are odd so the key and value tables never select the same + physical page for one slot, and the batch index shifts the mapping so no two + requests share a page sequence. + """ + return [ + (offset + stride * page + batch) % POOL_PAGES + for batch in range(BATCH) + for page in range(NUM_PAGES) + ] + + +def emit_page_table(name, pages): + """Emits one host-constant page table and exports it as a buffer view.""" + flat_type = f"tensor<{BATCH * NUM_PAGES}xi64>" + table_type = f"tensor<{BATCH}x{NUM_PAGES}xi64>" + values = ", ".join(str(page) for page in pages) + return f"""\ + %{name}_flat = arith.constant dense<[{values}]> : {flat_type} + %{name}_table = tensor.expand_shape %{name}_flat [[0, 1]] output_shape [{BATCH}, {NUM_PAGES}] : {flat_type} into {table_type} + %{name}_page_table = hal.tensor.export %{name}_table : {table_type} -> !hal.buffer_view +""" + + +def generate_calls_mlir(output_path): + key_pages = page_table_entries(3, 5) + value_pages = page_table_entries(6, 7) + text = f"""\ +builtin.module @calls {{ + func.func private @attention_test.generate_random_4d_tensor( + %device: !hal.device, %dim0: i64, %dim1: i64, %dim2: i64, + %dim3: i64, %element_type: i32, %seed: i32) -> !hal.buffer_view + func.func private @attention_test.check_paged_attention_results( + %device: !hal.device, %batch: i64, %num_pages: i64, %page_size: i64, + %head_dim: i64, %scale: f32, %query: !hal.buffer_view, + %kv_storage: !hal.buffer_view, %key_page_table: !hal.buffer_view, + %value_page_table: !hal.buffer_view, %result: !hal.buffer_view) + func.func private @module.paged_attention( + %query: !hal.buffer_view, %kv_storage: !hal.buffer_view, + %key_page_table: !hal.buffer_view, %value_page_table: !hal.buffer_view) + -> !hal.buffer_view + func.func @paged_attention_test() attributes {{ + iree.reflection = {{description = "Paged KV attention B{BATCH} pages{NUM_PAGES} page{PAGE_SIZE} head{HEAD_DIM} f16"}} + }} {{ + %device_index = arith.constant 0 : index + %device = hal.devices.get %device_index : !hal.device + %batch = arith.constant {BATCH} : i64 + %num_pages = arith.constant {NUM_PAGES} : i64 + %page_size = arith.constant {PAGE_SIZE} : i64 + %head_dim = arith.constant {HEAD_DIM} : i64 + %scale = arith.constant {SCALE!r} : f32 + %c1 = arith.constant 1 : i64 + %pool_pages = arith.constant {POOL_PAGES} : i64 + %query_seed = arith.constant 17 : i32 + %kv_seed = arith.constant 23 : i32 + %query_type = hal.element_type<f16> : i32 + %query = call @attention_test.generate_random_4d_tensor( + %device, %batch, %c1, %c1, %head_dim, %query_type, %query_seed) : (!hal.device, i64, i64, i64, i64, i32, i32) -> !hal.buffer_view + %kv = call @attention_test.generate_random_4d_tensor( + %device, %pool_pages, %c1, %page_size, %head_dim, %query_type, %kv_seed) : (!hal.device, i64, i64, i64, i64, i32, i32) -> !hal.buffer_view +{emit_page_table("key", key_pages)}\ +{emit_page_table("value", value_pages)}\ + %result = call @module.paged_attention(%query, %kv, %key_page_table, %value_page_table) : (!hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> !hal.buffer_view + call @attention_test.check_paged_attention_results(%device, %batch, %num_pages, %page_size, %head_dim, %scale, %query, %kv, %key_page_table, %value_page_table, %result) : (!hal.device, i64, i64, i64, i64, f32, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) -> () + return + }} +}} +""" + with open(output_path, "w") as file: + file.write(text) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output_attention_mlir", required=True) + parser.add_argument("--output_calls_mlir", required=True) + args = parser.parse_args() + generate_attention_mlir(args.output_attention_mlir) + generate_calls_mlir(args.output_calls_mlir) + + +if __name__ == "__main__": + main()
diff --git a/tools/testing/e2e/iree-e2e-attention-test.cc b/tools/testing/e2e/iree-e2e-attention-test.cc index 7dc96bd..aedff21 100644 --- a/tools/testing/e2e/iree-e2e-attention-test.cc +++ b/tools/testing/e2e/iree-e2e-attention-test.cc
@@ -5,11 +5,14 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include <float.h> +#include <inttypes.h> #include <math.h> #include <stdio.h> #include <stdlib.h> #include <string.h> +#include <vector> + #include "iree/base/api.h" #include "iree/base/internal/cpu.h" #include "iree/base/internal/math.h" @@ -50,9 +53,9 @@ static void reference_attention_f32_f32_f32_f32( iree_hal_dim_t M, iree_hal_dim_t K1, iree_hal_dim_t K2, iree_hal_dim_t N, - iree_hal_dim_t B, const float* query_data, const float* key_data, - const float* value_data, float* result_data, iree_hal_dim_t b, - float* Attention, const uint8_t* mask_data) { + iree_hal_dim_t B, float scale, const float* query_data, + const float* key_data, const float* value_data, float* result_data, + iree_hal_dim_t b, float* Attention, const uint8_t* mask_data) { // Compute Q * K^T for (int m = 0; m < M; ++m) { for (int k2 = 0; k2 < K2; ++k2) { @@ -64,7 +67,7 @@ sum += query_data[q_idx] * key_data[k_idx]; } int att_idx = index_3d(0, m, k2, M, K2); - Attention[att_idx] = sum / sqrt(K1); // Scale by sqrt(K1) + Attention[att_idx] = sum * scale; } } @@ -101,7 +104,7 @@ // Apply softmax for (int k2 = 0; k2 < K2; ++k2) { int att_idx = index_3d(0, m, k2, M, K2); - Attention[att_idx] = exp(Attention[att_idx]) / sum; + Attention[att_idx] = exp(Attention[att_idx] - max_val) / sum; } } @@ -122,17 +125,18 @@ static iree_status_t reference_attention_element( iree_hal_dim_t M, iree_hal_dim_t K1, iree_hal_dim_t K2, iree_hal_dim_t N, - iree_hal_dim_t B, iree_hal_element_type_t query_elem_type, + iree_hal_dim_t B, float scale, iree_hal_element_type_t query_elem_type, iree_hal_element_type_t key_elem_type, iree_hal_element_type_t value_elem_type, void* query_data, void* key_data, - void* value_data, void* actual_data, void* result_data, iree_hal_dim_t b, - float* Attention, const uint8_t* mask_data) { + void* value_data, void* result_data, iree_hal_dim_t b, float* Attention, + const uint8_t* mask_data) { if (query_elem_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32 && key_elem_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32 && value_elem_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) { reference_attention_f32_f32_f32_f32( - M, K1, K2, N, B, (const float*)query_data, (const float*)key_data, - (const float*)value_data, (float*)result_data, b, Attention, mask_data); + M, K1, K2, N, B, scale, (const float*)query_data, + (const float*)key_data, (const float*)value_data, (float*)result_data, + b, Attention, mask_data); } else { return iree_make_status( @@ -146,12 +150,12 @@ // against. static iree_status_t reference_attention( iree_hal_dim_t B, iree_hal_dim_t M, iree_hal_dim_t K1, iree_hal_dim_t K2, - iree_hal_dim_t N, iree_hal_element_type_t query_elem_type, + iree_hal_dim_t N, float scale, iree_hal_element_type_t query_elem_type, iree_hal_element_type_t key_elem_type, iree_hal_element_type_t value_elem_type, iree_byte_span_t query_contents, iree_byte_span_t key_contents, iree_byte_span_t value_contents, - iree_byte_span_t actual_contents, iree_byte_span_t result_contents, - iree_byte_span_t mask_contents, int compute_every) { + iree_byte_span_t result_contents, iree_byte_span_t mask_contents, + int compute_every) { IREE_TRACE_ZONE_BEGIN(z0); IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, B); IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, M); @@ -168,12 +172,15 @@ for (iree_hal_dim_t b = 0; b < B; ++b) { if (++count < compute_every) continue; count = 0; - IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, reference_attention_element( - M, K1, K2, N, B, query_elem_type, key_elem_type, - value_elem_type, query_contents.data, key_contents.data, - value_contents.data, actual_contents.data, result_contents.data, - b, Attention, mask_data)); + iree_status_t status = reference_attention_element( + M, K1, K2, N, B, scale, query_elem_type, key_elem_type, value_elem_type, + query_contents.data, key_contents.data, value_contents.data, + result_contents.data, b, Attention, mask_data); + if (!iree_status_is_ok(status)) { + free_tensor(Attention); + IREE_TRACE_ZONE_END(z0); + return status; + } } free_tensor(Attention); @@ -191,6 +198,9 @@ iree_hal_dim_t k1; iree_hal_dim_t k2; iree_hal_dim_t n; + // Scale applied to Q*K^T by the attention op under test. The reference must + // use the exact same value or the softmax distributions will not match. + float scale; iree_hal_element_type_t query_elem_type; iree_hal_element_type_t key_elem_type; iree_hal_element_type_t value_elem_type; @@ -208,7 +218,7 @@ static iree_status_t attention_results_initialize( iree_hal_device_t* device, iree_hal_dim_t b_size, iree_hal_dim_t m_size, iree_hal_dim_t k1_size, iree_hal_dim_t k2_size, iree_hal_dim_t n_size, - iree_hal_buffer_view_t* query, iree_hal_buffer_view_t* key, + float scale, iree_hal_buffer_view_t* query, iree_hal_buffer_view_t* key, iree_hal_buffer_view_t* value, iree_hal_buffer_view_t* result, iree_hal_buffer_view_t* mask, // Optional: can be nullptr. iree_allocator_t host_allocator, attention_results_t* out_results) { @@ -222,6 +232,7 @@ out_results->k1 = k1_size; out_results->k2 = k2_size; out_results->n = n_size; + out_results->scale = scale; out_results->query_elem_type = iree_hal_buffer_view_element_type(query); out_results->key_elem_type = iree_hal_buffer_view_element_type(key); @@ -338,13 +349,48 @@ IREE_TRACE_ZONE_BEGIN(z0); IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, reference_attention(results->b, results->m, results->k1, results->k2, - results->n, results->query_elem_type, - results->key_elem_type, results->value_elem_type, - results->query_contents, results->key_contents, - results->value_contents, results->actual_contents, - results->expected_contents, - results->mask_contents, check_every)); + z0, reference_attention( + results->b, results->m, results->k1, results->k2, results->n, + results->scale, results->query_elem_type, results->key_elem_type, + results->value_elem_type, results->query_contents, + results->key_contents, results->value_contents, + results->expected_contents, results->mask_contents, check_every)); + + iree_host_size_t count = 0; + for (iree_hal_dim_t b = 0; b < results->b; ++b) { + if (++count < check_every) continue; + count = 0; + for (iree_hal_dim_t m = 0; m < results->m; ++m) { + for (iree_hal_dim_t n = 0; n < results->n; ++n) { + iree_hal_dim_t idx = (b * results->m + m) * results->n + n; + iree_test_utils_e2e_value_t actual_value = + iree_test_utils_read_buffer_element(idx, results->result_elem_type, + results->actual_contents.data); + iree_test_utils_e2e_value_t expected_value = + iree_test_utils_read_buffer_element( + idx, results->result_elem_type, + results->expected_contents.data); + if (!iree_test_utils_result_elements_agree(expected_value, + actual_value)) { + if (file) { + char actual_str[32]; + char expected_str[32]; + iree_test_utils_snprintf_value(actual_str, sizeof(actual_str), + actual_value); + iree_test_utils_snprintf_value(expected_str, sizeof(expected_str), + expected_value); + fprintf(file, + "\n\nerror: the actual and expected attention results " + "disagree at b %" PRIdim ", m %" PRIdim ", n %" PRIdim + ": actual %s, expected %s\n\n", + b, m, n, actual_str, expected_str); + } + IREE_TRACE_ZONE_END(z0); + return iree_make_status(IREE_STATUS_ABORTED); + } + } + } + } IREE_TRACE_ZONE_END(z0); return iree_ok_status(); @@ -443,8 +489,197 @@ return std::move(result_view); } + StatusOr<vm::ref<iree_hal_buffer_view_t>> GenerateRandom4dTensor( + iree_hal_device_t* device, int64_t dim0, int64_t dim1, int64_t dim2, + int64_t dim3, iree_hal_element_type_t element_type, int32_t seed) { + iree_hal_dim_t dims[4] = { + (iree_hal_dim_t)dim0, + (iree_hal_dim_t)dim1, + (iree_hal_dim_t)dim2, + (iree_hal_dim_t)dim3, + }; + iree_hal_buffer_params_t buffer_params = {0}; + buffer_params.usage = IREE_HAL_BUFFER_USAGE_DEFAULT; + buffer_params.access = IREE_HAL_MEMORY_ACCESS_ALL; + buffer_params.type = IREE_HAL_MEMORY_TYPE_OPTIMAL_FOR_DEVICE; + vm::ref<iree_hal_buffer_view_t> result_view; + struct callback_state_t { + iree_hal_element_type_t element_type; + int32_t seed; + } callback_state = {element_type, seed}; + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_generate_buffer( + device, iree_hal_device_allocator(device), IREE_ARRAYSIZE(dims), dims, + element_type, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, buffer_params, + +[](iree_hal_buffer_mapping_t* mapping, void* user_data) { + callback_state_t callback_state = *(callback_state_t*)user_data; + iree_byte_span_t span = mapping->contents; + int32_t min = 0; + int32_t max = 0; + iree_test_utils_get_min_max_for_element_type( + callback_state.element_type, &min, &max); + uint32_t range = (max - min + 1); + iree_host_size_t element_byte_count = + iree_hal_element_dense_byte_count(callback_state.element_type); + uint8_t* data_end = span.data + span.data_length; + uint32_t state = callback_state.seed; + for (uint8_t* data = span.data; data < data_end; + data += element_byte_count) { + int32_t value = + (int32_t)iree_test_utils_pseudorandom_range(&state, range) + + min; + iree_test_utils_write_element(callback_state.element_type, value, + data); + } + return iree_ok_status(); + }, + &callback_state, &result_view)); + return std::move(result_view); + } + + // Gathers the paged KV cache into dense key/value tensors and then defers to + // the shared attention reference and result checker. + Status CheckPagedAttentionResults(iree_hal_device_t* device, int64_t batch, + int64_t num_pages, int64_t page_size, + int64_t head_dim, float scale, + iree_hal_buffer_view_t* query, + iree_hal_buffer_view_t* kv_storage, + iree_hal_buffer_view_t* key_page_table, + iree_hal_buffer_view_t* value_page_table, + iree_hal_buffer_view_t* actual_result) { + iree_hal_buffer_t* query_buffer = iree_hal_buffer_view_buffer(query); + iree_hal_buffer_t* kv_buffer = iree_hal_buffer_view_buffer(kv_storage); + iree_hal_buffer_t* key_table_buffer = + iree_hal_buffer_view_buffer(key_page_table); + iree_hal_buffer_t* value_table_buffer = + iree_hal_buffer_view_buffer(value_page_table); + iree_hal_buffer_t* result_buffer = + iree_hal_buffer_view_buffer(actual_result); + + const iree_host_size_t query_bytes = + iree_hal_buffer_byte_length(query_buffer); + const iree_host_size_t kv_bytes = iree_hal_buffer_byte_length(kv_buffer); + const iree_host_size_t key_table_bytes = + iree_hal_buffer_byte_length(key_table_buffer); + const iree_host_size_t value_table_bytes = + iree_hal_buffer_byte_length(value_table_buffer); + const iree_host_size_t result_bytes = + iree_hal_buffer_byte_length(result_buffer); + + std::vector<uint8_t> query_data(query_bytes); + std::vector<uint8_t> kv_data(kv_bytes); + std::vector<uint8_t> key_table_data(key_table_bytes); + std::vector<uint8_t> value_table_data(value_table_bytes); + std::vector<uint8_t> result_data(result_bytes); + + IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h( + device, query_buffer, 0, query_data.data(), query_data.size(), + IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout())); + IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h( + device, kv_buffer, 0, kv_data.data(), kv_data.size(), + IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout())); + IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h( + device, key_table_buffer, 0, key_table_data.data(), + key_table_data.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, + iree_infinite_timeout())); + IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h( + device, value_table_buffer, 0, value_table_data.data(), + value_table_data.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, + iree_infinite_timeout())); + IREE_RETURN_IF_ERROR(iree_hal_device_transfer_d2h( + device, result_buffer, 0, result_data.data(), result_data.size(), + IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout())); + + if (iree_hal_buffer_view_element_type(query) != + IREE_HAL_ELEMENT_TYPE_FLOAT_16 || + iree_hal_buffer_view_element_type(kv_storage) != + IREE_HAL_ELEMENT_TYPE_FLOAT_16 || + iree_hal_buffer_view_element_type(key_page_table) != + IREE_HAL_ELEMENT_TYPE_INT_64 || + iree_hal_buffer_view_element_type(value_page_table) != + IREE_HAL_ELEMENT_TYPE_INT_64 || + iree_hal_buffer_view_element_type(actual_result) != + IREE_HAL_ELEMENT_TYPE_FLOAT_32) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "unexpected paged attention element types"); + } + + const int64_t pool_pages = + kv_bytes / (sizeof(uint16_t) * page_size * head_dim); + const uint16_t* query_f16 = + reinterpret_cast<const uint16_t*>(query_data.data()); + const uint16_t* kv_f16 = reinterpret_cast<const uint16_t*>(kv_data.data()); + const int64_t* key_pages = + reinterpret_cast<const int64_t*>(key_table_data.data()); + const int64_t* value_pages = + reinterpret_cast<const int64_t*>(value_table_data.data()); + + if (query_bytes != batch * head_dim * sizeof(uint16_t) || + key_table_bytes != batch * num_pages * sizeof(int64_t) || + value_table_bytes != batch * num_pages * sizeof(int64_t) || + result_bytes != batch * head_dim * sizeof(float) || pool_pages <= 0) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "unexpected paged attention buffer sizes"); + } + + const int64_t sequence_length = num_pages * page_size; + std::vector<float> query_f32(batch * head_dim); + std::vector<float> key_f32(batch * sequence_length * head_dim); + std::vector<float> value_f32(batch * sequence_length * head_dim); + std::vector<float> expected(batch * head_dim); + for (int64_t b = 0; b < batch; ++b) { + for (int64_t d = 0; d < head_dim; ++d) { + query_f32[b * head_dim + d] = + iree_math_f16_to_f32(query_f16[b * head_dim + d]); + } + for (int64_t position = 0; position < sequence_length; ++position) { + const int64_t logical_page = position / page_size; + const int64_t page_offset = position % page_size; + const int64_t key_page = key_pages[b * num_pages + logical_page]; + const int64_t value_page = value_pages[b * num_pages + logical_page]; + if (key_page < 0 || key_page >= pool_pages || value_page < 0 || + value_page >= pool_pages) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "paged attention page is out of bounds"); + } + for (int64_t d = 0; d < head_dim; ++d) { + const int64_t dense_index = + (b * sequence_length + position) * head_dim + d; + key_f32[dense_index] = iree_math_f16_to_f32( + kv_f16[(key_page * page_size + page_offset) * head_dim + d]); + value_f32[dense_index] = iree_math_f16_to_f32( + kv_f16[(value_page * page_size + page_offset) * head_dim + d]); + } + } + } + + // All spans below point at buffers owned by this function, so this struct + // must not be passed to attention_results_deinitialize. + attention_results_t results = {}; + results.b = batch; + results.m = 1; + results.k1 = head_dim; + results.k2 = sequence_length; + results.n = head_dim; + results.scale = scale; + results.query_elem_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32; + results.key_elem_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32; + results.value_elem_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32; + results.result_elem_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32; + results.query_contents = + iree_make_byte_span(query_f32.data(), query_f32.size() * sizeof(float)); + results.key_contents = + iree_make_byte_span(key_f32.data(), key_f32.size() * sizeof(float)); + results.value_contents = + iree_make_byte_span(value_f32.data(), value_f32.size() * sizeof(float)); + results.actual_contents = + iree_make_byte_span(result_data.data(), result_data.size()); + results.expected_contents = + iree_make_byte_span(expected.data(), expected.size() * sizeof(float)); + return check_attention_results(stderr, &results); + } + Status CheckAttentionResults(iree_hal_device_t* device, int64_t b, int64_t m, - int64_t k1, int64_t k2, int64_t n, + int64_t k1, int64_t k2, int64_t n, float scale, iree_hal_buffer_view_t* query, iree_hal_buffer_view_t* key, iree_hal_buffer_view_t* value, @@ -452,7 +687,8 @@ attention_results_t results = {}; IREE_RETURN_IF_ERROR(attention_results_initialize( device, (iree_hal_dim_t)b, (iree_hal_dim_t)m, (iree_hal_dim_t)k1, - (iree_hal_dim_t)k2, (iree_hal_dim_t)n, query, key, value, actual_result, + (iree_hal_dim_t)k2, (iree_hal_dim_t)n, scale, query, key, value, + actual_result, nullptr, // No mask host_allocator_, &results)); iree_status_t status = check_attention_results(stderr, &results); @@ -460,18 +696,16 @@ return status; } - Status CheckAttentionResultsWithMask(iree_hal_device_t* device, int64_t b, - int64_t m, int64_t k1, int64_t k2, - int64_t n, iree_hal_buffer_view_t* query, - iree_hal_buffer_view_t* key, - iree_hal_buffer_view_t* value, - iree_hal_buffer_view_t* mask, - iree_hal_buffer_view_t* actual_result) { + Status CheckAttentionResultsWithMask( + iree_hal_device_t* device, int64_t b, int64_t m, int64_t k1, int64_t k2, + int64_t n, float scale, iree_hal_buffer_view_t* query, + iree_hal_buffer_view_t* key, iree_hal_buffer_view_t* value, + iree_hal_buffer_view_t* mask, iree_hal_buffer_view_t* actual_result) { attention_results_t results = {}; IREE_RETURN_IF_ERROR(attention_results_initialize( device, (iree_hal_dim_t)b, (iree_hal_dim_t)m, (iree_hal_dim_t)k1, - (iree_hal_dim_t)k2, (iree_hal_dim_t)n, query, key, value, actual_result, - mask, host_allocator_, &results)); + (iree_hal_dim_t)k2, (iree_hal_dim_t)n, scale, query, key, value, + actual_result, mask, host_allocator_, &results)); iree_status_t status = check_attention_results(stderr, &results); attention_results_deinitialize(&results); return status; @@ -487,6 +721,12 @@ "generate_random_tensor", &AttentionTestModuleState::GenerateRandom3dTensor), vm::MakeNativeFunction( + "generate_random_4d_tensor", + &AttentionTestModuleState::GenerateRandom4dTensor), + vm::MakeNativeFunction( + "check_paged_attention_results", + &AttentionTestModuleState::CheckPagedAttentionResults), + vm::MakeNativeFunction( "check_attention_results", &AttentionTestModuleState::CheckAttentionResults), vm::MakeNativeFunction(