Skip to content

Optimize rmsnorm/layernorm to get better performance than aiter/triton#610

Open
cschenjunlin wants to merge 11 commits into
mainfrom
cjl/norm_optimization
Open

Optimize rmsnorm/layernorm to get better performance than aiter/triton#610
cschenjunlin wants to merge 11 commits into
mainfrom
cjl/norm_optimization

Conversation

@cschenjunlin

Copy link
Copy Markdown
Contributor

Motivation

Optimize rmsnorm/layernorm to get better performance than aiter/triton

Technical Details

Test Plan

Test Result

Submission Checklist

@cschenjunlin cschenjunlin force-pushed the cjl/norm_optimization branch from 8d34dfc to 2b28750 Compare June 4, 2026 09:27
@coderfeli

Copy link
Copy Markdown
Collaborator

CI failed @cschenjunlin

@coderfeli

coderfeli commented Jun 17, 2026

Copy link
Copy Markdown
Collaborator

@cschenjunlin Test coverage: both _get_*_configs() now default to only (32768, 8192, "bf16"), with f32/f16, small-N, and unaligned-tail cases commented out. The generic/scalar path (which this PR also modified, e.g. xscale dtype) and f32/f16 are no longer exercised in CI. Suggest keeping at least one unaligned-N and one f32 case active.

Large default shape: 32768x8192 bf16 (~512MB/tensor, several for fused-add-quant) runs in the default non-large flow since these tests are only l2_device/rocm_lower, not large_shape (MI runs took 45-59 min). Consider a small fast-path shape for correctness + a large_shape-marked big shape for perf.

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.

2 participants