Skip to content

Fix CUDA RMSNorm small-row dispatch#3792

Open
dhiltgen wants to merge 2 commits into
ml-explore:mainfrom
dhiltgen:cuda-rmsnorm-fix
Open

Fix CUDA RMSNorm small-row dispatch#3792
dhiltgen wants to merge 2 commits into
ml-explore:mainfrom
dhiltgen:cuda-rmsnorm-fix

Conversation

@dhiltgen

@dhiltgen dhiltgen commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Proposed changes

CUDA RMSNorm used a launch shape that packed multiple rows into one block even when each row required more than one reduction group. BlockBroadcastReduce is block-wide, so those rows mixed partial reductions and produced incorrect results for small dimensions such as fp32 D=256 and bf16/fp16 D=512.

Fix the affected dispatch case by launching one row per block when two reduction groups are needed. Add CUDA regression coverage for forward and VJP fast RMSNorm on the affected small multi-row shapes.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

CUDA RMSNorm used a launch shape that packed multiple rows into one block even when each row required more than one reduction group. BlockBroadcastReduce is block-wide, so those rows mixed partial reductions and produced incorrect results for small dimensions such as fp32 D=256 and bf16/fp16 D=512.

Fix the affected dispatch case by launching one row per block when two reduction groups are needed. Add CUDA regression coverage for forward and VJP fast RMSNorm on the affected small multi-row shapes.

@zcbenz zcbenz left a comment

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.

Nice catch, thanks!

@zcbenz zcbenz force-pushed the cuda-rmsnorm-fix branch from 5ed47ed to 98323e8 Compare July 3, 2026 06:31
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