Repository navigation
[Feat] Align quant and fused rmsnorm kernels with aiter/triton - #481
Conversation
| y = (added * rrms) * g | ||
| _store_vec(_to_elem_vec(y), out_div, idx) | ||
|
|
||
| else: |
There was a problem hiding this comment.
no need to copy all the codes in else?
|
@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. If further reduction of code duplication is needed, I can attempt to converge the |
|
@coderfeli Sorry for the delay, please check this PR again. Thanks! |
| return allocator, red_offset, red2_offset | ||
|
|
||
|
|
||
| def _load_scalar(copy_atom, scalar_reg_ty, scalar_reg_lay, divided_tensor, index): |
There was a problem hiding this comment.
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(): |
There was a problem hiding this comment.
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>
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>
…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.
Motivation
Align quant and fused kernels with aiter/triton
Technical Details
Test Plan
Test Result
Tested on MI308+ROCm7.1:
quant rmsnorm performance compare:
fused_add rmsnorm performance compare:
fused add quant rmsnorm performance compare:
Submission Checklist