[GLM5.2 FP8/MXFP4] optimize GLM5.2 by zejunchen-zejun · Pull Request #1458 · ROCm/ATOM

zejunchen-zejun · 2026-07-03T09:47:06Z

indexer fusion
online quant support quant moe from 1x128+FP32 scale to 1x32+E8M0 scale
update the recipe for specifying the online quant and MTP3
enable ar+rmsnorm+quant fusion

Copilot

Pull request overview

Updates ATOM’s DeepSeek/GLM shared codepaths to enable GLM‑5.2 performance optimizations (fused indexer path + MXFP8 online quant behavior) and documents server launch recipes for GLM‑5.2 FP8/MXFP4.

Changes:

Add a GLM-specific env toggle (ATOM_ENABLE_GLM_FUSED_INDEXER) and use it to allow GLM‑5.2 (glm_moe_dsa) to take the fused indexer kernel path.
Force scale_type=fp8_e8m0 for MXFP8 (per_1x32 + fp8) online weight quant to match expected scale buffer/kernel consumption.
Document GLM‑5.2 FP8/MXFP4 server and MTP server launch recipes (docker image + example commands).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
recipes/GLM-5.md	Adds GLM‑5.2 FP8/MXFP4 server launch recipes and environment hints.
atom/utils/envs.py	Introduces `ATOM_ENABLE_GLM_FUSED_INDEXER` env toggle for GLM fused indexer enablement.
atom/quantization/quark/utils.py	Adjusts MXFP8 online weight quant to force e8m0 scale output.
atom/models/deepseek_v2.py	Allows `glm_moe_dsa` configs to use fused indexer kernel path when env-gated.

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

Copilot

Pull request overview

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

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

GLM-5.2 (glm_moe_dsa) was hard-excluded from the fused indexer path added for DeepSeek-V3.2 in #788, forcing the indexer's rope + fp8-quant + kv-cache write to run as 5-6 separate ops per layer per token. GLM's indexer is structurally identical to V3.2 (index_head_dim=128, qk_rope_head_dim=64, per_1x128 fp8 quant, always-neox indexer rope), so the fused indexer_qk_rope_quant_and_cache kernel is math-equivalent to the per-op path. Notably the recent GLM acc fix (da360a2) fixed a bug that only exists in the per-op path (dropping q_nope on recombine); the fused kernel never had it. - Allow glm_moe_dsa through _supports_fused_indexer_kernel_config, gated by a new env ATOM_ENABLE_GLM_FUSED_INDEXER (default on) for easy A/B + rollback. - Keep the wk+weights_proj GEMM merge OFF for GLM (decoupled in _can_fuse_indexer_wk_weights_proj): that fusion requires merging two checkpoint tensors and GLM names weights_proj "indexers_proj" without a pre-mergeable layout. The dominant win (rope+quant+cache) needs no merge. DeepSeek-V3.2 behavior is unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Remove the guard that kept the indexer wk+weights_proj GEMM-merge off for glm_moe_dsa. The guard was overly conservative: it assumed GLM checkpoints ship the projection under a non-standard "indexers_proj" name that couldn't feed the packed_modules_mapping merge. In fact "indexers_proj" only appears in the HF quant config's modules_to_not_convert (remapped via quant_exclude_name_mapping); the actual checkpoint tensors use the standard indexer.wk (fp8 block-scale) + indexer.weights_proj (bf16) paths — exactly the layout IndexerWkWeightsProjLinear's fp8-wk load and the merge expect. GLM now takes the same path as DeepSeek-V3.2: wk + weights_proj collapse into a single BF16 GEMM (IndexerWkWeightsProjLinear), on top of the already-enabled qk-rope+quant+cache fusion. Still gated by ATOM_ENABLE_GLM_FUSED_INDEXER for rollback. DeepSeek-V3.2 behavior unchanged. Verify on GPU: weights load without missing indexer.wk / indexer.weights_proj, accuracy unchanged, and the two small indexer GEMMs collapse to one in the decode breakdown. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

