Skip to content

[Feat] Align quant and fused rmsnorm kernels with aiter/triton - #481

Merged
coderfeli merged 13 commits into
mainfrom
cjl/fused_quant_rmsnorm
May 16, 2026
Merged

coderfeli merged 13 commits into
mainfrom
cjl/fused_quant_rmsnorm

Conversation

@cschenjunlin

@cschenjunlin cschenjunlin commented May 8, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

Align quant and fused kernels with aiter/triton

Technical Details

Test Plan

Test Result

Tested on MI308+ROCm7.1:

quant rmsnorm performance compare:

====================================================================================================
Perf Compare (gpu us): FlyDSL vs AIter
====================================================================================================
op         shape              dtype  FlyDSL(gpu us)  AIter(gpu us)    speedup
rmsnorm_dq 64x256             f32              28.6           37.0      1.29x
rmsnorm_dq 128x1024           f32              28.1           37.5      1.33x
rmsnorm_dq 32x128             f16              29.5           36.9      1.25x
rmsnorm_dq 64x2000            f32              29.3           38.7      1.32x
rmsnorm_dq 16x512             bf16             29.1           37.1      1.27x
rmsnorm_dq 1024x8192          bf16             28.9           37.4      1.30x
rmsnorm_dq 32768x8192         bf16            400.8        1,089.9      2.72x
rmsnorm_sq 64x256             f32              30.9           38.8      1.25x
rmsnorm_sq 128x1024           f32              30.7           38.0      1.24x
rmsnorm_sq 32x128             f16              30.1           38.7      1.29x
rmsnorm_sq 64x2000            f32              30.3           38.4      1.27x
rmsnorm_sq 16x512             bf16             30.5           39.7      1.30x
rmsnorm_sq 1024x8192          bf16             30.6           46.6      1.52x
rmsnorm_sq 32768x8192         bf16            535.8        1,476.0      2.75x
====================================================================================================

fused_add rmsnorm performance compare:

====================================================================================================
Perf Compare (gpu us): FlyDSL vs AIter
====================================================================================================
op         shape              dtype  FlyDSL(gpu us)  AIter(gpu us)    speedup
rmsnorm_add 64x256             f32              31.5           52.1      1.65x
rmsnorm_add 128x1024           f32              31.5           51.5      1.63x
rmsnorm_add 32x128             f16              31.3           50.6      1.62x
rmsnorm_add 64x2000            f32              30.8           51.4      1.67x
rmsnorm_add 16x512             bf16             31.1           51.3      1.65x
rmsnorm_add 1024x8192          bf16             31.2           55.7      1.79x
rmsnorm_add 32768x8192         bf16            814.7        1,661.4      2.04x
====================================================================================================

fused add quant rmsnorm performance compare:

====================================================================================================
Perf Compare (gpu us): FlyDSL vs AIter
====================================================================================================
op         shape              dtype  FlyDSL(gpu us)  AIter(gpu us)    speedup
rmsnorm_add_dq 64x256             f32              32.8           38.8      1.18x
rmsnorm_add_dq 128x1024           f32              33.3           37.6      1.13x
rmsnorm_add_dq 32x128             f16              33.2           37.6      1.13x
rmsnorm_add_dq 64x2000            f32              33.5           39.3      1.17x
rmsnorm_add_dq 16x512             bf16             32.8           39.2      1.20x
rmsnorm_add_dq 1024x8192          bf16             33.3           54.6      1.64x
rmsnorm_add_dq 32768x8192         bf16            731.4        1,562.0      2.14x
rmsnorm_add_sq 64x256             f32              35.2           42.2      1.20x
rmsnorm_add_sq 128x1024           f32              35.1           41.9      1.19x
rmsnorm_add_sq 32x128             f16              35.5           39.8      1.12x
rmsnorm_add_sq 64x2000            f32              36.4           40.8      1.12x
rmsnorm_add_sq 16x512             bf16             35.1           43.1      1.23x
rmsnorm_add_sq 1024x8192          bf16             35.8           64.1      1.79x
rmsnorm_add_sq 32768x8192         bf16            888.1        1,847.8      2.08x
====================================================================================================

