`bridge.models.glm_moe_dsa.glm5_hybrid`#

GLM adaptations for Core’s split attention/MLP Hybrid stack.

Module Contents#

Classes#

`GLMDSAttention`	Read index-sharing state from the enclosing GLM stage forward instead of carrier attributes.
`GLMHybridStack`	Isolate top-k state by forward and replay a complete stage on recomputation.
`GLMHybridMTPLayer`	Use the same forward-scoped state for Core’s nested MTP HybridStack.

Functions#

`_forward_glm_stack`
`glm_hybrid_stack_spec`	Return a private GLM spec without mutating Core’s shared specification.

Data#

_FORWARD_STATE

API#

bridge.models.glm_moe_dsa.glm5_hybrid._FORWARD_STATE: contextvars.ContextVar[tuple[dict[int, torch.Tensor], dict[int, torch.Tensor]] | None]#: ‘ContextVar(…)’

class bridge.models.glm_moe_dsa.glm5_hybrid.GLMDSAttention#

Bases: megatron.core.transformer.experimental_attention_variant.dsa.DSAttention

Read index-sharing state from the enclosing GLM stage forward instead of carrier attributes.

_get_index_share_topk_holder( packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None, attention_mask: torch.Tensor | None = None, ) → dict[int, torch.Tensor]#

_get_index_share_topk_length_holder( packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None, attention_mask: torch.Tensor | None = None, ) → dict[int, torch.Tensor]#

bridge.models.glm_moe_dsa.glm5_hybrid._forward_glm_stack(

stack: megatron.core.models.hybrid.hybrid_block.HybridStack,

hidden_states: Any,

attention_mask: Any,

**kwargs: Any,

) → torch.Tensor#

class bridge.models.glm_moe_dsa.glm5_hybrid.GLMHybridStack(

config: Any,

submodules: Any,

*,

layer_config_list: Any = None,

layer_type_list: list[str] | None = None,

**kwargs: Any,

)#

Bases: megatron.core.models.hybrid.hybrid_block.HybridStack

Isolate top-k state by forward and replay a complete stage on recomputation.

Initialization

forward(

hidden_states: Any,

attention_mask: Any,

**kwargs: Any,

) → torch.Tensor#: Run the stage with fresh state, including during backward recomputation.

class bridge.models.glm_moe_dsa.glm5_hybrid.GLMHybridMTPLayer#

Bases: megatron.core.transformer.multi_token_prediction.MultiTokenPredictionLayer

Use the same forward-scoped state for Core’s nested MTP HybridStack.

_proj_and_transformer_layer(

hidden_states: torch.Tensor,

decoder_input: torch.Tensor,

**kwargs: Any,

) → torch.Tensor#

bridge.models.glm_moe_dsa.glm5_hybrid.glm_hybrid_stack_spec( config: Any, ) → megatron.core.transformer.spec_utils.ModuleSpec#: Return a private GLM spec without mutating Core’s shared specification.

bridge.models.glm_moe_dsa.glm5_hybrid#