Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Transformer in FINN: Scaled Dot-Product Attention #13

Draft
wants to merge 94 commits into
base: dev
Choose a base branch
from

Conversation

iksnagreb
Copy link

@iksnagreb iksnagreb commented Jan 20, 2025

Adds support for multi-head scaled dot-product attention, i.e., the core operation of a Transformer, to FINN. This includes compiler integration of hardware operators for the attention mechanism and multi-head splitting/and merging as well as related graph transformations. Heavily depends on the related streamlining of scaled dot product attention: #12

  • Add attention-hlslib dependency to fetch-repos.sh, see https://github.com/iksnagreb/attention-hlslib
  • Figure out how to integrate the Brevitas modifications...
  • There are probably some undocumented fixes/modification lying around on some other branches...

To support a complete Transformer, the following PRs must be meged:

WIP: Merge branch for testing the integration of all the Transformer related PRs until they are fully merged into dev: https://github.com/eki-project/finn-plus/tree/transformer

iksnagreb added 30 commits April 3, 2024 15:21
Currently this is not a HLSCustomOp, but a QONNX CustomOp.

Implemented are first operator attributes, ONNX graph/model construction
and a rather improvised python mode node execution for debugging.
This causes the C++ simulation to fail as multithreshold activations are
not implemented on the HLS side yet.
Note: The threshold parameters are generated and included but not
connected to the attention operator yet. The attention operator uses
uninitialized thresholds of the same type and shape.
Note: Currently there is no method for optimizing the accumulator width
of both, the HLSCustomOp and the python simulation. Thus, to make the
tests pass, both must be specified manually to the maximum possible
accumulator bitwidth. Doing the MinimizeAccumulatorWidth transform would
cause the HLS and python operator behavior to diverge.
iksnagreb and others added 22 commits April 26, 2024 15:54
Note: This is currently not controlling the memory used by the internal
threshold operations and also not controlling the resoruce type used for
implementing the floating-point operations within the softmax. These are
all still handled by the tools' automatic strategy.
This is a temporary solution to get at least node-by-node RTL simulation
of models working by simply skipping the attention operator.
The inferred shape is not taken from the model graph but from the node
attributes specifying the shape.
Instead of manually squeezing all shapes, explicit Squeeze and Unsqueeze
operations are inserted into the graph before deleting and redoing all
shape annotations from scratch. This should be more robust and keeps the
interface (data layout) the model exposes to the outside.

Wraps Im2Col operations in Unsqueeze-Squeeze operators to shield it from
squeezing as Im2Col always operates on 4-dimensional layouts.
@iksnagreb iksnagreb self-assigned this Jan 28, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
Status: In Progress
Development

Successfully merging this pull request may close these issues.

3 participants