Repository navigation
Conversation
…un.py - Create .github/benchmark/models.json as single source of truth for model configs, shared between benchmark matrix and regression re-run - Create .github/scripts/regression_rerun.py to parse regression report and generate action file (DOWNLOAD/LAUNCH/BENCHMARK/STOP), matching by model_id + variant suffix instead of fragile display name / org guessing - Add model_id field to regression report entries in summarize.py - Add NUM_PROMPTS_OVERRIDE and RESULT_FILENAME fallback to atom_test.sh - Add load-models job to read models.json for benchmark matrix via fromJson - Replace ~130 lines of inline shell+python in regression-analysis with ~40 lines that delegate to regression_rerun.py + atom_test.sh
…heck Boolean inputs.xxx is falsy when false, so `false || 'true'` always evaluates to 'true', making it impossible to disable a model via workflow_dispatch. Input defaults already provide the true fallback.
- atom_test.sh: add `stop` subcommand, add `--mark-trace` to launch when ENABLE_TORCH_PROFILER=1 (profiler logic fully in atom_test.sh) - regression_rerun.py: output flat config list (model|args|bench_args| isl|osl|conc) instead of DOWNLOAD/LAUNCH/STOP lifecycle actions; sorted by model+args for server grouping - YAML: split monolithic step into 4 independent steps: Prepare configs → Download models → Run with profiler → Wait traces "Run with profiler" uses same atom_test.sh launch/benchmark/stop pattern as the benchmark job
…analysis - regression-rerun (GPU): only on regression, re-runs configs with profiler - profiler-analysis (ubuntu-latest): runs on regression OR manual enable_profiler - Regression mode: downloads traces from regression-rerun - Manual mode: downloads traces from benchmark jobs - Issue creation only for regression mode This allows manual profiler data collection via workflow_dispatch with enable_profiler=true, without requiring an actual regression.
Stale generated code in ~/.cache/atom/ may reference record_function without importing it, causing NameError when --mark-trace is enabled. Clear cache on every launch to ensure fresh compilation with correct instrumentation.
…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.
Contributor
There was a problem hiding this comment.
Pull request overview
Refactors RMSNorm+quantization fusion support to be reusable across models (rather than DeepSeek-local), and modernizes the benchmark CI workflow by centralizing model configuration and improving regression re-run + profiler trace handling.
Changes:
- Add a master env toggle for RMSNorm+quant fusion and wire deprecated DeepSeek/Llama toggles to fall back to it.
- Move DeepSeek’s fused RMSNorm+group-quant helpers into
atom/model_ops/layernorm.py, addDualRMSNorm, and route DeepSeek paths through the shared APIs. - Centralize benchmark model definitions in
.github/benchmark/models.json, and update the benchmark workflow to load configs from it and to generate regression re-run configs via a new script.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
atom/utils/envs.py |
Introduces a global RMSNorm+quant fusion switch and deprecates older per-model toggles via fallback defaults. |
atom/models/deepseek_v2.py |
Switches DeepSeek fused QK-norm+quant and input-norm+quant paths to shared RMSNorm/DualRMSNorm implementations. |
atom/model_ops/layernorm.py |
Adds shared fused RMSNorm+group-quant dispatch and a DualRMSNorm helper; extends RMSNorm to support additional fused-quant modes. |
.github/workflows/atom-benchmark.yaml |
Loads benchmark models from JSON, streamlines regression re-run/profiler artifact flow, and adds a dedicated profiler analysis job. |
.github/scripts/summarize.py |
Includes model_id in regression entries to support stable model resolution. |
.github/scripts/regression_rerun.py |
New helper to map regression report entries back to models.json and emit a flat rerun config list. |
.github/scripts/atom_test.sh |
Adds cache cleanup on launch, supports stop, and improves benchmark result naming / prompt count overrides. |
.github/benchmark/models.json |
New single source of truth for benchmark model registry used by CI. |
Comments suppressed due to low confidence (1)
atom/model_ops/layernorm.py:448
RMSNorm.forward()’s return type annotation doesn’t match what the function can actually return whenresidualis provided and fused quantization is active (e.g., it can return((x_quant, x_scale), residual_out)). This makes the API contract misleading for callers/type-checkers; please update the annotation (and/or docstring) to reflect the nested-tuple shape, or normalize the return shape across branches.
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.
- 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.
Copilot wasn't able to review any files in this pull request.
💡 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.
This was referenced Mar 16, 2026
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
RMSNormclass withquant_config-driven auto-routing. AddDualRMSNormfor MLA dual-norm fusion. Delete ~200 lines of dead code from deepseek_v2.py.models.json, addregression_rerun.pyfor profiler re-run, splitregression-analysisinto GPU + lightweight jobs, fix boolean input bug and compile cache NameError.ATOM_ENABLE_RMSNORM_QUANT_FUSIONenv var to control all RMSNorm fusion paths.RMSNorm changes
atom/model_ops/layernorm.pyRMSNormwith auto-routing, addDualRMSNormatom/models/deepseek_v2.pyatom/utils/envs.pyATOM_ENABLE_RMSNORM_QUANT_FUSIONmaster switchCI changes
.github/benchmark/models.json.github/scripts/regression_rerun.py.github/scripts/atom_test.shstopcommand, compile cache clear, profiler args.github/workflows/atom-benchmark.yamlVerification
Test plan
ATOM_ENABLE_RMSNORM_QUANT_FUSION=0disables all fusion without errorsSupersedes #330.
🤖 Generated with Claude Code