Skip to content

[Mthreads] Optimized mm addmm bmm with FlagOSTune - #2643

Merged
huangyiqun merged 14 commits into
flagos-ai:masterfrom
Fattysand:feature/mm-addmm-bmm-flagtune-support
Apr 27, 2026
Merged

[Mthreads] Optimized mm addmm bmm with FlagOSTune#2643
huangyiqun merged 14 commits into
flagos-ai:masterfrom
Fattysand:feature/mm-addmm-bmm-flagtune-support

Conversation

@Fattysand

@Fattysand Fattysand commented Apr 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

PR Category

Operator

Type of Change

Performance Optimization

Description

  1. USE FlagOSTune to optimized mthreads mm/addmm.bmm performance
  2. Achieved 1.02x to 5.6x speedup for mm operator on MTT S5000;
  3. Achieved 1x to 2.82x speedup for addmm operator on MTT S5000;
  4. Achieved 1x to 1.19x speedup for bmm operator on MTT S5000;

Issue

Progress

  • Change is properly reviewed (1 reviewer required, 2 recommended).
  • Change is responded to an issue.
  • Change is fully covered by a UT.

Performance

tengqm
tengqm reviewed Apr 24, 2026
View reviewed changes
Comment thread src/flag_gems/runtime/backend/_mthreads/addmm_mthreads_expand.yaml
tengqm
tengqm reviewed Apr 24, 2026
View reviewed changes
Comment thread src/flag_gems/runtime/backend/_mthreads/ops/mm.py
@github-actions github-actions Bot added the tests label Apr 24, 2026
Fattysand added 10 commits April 27, 2026 10:05
@Fattysand
add flagtune support for mm and addmm operators on mthreads backend
b00c233 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
feat: add FlagTune auto-tuning support for mm and addmm operators
3f54d09 
- Add get_expand_config() to read YAML tuning configs
- Update @libtuner decorators with dynamic strategy and warmup/rep params
- Add SQMMA kernel pre-hooks using cached TMA descriptors
- Preserve master's fp32 support and transpose logic
- Add SQMMA_SWITCH control for mm (False) and addmm (True)

Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
add flagtune for bmm.py and bmm_mthreads_expand.yaml
d7e9f4e 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
modify mm_hopper_expand.yaml
a98db0e 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
modify SQMMA_ON switch for mthreads mm addmm bmm
cbf288f 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
fix mthreads ops and use configloader logic
79b2611 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
sync utils.py
37e129b 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
fix bug in fp32 for mm/addmm/bmm
3953ba7 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
modify test_blas_perf_parallel func: test_blas_benchmark
c6605d6 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
Modify as required
0dcac0b 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
Fattysand force-pushed the feature/mm-addmm-bmm-flagtune-support branch from 06f63f4 to 0dcac0b Compare April 27, 2026 02:12
@Fattysand
fix test_blas_perf_parallel
1939717 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand Fattysand changed the title [Feature & Bugfix] mthreads mm addmm bmm flagtune support [Mthreads] mm addmm bmm flagtune support Apr 27, 2026
@Fattysand Fattysand changed the title [Mthreads] mm addmm bmm flagtune support [Mthreads] Optimized mm addmm bmm with FlagOSTune Apr 27, 2026
@Fattysand
precommit test_blas_perf_parallel
1dc8970 
Signed-off-by: fattysand <44150064+Fattysand@users.noreply.github.com>
@Fattysand
Fattysand marked this pull request as ready for review April 27, 2026 04:37
tengqm
tengqm reviewed Apr 27, 2026
View reviewed changes
Comment thread benchmark/test_blas_perf_parallel.py
@huangyiqun
huangyiqun merged commit 6fed5de into flagos-ai:master Apr 27, 2026
47 of 50 checks passed
@Vincent-Xiao Vincent-Xiao mentioned this pull request Apr 28, 2026
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tengqm tengqm tengqm left review comments
@Vincent-Xiao Vincent-Xiao Vincent-Xiao approved these changes
@huangyiqun huangyiqun huangyiqun approved these changes
@0x45f 0x45f Awaiting requested review from 0x45f 0x45f is a code owner
@douxetpur douxetpur Awaiting requested review from douxetpur douxetpur is a code owner
@bin913 bin913 Awaiting requested review from bin913 bin913 is a code owner
@w1120029931-bit w1120029931-bit Awaiting requested review from w1120029931-bit w1120029931-bit is a code owner

Assignees

No one assigned

Labels

benchmark core size/XLarge tests vendor/MooreThreads

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@Fattysand @tengqm @Vincent-Xiao @huangyiqun