[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(
