From cfc634b4a7baab626102c10fc66d44983a37aad5 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Tue, 23 Jun 2026 20:39:20 +0800 Subject: [PATCH 1/5] update --- src/mcore_bridge/bridge/gpt_bridge.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 7f7ea5d..1b90308 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -1611,8 +1611,11 @@ def _set_mla_attn_state( self._set_state_dict(mg_attn, 'linear_kv_up_proj.layer_norm_weight', hf_state_dict, 'kv_a_layernorm.weight', to_mcore) if self.config.experimental_attention_variant == 'dsa': - indexer = None if mg_attn is None else mg_attn.core_attention.indexer - hf_state_dict.update(self._set_indexer(indexer, hf_state_dict, 'indexer.', to_mcore)) + has_indexer = False if mg_attn is None else mg_attn.core_attention.indexer is not None + has_indexer = self._reduce_tensor_pp_group(has_indexer, to_mcore) + if has_indexer: + indexer = None if mg_attn is None else mg_attn.core_attention.indexer + hf_state_dict.update(self._set_indexer(indexer, hf_state_dict, 'indexer.', to_mcore)) if to_mcore: hf_state_dict = {} else: From b13eb6f9fe7ad385399b9e2f9c61b7a97da6ee2d Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Tue, 23 Jun 2026 20:44:15 +0800 Subject: [PATCH 2/5] fix glm5.2 indexer shared --- src/mcore_bridge/model/constant.py | 1 + src/mcore_bridge/model/gpts/__init__.py | 3 +- src/mcore_bridge/model/gpts/glm_moe_dsa.py | 180 +++++++++++++++++++++ src/mcore_bridge/model/gpts/llm.py | 2 +- 4 files changed, 184 insertions(+), 2 deletions(-) create mode 100644 src/mcore_bridge/model/gpts/glm_moe_dsa.py diff --git a/src/mcore_bridge/model/constant.py b/src/mcore_bridge/model/constant.py index 1c3ae97..02ce638 100644 --- a/src/mcore_bridge/model/constant.py +++ b/src/mcore_bridge/model/constant.py @@ -11,6 +11,7 @@ class LLMModelType: bailing_moe = 'bailing_moe' bailing_hybrid = 'bailing_hybrid' deepseek_v4 = 'deepseek_v4' + glm_moe_dsa = 'glm_moe_dsa' qwen3_emb = 'qwen3_emb' diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py index 52025b8..e79d01a 100644 --- a/src/mcore_bridge/model/gpts/__init__.py +++ b/src/mcore_bridge/model/gpts/__init__.py @@ -1,2 +1,3 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from . import bailing_hybrid, bailing_moe, deepseek_v4, glm4, hunyuan, llm, minimax_m2, olmoe, qwen3_emb, qwen3_next +from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, olmoe, qwen3_emb, + qwen3_next) diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py new file mode 100644 index 0000000..8ba0c96 --- /dev/null +++ b/src/mcore_bridge/model/gpts/glm_moe_dsa.py @@ -0,0 +1,180 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from typing import Optional + +from ..constant import ModelType +from ..gpt_model import GPTModel +from ..modules import TransformerBlock +from ..register import ModelLoader, ModelMeta, register_model + +try: + from megatron.core.transformer.experimental_attention_variant.dsa import (DSAIndexerLossAutoScaler, + DSAIndexerLossLoggingHelper, DSAttention, + DSAttentionSubmodules, + FusedDSAIndexerLoss, unfused_dsa_fn) +except ImportError: + DSAttention = object + + +class GlmMoeDsaDSAttention(DSAttention): + """DSAttention with shared indexer support for GLM 5.2. + + For "full" layers: computes topk_indices via the indexer and stores them + in ``shared_topk_indices`` for subsequent shared layers. + For "shared" layers: skips the indexer and reuses topk_indices from the + most recent "full" layer, analogous to gemma4's ``shared_kv_states``. + + Refer: https://arxiv.org/abs/2603.12201 for more details. + """ + + def __init__(self, config, submodules, layer_number, *args, **kwargs): + indexer_types = getattr(config.hf_config, 'indexer_types', None) + self.skip_topk = False + if indexer_types is not None: + layer_idx = layer_number - 1 + if layer_idx < len(indexer_types): + self.skip_topk = indexer_types[layer_idx] == 'shared' + + if self.skip_topk: + # Don't create indexer for shared layers to save memory + submodules = DSAttentionSubmodules(indexer=None) + + super().__init__(config, submodules, layer_number, *args, **kwargs) + + def _get_float_mask(self, query, key, attention_mask, x, attn_mask_type): + """Build a FP32 mask with -inf for masked positions.""" + sq = query.size(0) + skv = key.size(0) + if attn_mask_type is not None: + from megatron.core.transformer.enums import AttnMaskType + assert attn_mask_type == AttnMaskType.causal + float_mask = torch.triu( + torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=x.device), + diagonal=1, + ) + else: + b = query.size(1) + assert attention_mask.shape == (b, 1, sq, skv) + mask = attention_mask.squeeze() + float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill(mask, float('-inf')) + return float_mask + + def forward( + self, + query, + key, + value, + attention_mask, + x, + qr, + attn_mask_type=None, + attention_bias=None, + packed_seq_params=None, + ): + shared_topk_indices = getattr(self, '_shared_topk_indices', None) + + if self.skip_topk: + # Shared layer: reuse topk_indices from previous full layer + assert shared_topk_indices is not None and 'topk_indices' in shared_topk_indices, ( + 'Shared DSA layers require topk_indices from a previous full indexer layer.') + topk_indices = shared_topk_indices['topk_indices'] + output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) + return output + + # Full layer: compute topk_indices, store for shared layers, then run sparse attention. + # We override the full forward to avoid double-computing topk_indices. + x = x.detach() + qr = qr.detach() + float_mask = self._get_float_mask(query, key, attention_mask, x, attn_mask_type) + + if self.training and torch.is_grad_enabled(): + q, k, weights = self.indexer.forward_before_topk(x, qr, packed_seq_params) + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + + topk_indices, indexer_loss = FusedDSAIndexerLoss.apply( + q, weights, k, query.detach(), key.detach(), self.softmax_scale, + self.indexer.index_topk, indexer_loss_coeff, float_mask, + getattr(self.config, 'dsa_indexer_use_sparse_loss', + False), self.indexer.pg_collection, self.config.calculate_per_token_loss) + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=max( + self.layer_number, + self.config.num_layers + (self.config.mtp_num_layers or 0), + ), + ) + output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + else: + _, topk_indices = self.indexer.forward_with_scores( + x, qr, mask=float_mask, packed_seq_params=packed_seq_params) + output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) + + # Store topk_indices for subsequent shared layers (in-place dict mutation) + if shared_topk_indices is not None: + shared_topk_indices['topk_indices'] = topk_indices.detach() + + return output + + +class GlmMoeDsaGPTModel(GPTModel): + """GPT model for GLM 5.2 with shared DSA indexer support. + + Creates a ``shared_topk_indices`` dict and passes it through + ``extra_block_kwargs`` so that "full" DSA layers can store their + topk_indices for reuse by subsequent "shared" layers. + """ + + def forward(self, *args, **kwargs): + extra_block_kwargs = kwargs.get('extra_block_kwargs') or {} + extra_block_kwargs['shared_topk_indices'] = {} + kwargs['extra_block_kwargs'] = extra_block_kwargs + return super().forward(*args, **kwargs) + + +class GlmMoeDsaTransformerBlock(TransformerBlock): + """TransformerBlock that routes ``shared_topk_indices`` to DSAttention. + + Pops ``shared_topk_indices`` from kwargs before calling the layer + (so it doesn't hit ``_forward_attention``'s fixed signature), sets + it as an attribute on the core attention module, and restores it + afterward for subsequent layers. + """ + + def _layer_forward(self, layer, hidden_states, **kwargs): + shared_topk_indices = kwargs.pop('shared_topk_indices', None) + if shared_topk_indices is not None and hasattr(layer, 'self_attention'): + core_attn = getattr(layer.self_attention, 'core_attention', None) + if isinstance(core_attn, GlmMoeDsaDSAttention): + core_attn._shared_topk_indices = shared_topk_indices + result = super()._layer_forward(layer, hidden_states, **kwargs) + # Restore for subsequent layers + if shared_topk_indices is not None: + kwargs['shared_topk_indices'] = shared_topk_indices + return result + + +class GlmMoeDsaLoader(ModelLoader): + model_cls = GlmMoeDsaGPTModel + transformer_block = GlmMoeDsaTransformerBlock + + def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): + transformer_layer_spec = super().get_transformer_layer_spec(vp_stage) + + indexer_types = getattr(self.config.hf_config, 'indexer_types', None) + if indexer_types is not None: + for i, layer_spec in enumerate(transformer_layer_spec.layer_specs): + core_attn = layer_spec.submodules.self_attention.submodules.core_attention + if hasattr(core_attn, 'module') and issubclass(core_attn.module, DSAttention): + core_attn.module = GlmMoeDsaDSAttention + + return transformer_layer_spec + + +register_model(ModelMeta( + ModelType.glm_moe_dsa, + ['glm_moe_dsa'], + loader=GlmMoeDsaLoader, +)) diff --git a/src/mcore_bridge/model/gpts/llm.py b/src/mcore_bridge/model/gpts/llm.py index 424ecc9..b072f77 100644 --- a/src/mcore_bridge/model/gpts/llm.py +++ b/src/mcore_bridge/model/gpts/llm.py @@ -8,6 +8,6 @@ [ 'qwen2', 'llama', 'qwen3', 'qwen2_moe', 'qwen3_moe', 'internlm3', 'mimo', 'deepseek', 'deepseek_v2', 'deepseek_v3', 'deepseek_v32', 'kimi_k2', 'dots1', 'ernie4_5', 'ernie4_5_moe', 'glm4_moe', 'glm4_moe_lite', - 'glm_moe_dsa', 'gpt_oss' + 'gpt_oss' ], )) From 600ee06df3f2dc8c684f3eefcd4b25f83c75329c Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Tue, 23 Jun 2026 21:03:39 +0800 Subject: [PATCH 3/5] update --- src/mcore_bridge/model/gpts/glm_moe_dsa.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py index 8ba0c96..94b4c8c 100644 --- a/src/mcore_bridge/model/gpts/glm_moe_dsa.py +++ b/src/mcore_bridge/model/gpts/glm_moe_dsa.py @@ -19,16 +19,12 @@ class GlmMoeDsaDSAttention(DSAttention): """DSAttention with shared indexer support for GLM 5.2. - For "full" layers: computes topk_indices via the indexer and stores them - in ``shared_topk_indices`` for subsequent shared layers. - For "shared" layers: skips the indexer and reuses topk_indices from the - most recent "full" layer, analogous to gemma4's ``shared_kv_states``. - Refer: https://arxiv.org/abs/2603.12201 for more details. """ def __init__(self, config, submodules, layer_number, *args, **kwargs): - indexer_types = getattr(config.hf_config, 'indexer_types', None) + super().__init__(config, submodules, layer_number, *args, **kwargs) + indexer_types = config.hf_config.indexer_types self.skip_topk = False if indexer_types is not None: layer_idx = layer_number - 1 @@ -36,10 +32,7 @@ def __init__(self, config, submodules, layer_number, *args, **kwargs): self.skip_topk = indexer_types[layer_idx] == 'shared' if self.skip_topk: - # Don't create indexer for shared layers to save memory - submodules = DSAttentionSubmodules(indexer=None) - - super().__init__(config, submodules, layer_number, *args, **kwargs) + self.indexer = None def _get_float_mask(self, query, key, attention_mask, x, attn_mask_type): """Build a FP32 mask with -inf for masked positions.""" @@ -145,8 +138,8 @@ class GlmMoeDsaTransformerBlock(TransformerBlock): def _layer_forward(self, layer, hidden_states, **kwargs): shared_topk_indices = kwargs.pop('shared_topk_indices', None) - if shared_topk_indices is not None and hasattr(layer, 'self_attention'): - core_attn = getattr(layer.self_attention, 'core_attention', None) + if shared_topk_indices is not None: + core_attn = layer.self_attention.core_attention if isinstance(core_attn, GlmMoeDsaDSAttention): core_attn._shared_topk_indices = shared_topk_indices result = super()._layer_forward(layer, hidden_states, **kwargs) From e618f2961d0b952dc1d7f9ce582c5a2cc7b2e20c Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Tue, 23 Jun 2026 21:30:18 +0800 Subject: [PATCH 4/5] update --- src/mcore_bridge/model/gpts/glm_moe_dsa.py | 38 +++++++++++++++------- 1 file changed, 26 insertions(+), 12 deletions(-) diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py index 94b4c8c..5153498 100644 --- a/src/mcore_bridge/model/gpts/glm_moe_dsa.py +++ b/src/mcore_bridge/model/gpts/glm_moe_dsa.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import torch +from contextlib import contextmanager from typing import Optional from ..constant import ModelType @@ -10,7 +11,6 @@ try: from megatron.core.transformer.experimental_attention_variant.dsa import (DSAIndexerLossAutoScaler, DSAIndexerLossLoggingHelper, DSAttention, - DSAttentionSubmodules, FusedDSAIndexerLoss, unfused_dsa_fn) except ImportError: DSAttention = object @@ -69,7 +69,9 @@ def forward( if self.skip_topk: # Shared layer: reuse topk_indices from previous full layer assert shared_topk_indices is not None and 'topk_indices' in shared_topk_indices, ( - 'Shared DSA layers require topk_indices from a previous full indexer layer.') + f'GLM 5.2 DSA: Layer {self.layer_number} is a "shared" indexer layer but no ' + f'"full" layer precedes it in this PP stage. Please adjust ' + f'`pipeline_model_parallel_layout` to ensure each PP stage starts with a "full" indexer layer.') topk_indices = shared_topk_indices['topk_indices'] output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) return output @@ -127,22 +129,34 @@ def forward(self, *args, **kwargs): return super().forward(*args, **kwargs) +@contextmanager +def _shared_topk_indices_context(layer, shared_topk_indices): + """Temporarily inject shared_topk_indices into the core attention module.""" + core_attn = None + if shared_topk_indices is not None and hasattr(layer, 'self_attention'): + _attn = getattr(layer.self_attention, 'core_attention', None) + if isinstance(_attn, GlmMoeDsaDSAttention): + core_attn = _attn + core_attn._shared_topk_indices = shared_topk_indices + try: + yield + finally: + if core_attn is not None: + core_attn._shared_topk_indices = None + + class GlmMoeDsaTransformerBlock(TransformerBlock): """TransformerBlock that routes ``shared_topk_indices`` to DSAttention. Pops ``shared_topk_indices`` from kwargs before calling the layer - (so it doesn't hit ``_forward_attention``'s fixed signature), sets - it as an attribute on the core attention module, and restores it - afterward for subsequent layers. + (so it doesn't hit ``_forward_attention``'s fixed signature), injects + it via context manager, and restores it afterward for subsequent layers. """ def _layer_forward(self, layer, hidden_states, **kwargs): shared_topk_indices = kwargs.pop('shared_topk_indices', None) - if shared_topk_indices is not None: - core_attn = layer.self_attention.core_attention - if isinstance(core_attn, GlmMoeDsaDSAttention): - core_attn._shared_topk_indices = shared_topk_indices - result = super()._layer_forward(layer, hidden_states, **kwargs) + with _shared_topk_indices_context(layer, shared_topk_indices): + result = super()._layer_forward(layer, hidden_states, **kwargs) # Restore for subsequent layers if shared_topk_indices is not None: kwargs['shared_topk_indices'] = shared_topk_indices @@ -156,9 +170,9 @@ class GlmMoeDsaLoader(ModelLoader): def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): transformer_layer_spec = super().get_transformer_layer_spec(vp_stage) - indexer_types = getattr(self.config.hf_config, 'indexer_types', None) + indexer_types = self.config.hf_config.indexer_types if indexer_types is not None: - for i, layer_spec in enumerate(transformer_layer_spec.layer_specs): + for layer_spec in transformer_layer_spec.layer_specs: core_attn = layer_spec.submodules.self_attention.submodules.core_attention if hasattr(core_attn, 'module') and issubclass(core_attn.module, DSAttention): core_attn.module = GlmMoeDsaDSAttention From c40515f83773460364c60b71332ea22117d753e4 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Tue, 23 Jun 2026 23:51:04 +0800 Subject: [PATCH 5/5] fix --- src/mcore_bridge/model/gpts/glm_moe_dsa.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py index 5153498..01786a4 100644 --- a/src/mcore_bridge/model/gpts/glm_moe_dsa.py +++ b/src/mcore_bridge/model/gpts/glm_moe_dsa.py @@ -71,7 +71,8 @@ def forward( assert shared_topk_indices is not None and 'topk_indices' in shared_topk_indices, ( f'GLM 5.2 DSA: Layer {self.layer_number} is a "shared" indexer layer but no ' f'"full" layer precedes it in this PP stage. Please adjust ' - f'`pipeline_model_parallel_layout` to ensure each PP stage starts with a "full" indexer layer.') + f'`--pipeline_model_parallel_layout` to ensure each PP stage starts with a "full" indexer layer. ' + f'indexer_types: {self.config.hf_config.indexer_types}.') topk_indices = shared_topk_indices['topk_indices'] output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) return output