refactor: unify RMSNorm fusion with DualRMSNorm + master switch

[valarLip]

Summary

Unify RMSNorm + quantization fusion into DualRMSNorm class with ATOM_ENABLE_RMSNORM_QUANT_FUSION master switch.

Changes

• layernorm.py — add DualRMSNorm, move group-quant kernel wrappers
• deepseek_v2.py — migrate to RMSNorm/DualRMSNorm, remove 300+ lines
• envs.py — add ATOM_ENABLE_RMSNORM_QUANT_FUSION
• docs/ — update environment_variables.md + model_ops_guide.md

5 commits, 5 files. No CI/workflow changes.

Split from #334 (model code part only).

[Copilot]

Pull request overview

This PR consolidates DeepSeek-specific RMSNorm+quant fusion logic into shared atom/model_ops/layernorm.py abstractions (DualRMSNorm + shared group-quant dispatch) and introduces a repo-wide master env switch (ATOM_ENABLE_RMSNORM_QUANT_FUSION) to control all RMSNorm+quant fusion paths.

Changes:

• Add DualRMSNorm and move FP8/FP4 group-quant fused RMSNorm dispatch/wrappers into layernorm.py, extending RMSNorm quant routing.
• Refactor deepseek_v2.py to use DualRMSNorm / unified RMSNorm.forward() paths, removing large blocks of model-local fusion code.
• Add ATOM_ENABLE_RMSNORM_QUANT_FUSION and deprecate older per-model toggles with fallback behavior; update docs accordingly.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
atom/model_ops/layernorm.py	Introduces DualRMSNorm, adds group-quant fusion dispatcher, and refactors RMSNorm fused-quant routing.
atom/models/deepseek_v2.py	Migrates DeepSeek MLA QK-norm and input-layernorm fusion to DualRMSNorm/RMSNorm APIs and deletes model-local wrappers.
atom/utils/envs.py	Adds ATOM_ENABLE_RMSNORM_QUANT_FUSION master switch and makes older DS/Llama switches deprecated overrides.
docs/environment_variables.md	Documents the new master switch and DS deprecated overrides.
docs/model_ops_guide.md	Adds DualRMSNorm to operator docs and fusion table.

Comments suppressed due to low confidence (2)

docs/environment_variables.md:72

• The Llama RMSNorm fusion env var docs still show ATOM_LLAMA_ENABLE_AITER_TRITON_FUSED_RMSNORM_QUANT as non-deprecated with a fixed default of 1 (true), but atom/utils/envs.py now makes it a deprecated override that falls back to ATOM_ENABLE_RMSNORM_QUANT_FUSION when unset. Please update this section to reflect the new defaulting/deprecation behavior so users don’t get surprised when the master switch disables Llama fusion too.

| **ATOM_ENABLE_RMSNORM_QUANT_FUSION** | bool | 1 (true) | Master switch for all RMSNorm + quantization fusion paths (all models). |
| **ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION** | bool | (master switch) | *Deprecated.* Override for DeepSeek input layernorm fusion. Falls back to `ATOM_ENABLE_RMSNORM_QUANT_FUSION` when unset. |
| **ATOM_ENABLE_DS_QKNORM_QUANT_FUSION** | bool | (master switch) | *Deprecated.* Override for DeepSeek QK-norm fusion. Falls back to `ATOM_ENABLE_RMSNORM_QUANT_FUSION` when unset. |

### Qwen3-MoE style

| Variable | Type | Default | Description |
|----------|------|---------|-------------|
| **ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION** | bool | 0 (false) | If set to `1`, fuse QK norm, RoPE, and cache quantization into one kernel. **Enable this for Qwen3-MoE models for better performance.** |

### Llama-style

| Variable | Type | Default | Description |
|----------|------|---------|-------------|
| **ATOM_LLAMA_ENABLE_AITER_TRITON_FUSED_RMSNORM_QUANT** | bool | 1 (true) | If set to `1`, use Triton kernel to fuse RMSNorm with quantization. |
| **ATOM_LLAMA_ENABLE_AITER_TRITON_FUSED_SILU_MUL_QUANT** | bool | 1 (true) | If set to `1`, use Triton kernel to fuse SiLU and mul with quantization in MLP module. |

atom/model_ops/layernorm.py:448

• RMSNorm.forward() can return (x, residual) (plain/allreduce/pad paths) and ((x_quant, x_scale), residual) (fused quant + residual), but the return type annotation only allows Tensor | tuple[Tensor, Tensor]. Update the annotation to reflect the actual nested-tuple variants (or refactor to a consistent return type) so callers and type-checkers don't get misled.

    @mark_trace(prefix="rmsnorm", torch_compile=True)
    def forward(
        self,
        x: torch.Tensor,
        residual: torch.Tensor | None = None,
        x_scale: Optional[torch.Tensor] = None,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.