Skip to content

Conversation

@andrea-fasoli
Copy link
Collaborator

@andrea-fasoli andrea-fasoli commented Jun 25, 2025

Description of the change

Implementation of addons FP8 attention BMMs for AIU. Similarly to GPTQ and INT8 addons, the functions build on FMS and expand its functionalities.

Original implementation by @ani300

Components:

  • scaled BMM operation registered for torch.compile
  • custom scaled attention BMM operation registered with FMS

Related issues or PRs

n/a

How to verify the PR

Code review (including math) is required.
cc: @ani300 @nwang-ibm @chichun-charlie-liu

Was the PR tested

  • I have added >=1 unit test(s) for every new method I have added (if that coverage is difficult, please briefly explain the reason)
  • I have ensured all unit tests pass

Checklist for passing CI/CD:

  • All commits are signed showing "Signed-off-by: Name <email@domain.com>" with git commit -signoff or equivalent
  • PR title and commit messages adhere to Conventional Commits
  • Contribution is formatted with tox -e fix
  • Contribution passes linting with tox -e lint
  • Contribution passes spellcheck with tox -e spellcheck
  • Contribution passes all unit tests with tox -e unit

Note: CI/CD performs unit tests on multiple versions of Python from a fresh install. There may be differences with your local environment and the test environment.

Signed-off-by: Andrea Fasoli <andrea.fasoli@ibm.com>
@andrea-fasoli andrea-fasoli changed the title Addons for FP8 attention bmm in FMS feat: addons for FP8 attention bmm in FMS Jun 25, 2025
@github-actions github-actions bot added the feat label Jun 25, 2025
Signed-off-by: Andrea Fasoli <andrea.fasoli@ibm.com>
@andrea-fasoli andrea-fasoli changed the title feat: addons for FP8 attention bmm in FMS feat: addons for FP8 attention bmm and linear in FMS Jun 27, 2025
Signed-off-by: Andrea Fasoli <andrea.fasoli@ibm.com>
Signed-off-by: Andrea Fasoli <andrea.fasoli@ibm.com>
@andrea-fasoli
Copy link
Collaborator Author

Bringing in FP8 attention and linear modules into FMS-MO.

CI/CD only fails lint at this time. All errors current errors are import errors. They come in two kinds:

  • torchao import error -> need to add this to our installed packages
  • import of some fms functions -> I'm investigating this, same functions are loaded fine from other addons...

cc: @chichun-charlie-liu @BrandonGroth @tharapalanivel @ani300

@andrea-fasoli
Copy link
Collaborator Author

tox -e lint run locally does not report fms import errors, only torchao errors

Signed-off-by: Andrea Fasoli <andrea.fasoli@ibm.com>
@ani300 ani300 merged commit c931ad7 into foundation-model-stack:main Jul 3, 2025
10 of 11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants