Skip to content

Implement RMSNORM in cuTile and cuTeDSL#1280

Open
Charlesrizzly wants to merge 2 commits into
linkedin:mainfrom
Charlesrizzly:kenemuo/rmsnorm-pr1-core
Open

Implement RMSNORM in cuTile and cuTeDSL#1280
Charlesrizzly wants to merge 2 commits into
linkedin:mainfrom
Charlesrizzly:kenemuo/rmsnorm-pr1-core

Conversation

@Charlesrizzly

Copy link
Copy Markdown
Contributor

Summary

This PR adds the core cuTile/cuTeDSL RMSNorm backend work and related backend/runtime updates

Main changes:

  • Added cuTile RMSNorm implementation and export wiring.
  • Added cuTeDSL backend registration and RMSNorm wiring.
  • Standardized DTensor handling by introducing a shared DTensor helper in ops utils and applying it in RMSNorm,
  • ModulatedRMSNorm, and SwiGLU.
  • Updated cuTile backend routing coverage to include RMSNorm.
  • Includes the current benchmark data CSV update that was part of the original branch split.

Testing Done

Implemented/updated backend-related coverage in:
test_cutile_backend.py

Full GPU validation and broader test execution were not completed as part of preparing this stacked PR split.

Hardware Type: H100 and B200

  • Hardware Type:
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

@Charlesrizzly
Charlesrizzly marked this pull request as ready for review July 1, 2026 04:51
@vaibhavjindal

Copy link
Copy Markdown
Collaborator

@Charlesrizzly let's split it into 2 PRs - for cutile and cutedsl separately. Would be easier to track things that way.

@Charlesrizzly

Copy link
Copy Markdown
Contributor Author

hey yes I split it already the new one is here: #1285

@vaibhavjindal

vaibhavjindal commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator

hey yes I split it already the new one is here: #1285

@Charlesrizzly a few things to do here:

  • let's remove the cutile code from this PR and have only cutedsl logic in this one
  • Rebase on top of the latest changes in the main branch (see [CuteDSL] Add CuTe DSL cross-entropy kernel and CuteDSL integration scaffolding (B200) #1279)
  • Let's add some benchmarking results. We want to see how this kernel performs in terms of memory and speed when compared with existing triton kernels. Both on H100 and B200s.
  • Let's also add tests for the cutedsl implementation.

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