[Bug] Using q instead of normalized q in Megatron's DSA MLA indexer for GLM 5 models

Bug Description

In slime_plugins/models/glm5/glm5.py, the DSA lightning-indexer projects its query from the raw q_a_proj output (q_compressed), i.e. it skips the q_a_layernorm RMSNorm. The reference GLM-5 / DeepSeek-V3.2 implementations (HF transformers and vLLM) feed the indexer the q_a_layernorm-normalized q.

I've run an experiment to see the difference of doing this or not and couldn't see a significant difference in terms of abs diff between the logprobs of vLLM and megatron, but wanted to raise this still, since I believe this should actually change the top k and thus make the model a little bit off policy.

Locations in the code

I've let AI generate views of where this happens (or rather not) in slime vs HF reference to find it easily in the code:

Where slime skips the norm

q_layernorm is IdentityOp, so the q-a-layernorm is fused only into the main linear_q_up_proj; the indexer reads the pre-norm q_compressed:

# slime_plugins/models/glm5/glm5.py
q_compressed, _ = self.linear_q_down_proj(hidden_states)   # L537: raw q_a_proj output
...
q_compressed = self.q_layernorm(q_compressed)              # L551: q_layernorm = IdentityOp -> no-op
q, _ = self.linear_q_up_proj(q_compressed)                 # L552: main path normalizes internally (fused LN)
...
q_compressed = q_compressed.detach()                       # L616: still un-normalized
index_q, _ = self.wq_b(q_compressed)                       # L620: indexer query from RAW q

• https://github.com/THUDM/slime/blob/fa3c990af6f18efd3fd9922698bf4bf4048d1263/slime_plugins/models/glm5/glm5.py#L616-L620
• spec sets q_layernorm=IdentityOp: https://github.com/THUDM/slime/blob/fa3c990af6f18efd3fd9922698bf4bf4048d1263/slime_plugins/models/glm5/glm5.py#L764

Reference (HF transformers)

GlmMoeDsaAttention computes q_resid = q_a_layernorm(q_a_proj(hidden)) and passes that to the indexer, whose wq_b consumes it:

# transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py
q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states))   # L423
...
topk_indices = self.indexer(hidden_states, q_resid, ...)     # L446

# class GlmMoeDsaIndexer.forward:
#   q_resid: Query residual from `q_a_layernorm(q_a_proj(x))`  (L220, docstring)
q = self.wq_b(q_resid)                                       # L231

• https://github.com/huggingface/transformers/blob/8edd87b6331bf4f6fb290a6ee302637574ec2e7d/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py#L423
• https://github.com/huggingface/transformers/blob/8edd87b6331bf4f6fb290a6ee302637574ec2e7d/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py#L446
• https://github.com/huggingface/transformers/blob/8edd87b6331bf4f6fb290a6ee302637574ec2e7d/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py#L220-L231

Steps to Reproduce

Just any normal GLM 5 training.

Expected Behavior

We are very close to expected behavior, this is more of a slight imprecision.

Actual Behavior

No normalization, where HF normalizes.

Environment

none

Logs

Additional Context

No response

Pre-submission Checklist

• I have read the CONTRIBUTING.md and understand the collaboration scope.
• I have read the documentation and my issue is not addressed there.
• I have searched for existing issues and this is not a duplicate.
• I have provided a minimal, reproducible example.

Read slime_plugins/models/glm5/glm5.py around lines 537-620, especially q_layernorm, q_compressed, and the indexer call. Compare the data flow with the referenced Hugging Face transformers implementation around lines 423, 446, and 220-231; done when the indexer receives the normalized query and normal GLM 5 training can be checked for the expected behavior.