Submission Checklist

  • quant_rms_norm_kernel
  • fused_add_rmsnorm_kernel
  • quant_fused_add_rmsnorm_kernel
  • rmsnorm_kernel_large_m_small_n

Comment thread kernels/rmsnorm_kernel.py
y = (added * rrms) * g
_store_vec(_to_elem_vec(y), out_div, idx)

else:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no need to copy all the codes in else?

@coderfeli

Copy link
Copy Markdown
Collaborator

@cschenjunlin any update?

@cschenjunlin

Copy link
Copy Markdown
Contributor Author

@cschenjunlin any update?

I have pushed new commits to solve the duplication issue mentioned above.

Some common functions are added, to reduce code duplication within the variant kernels. However, the scope of refactoring here is limited, as the variants have introduced some new data flow logic.
For each variant kernel, a standalone build_xx_module function is retained. This is to align with the implementation of aiter/triton, where each variant kernel is also implemented as an standalone kernel separate from the base version.

If further reduction of code duplication is needed, I can attempt to converge the build_xx_module functions of all variants into a unified implementation, and use flag parameters to differentiate branch logic internally.

@i-chaochen

Copy link
Copy Markdown

@coderfeli Sorry for the delay, please check this PR again.

Thanks!

Comment thread kernels/rmsnorm_kernel.py Outdated
return allocator, red_offset, red2_offset


def _load_scalar(copy_atom, scalar_reg_ty, scalar_reg_lay, divided_tensor, index):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This now conflicts with main, which moved these register temporaries to fx.make_rmem_tensor/internal types. Please rebase and keep the shared helpers on the new API instead of reintroducing MemRefType + memref_alloca.

return ok, flydsl_gpu_us


def test_rmsnorm_fused_add_dynamicquant():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These new variant tests are pytest-only today. run_benchmark.sh executes this file as a script, but main only calls test_all(), so the fused/quant variants are not exercised in that benchmark path.

Resolve RMSNorm conflicts by keeping the branch variants on the current make_rmem_tensor-based register helper API.

Co-authored-by: Cursor <cursoragent@cursor.com>
@coderfeli
coderfeli merged commit ba2b1c4 into main May 16, 2026
12 checks passed
@coderfeli
coderfeli deleted the cjl/fused_quant_rmsnorm branch May 16, 2026 06:47
jhinpan added a commit to jhinpan/FlyDSL-lab that referenced this pull request Jun 15, 2026
The Python style pre-check (ruff F841) lints every file the PR touches.
Merging main into this branch pulled kernels/rmsnorm_kernel.py into the
changed-file set, surfacing a latent dead assignment from ROCm#481:

  F841 Local variable `arch` is assigned to but never used
     --> kernels/rmsnorm_kernel.py:631

`_build_rmsnorm_quant_module` called get_hip_arch() but never consumed the
result (unlike the three other call sites, which gate USE_HW_CVT_PK_BF16_F32
on it). Remove the dead assignment. This failure was blocking the GPU "Fly
DSL test" job, which is gated on the Checks workflow and skipped with exit
code 78 once style failed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jhinpan added a commit that referenced this pull request Aug 21, 2026
…ged an outcome

Dropping F5 with the other four was wrong. It is the single family review-pr
alone missed in the seeded test (4/5 -> 5/5), so it is the only one with direct
evidence that it adds something, and it was cut for tidiness rather than for a
reason.

It returns as a scanner rather than prose, which is what the evidence supports.
scan_unreachable_tests.py parses the test file, resolves what `__main__` reaches
transitively, and reports only tests the diff ADDS. Validated three ways:

  - replaying #481 it names exactly the three tests coderfeli named
    (test_rmsnorm_fused_add and its dynamicquant/smoothquant variants -- "the
    fused/quant variants are not exercised in that path")
  - on a seeded PR adding one pytest-only test, it flags that test and nothing else
  - on a comment-only control it reports nothing

Only added tests count. A first version reported absolute state and fired on 6 of
7 untouched tests in one file, because pytest-first files with a one-line
__main__ are the norm here; that version would have been noise.

F5 pairs with validate-kernel-pr from the other side: this catches a test that
will not run, the shape grid catches shapes a running test never reaches.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants