Repository navigation
Conversation
…MSNorm Phase 1 of RMSNorm fusion unification: - Move _fuse_rmsnorm_fp4_quant, _fused_rms_fp8_group_quant, and their fake-tensor versions from deepseek_v2.py to model_ops/layernorm.py - Add public fuse_rmsnorm_group_quant() dispatcher - Extend RMSNorm with quant_config-based auto-routing: group-quant path selected automatically from quant_type + params_dtype - Add transpose_scale and shuffle constructor parameters - Add DualRMSNorm class for fused dual-norm (q_a + kv_a in MLA)
Phase 2 of RMSNorm fusion unification: - DecoderLayer.forward(): replace 44-line _fuse_rmsnorm_quant bypass with 7-line self.input_layernorm() call (RMSNorm handles routing internally) - DecoderLayer.__init__(): pass fused_quant + quant_config to RMSNorm - MLAAttention: add DualRMSNorm for non-triton-GEMM QK-norm fusion, replacing 15-line _fuse_rmsnorm_quant(x2=kv_c) with 3-line call - Triton GEMM QKV-projection fusion path unchanged (MLA-specific)
…dead code Phase 3+4 of RMSNorm fusion unification: - envs.py: add ATOM_ENABLE_RMSNORM_QUANT_FUSION master switch (default ON) Old per-model vars fallback to master switch when not explicitly set - deepseek_v2.py: delete ~210 lines of private _fuse_rmsnorm_quant functions and fake-tensor versions (now in layernorm.py) - Remove unused self.fuse_rmsnorm_quant flag from DecoderLayer
…ights DualRMSNorm previously created weight1/weight2 parameters that didn't match checkpoint keys (qk_layernorm.weight1 vs q_a_layernorm.weight). This caused weights to stay at initial 1.0 values, producing garbage output (GSM8K 0.0). Fix: DualRMSNorm now takes existing RMSNorm modules as constructor arguments and uses their .weight and .eps directly. No duplicate parameters, checkpoint loading works correctly.
- Remove use_triton_gemm() guard for DualRMSNorm: FP4 models now use fused QK-norm + quant via DualRMSNorm regardless of GEMM backend - Update comments/docs to match: clarify input_layernorm still requires triton GEMM while QK-norm DualRMSNorm does not - Add DualRMSNorm to model_ops_guide (source table, normalization section, fused kernel chains table) - Add ATOM_ENABLE_RMSNORM_QUANT_FUSION master switch to env var docs, mark old per-model vars as deprecated - Fix DualRMSNorm docstring terminology consistency Verified: DeepSeek BF16 GSM8K 0.957, DeepSeek MXFP4 GSM8K 0.948
Contributor
There was a problem hiding this comment.
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
DualRMSNormand move FP8/FP4 group-quant fused RMSNorm dispatch/wrappers intolayernorm.py, extendingRMSNormquant routing. - Refactor
deepseek_v2.pyto useDualRMSNorm/ unifiedRMSNorm.forward()paths, removing large blocks of model-local fusion code. - Add
ATOM_ENABLE_RMSNORM_QUANT_FUSIONand 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_QUANTas non-deprecated with a fixed default of1 (true), butatom/utils/envs.pynow makes it a deprecated override that falls back toATOM_ENABLE_RMSNORM_QUANT_FUSIONwhen 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 allowsTensor | 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.
Comment on lines
333
to
354
| `RMSNorm` supports multiple forward paths depending on configuration flags: | ||
|
|
||
| | Condition | Kernel / Path | Returns | | ||
| |---|---|---| | ||
| | `x_pad_to_multiple > 0`, no residual | `fused_rmsnorm_pad_` (Triton `fused_add_rmsnorm_pad`) | Padded output | | ||
| | `x_pad_to_multiple > 0`, with residual | `fused_add_rmsnorm_pad_` | (output, residual) | | ||
| | `fused_allreduce=True` and `tp_size > 1` | `tensor_model_parallel_fused_allreduce_rmsnorm` | (output, residual) | | ||
| | `fused_quant=True` and `x_scale` provided | `fused_rms_fp8_per_tensor_static_quant` | (FP8 output, scale) | | ||
| | `fused_quant=True` and `per_1x32` | `fused_rms_mxfp4_quant` | (MXFP4 output, scale) | | ||
| | Default, no residual | `rmsnorm2d_fwd` | Output | | ||
| | Default, with residual | `rmsnorm2d_fwd_with_add` | (output, residual) | | ||
|
|
||
| Constructor parameters: | ||
| ```python | ||
| RMSNorm( | ||
| dim: int, | ||
| eps: float = 1e-6, | ||
| x_pad_to_multiple: int = 0, | ||
| fused_allreduce: bool = False, | ||
| fused_quant: bool = False, | ||
| quant_config: Optional[QuantizationConfig] = None, | ||
| ) |
| from atom.model_ops.base_attention import Attention | ||
| from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding | ||
| from atom.model_ops.layernorm import LayerNorm, RMSNorm | ||
| from atom.model_ops.layernorm import DualRMSNorm, LayerNorm, RMSNorm # noqa: F401 |
Comment on lines
1213
to
+1221
| self.fuse_qknorm_quant = True | ||
| # DualRMSNorm: fused dual norm + quant for both FP8 and FP4 paths | ||
| self.qk_layernorm = DualRMSNorm( | ||
| self.q_a_layernorm, | ||
| self.kv_a_layernorm, | ||
| quant_config=quant_config, | ||
| transpose_scale=True, | ||
| shuffle=False, | ||
| ) |
| | `MLAAttention` | `attention_mla.py` | `mla_decode_fwd`, `mla_prefill_fwd`, `concat_and_cache_mla`, `fused_qk_rope_concat_and_cache_mla` | Multi-head latent attention | | ||
| | `FusedMoE` | `moe.py` | `aiter.fused_moe.fused_moe`, `asm_moe` | Mixture of experts | | ||
| | `RMSNorm` | `layernorm.py` | `rmsnorm2d_fwd`, `rmsnorm2d_fwd_with_add`, `fused_add_rmsnorm_pad` | RMS normalization | | ||
| | `DualRMSNorm` | `layernorm.py` | `fuse_rmsnorm_group_quant` | Fused dual RMSNorm + quant (MLA q/kv norms) | |
1 task
Contributor
|
Closing due to inactivity: no pushes, reviews, or comments for 88 days. This is housekeeping, not a judgement on the change — if you want to continue this work, comment here and we'll reopen, or open a fresh PR rebased on the latest main. |
qichu-yun
added a commit
that referenced
this pull request
Oct 29, 2026
The previous commit planned an MTP verify step's sparse per-token work over the real tokens and attributed the GPU memory fault to the padded tokens' empty work items reading the sparse-index buffer past what the step wrote. That was a guess, and it was wrong. What actually runs past an end: the MTP drafter's step 0 runs the target's token stream through its own sparse MLA layer reusing the target's verify metadata, and that stream is not padded (eagle_proposer: "Step 0 carries the target's stream and pads nothing"), so its q and o hold scheduled_tokens rows. The persistent MLA kernel reads q and writes o for every work item, zero-KV padding included, so a plan over the cudagraph-padded running_tokens read q and wrote o past their ends. Whether that faulted depended on whether the page after q happened to be mapped, so the fault was intermittent; when it did not fault, it silently overwrote whatever followed o. Measured on MI455 (gfx1250, FlyDSL PS1 decode), GLM-5.2 MTP K3, 39 requests padded to 40, LLM-Booster #287 / #342: - the four GPUs' fault addresses all lay 20-24 KB past the end of a 140-row q (35 x 4 tokens), with the next page unmapped; - an eager-call check saw q with 156 rows against a plan reaching row 158; - a model-free test with q/o sized to the real tokens and a sentinel tail shows 3 work items past q and 3 rows written past o at step 0, with the original aiter kernel and with ROCm/aiter#7089 alike, and none when planned over the real tokens. So the code change of the previous commit is the fix; this commit corrects its rationale, records the planned row count in the metadata (sparse_mtp_num_tokens), and has MLAAttention._forward_decode refuse a sparse MTP plan larger than its q -- host ints only, no sync, skipped under DCP, which rebuilds the plan per layer. Tests: tests/model_ops/test_mla_sparse_mtp_plan_rows.py (the refusal, and a plan within q reaching the kernel for both the drafter's unpadded and the target's padded rows); the builder test now also checks the recorded row count.
qichu-yun
added a commit
that referenced
this pull request
Nov 5, 2026
The previous commit planned an MTP verify step's sparse per-token work over the real tokens and attributed the GPU memory fault to the padded tokens' empty work items reading the sparse-index buffer past what the step wrote. That was a guess, and it was wrong. What actually runs past an end: the MTP drafter's step 0 runs the target's token stream through its own sparse MLA layer reusing the target's verify metadata, and that stream is not padded (eagle_proposer: "Step 0 carries the target's stream and pads nothing"), so its q and o hold scheduled_tokens rows. The persistent MLA kernel reads q and writes o for every work item, zero-KV padding included, so a plan over the cudagraph-padded running_tokens read q and wrote o past their ends. Whether that faulted depended on whether the page after q happened to be mapped, so the fault was intermittent; when it did not fault, it silently overwrote whatever followed o. Measured on MI455 (gfx1250, FlyDSL PS1 decode), GLM-5.2 MTP K3, 39 requests padded to 40, LLM-Booster #287 / #342: - the four GPUs' fault addresses all lay 20-24 KB past the end of a 140-row q (35 x 4 tokens), with the next page unmapped; - an eager-call check saw q with 156 rows against a plan reaching row 158; - a model-free test with q/o sized to the real tokens and a sentinel tail shows 3 work items past q and 3 rows written past o at step 0, with the original aiter kernel and with ROCm/aiter#7089 alike, and none when planned over the real tokens. So the code change of the previous commit is the fix; this commit corrects its rationale, records the planned row count in the metadata (sparse_mtp_num_tokens), and has MLAAttention._forward_decode refuse a sparse MTP plan larger than its q -- host ints only, no sync, skipped under DCP, which rebuilds the plan per layer. Tests: tests/model_ops/test_mla_sparse_mtp_plan_rows.py (the refusal, and a plan within q reaching the kernel for both the drafter's unpadded and the target's padded rows); the builder test now also checks the recorded row count.
valarLip
pushed a commit
that referenced
this pull request
Nov 9, 2026
#2715) * [MLA] Plan a padded sparse MTP verify's work over the real tokens only prepare_decode built the sparse per-token persistent metadata of an MTP verify step over running_tokens, so a cudagraph-padded batch planned its padded tail tokens too. Those have zero sparse KV, but get_mla_metadata_v1 still gives every row a work item, starting at sparse_kv_indptr[-1] -- the end of what the step wrote to the sparse-index buffer. With those items planned, GLM-5.2 MXFP4 with MTP K=3 on MI455 (gfx1250, FlyDSL PS1 MLA decode) memory-faults in the decode kernel some tens of steps into any padded batch -- presumably by reading indices past that end -- which ends every GSM8K full run near its tail. Plan over scheduled_tokens instead, as the cprr branch already does for its requests and as the capture path does. The persistent kernel is work_indptr-driven, so the padded rows are left unplanned. Adds a regression test: a padded sparse MTP verify plans 2*q tokens for 2 real requests padded to 4. * [MLA] Explain the padded sparse MTP plan fix: the drafter's step 0 rows The previous commit planned an MTP verify step's sparse per-token work over the real tokens and attributed the GPU memory fault to the padded tokens' empty work items reading the sparse-index buffer past what the step wrote. That was a guess, and it was wrong. What actually runs past an end: the MTP drafter's step 0 runs the target's token stream through its own sparse MLA layer reusing the target's verify metadata, and that stream is not padded (eagle_proposer: "Step 0 carries the target's stream and pads nothing"), so its q and o hold scheduled_tokens rows. The persistent MLA kernel reads q and writes o for every work item, zero-KV padding included, so a plan over the cudagraph-padded running_tokens read q and wrote o past their ends. Whether that faulted depended on whether the page after q happened to be mapped, so the fault was intermittent; when it did not fault, it silently overwrote whatever followed o. Measured on MI455 (gfx1250, FlyDSL PS1 decode), GLM-5.2 MTP K3, 39 requests padded to 40, LLM-Booster #287 / #342: - the four GPUs' fault addresses all lay 20-24 KB past the end of a 140-row q (35 x 4 tokens), with the next page unmapped; - an eager-call check saw q with 156 rows against a plan reaching row 158; - a model-free test with q/o sized to the real tokens and a sentinel tail shows 3 work items past q and 3 rows written past o at step 0, with the original aiter kernel and with ROCm/aiter#7089 alike, and none when planned over the real tokens. So the code change of the previous commit is the fix; this commit corrects its rationale, records the planned row count in the metadata (sparse_mtp_num_tokens), and has MLAAttention._forward_decode refuse a sparse MTP plan larger than its q -- host ints only, no sync, skipped under DCP, which rebuilds the plan per layer. Tests: tests/model_ops/test_mla_sparse_mtp_plan_rows.py (the refusal, and a plan within q reaching the kernel for both the drafter's unpadded and the target's padded rows); the builder test now also checks the recorded row count. * [MLA] Cut the sparse MTP verify inputs to q's rows for the DCP rebuild Under DCP the sparse_mtp_ plan is rebuilt per layer from paged_cu_seqlens_q, paged_kv_indptr and paged_kv_last_page_lens. paged_kv_indptr was already cut to q's rows (B + 1), but the other two kept the padded running_tokens width, so the drafter's unpadded step 0 could still be planned past q and o. Cut them to B + 1 / B as the non-MTP branch does; the target's padded verify, where B == running_tokens, is unchanged. * [MLA] Scope the sparse_mtp_num_tokens comment to where the check runs The assertion in MLAAttention._forward_decode is skipped under DCP and, being a host check, runs on eager steps and at capture but not on graph replay.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Unify RMSNorm + quantization fusion into
DualRMSNormclass withATOM_ENABLE_RMSNORM_QUANT_FUSIONmaster switch.Changes
layernorm.py— add DualRMSNorm, move group-quant kernel wrappersdeepseek_v2.py— migrate to RMSNorm/DualRMSNorm, remove 300+ linesenvs.py— add ATOM_ENABLE_RMSNORM_QUANT_FUSIONdocs/— update environment_variables.md + model_ops_guide.md5 commits, 5 files. No CI/workflow changes.
Split from #334 (model code part only).