support absorbed_mla & support glm_moe_dsa CP/packing#144
Conversation
Jintao-Huang
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the AbsorbedMLASelfAttention module to support multi-latent attention with matrix absorption, registering it within the model's replacement logic. Feedback highlights several critical robustness issues: inheriting directly from McoreAbsorbedMLASelfAttention when it is None will raise a TypeError at import time, which can be resolved with a dummy fallback class. Additionally, accessing __name__ on self_attention.module and accessing self.pg_collection.cp are prone to AttributeErrors if those objects are None or of unexpected types, and should be safely guarded.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| elif self_attention.module.__name__ == 'AbsorbedMLASelfAttention': | ||
| self_attention.module = AbsorbedMLASelfAttention |
There was a problem hiding this comment.
Accessing __name__ directly on self_attention.module can raise an AttributeError if the module is None or a partial object (which does not have a __name__ attribute). Use getattr to safely retrieve the attribute and avoid potential crashes during model registration.
| elif self_attention.module.__name__ == 'AbsorbedMLASelfAttention': | |
| self_attention.module = AbsorbedMLASelfAttention | |
| elif getattr(self_attention.module, '__name__', None) == 'AbsorbedMLASelfAttention': | |
| self_attention.module = AbsorbedMLASelfAttention |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for AbsorbedMLASelfAttention and refactors the GLM MoE DSA implementation to align with Megatron-Core, including updates to configuration parsing and the DSA indexer. The review feedback highlights several critical opportunities to improve robustness and prevent runtime crashes, particularly by defensively handling cases where packed_seq_params is empty or passed as a tuple, ensuring sequence packing is active before reshaping tensors, and safely accessing module attributes using getattr.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
3 participants
No description provided.