[FlyDSL] flydsl implementation of a16w16 gemm by omuhamma · Pull Request #2725 · ROCm/aiter

omuhamma · 2026-04-13T21:10:32Z

Motivation

To write the gemm for gfx1250 in flydsl

FlyDsl Pr: ROCm/FlyDSL#401

Test Plan

Unit tests & performance taken on HW

e2e testing on HW

Benchmarked with ISL/OSL 1k, 64 CONC, 640 num prompts

Test Result

All tests pass

server and benchmark work, accuracy remains the same (90%) before and after the change as well (with lm_eval)

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

github-actions · 2026-04-13T21:11:29Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-355`	Run Triton tests on MI355 in addition to MI325
`ci:sglang`	SGLang integration tests
`ci:atom`	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
`ci:vllm`	vLLM benchmark
`ci:all`	All of the above

Add labels via the sidebar or gh pr edit 2725 --add-label <label>

coderfeli · 2026-04-14T03:46:36Z

@omuhamma is this validated on silicon? could you plz also create a pr to flydsl if ready? Also, @yadaish is also trying naive gemm here.

omuhamma · 2026-04-14T16:27:53Z

@coderfeli Its only been tested on AM workflows as of now, I will further optimize the code and test it on hardware soon.

Added FlyDSL PR to the description

Co-authored-by: Cursor <cursoragent@cursor.com>

valarLip

all changes should under "aiter/ops/flydsl/"

Updated

* flydsl implementation of a16w16 gemm, WIP * Code styling and unit test fixes * UT update * code updates, best runs * Kernel updates,split k for smaller configs and unroll variable * Update shuffle.py * Added compute bound variant + black/ruff * black/ruff * black/ruff * new api * Cleanup pt2 * Reverting UT changes that existed for AM * Revert prev commit * small refactor * Splitting up kernel / wrapper * black/ruff * refactor with claude skill * black/ruff * kernel updates based off comments * black/ruff * Update gemm_a16w16_gfx1250.py * oob edge cases * Revert "Update gemm_a16w16_gfx1250.py" This reverts commit da29bb5. * e2e testing and fixes * Wire flydsl_gfx1250 a16w16 dispatch into tuned_gemm solMap The tuned CSV rows added for gfx1250 (libtype=flydsl_gfx1250) had no matching solMap entry, so those bf16 shapes would KeyError at dispatch. Add the kernelName parser (tile/warp/buffer/split_k/unroll encoding) and the flydsl_gfx1250_gemm dispatch wrapper, and register it in solMap. Addresses PR review feedback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Update aiter/tuned_gemm.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Address review nits: fx-native LDS pointer, single typed casts - _frag transposed load: replace raw llvm_dialect.inttoptr + ir.Type.parse("!llvm.ptr<3>") with fx.PointerType/fx.inttoptr/ fx.to_llvm_ptr (matches chunk_gated_delta_h); drops the flydsl._mlir imports entirely. - Collapse redundant fx wraps in _f32/_imm64/_slot_off to single typed casts. No functional change: 172/172 op tests pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Support FlyDSL >=0.3.1: version-gate raw buffer op aux argument FlyDSL 0.3.1 changed raw_ptr_buffer_store/atomic_fadd's trailing aux from a required SSA operand (i32 Value) to an optional MLIR attribute; the old Value form raises on store and is silently dropped on atomic_fadd (losing the split-k device-scope bits). Gate the four call sites on the installed version at compile time: <=0.3.0 keeps the operand form, >=0.3.1 passes aux=SCOPE_DEV as an int and omits the zero aux on stores. Tested: 172/172 op tests on both flydsl 0.3.0 and 0.3.1. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * resolving comments * removing bad copies, and unncessary padding * config picker, refactor, kernel change to prevent copies with split k * Fix ruff I001 and drop deleted is_flydsl_available import Co-authored-by: Cursor <cursoragent@cursor.com> * Addressing comments * CI fix --------- Co-authored-by: Satya Nikhil <nikhil.kodukula@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Felix Li <felix.li@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Xiake Sun <xiake.sun@amd.com>

* flydsl implementation of a16w16 gemm, WIP * Code styling and unit test fixes * UT update * code updates, best runs * Kernel updates,split k for smaller configs and unroll variable * Update shuffle.py * Added compute bound variant + black/ruff * black/ruff * black/ruff * new api * Cleanup pt2 * Reverting UT changes that existed for AM * Revert prev commit * small refactor * Splitting up kernel / wrapper * black/ruff * refactor with claude skill * black/ruff * kernel updates based off comments * black/ruff * Update gemm_a16w16_gfx1250.py * oob edge cases * Revert "Update gemm_a16w16_gfx1250.py" This reverts commit da29bb5. * e2e testing and fixes * Wire flydsl_gfx1250 a16w16 dispatch into tuned_gemm solMap The tuned CSV rows added for gfx1250 (libtype=flydsl_gfx1250) had no matching solMap entry, so those bf16 shapes would KeyError at dispatch. Add the kernelName parser (tile/warp/buffer/split_k/unroll encoding) and the flydsl_gfx1250_gemm dispatch wrapper, and register it in solMap. Addresses PR review feedback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Update aiter/tuned_gemm.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Address review nits: fx-native LDS pointer, single typed casts - _frag transposed load: replace raw llvm_dialect.inttoptr + ir.Type.parse("!llvm.ptr<3>") with fx.PointerType/fx.inttoptr/ fx.to_llvm_ptr (matches chunk_gated_delta_h); drops the flydsl._mlir imports entirely. - Collapse redundant fx wraps in _f32/_imm64/_slot_off to single typed casts. No functional change: 172/172 op tests pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Support FlyDSL >=0.3.1: version-gate raw buffer op aux argument FlyDSL 0.3.1 changed raw_ptr_buffer_store/atomic_fadd's trailing aux from a required SSA operand (i32 Value) to an optional MLIR attribute; the old Value form raises on store and is silently dropped on atomic_fadd (losing the split-k device-scope bits). Gate the four call sites on the installed version at compile time: <=0.3.0 keeps the operand form, >=0.3.1 passes aux=SCOPE_DEV as an int and omits the zero aux on stores. Tested: 172/172 op tests on both flydsl 0.3.0 and 0.3.1. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * resolving comments * removing bad copies, and unncessary padding * config picker, refactor, kernel change to prevent copies with split k * Fix ruff I001 and drop deleted is_flydsl_available import Co-authored-by: Cursor <cursoragent@cursor.com> * Addressing comments * CI fix --------- Co-authored-by: Satya Nikhil <nikhil.kodukula@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Felix Li <felix.li@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

flydsl implementation of a16w16 gemm, WIP

727dd89

omuhamma self-assigned this Apr 13, 2026

omuhamma mentioned this pull request Apr 14, 2026

gemm a16w16 flydsl implementation (WIP) ROCm/FlyDSL#401

Closed

1 task

apicciau mentioned this pull request Apr 20, 2026

Add FlyDSL GemmTuner integration #2816

Closed

omuhamma added 4 commits April 20, 2026 19:00

Code styling and unit test fixes

aa50b0e

UT update

eb82c24

Merge branch 'main' into omuhamma/fly_a16w16

d8ebec7

Merge branch 'main' into omuhamma/fly_a16w16

0166ffe

zufayu requested review from yzhou103 and removed request for yzhou103 June 26, 2026 08:57

code updates, best runs

6e04e37

zufayu removed the request for review from yzhou103 July 13, 2026 02:23

omuhamma added 3 commits July 16, 2026 20:35

Kernel updates,split k for smaller configs and unroll variable

a356320

Update shuffle.py

c61cee8

Merge branch 'main' into omuhamma/fly_a16w16

9095aeb

omuhamma marked this pull request as ready for review July 20, 2026 15:18

omuhamma requested a review from a team July 20, 2026 15:18

Added compute bound variant + black/ruff

3ae9e68

omuhamma mentioned this pull request Jul 20, 2026

a16w16 for gfx1250 on flydsl ROCm/FlyDSL#875

Open

zufayu requested a review from coderfeli July 21, 2026 01:40

omuhamma mentioned this pull request Jul 24, 2026

Enabling coexec llvm for Flydsl ROCm/FlyDSL#900

Closed

omuhamma added 5 commits July 24, 2026 13:08

Merge branch 'main' into omuhamma/fly_a16w16

367e648

black/ruff

56f7c17

black/ruff

3244365

Merge branch 'main' into omuhamma/fly_a16w16

6560453

Merge branch 'main' into omuhamma/fly_a16w16

b453d37

omuhamma added 2 commits August 31, 2026 18:22

removing bad copies, and unncessary padding

a94e851

Merge branch 'main' into omuhamma/fly_a16w16

ceed330

coderfeli previously approved these changes Sep 4, 2026

View reviewed changes

valarLip requested changes Sep 4, 2026

View reviewed changes

Comment thread aiter/tuned_gemm.py Resolved

Comment thread aiter/tuned_gemm.py Resolved

coderfeli reviewed Sep 4, 2026

View reviewed changes

Comment thread op_tests/flydsl_tests/test_gemm_a16w16.py Resolved

config picker, refactor, kernel change to prevent copies with split k

e659da7

omuhamma dismissed coderfeli’s stale review via e659da7 September 8, 2026 22:24

coderfeli and others added 2 commits September 9, 2026 05:56

Merge branch 'main' into omuhamma/fly_a16w16

3fcfd9b

Fix ruff I001 and drop deleted is_flydsl_available import

c56625e

Co-authored-by: Cursor <cursoragent@cursor.com>

coderfeli previously approved these changes Sep 9, 2026

View reviewed changes

Merge branch 'main' into omuhamma/fly_a16w16

44296e3

valarLip previously requested changes Sep 9, 2026

View reviewed changes

Comment thread op_tests/flydsl_tests/test_gemm_a16w16.py Resolved

Addressing comments

8d828df

omuhamma dismissed coderfeli’s stale review via 8d828df September 10, 2026 01:16

omuhamma added 6 commits September 9, 2026 21:48

Merge branch 'main' into omuhamma/fly_a16w16

8256d65

Merge branch 'main' into omuhamma/fly_a16w16

8adc7ac

CI fix

2101f4b

Merge branch 'main' into omuhamma/fly_a16w16

f83cde9

Merge branch 'main' into omuhamma/fly_a16w16

f6f0dc9

Merge branch 'main' into omuhamma/fly_a16w16

eea9470

coderfeli approved these changes Sep 14, 2026

View reviewed changes

omuhamma merged commit acc17ad into main Sep 14, 2026
56 checks passed

omuhamma deleted the omuhamma/fly_a16w16 branch September 14, 2026 14:30

JiaoliangYu mentioned this pull request Oct 27, 2026

[HIP] fix(fused_moe): give non-persistent per_1x128 1-stage kernels the full-height scale under num_local_tokens #7010

Open

gbyu-amd mentioned this pull request Oct 27, 2026

[HIP] [MoE] MXFP4 inline-sort: accept num_local_tokens, skip ids outside [0, E) #7020

Merged

ZLkanyo009 mentioned this pull request Oct 28, 2026

[HIP] [Bugfix] Honor num_local_tokens on the MXFP4 inline-sort path #7083

Closed

4 tasks done

Navigation Menu

Repository navigation

[FlyDSL] flydsl implementation of a16w16 gemm - #2725

[FlyDSL] flydsl implementation of a16w16 gemm#2725
omuhamma merged 63 commits into
mainfrom
omuhamma/fly_a16w16

omuhamma commented Apr 13, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Apr 13, 2026

Uh oh!

coderfeli commented Apr 14, 2026

Uh oh!

omuhamma commented Apr 14, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

valarLip left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

omuhamma commented Apr 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Apr 13, 2026

🏷️ CI Guide

Uh oh!

coderfeli commented Apr 14, 2026

Uh oh!

omuhamma commented Apr 14, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

valarLip left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

omuhamma commented Apr 13, 2026 •

edited

Loading