From efce593f613abf3a35aeb5d2ace3c3486a9cd753 Mon Sep 17 00:00:00 2001 From: Ruize Yu Date: Fri, 6 Mar 2026 11:09:20 -0500 Subject: [PATCH 1/6] feat: add flashattention in example and modify ci test --- .github/workflows/config.yml | 7 ++ examples/flashattention/flash_Atten.py | 156 +++++++++++++++++++++++++ examples/flashattention/test_flash.py | 101 ++++++++++++++++ 3 files changed, 264 insertions(+) create mode 100644 examples/flashattention/flash_Atten.py create mode 100644 examples/flashattention/test_flash.py diff --git a/.github/workflows/config.yml b/.github/workflows/config.yml index 8f1a3dcbf4..2f811c817c 100644 --- a/.github/workflows/config.yml +++ b/.github/workflows/config.yml @@ -84,3 +84,10 @@ jobs: python3 -m pip install torch==2.5.1 python3 examples/torch/toy.py python3 examples/torch/mlp.py + - name: flashAttention + shell: bash + run: | + source activate allo + export PATH=/root/llvm-project/build/bin:${PATH} + export LLVM_BUILD_DIR=/root/llvm-project/build + python3 examples/flashattention/test_flash.py diff --git a/examples/flashattention/flash_Atten.py b/examples/flashattention/flash_Atten.py new file mode 100644 index 0000000000..4739953752 --- /dev/null +++ b/examples/flashattention/flash_Atten.py @@ -0,0 +1,156 @@ +# Copyright Allo authors. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import allo +from allo.ir.types import float32, int32, bool, index + +def get_scheduled_flash_attention( + BATCH_SIZE: int, + CONTEXT_LENGTH: int, + HIDDEN_SIZE: int, + NUM_HEADS: int, + BLOCK_T: int = 4 +): + + HEAD_DIM = HIDDEN_SIZE // NUM_HEADS + D_SQRT = float(HEAD_DIM ** 0.5) + THREE_H = 3 * HIDDEN_SIZE + IN_ELEMS = BATCH_SIZE * CONTEXT_LENGTH * THREE_H + OUT_ELEMS = BATCH_SIZE * CONTEXT_LENGTH * NUM_HEADS * HEAD_DIM + + + def load_tile(global_mem: float32[IN_ELEMS], local_tile: float32[BLOCK_T, HEAD_DIM], b: index, h: index, t_start: index, type_offset: int32): + for t, d in allo.grid(BLOCK_T, HEAD_DIM): + t_global: index = t_start + t + global_c: index = type_offset * HIDDEN_SIZE + h * HEAD_DIM + d + idx: index = b * (CONTEXT_LENGTH * THREE_H) + t_global * THREE_H + global_c + local_tile[t, d] = global_mem[idx] + + def store_tile(global_mem: float32[OUT_ELEMS], local_tile: float32[BLOCK_T, HEAD_DIM], b: index, h: index, t_start: index): + for t, d in allo.grid(BLOCK_T, HEAD_DIM): + t_global: index = t_start + t + idx: index = (((b * CONTEXT_LENGTH + t_global) * NUM_HEADS + h) * HEAD_DIM + d) + global_mem[idx] = local_tile[t, d] + + def compute_block_attention(Q_tile: float32[BLOCK_T, HEAD_DIM], K_tile: float32[BLOCK_T, HEAD_DIM], V_tile: float32[BLOCK_T, HEAD_DIM], O_tile: float32[BLOCK_T, HEAD_DIM], m_vec: float32[BLOCK_T], l_vec: float32[BLOCK_T], scale: float32, is_first_block: bool): + S_tile: float32[BLOCK_T, BLOCK_T] = 0.0 + + for i, j in allo.grid(BLOCK_T, BLOCK_T): + for d in range(HEAD_DIM): + S_tile[i, j] += Q_tile[i, d] * K_tile[j, d] + + for i1, j1 in allo.grid(BLOCK_T, BLOCK_T): + S_tile[i1, j1] = S_tile[i1, j1] * scale + + for m in range(BLOCK_T): + row_max_val: float32 = -1e30 + + for n in range(BLOCK_T): + if S_tile[m, n] > row_max_val: + row_max_val = S_tile[m, n] + else: + row_max_val = row_max_val + + m_prev: float32 = m_vec[m] + m_new: float32 = 0.0 + if row_max_val > m_prev: + m_new = row_max_val + else: + m_new = m_prev + + alpha: float32 = 0.0 + if is_first_block: + alpha = 0.0 + else: + alpha = allo.exp(m_prev - m_new) + beta: float32 = allo.exp(row_max_val - m_new) + row_sum_exp: float32 = 0.0 + + for k in range(BLOCK_T): + row_sum_exp = row_sum_exp + allo.exp(S_tile[m, k] - row_max_val) + + l_new: float32 = (l_vec[m] * alpha) + (row_sum_exp * beta) + + p_val: float32[BLOCK_T] = 0.0 + for z in range(BLOCK_T): + p_val[z] = allo.exp(S_tile[m, z] - row_max_val) + + for h in range(HEAD_DIM): + pv_sum: float32 = 0.0 + for l in range(BLOCK_T): + pv_sum = pv_sum + p_val[l] * V_tile[l, h] + O_tile[m, h] = O_tile[m, h] * alpha + pv_sum * beta + + m_vec[m] = m_new + l_vec[m] = l_new + + def compute_engine(input_mem: float32[IN_ELEMS], output_mem: float32[OUT_ELEMS]): + scale: float32 = 1.0 / D_SQRT + Q_sram: float32[BLOCK_T, HEAD_DIM] + K_sram: float32[BLOCK_T, HEAD_DIM] + V_sram: float32[BLOCK_T, HEAD_DIM] + O_sram: float32[BLOCK_T, HEAD_DIM] + + m_sram: float32[BLOCK_T] + l_sram: float32[BLOCK_T] + + for b, h in allo.grid(BATCH_SIZE, NUM_HEADS): + for tr in range(0, CONTEXT_LENGTH, BLOCK_T): + load_tile(input_mem, Q_sram, b, h, tr, 0) + for m in range(BLOCK_T): + m_sram[m] = -1e30 + l_sram[m] = 0.0 + for n in range(HEAD_DIM): + O_sram[m, n] = 0.0 + + for tc in range(0, CONTEXT_LENGTH, BLOCK_T): + load_tile(input_mem, K_sram, b, h, tc, 1) + load_tile(input_mem, V_sram, b, h, tc, 2) + compute_block_attention(Q_sram, K_sram, V_sram, O_sram, m_sram, l_sram, scale, (tc == 0)) + + for i in range(BLOCK_T): + inv_l: float32 = 1.0 / (l_sram[i] + 1e-9) + for d in range(HEAD_DIM): + O_sram[i, d] = O_sram[i, d] * inv_l + + store_tile(output_mem, O_sram, b, h, tr) + + s1 = allo.customize(load_tile) + s1.pipeline("d") + s1.pipeline("t") + + s2 = allo.customize(store_tile) + s2.pipeline("d") + s2.pipeline("t") + + s3 = allo.customize(compute_block_attention) + s3.partition(s3.Q_tile, partition_type=0, dim=0) + s3.partition(s3.K_tile, partition_type=0, dim=0) + s3.unroll("d") + s3.unroll("j", factor=4) + s3.unroll("i", factor=4) + + s3.partition(s3.S_tile, partition_type=0, dim=0) + s3.partition(s3.O_tile, partition_type=0, dim=0) + s3.partition(s3.V_tile, partition_type=0, dim=0) + s3.partition(s3.p_val, partition_type=0, dim=0) + s3.unroll("z") + s3.unroll("j1") + s3.unroll("i1") + s3.pipeline("n") + s3.unroll("k") + s3.unroll("l") + s3.unroll("h") + s3.unroll("m") + + s = allo.customize(compute_engine) + s.partition(s.O_sram, partition_type=0, dim=0) + s.partition(s.m_sram, partition_type=0, dim=0) + s.partition(s.l_sram, partition_type=0, dim=0) + s.unroll("n") + s.unroll("m") + s.pipeline("d") + + s.compose([s1, s2, s3]) + + return s \ No newline at end of file diff --git a/examples/flashattention/test_flash.py b/examples/flashattention/test_flash.py new file mode 100644 index 0000000000..1790dd34a2 --- /dev/null +++ b/examples/flashattention/test_flash.py @@ -0,0 +1,101 @@ +# Copyright Allo authors. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import tempfile +import pytest +import numpy as np +import scipy.special +import allo.backend.hls as hls + +from flash_Atten import get_scheduled_flash_attention + +def run_test_with_params(BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLOCK_T): + print("=" * 60) + print(f"Testing FlashAttention: BATCH={BATCH_SIZE}, SEQ_LEN={CONTEXT_LENGTH}, " + f"HIDDEN={HIDDEN_SIZE}, HEADS={NUM_HEADS}, BLOCK_T={BLOCK_T}") + print("=" * 60) + + HEAD_DIM = HIDDEN_SIZE // NUM_HEADS + D_SQRT = HEAD_DIM ** 0.5 + THREE_H = 3 * HIDDEN_SIZE + IN_ELEMS = BATCH_SIZE * CONTEXT_LENGTH * THREE_H + OUT_ELEMS = BATCH_SIZE * CONTEXT_LENGTH * NUM_HEADS * HEAD_DIM + + A = np.random.rand(IN_ELEMS).astype(np.float32) + B_out = np.zeros(OUT_ELEMS, dtype=np.float32) + + A_reshaped = A.reshape((BATCH_SIZE, CONTEXT_LENGTH, 3, NUM_HEADS, HEAD_DIM)) + Q_np = A_reshaped[:, :, 0, :, :].transpose((0, 2, 1, 3)) + K_np = A_reshaped[:, :, 1, :, :].transpose((0, 2, 1, 3)) + V_np = A_reshaped[:, :, 2, :, :].transpose((0, 2, 1, 3)) + + scores = np.matmul(Q_np, K_np.transpose((0, 1, 3, 2))) + scores = scores * (1.0 / D_SQRT) + attn_weights = scipy.special.softmax(scores, axis=-1) + + out_np = np.matmul(attn_weights, V_np) + B_golden = out_np.transpose((0, 2, 1, 3)).flatten() + + s = get_scheduled_flash_attention( + BATCH_SIZE=BATCH_SIZE, + CONTEXT_LENGTH=CONTEXT_LENGTH, + HIDDEN_SIZE=HIDDEN_SIZE, + NUM_HEADS=NUM_HEADS, + BLOCK_T=BLOCK_T + ) + + print("Running Software Simulator for numerical correctness...") + sim_mod = s.build(target="llvm") + sim_mod(A, B_out) + + try: + np.testing.assert_allclose(B_out, B_golden, rtol=1e-4, atol=1e-4) + print("✅ Simulator Test Passed: Outputs match Golden Reference!") + except AssertionError as e: + print("❌ Simulator Test Failed!") + raise e + + if hls.is_available("vitis_hls"): + print("Running Vitis HLS Synthesis") + with tempfile.TemporaryDirectory() as tmpdir: + hls_mod = s.build( + target="vitis_hls", + mode="csyn", + project=tmpdir + ) + hls_mod() + print("✅ HLS Synthesis Passed!") + else: + print("⚠️ Vitis HLS not available, skipping C synthesis.") + + +def test_flashattention(): + run_test_with_params( + BATCH_SIZE=4, + CONTEXT_LENGTH=16, + HIDDEN_SIZE=64, + NUM_HEADS=4, + BLOCK_T=4 + ) + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser(description="Allo FlashAttention Testbench") + + parser.add_argument("--BATCH_SIZE", type=int, default=4, required=False, help="Batch size of input data") + parser.add_argument("--CONTEXT_LENGTH", type=int, default=16, required=False, help="Context length of input data") + parser.add_argument("--HIDDEN_SIZE", type=int, default=64, required=False, help="Hidden size of input data") + parser.add_argument("--NUM_HEADS", type=int, default=4, required=False, help="Number of heads of input data") + parser.add_argument("--BLOCK_T", type=int, default=4, required=False, help="Size of tiles") + + args = parser.parse_args() + + run_test_with_params( + BATCH_SIZE=args.BATCH_SIZE, + CONTEXT_LENGTH=args.CONTEXT_LENGTH, + HIDDEN_SIZE=args.HIDDEN_SIZE, + NUM_HEADS=args.NUM_HEADS, + BLOCK_T=args.BLOCK_T + ) \ No newline at end of file From 14e26098c409ee090ba4edba2940892107a50575 Mon Sep 17 00:00:00 2001 From: Ruize Yu Date: Fri, 6 Mar 2026 13:52:34 -0500 Subject: [PATCH 2/6] [feat] implementing flash attention in allo and add it to the example --- examples/flashattention/flash_Atten.py | 64 ++++++++++++++------ examples/flashattention/test_flash.py | 81 ++++++++++++++++---------- 2 files changed, 96 insertions(+), 49 deletions(-) diff --git a/examples/flashattention/flash_Atten.py b/examples/flashattention/flash_Atten.py index 4739953752..7fc21b1034 100644 --- a/examples/flashattention/flash_Atten.py +++ b/examples/flashattention/flash_Atten.py @@ -4,41 +4,65 @@ import allo from allo.ir.types import float32, int32, bool, index + def get_scheduled_flash_attention( - BATCH_SIZE: int, - CONTEXT_LENGTH: int, - HIDDEN_SIZE: int, - NUM_HEADS: int, - BLOCK_T: int = 4 + BATCH_SIZE: int, + CONTEXT_LENGTH: int, + HIDDEN_SIZE: int, + NUM_HEADS: int, + BLOCK_T: int = 4, ): - + HEAD_DIM = HIDDEN_SIZE // NUM_HEADS - D_SQRT = float(HEAD_DIM ** 0.5) + D_SQRT = float(HEAD_DIM**0.5) THREE_H = 3 * HIDDEN_SIZE IN_ELEMS = BATCH_SIZE * CONTEXT_LENGTH * THREE_H OUT_ELEMS = BATCH_SIZE * CONTEXT_LENGTH * NUM_HEADS * HEAD_DIM - - def load_tile(global_mem: float32[IN_ELEMS], local_tile: float32[BLOCK_T, HEAD_DIM], b: index, h: index, t_start: index, type_offset: int32): + def load_tile( + global_mem: float32[IN_ELEMS], + local_tile: float32[BLOCK_T, HEAD_DIM], + b: index, + h: index, + t_start: index, + type_offset: int32, + ): for t, d in allo.grid(BLOCK_T, HEAD_DIM): t_global: index = t_start + t global_c: index = type_offset * HIDDEN_SIZE + h * HEAD_DIM + d idx: index = b * (CONTEXT_LENGTH * THREE_H) + t_global * THREE_H + global_c local_tile[t, d] = global_mem[idx] - def store_tile(global_mem: float32[OUT_ELEMS], local_tile: float32[BLOCK_T, HEAD_DIM], b: index, h: index, t_start: index): + def store_tile( + global_mem: float32[OUT_ELEMS], + local_tile: float32[BLOCK_T, HEAD_DIM], + b: index, + h: index, + t_start: index, + ): for t, d in allo.grid(BLOCK_T, HEAD_DIM): t_global: index = t_start + t - idx: index = (((b * CONTEXT_LENGTH + t_global) * NUM_HEADS + h) * HEAD_DIM + d) + idx: index = ( + (b * CONTEXT_LENGTH + t_global) * NUM_HEADS + h + ) * HEAD_DIM + d global_mem[idx] = local_tile[t, d] - def compute_block_attention(Q_tile: float32[BLOCK_T, HEAD_DIM], K_tile: float32[BLOCK_T, HEAD_DIM], V_tile: float32[BLOCK_T, HEAD_DIM], O_tile: float32[BLOCK_T, HEAD_DIM], m_vec: float32[BLOCK_T], l_vec: float32[BLOCK_T], scale: float32, is_first_block: bool): + def compute_block_attention( + Q_tile: float32[BLOCK_T, HEAD_DIM], + K_tile: float32[BLOCK_T, HEAD_DIM], + V_tile: float32[BLOCK_T, HEAD_DIM], + O_tile: float32[BLOCK_T, HEAD_DIM], + m_vec: float32[BLOCK_T], + l_vec: float32[BLOCK_T], + scale: float32, + is_first_block: bool, + ): S_tile: float32[BLOCK_T, BLOCK_T] = 0.0 for i, j in allo.grid(BLOCK_T, BLOCK_T): for d in range(HEAD_DIM): S_tile[i, j] += Q_tile[i, d] * K_tile[j, d] - + for i1, j1 in allo.grid(BLOCK_T, BLOCK_T): S_tile[i1, j1] = S_tile[i1, j1] * scale @@ -57,7 +81,7 @@ def compute_block_attention(Q_tile: float32[BLOCK_T, HEAD_DIM], K_tile: float32[ m_new = row_max_val else: m_new = m_prev - + alpha: float32 = 0.0 if is_first_block: alpha = 0.0 @@ -102,12 +126,14 @@ def compute_engine(input_mem: float32[IN_ELEMS], output_mem: float32[OUT_ELEMS]) l_sram[m] = 0.0 for n in range(HEAD_DIM): O_sram[m, n] = 0.0 - + for tc in range(0, CONTEXT_LENGTH, BLOCK_T): load_tile(input_mem, K_sram, b, h, tc, 1) load_tile(input_mem, V_sram, b, h, tc, 2) - compute_block_attention(Q_sram, K_sram, V_sram, O_sram, m_sram, l_sram, scale, (tc == 0)) - + compute_block_attention( + Q_sram, K_sram, V_sram, O_sram, m_sram, l_sram, scale, (tc == 0) + ) + for i in range(BLOCK_T): inv_l: float32 = 1.0 / (l_sram[i] + 1e-9) for d in range(HEAD_DIM): @@ -152,5 +178,5 @@ def compute_engine(input_mem: float32[IN_ELEMS], output_mem: float32[OUT_ELEMS]) s.pipeline("d") s.compose([s1, s2, s3]) - - return s \ No newline at end of file + + return s diff --git a/examples/flashattention/test_flash.py b/examples/flashattention/test_flash.py index 1790dd34a2..66dda4c028 100644 --- a/examples/flashattention/test_flash.py +++ b/examples/flashattention/test_flash.py @@ -10,14 +10,17 @@ from flash_Atten import get_scheduled_flash_attention + def run_test_with_params(BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLOCK_T): print("=" * 60) - print(f"Testing FlashAttention: BATCH={BATCH_SIZE}, SEQ_LEN={CONTEXT_LENGTH}, " - f"HIDDEN={HIDDEN_SIZE}, HEADS={NUM_HEADS}, BLOCK_T={BLOCK_T}") + print( + f"Testing FlashAttention: BATCH={BATCH_SIZE}, SEQ_LEN={CONTEXT_LENGTH}, " + f"HIDDEN={HIDDEN_SIZE}, HEADS={NUM_HEADS}, BLOCK_T={BLOCK_T}" + ) print("=" * 60) HEAD_DIM = HIDDEN_SIZE // NUM_HEADS - D_SQRT = HEAD_DIM ** 0.5 + D_SQRT = HEAD_DIM**0.5 THREE_H = 3 * HIDDEN_SIZE IN_ELEMS = BATCH_SIZE * CONTEXT_LENGTH * THREE_H OUT_ELEMS = BATCH_SIZE * CONTEXT_LENGTH * NUM_HEADS * HEAD_DIM @@ -33,16 +36,16 @@ def run_test_with_params(BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLO scores = np.matmul(Q_np, K_np.transpose((0, 1, 3, 2))) scores = scores * (1.0 / D_SQRT) attn_weights = scipy.special.softmax(scores, axis=-1) - + out_np = np.matmul(attn_weights, V_np) B_golden = out_np.transpose((0, 2, 1, 3)).flatten() s = get_scheduled_flash_attention( - BATCH_SIZE=BATCH_SIZE, - CONTEXT_LENGTH=CONTEXT_LENGTH, - HIDDEN_SIZE=HIDDEN_SIZE, - NUM_HEADS=NUM_HEADS, - BLOCK_T=BLOCK_T + BATCH_SIZE=BATCH_SIZE, + CONTEXT_LENGTH=CONTEXT_LENGTH, + HIDDEN_SIZE=HIDDEN_SIZE, + NUM_HEADS=NUM_HEADS, + BLOCK_T=BLOCK_T, ) print("Running Software Simulator for numerical correctness...") @@ -59,11 +62,7 @@ def run_test_with_params(BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLO if hls.is_available("vitis_hls"): print("Running Vitis HLS Synthesis") with tempfile.TemporaryDirectory() as tmpdir: - hls_mod = s.build( - target="vitis_hls", - mode="csyn", - project=tmpdir - ) + hls_mod = s.build(target="vitis_hls", mode="csyn", project=tmpdir) hls_mod() print("✅ HLS Synthesis Passed!") else: @@ -72,30 +71,52 @@ def run_test_with_params(BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLO def test_flashattention(): run_test_with_params( - BATCH_SIZE=4, - CONTEXT_LENGTH=16, - HIDDEN_SIZE=64, - NUM_HEADS=4, - BLOCK_T=4 + BATCH_SIZE=4, CONTEXT_LENGTH=16, HIDDEN_SIZE=64, NUM_HEADS=4, BLOCK_T=4 ) if __name__ == "__main__": - + parser = argparse.ArgumentParser(description="Allo FlashAttention Testbench") - - parser.add_argument("--BATCH_SIZE", type=int, default=4, required=False, help="Batch size of input data") - parser.add_argument("--CONTEXT_LENGTH", type=int, default=16, required=False, help="Context length of input data") - parser.add_argument("--HIDDEN_SIZE", type=int, default=64, required=False, help="Hidden size of input data") - parser.add_argument("--NUM_HEADS", type=int, default=4, required=False, help="Number of heads of input data") - parser.add_argument("--BLOCK_T", type=int, default=4, required=False, help="Size of tiles") - + + parser.add_argument( + "--BATCH_SIZE", + type=int, + default=4, + required=False, + help="Batch size of input data", + ) + parser.add_argument( + "--CONTEXT_LENGTH", + type=int, + default=16, + required=False, + help="Context length of input data", + ) + parser.add_argument( + "--HIDDEN_SIZE", + type=int, + default=64, + required=False, + help="Hidden size of input data", + ) + parser.add_argument( + "--NUM_HEADS", + type=int, + default=4, + required=False, + help="Number of heads of input data", + ) + parser.add_argument( + "--BLOCK_T", type=int, default=4, required=False, help="Size of tiles" + ) + args = parser.parse_args() - + run_test_with_params( BATCH_SIZE=args.BATCH_SIZE, CONTEXT_LENGTH=args.CONTEXT_LENGTH, HIDDEN_SIZE=args.HIDDEN_SIZE, NUM_HEADS=args.NUM_HEADS, - BLOCK_T=args.BLOCK_T - ) \ No newline at end of file + BLOCK_T=args.BLOCK_T, + ) From 2a034292b5e8c22ffbb85fdfbb19d5a522999f43 Mon Sep 17 00:00:00 2001 From: Ruize Yu Date: Fri, 6 Mar 2026 15:55:56 -0500 Subject: [PATCH 3/6] fix: replace scipy with pure numpy to fix CI dependency error --- examples/flashattention/test_flash.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/flashattention/test_flash.py b/examples/flashattention/test_flash.py index 66dda4c028..dbb3fb7bdf 100644 --- a/examples/flashattention/test_flash.py +++ b/examples/flashattention/test_flash.py @@ -5,7 +5,6 @@ import tempfile import pytest import numpy as np -import scipy.special import allo.backend.hls as hls from flash_Atten import get_scheduled_flash_attention @@ -35,7 +34,9 @@ def run_test_with_params(BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLO scores = np.matmul(Q_np, K_np.transpose((0, 1, 3, 2))) scores = scores * (1.0 / D_SQRT) - attn_weights = scipy.special.softmax(scores, axis=-1) + max_scores = np.max(scores, axis=-1, keepdims=True) + exp_scores = np.exp(scores - max_scores) + attn_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True) out_np = np.matmul(attn_weights, V_np) B_golden = out_np.transpose((0, 2, 1, 3)).flatten() From b80b1d061484a0ad8479d187d1dccaa591bf03c0 Mon Sep 17 00:00:00 2001 From: Ruize Yu Date: Sun, 8 Mar 2026 19:18:20 -0400 Subject: [PATCH 4/6] feat: fix package import issue and divisibility check --- examples/flashattention/__init__.py | 0 examples/flashattention/flash_Atten.py | 7 +++---- examples/flashattention/test_flash.py | 7 +++++-- 3 files changed, 8 insertions(+), 6 deletions(-) create mode 100644 examples/flashattention/__init__.py diff --git a/examples/flashattention/__init__.py b/examples/flashattention/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/flashattention/flash_Atten.py b/examples/flashattention/flash_Atten.py index 7fc21b1034..4f416da3e1 100644 --- a/examples/flashattention/flash_Atten.py +++ b/examples/flashattention/flash_Atten.py @@ -12,7 +12,6 @@ def get_scheduled_flash_attention( NUM_HEADS: int, BLOCK_T: int = 4, ): - HEAD_DIM = HIDDEN_SIZE // NUM_HEADS D_SQRT = float(HEAD_DIM**0.5) THREE_H = 3 * HIDDEN_SIZE @@ -99,11 +98,11 @@ def compute_block_attention( for z in range(BLOCK_T): p_val[z] = allo.exp(S_tile[m, z] - row_max_val) - for h in range(HEAD_DIM): + for g in range(HEAD_DIM): pv_sum: float32 = 0.0 for l in range(BLOCK_T): - pv_sum = pv_sum + p_val[l] * V_tile[l, h] - O_tile[m, h] = O_tile[m, h] * alpha + pv_sum * beta + pv_sum = pv_sum + p_val[l] * V_tile[l, g] + O_tile[m, g] = O_tile[m, g] * alpha + pv_sum * beta m_vec[m] = m_new l_vec[m] = l_new diff --git a/examples/flashattention/test_flash.py b/examples/flashattention/test_flash.py index dbb3fb7bdf..2121bd8a32 100644 --- a/examples/flashattention/test_flash.py +++ b/examples/flashattention/test_flash.py @@ -3,14 +3,17 @@ import argparse import tempfile -import pytest import numpy as np import allo.backend.hls as hls -from flash_Atten import get_scheduled_flash_attention +from examples.flashattention.flash_Atten import get_scheduled_flash_attention def run_test_with_params(BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLOCK_T): + + assert HIDDEN_SIZE % NUM_HEADS == 0, f"HIDDEN_SIZE ({HIDDEN_SIZE}) must be exactly divisible by NUM_HEADS ({NUM_HEADS})" + assert CONTEXT_LENGTH % BLOCK_T == 0, f"CONTEXT_LENGTH ({CONTEXT_LENGTH}) must be exactly divisible by BLOCK_T ({BLOCK_T})" + print("=" * 60) print( f"Testing FlashAttention: BATCH={BATCH_SIZE}, SEQ_LEN={CONTEXT_LENGTH}, " From e1cc1f5502b59b3f18e7e43e0a37e11e2b33812b Mon Sep 17 00:00:00 2001 From: Ruize Yu Date: Sun, 8 Mar 2026 19:22:30 -0400 Subject: [PATCH 5/6] feat: fix schedule error --- examples/flashattention/flash_Atten.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/flashattention/flash_Atten.py b/examples/flashattention/flash_Atten.py index 4f416da3e1..b6a87063f5 100644 --- a/examples/flashattention/flash_Atten.py +++ b/examples/flashattention/flash_Atten.py @@ -165,7 +165,7 @@ def compute_engine(input_mem: float32[IN_ELEMS], output_mem: float32[OUT_ELEMS]) s3.pipeline("n") s3.unroll("k") s3.unroll("l") - s3.unroll("h") + s3.unroll("g") s3.unroll("m") s = allo.customize(compute_engine) From 11738e5a9257f657dbea673c6f992fd668c4ce6c Mon Sep 17 00:00:00 2001 From: Ruize Yu Date: Sun, 8 Mar 2026 19:29:53 -0400 Subject: [PATCH 6/6] feat: fix license and format error --- examples/flashattention/__init__.py | 2 ++ examples/flashattention/test_flash.py | 8 ++++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/examples/flashattention/__init__.py b/examples/flashattention/__init__.py index e69de29bb2..1ab7e98c71 100644 --- a/examples/flashattention/__init__.py +++ b/examples/flashattention/__init__.py @@ -0,0 +1,2 @@ +# Copyright Allo authors. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/examples/flashattention/test_flash.py b/examples/flashattention/test_flash.py index 2121bd8a32..03a0bb993b 100644 --- a/examples/flashattention/test_flash.py +++ b/examples/flashattention/test_flash.py @@ -11,8 +11,12 @@ def run_test_with_params(BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLOCK_T): - assert HIDDEN_SIZE % NUM_HEADS == 0, f"HIDDEN_SIZE ({HIDDEN_SIZE}) must be exactly divisible by NUM_HEADS ({NUM_HEADS})" - assert CONTEXT_LENGTH % BLOCK_T == 0, f"CONTEXT_LENGTH ({CONTEXT_LENGTH}) must be exactly divisible by BLOCK_T ({BLOCK_T})" + assert ( + HIDDEN_SIZE % NUM_HEADS == 0 + ), f"HIDDEN_SIZE ({HIDDEN_SIZE}) must be exactly divisible by NUM_HEADS ({NUM_HEADS})" + assert ( + CONTEXT_LENGTH % BLOCK_T == 0 + ), f"CONTEXT_LENGTH ({CONTEXT_LENGTH}) must be exactly divisible by BLOCK_T ({BLOCK_T})" print("=" * 60) print(