[Example] Add two Flashattention accelerator example by RuizeYu05 · Pull Request #615 · cornell-zhang/allo

RuizeYu05 · 2026-09-30T01:25:43Z

Description

I merged the two PR (#564 and #581) submitted before about Flashattention examples

Problems

Add example for Allo

Proposed Solutions

Implement Flashattention accelerator in allo

Examples

It's an example

Checklist

Please make sure to review and check all of these items:

PR's title starts with a category (e.g. [Bugfix], [IR], [Builder], etc)
All changes have test coverage (It would be good to provide ~2 different test cases to test the robustness of your code)
Pass the formatting check locally
Code is well-documented

Fangtangtang · 2026-09-30T02:25:44Z

Thanks! Does this include PRs #564 and #581? If so, I'll close them.

RuizeYu05 · 2026-09-30T02:27:48Z

Yes you could close them

…

On Tue, Sep 29, 2026 at 10:26 PM Shihan Fang ***@***.***> wrote: *Fangtangtang* left a comment (cornell-zhang/allo#615) <#615 (comment)> Thanks! Does this include PRs #564 <#564> and #581 <#581>? If so, I'll close them. — Reply to this email directly, view it on GitHub <#615?email_source=notifications&email_token=BPK7XM6CFLXQCMQ5FTK2SM35RRVL5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKOJQGI4DINBVHA32M4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2KYZTPN52GK4S7MNWGSY3L#issuecomment-5902844587>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/BPK7XM5STXXA25U2GVOVLI35RRVL5AVCNFSNUABFKJSXA33TNF2G64TZHM3DMOJSHEZTSNRVHNEXG43VMU5TKNRTHE2DAMJYHAYKC5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/BPK7XMYWJVOHGUZBCKTZK2D5RRVL5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKOJQGI4DINBVHA32M4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2KUZTPN52GK4S7NFXXG> and Android <https://github.com/notifications/mobile/android/BPK7XM5ALG4FWX47DK6CBI35RRVL5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKOJQGI4DINBVHA32M4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2K4ZTPN52GK4S7MFXGI4TPNFSA>. Download it today! You are receiving this because you authored the thread.Message ID: ***@***.***>

Copilot

Copilot review overview

🟡 Changes recommended

Invalid shape handling, softmax initialization, and silently skipped systolic correctness testing can produce incorrect results or ineffective CI coverage.

Review effort: Balanced
Findings: 3 · 1

Open (4)

Reject non-divisible context and hidden-size shapes · New
Validate exact context and head divisibility in factory · New
Handle empty online-softmax state before exponentiation · New
Run simulator comparison independently of HLS synthesis · New

What changed in this PR

Adds two FPGA attention accelerator examples: tiled FlashAttention and a quantized systolic MHA design.

Changes:

Adds both accelerator implementations and numerical/HLS test drivers.
Documents their algorithms and hardware organization.
Integrates checks into standard and weekly FPGA workflows.

File	Description
`examples/attention/README.md`	Documents both attention accelerators.
`examples/attention/fused_MHA_systolic/fused_MHA_systolic.py`	Implements quantized systolic MHA.
`examples/attention/fused_MHA_systolic/test_systolic.py`	Adds reference and HLS testing.
`examples/attention/fused_MHA_systolic/__init__.py`	Initializes the example package.
`examples/attention/flashattention/flash_Atten.py`	Implements tiled FlashAttention.
`examples/attention/flashattention/test_flash.py`	Adds LLVM and HLS testing.
`examples/attention/flashattention/__init__.py`	Initializes the example package.
`.github/workflows/config.yml`	Adds both examples to PR CI.
`.github/workflows/fpga_weekly.yml`	Adds weekly synthesis coverage.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+    NUM_HEADS: int,
+    BLOCK_T: int = 4,
+):
+    HEAD_DIM = HIDDEN_SIZE // NUM_HEADS


+    HEAD_DIM = HIDDEN_SIZE // NUM_HEADS
+    NUM_TC = CONTEXT_LENGTH // BLOCK_T
+
+    assert NUM_TC == BLOCK_T, "This design requires NUM_TC == BLOCK_T"


+                            m_new: float32 = m_cur
+                            if x > m_cur:
+                                m_new = x
+                            ep: float32 = allo.exp(m_cur - m_new)
+                            ex: float32 = allo.exp(x - m_new)


Fangtangtang

Could you please clean up some of the structural issues first? and check copilot's suggestions

RuizeYu05 · 2026-10-01T01:36:04Z

Could you please clean up some of the structural issues first? and check copilot's suggestions

All done as you wish, my lord.

Fangtangtang

Thanks for your patience! I have a few additional suggestions.

Fangtangtang · 2026-10-05T13:25:52Z

+def test_flashattention():
+    run_test_with_params(
+        BATCH_SIZE=4, CONTEXT_LENGTH=16, HIDDEN_SIZE=64, NUM_HEADS=4, BLOCK_T=4
+    )


This wrapper seems to exist only for pytest test discovery?
We can make run_test_with_params the pytest test directly and pass these arguments using @pytest.mark.parametrize. Something like

@pytest.mark.parametrize( "BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLOCK_T", [ (4, 16, 64, 4, 4), ], ) def test_flashattention( BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLOCK_T, ):

Actually I was thinking of removing the wrapper: rename run_test_with_params to test_flashattention so it can be triggered by pytest. The goal is to avoid introducing unnecessary functions.

Fangtangtang · 2026-10-05T13:28:44Z

+import allo.dataflow as df
+
+int8 = Int(8)
+int32 = Int(32)


why not import int8, int32directly from allo.ir.types

We are testing quantization method then. It's easier for us to change data type.

If you'd like in this way, I will modify it.

I see, but I'm not sure how this makes changing data types easier. I'd prefer directly importing the predefined types

Fangtangtang · 2026-10-05T13:33:22Z

+def test_fused_MHA_systolic():
+    run_test_with_params(
+        BATCH_SIZE=4, CONTEXT_LENGTH=16, HIDDEN_SIZE=16, NUM_HEADS=4, BLOCK_T=4
+    )


Similar to test_flash_attention.py above.

Fangtangtang · 2026-10-05T13:41:58Z

you may need to update file names in this README. It would be better to make them relative links to the corresponding files.
Also, shall we simplify this and make it more technical? I think a short description of each design plus a concrete result table would be clearer.

Fangtangtang · 2026-10-08T03:29:06Z

+def test_flashattention():
+    run_test_with_params(
+        BATCH_SIZE=4, CONTEXT_LENGTH=16, HIDDEN_SIZE=64, NUM_HEADS=4, BLOCK_T=4
+    )


Actually I was thinking of removing the wrapper: rename run_test_with_params to test_flashattention so it can be triggered by pytest. The goal is to avoid introducing unnecessary functions.

Fangtangtang · 2026-10-08T03:31:43Z

+@pytest.mark.parametrize(
+    "BATCH_SIZE, CONTEXT_LENGTH, HIDDEN_SIZE, NUM_HEADS, BLOCK_T",
+    [
+        (4, 16, 16, 4, 4),
+    ],
+)
+def test_fused_MHA_systolic(
+    BATCH_SIZE,
+    CONTEXT_LENGTH,
+    HIDDEN_SIZE,
+    NUM_HEADS,
+    BLOCK_T,
+):
+    run_test_with_params(
+        BATCH_SIZE=BATCH_SIZE,
+        CONTEXT_LENGTH=CONTEXT_LENGTH,
+        HIDDEN_SIZE=HIDDEN_SIZE,
+        NUM_HEADS=NUM_HEADS,
+        BLOCK_T=BLOCK_T,
+    )


Similar to test_flash_attention.py above. Please avoid introducing unnecessary functions.

RuizeYu05 · 2026-10-08T14:24:53Z

All changes applied now.

Ruize Yu added 2 commits September 29, 2026 21:15

merge two pr into one

328e0e9

syntx error

4ef0826

RuizeYu05 mentioned this pull request Sep 30, 2026

[Example] add flashattention in example and modify ci test #564

Closed

4 tasks done

RuizeYu05 changed the title ~~Add two Flashattention accelerator example~~ [Example] Add two Flashattention accelerator example Sep 30, 2026

fix file

135dad2

Fangtangtang requested a balanced review from Copilot September 30, 2026 13:18

Copilot started reviewing on behalf of Fangtangtang September 30, 2026 13:19 View session

Copilot AI reviewed Sep 30, 2026

View reviewed changes

Fangtangtang requested changes Sep 30, 2026

View reviewed changes

Ruize Yu added 3 commits September 30, 2026 20:51

fix structural issues

b6be897

format change

b986327

format change

00088bd

Fangtangtang reviewed Oct 5, 2026

View reviewed changes

Ruize Yu added 2 commits October 7, 2026 21:07

format polishment

99af013

format change

3efc642

Fangtangtang reviewed Oct 8, 2026

View reviewed changes

change test, readme

98c5d83

Conversation

RuizeYu05 commented Sep 30, 2026 • edited by Fangtangtang Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Problems

Proposed Solutions

Examples

Checklist

Uh oh!

Fangtangtang commented Sep 30, 2026

Uh oh!

RuizeYu05 commented Sep 30, 2026 via email

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Fangtangtang left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

RuizeYu05 commented Oct 1, 2026

Uh oh!

Fangtangtang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Fangtangtang Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

RuizeYu05 commented Oct 8, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

RuizeYu05 commented Sep 30, 2026 •

edited by Fangtangtang

Loading

Fangtangtang left a comment •

edited

Loading

Fangtangtang Oct 8, 2026 •

edited

Loading