feat: extend QK-Clip to support MLA (MuonClip Algorithm 1) [skip-build] (#28)

Revert "fix: disable CUDA graphs in Newton-Schulz for cpu_offload compatibility" (#29)

fix: disable CUDA graphs in Newton-Schulz for cpu_offload compatibility

Replace cpu_offload constructor param with turn_on/turn_off API (#26)

Invalidate AdamW tensor caches on load_state_dict [skip-build]

draft commit for cpu_offload (#23)

Update fast path comment to reflect current behavior [skip-build]

Update comment to reflect use_local_synchronization behavior [skip-build]

Fix deadlock in construct_shard_mesh with PP + dp_replicate > 1

Muon optimizer: expert batching, parallel caching, A2A overlap [skip-build]

Optimize pipeline: batched update, zero-copy scatter, prelaunch gather [skip-build]

Cache AdamW placement grouping and tensor lists [skip-build]

Add torch.compile, CUDA graph, and compiled momentum [skip-build]

Add mhc_attn, mhc_ffn, lambda_proj to skip_keys

Remove verbose param_groups summary logging

Support multi-component expert_keys (e.g. "experts.w1")

Extract is_expert_param() helper to consolidate expert key matching

Include original (pre-normalize) FQN in is_muon logging

Add info-level logging for param group classification (Muon vs AdamW)

Use component-level matching for expert_keys to avoid shared_experts collision

Normalize parameter FQNs to handle torch.compile / checkpoint wrappers

Apply pre-commit formatting (yapf) [skip-build]

Add max_iter cap and non-finite checks to _optimal_quintic [skip-build]

Apply pre-commit formatting (yapf, isort) [skip-build]

Add comment explaining _coeffs_list and Polar Express vs former NS [skip-build]

Replace hardcoded NS coefficients with analytically optimal ones [skip-build]

Refactor pipeline to async generator pattern (#16)

Support mHC (#15)

Support param group with various placements (#13)

fix bug in fsdp

Update torch-ext/optimizer/muon.py

Update torch-ext/optimizer/muon.py

Update muon.py

fix assert in a2a gather scatter

delete state in split_func

change owner_params to owned_params

modify pre step (overlap step) can get from arsgs

add doc strings + init self rank on init_assign_params

license added for flash_muon

apply pre-commit hook

consider multi node

misc

use inpalce op in update_g

use COMM_DTYPE instead of hardcoded dtype

apply all2all scatter gather

feat(muon_clip) : add muon clip (#6)

feat(muon) : add tuned-abc-values & blfoat16 communication

feat: update muon to receive paramgroups, not model (#4)

fix(muon): add update_p stage and dealloc tensors properly

chore: add .gitignore