feat: add fused_MHA_systolic example

[RuizeYu05]

Description

This request's central contribution is a fused MHA systolic array. Our design is a pure dataflow description in Allo dataflow programming. We re-architect the algorithm-to-hardware mapping so that the entire FlashAttention loop fits a strictly unidirectional, feedback-free dataflow with one tc block per PE column. We additionally apply SageAttention( https://arxiv.org/abs/2410.02367), a quantization method that preserves precision under 8-bit quantization, to save resources.

Problems

Using online softmax we can fuse the self-attention calculation into one single pass. Is it possible to implement it in a pure dataflow architecture? Cascading two systolic arrays does not seem to be the optimal implementation.

Proposed Solutions

The kernel routes three concurrent streams through a block_T * block_T PE array. The query $Q$ and the FlashAttention state $(m, d, \mathbf{O})$ enter from the left boundary and propagate strictly rightward in lockstep, with each PE forwarding both to its right neighbour after one use. The key and value blocks $K$ and $V$ enter from the top and propagate downward, where column $j$ is statically bound to tc block j-1: column1 owns the first block_T keys, column2 the next block_T, and so on. Each compute PE iterates over its assigned $tc$ block, computing $\mathbf{q}\cdot\mathbf{k}$, applying the online softmax update against the incoming state, and accumulating $V$ into $\mathbf{O}$ before passing the updated state right and $K{,}V$ down. State traverses the array exactly once, so the channel graph is acyclic and no feedback FIFOs are required.

Examples

This is 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

[Copilot]

Pull request overview

Adds a new example implementation and testbench for a fused MHA systolic-array dataflow design (with int8 quantization and online softmax) in Allo.

Changes:

• Introduces get_systolic_top(...) dataflow region implementing the fused MHA systolic kernel.
• Adds a numerical-correctness testbench comparing simulator output against a NumPy golden reference and optionally running Vitis HLS synthesis.
• Adds example package scaffolding for examples/fused_MHA_systolic.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

File	Description
examples/fused_MHA_systolic/fused_MHA_systolic.py	New fused MHA systolic dataflow kernel generator (get_systolic_top).
examples/fused_MHA_systolic/test_systolic.py	New testbench for simulator correctness + optional HLS synthesis.
examples/fused_MHA_systolic/init.py	Package init with license header.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.