Is your feature request related to a problem? Please describe.
Further consolidate current Muon implementation with recent optimizer refactor. Discussion: #2241 (comment)

Tag the @mcore-oncall
to get oncall's attention to this issue.

Describe the solution you'd like
A clear and concise description of what you want to happen.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context or screenshots about the feature request here.