* enable ar+norm+quant fusion * fix fusion with online quant

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Copilot

Pull request overview

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

+    # GLM-5.2 (glm_moe_dsa): enable the fused indexer qk-rope + fp8-quant + kv-cache
+    # kernel (indexer_qk_rope_quant_and_cache), same path DeepSeek-V3.2 uses. GLM's
+    # indexer dims (index_head_dim=128, qk_rope_head_dim=64, per_1x128, neox rope) are
+    # identical to V3.2, so the fusion is math-equivalent to the unfused path. Set to
+    # "0" to fall back to the per-op Python path if a regression is suspected.


+    # GLM-5.2 (glm_moe_dsa) shares DeepSeek-V3.2's sparse-MLA indexer: same dims
+    # (index_head_dim=128, qk_rope_head_dim=64), same per_1x128 fp8 quant, and the
+    # indexer rope is always neox for both. The fused kernel path is therefore
+    # math-equivalent to the per-op path, so allow it here (gated by an env flag for
+    # easy rollback). This also enables the wk+weights_proj GEMM-merge for GLM — its


Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

@@ -2006,8 +2034,16 @@ def forward(
        hidden_states: torch.Tensor,


Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

quant draft model to fp8 Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

    "model_name": "GLM-5.2-MXFP4",
    "model_path": "amd/GLM-5.2-MXFP4",
-    "extraArgs": "--kv_cache_dtype fp8 -tp 4 --default-chat-template-kwargs '{\"enable_thinking\":false}' --online_quant_config '{\"global_quant_config\": \"ptpc_fp8\", \"exclude_layer\": [\"lm_head\", \"model.embed_tokens\", \"*.mlp.gate\", \"*expert*\"]}'",
+    "extraArgs": "--kv_cache_dtype fp8 -tp 4 --online_quant_config '{\"global_quant_config\": \"ptpc_fp8\", \"exclude_layer\": [\"lm_head\", \"model.embed_tokens\", \"*.mlp.gate\", \"*expert*\"]}'",
    "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nAITER_USE_FLYDSL_MOE_SORTING=1",


    "model_name": "GLM-5.2-MXFP4 MTP",
    "model_path": "amd/GLM-5.2-MXFP4",
-    "extraArgs": "--kv_cache_dtype fp8 -tp 4 --default-chat-template-kwargs '{\"enable_thinking\":false}' --online_quant_config '{\"global_quant_config\": \"ptpc_fp8\", \"exclude_layer\": [\"lm_head\", \"model.embed_tokens\", \"*.mlp.gate\", \"*expert*\"]}' --method mtp --num-speculative-tokens 3",
+    "extraArgs": "--kv_cache_dtype fp8 -tp 4 --online_quant_config '{\"global_quant_config\": \"ptpc_fp8\", \"exclude_layer\": [\"lm_head\", \"model.embed_tokens\", \"*.mlp.gate\", \"model.layers.[0-9].mlp.*expert*\", \"model.layers.[1-6][0-9].mlp.*expert*\", \"model.layers.7[0-7].mlp.*expert*\"]}' --method mtp --num-speculative-tokens 3",
    "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nAITER_USE_FLYDSL_MOE_SORTING=1",


    if getattr(config, "model_type", None) == "glm_moe_dsa":
-        return False
+        if not ENABLE_GLM_FUSED_INDEXER:
+            return False
    return (
        getattr(config, "index_head_dim", None) == 128


Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

 Tips on server configuration:
- Use `--kv_cache_dtype bf16` for the DSA sparse-attention path on CDNA4 (gfx950).
- `--gpu-memory-utilization 0.8` leaves headroom for the per-layer DSA index cache; higher values may OOM during KV-cache allocation.
+- Use the FP8, MXFP4, or MXFP4 MTP server recipes above for GLM-5.2.
+- Use `--kv_cache_dtype fp8` with the optimized GLM-5.2 server recipes unless you are intentionally comparing against the older bf16 KV-cache baseline.
 - No `--trust-remote-code` is needed — ATOM has built-in support for `GlmMoeDsaForCausalLM`.


Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

+    # GLM-5.2 (glm_moe_dsa): enable the fused indexer qk-rope + fp8-quant + kv-cache
+    # kernel (indexer_qk_rope_quant_and_cache), same path DeepSeek-V3.2 uses. GLM's
+    # indexer dims (index_head_dim=128, qk_rope_head_dim=64, per_1x128, neox rope) are
+    # identical to V3.2, so the fusion is math-equivalent to the unfused path. Set to
+    # "0" to fall back to the per-op Python path if a regression is suspected.


+            if self.use_fused_quant and self.quant_type.value in (
+                _QV_PER_1X128,
+                _QV_PER_TOKEN,
+            ):
+                # Combined AllReduce + RMSNorm + FP8 quant: the downstream GEMM


Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

    "model_name": "GLM-5.2-MXFP4 MTP",
    "model_path": "amd/GLM-5.2-MXFP4",
-    "extraArgs": "--kv_cache_dtype fp8 -tp 4 --default-chat-template-kwargs '{\"enable_thinking\":false}' --online_quant_config '{\"global_quant_config\": \"ptpc_fp8\", \"exclude_layer\": [\"lm_head\", \"model.embed_tokens\", \"*.mlp.gate\", \"*expert*\"]}' --method mtp --num-speculative-tokens 3",
+    "extraArgs": "-tp 4 --online_quant_config '{\"global_quant_config\": \"ptpc_fp8\", \"exclude_layer\": [\"lm_head\", \"model.embed_tokens\", \"*.mlp.gate\", \"model.layers.[0-9].mlp.*expert*\", \"model.layers.[1-6][0-9].mlp.*expert*\", \"model.layers.7[0-7].mlp.*expert*\"]}' --method mtp --num-speculative-tokens 3",


+python -m atom.entrypoints.openai_server \
+  --model "$model_path" \
+  --server-port 7777 \
+  --kv_cache_dtype fp8 \
+  --no-enable_prefix_caching \
+  -tp $TP 2>&1 | tee server.log &


zejunchen-zejun changed the title ~~[GLM5.2 FP8/MXFP4] optimization atom native GLM5.2~~ [GLM5.2 FP8/MXFP4] optimize atom native GLM5.2 Jul 3, 2026

zufayu requested a review from valarLip July 6, 2026 02:01

zejunchen-zejun changed the title ~~[GLM5.2 FP8/MXFP4] optimize atom native GLM5.2~~ [NOT READY] [GLM5.2 FP8/MXFP4] optimize atom native GLM5.2 Jul 6, 2026

zejunchen-zejun changed the title ~~[NOT READY] [GLM5.2 FP8/MXFP4] optimize atom native GLM5.2~~ [NOT READY][GLM5.2 FP8/MXFP4] optimize atom native GLM5.2 Jul 6, 2026

zejunchen-zejun marked this pull request as ready for review July 7, 2026 08:14

Copilot AI review requested due to automatic review settings July 7, 2026 08:14

Copilot started reviewing on behalf of zejunchen-zejun July 7, 2026 08:15 View session

Copilot AI reviewed Jul 7, 2026

View reviewed changes

Comment thread recipes/GLM-5.md Outdated

zejunchen-zejun changed the title ~~[NOT READY][GLM5.2 FP8/MXFP4] optimize atom native GLM5.2~~ [GLM5.2 FP8/MXFP4] optimize GLM5.2 Jul 7, 2026

Copilot AI review requested due to automatic review settings July 7, 2026 12:19

zejunchen-zejun force-pushed the zejun/opt_GLM5.2_0701 branch from f4c157f to 18eaab4 Compare July 7, 2026 12:19

Copilot started reviewing on behalf of zejunchen-zejun July 7, 2026 12:21 View session

Copilot AI reviewed Jul 7, 2026

View reviewed changes

Comment thread atom/model_ops/layernorm.py

Comment thread atom/model_ops/layernorm.py

Comment thread recipes/GLM-5.md Outdated

Copilot AI review requested due to automatic review settings July 8, 2026 02:42

Copilot started reviewing on behalf of zejunchen-zejun July 8, 2026 02:43 View session

zejunchen-zejun marked this pull request as draft July 8, 2026 02:43

Copilot AI reviewed Jul 8, 2026

View reviewed changes

Comment thread recipes/GLM-5.md Outdated

zejunchen-zejun and others added 12 commits July 8, 2026 10:49

fix GLM5.2 acc drop

e4b3f5f

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

retrieve the no-op for ROPE

27ad7b9

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

retrieve outplace return

03e3db2

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

add online quant for mxfp8

c963726

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

update recipe

2210eb2

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

add mtp recipe

6586895

Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>

add online quant for glm5.2 mxfp4 recipe

9872231

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

enable ar+norm+quant fusion (#1494)

e84e8d5

* enable ar+norm+quant fusion * fix fusion with online quant

add AITER_USE_FLYDSL_MOE_SORTING=1 into recipe

b80b01b

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

update recipe

5aacdc6

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

zejunchen-zejun force-pushed the zejun/opt_GLM5.2_0701 branch from 3fe3c11 to 5aacdc6 Compare July 8, 2026 02:52

Copilot started reviewing on behalf of zejunchen-zejun July 8, 2026 13:09 View session

Copilot AI reviewed Jul 8, 2026

View reviewed changes

add ATOM_ENABLE_RELAXED_MTP=1 for mtp case

c2e87f5

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Copilot AI review requested due to automatic review settings July 9, 2026 00:57

Copilot started reviewing on behalf of zejunchen-zejun July 9, 2026 00:58 View session

Copilot AI reviewed Jul 9, 2026

View reviewed changes

Comment thread atom/models/deepseek_v2.py

@@ -2006,8 +2034,16 @@ def forward(

hidden_states: torch.Tensor,

Merge branch 'main' into zejun/opt_GLM5.2_0701

94bdfa5

Copilot AI review requested due to automatic review settings July 9, 2026 03:02

Copilot started reviewing on behalf of zejunchen-zejun July 9, 2026 03:03 View session

Copilot AI reviewed Jul 9, 2026

View reviewed changes

Comment thread atom/models/deepseek_v2.py

update recipe and command for mtp case

2d2899d

quant draft model to fp8 Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Copilot AI review requested due to automatic review settings July 9, 2026 07:20

Copilot started reviewing on behalf of zejunchen-zejun July 9, 2026 07:22 View session

Copilot AI reviewed Jul 9, 2026

View reviewed changes

code refactor (#1544)

6908569

Copilot AI review requested due to automatic review settings July 9, 2026 15:13

Copilot started reviewing on behalf of zejunchen-zejun July 9, 2026 15:14 View session

Copilot AI reviewed Jul 9, 2026

View reviewed changes

Merge branch 'main' into zejun/opt_GLM5.2_0701

22df50a

Copilot AI review requested due to automatic review settings July 10, 2026 00:49

Copilot started reviewing on behalf of zejunchen-zejun July 10, 2026 00:50 View session

Copilot AI reviewed Jul 10, 2026

View reviewed changes

zejunchen-zejun added the ci:atom label Jul 10, 2026

update command

80b632f

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Copilot AI review requested due to automatic review settings July 10, 2026 13:48

Copilot started reviewing on behalf of zejunchen-zejun July 10, 2026 13:49 View session

Copilot AI reviewed Jul 10, 2026

View reviewed changes

valarLip approved these changes Jul 11, 2026

View reviewed changes

zejunchen-zejun merged commit 47d7d46 into main Jul 11, 2026
102 of 107 checks passed

zejunchen-zejun deleted the zejun/opt_GLM5.2_0701 branch July 11, 2026 09:44

		@@ -2006,8 +2034,16 @@ def forward(
		hidden_states: torch.Tensor,

Uh oh!

Conversation

zejunchen-zejun commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

zejunchen-zejun commented Jul 3, 2026 •

edited

Loading