Skip to content

Fix: interleaved RoPE application for MLA and Support Index Cache DSA indexer skip-topk sharing for GLM5 #46372

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ArthurZucker merged 12 commits into from
Jun 8, 2026
Merged

Fix: interleaved RoPE application for MLA and Support Index Cache DSA indexer skip-topk sharing for GLM5 #46372

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
67a9994
Fix error attributes mapping
Dovis01 Jun 2, 2026
3f02fef
Fix tests
Dovis01 Jun 2, 2026
3841bbc
Merge branch 'main' into fix-glm-moe-dsa
Dovis01 Jun 2, 2026
386e43f
Merge branch 'main' into fix-glm-moe-dsa
Dovis01 Jun 2, 2026
e59ee7c
Merge branch 'main' into fix-glm-moe-dsa
Dovis01 Jun 3, 2026
23d7895
Fix: Support interleaved RoPE for MLA
Dovis01 Jun 3, 2026
8eab52b
Fix tests
Dovis01 Jun 3, 2026
aa8928c
Remove redundant fields
Dovis01 Jun 4, 2026
45f7dcd
Fix tests
Dovis01 Jun 4, 2026
7da480e
Remove portentially codepath
Dovis01 Jun 4, 2026
f225757
Optimize code
Dovis01 Jun 5, 2026
d066cd3
Merge branch 'main' into fix-glm-moe-dsa
Dovis01 Jun 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 11 additions & 12 deletions src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,15 +32,15 @@ class GlmMoeDsaConfig(PreTrainedConfig):
n_group (`int`, *optional*, defaults to 1):
Number of groups for routed experts.
mlp_layer_types (`list`, *optional*):
MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to 3 dense + rest sparse.
MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to `3` dense layers and then every `moe_layer_freq`-th layer sparse.
index_topk (`int`, *optional*, defaults to 2048):
Number of top tokens selected by the indexer for sparse attention.
index_head_dim (`int`, *optional*, defaults to 128):
Head dimension for the indexer projections (DSA).
index_n_heads (`int | None`, *optional*, defaults to 32):
Number of heads for the indexer projections (DSA).
indexer_types (`list[str]`, *optional*):
Indexer mode for each layer (`"full"` or `"shared"`). Defaults to first layer full, then every `index_topk_freq`-th layer full, rest shared.
Indexer mode for each layer (`"full"` or `"shared"`). Defaults to the pattern derived from `index_topk_freq` and `index_skip_topk_offset`.

```python
>>> from transformers import GlmMoeDsaConfig, GlmMoeDsaModel
Expand Down Expand Up @@ -78,6 +78,7 @@ class GlmMoeDsaConfig(PreTrainedConfig):
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
"norm": (["hidden_states"], ["hidden_states"]),
}

attribute_map = {
"num_local_experts": "n_routed_experts",
}
Expand Down Expand Up @@ -122,25 +123,23 @@ class GlmMoeDsaConfig(PreTrainedConfig):

def __post_init__(self, **kwargs):
self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

# MLP layer types: first 3 dense, rest sparse
if self.mlp_layer_types is None:
self.mlp_layer_types = ["dense"] * min(3, self.num_hidden_layers) + ["sparse"] * (
self.num_hidden_layers - 3
)
moe_layer_freq = kwargs.get("moe_layer_freq", 1)
self.mlp_layer_types = [
"sparse" if i >= 3 and i % moe_layer_freq == 0 else "dense" for i in range(self.num_hidden_layers)
]

# Indexer layer types
if self.indexer_types is None:
pattern = kwargs.pop("index_topk_pattern", None)
freq = kwargs.pop("index_topk_freq", 1)
pattern = kwargs.get("index_topk_pattern")
if pattern is not None:
self.indexer_types = (
[{"F": "full", "S": "shared"}[c] for c in pattern] if isinstance(pattern, str) else list(pattern)
)
else:
# First layer full, then every freq-th layer full, rest shared
freq = max(kwargs.get("index_topk_freq", 1), 1)
offset = kwargs.get("index_skip_topk_offset", 2)
self.indexer_types = [
"full" if (max(i - 1, 0) % freq) == 0 else "shared" for i in range(self.num_hidden_layers)
"full" if (max(i - offset + 1, 0) % freq) == 0 else "shared" for i in range(self.num_hidden_layers)
]
super().__post_init__(**kwargs)

Expand Down
39 changes: 22 additions & 17 deletions src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,10 +77,12 @@ def apply_rotary_pos_emb(
unsqueeze_dim: int = 1,
) -> torch.Tensor:
"""
Applies Rotary Position Embedding to a single tensor.
Applies Rotary Position Embedding to a single tensor (query, key, or the indexer's q/k stream).

This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`.
Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding.
This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved=True)`.
Instead of complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding.
Rotary pairs are always interpreted as adjacent elements (GPT-J style), so we first de-interleave `x` into
halves before applying the standard `rotate_half`.

Args:
x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`.
Expand All @@ -94,20 +96,18 @@ def apply_rotary_pos_emb(
"""
cos = cos.unsqueeze(unsqueeze_dim)
sin = sin.unsqueeze(unsqueeze_dim)

# Split-half (NeoX/Llama style): (x[:d/2], x[d/2:])
# This matches llama's apply_rotary_pos_emb logic.
x_rotated = (x * cos) + (rotate_half(x) * sin)
return x_rotated
*leading_dims, head_dim = x.shape
x = x.view(*leading_dims, head_dim // 2, 2).transpose(-1, -2).reshape(*leading_dims, head_dim)
Comment thread
Dovis01 marked this conversation as resolved.
return (x * cos) + (rotate_half(x) * sin)


class GlmMoeDsaIndexer(nn.Module):
"""
DeepSeek Sparse Attention (DSA) indexer for selecting top-k tokens.

The Indexer has its own lightweight projections (wq_b, wk) separate from the
main MLA attention. It uses non-interleaved (NeoX/Llama) RoPE, unlike the main attention
which uses interleaved RoPE.
main MLA attention. RoPE uses the same interleaved pair layout as main MLA
attention.

**Cache strategy**: The Indexer manages its own key cache (`_cached_keys`) separately
from the DynamicCache used by MLA attention, since DynamicCache is sized for exactly
Expand Down Expand Up @@ -281,6 +281,9 @@ class GlmMoeDsaAttention(nn.Module):
and SDPA backends. The reference's compressed-cache decode path (which avoids the kv_b_proj
expansion at decode time) is a future optimization that would require a dedicated MLA cache class.

**DSA layer sharing**: full layers run the indexer; shared layers reuse the previous full
layer's top-k indices (`skip_topk` / `next_skip_topk`, derived from `config.indexer_types`).

**FP8 compatibility**: all weight accesses use standard nn.Linear forward calls (never
raw `.weight` access), so FP8-quantized checkpoints work transparently.
"""
Expand Down Expand Up @@ -332,15 +335,15 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int):

self.scaling = self.qk_head_dim ** (-0.5)

self.indexer = GlmMoeDsaIndexer(config, layer_idx)

# Refer: https://arxiv.org/abs/2603.12201 for more details.
# skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices.
# next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices.
# skip_topk: when True, this layer reuses the previous full indexer layer's top-k indices.
# next_skip_topk: when True, the next layer reuses this layer's top-k indices.
# Shared layers have no indexer of their own.
self.skip_topk = config.indexer_types[layer_idx] == "shared"
self.next_skip_topk = (
config.indexer_types[layer_idx + 1] == "shared" if layer_idx < len(config.indexer_types) - 1 else False
)
self.indexer = None if self.skip_topk else GlmMoeDsaIndexer(config, layer_idx)

def forward(
self,
Expand All @@ -364,7 +367,6 @@ def forward(
query_states = query_states.view(batch_size, seq_length, -1, self.qk_head_dim).transpose(1, 2)
# Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D]
q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format

# ===== KV path =====
compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D]
Expand All @@ -378,9 +380,10 @@ def forward(
k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D]
value_states = value_states.transpose(1, 2) # [B, H, S, v_D]

# RoPE on k_pe (single-head rope stream)
# RoPE on q_pe / k_pe (single-head rope stream for k), using interleaved pair layout.
k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D]
k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format
q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1)
k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1)
k_pe = k_pe.expand(-1, k_nope.shape[1], -1, -1) # [B, H, S, rope_D]

# Assemble full Q and K
Expand All @@ -394,6 +397,8 @@ def forward(
# ===== Indexer (DSA sparse mask) =====
# attention_mask is [B, 1, S, T] (4D) for eager and (2D) otherwise but indexer works with [B, S, T] (3D)
if not self.skip_topk or prev_topk_indices is None:
if self.indexer is None:
raise ValueError("Shared DSA layers require top-k indices from a previous full indexer layer.")
indexer_mask = (
attention_mask[:, 0, :, :]
if attention_mask is not None and attention_mask.dim() == 4
Expand Down
70 changes: 37 additions & 33 deletions src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,10 +53,12 @@ def apply_rotary_pos_emb(
unsqueeze_dim: int = 1,
) -> torch.Tensor:
"""
Applies Rotary Position Embedding to a single tensor.
Applies Rotary Position Embedding to a single tensor (query, key, or the indexer's q/k stream).

This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`.
Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding.
This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved=True)`.
Instead of complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding.
Rotary pairs are always interpreted as adjacent elements (GPT-J style), so we first de-interleave `x` into
halves before applying the standard `rotate_half`.

Args:
x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`.
Expand All @@ -70,11 +72,9 @@ def apply_rotary_pos_emb(
"""
cos = cos.unsqueeze(unsqueeze_dim)
sin = sin.unsqueeze(unsqueeze_dim)

# Split-half (NeoX/Llama style): (x[:d/2], x[d/2:])
# This matches llama's apply_rotary_pos_emb logic.
x_rotated = (x * cos) + (rotate_half(x) * sin)
return x_rotated
*leading_dims, head_dim = x.shape
x = x.view(*leading_dims, head_dim // 2, 2).transpose(-1, -2).reshape(*leading_dims, head_dim)
return (x * cos) + (rotate_half(x) * sin)


@auto_docstring(checkpoint="zai-org/GLM-5")
Expand All @@ -84,15 +84,15 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig):
n_group (`int`, *optional*, defaults to 1):
Number of groups for routed experts.
mlp_layer_types (`list`, *optional*):
MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to 3 dense + rest sparse.
MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to `3` dense layers and then every `moe_layer_freq`-th layer sparse.
index_topk (`int`, *optional*, defaults to 2048):
Number of top tokens selected by the indexer for sparse attention.
index_head_dim (`int`, *optional*, defaults to 128):
Head dimension for the indexer projections (DSA).
index_n_heads (`int | None`, *optional*, defaults to 32):
Number of heads for the indexer projections (DSA).
indexer_types (`list[str]`, *optional*):
Indexer mode for each layer (`"full"` or `"shared"`). Defaults to first layer full, then every `index_topk_freq`-th layer full, rest shared.
Indexer mode for each layer (`"full"` or `"shared"`). Defaults to the pattern derived from `index_topk_freq` and `index_skip_topk_offset`.

```python
>>> from transformers import GlmMoeDsaConfig, GlmMoeDsaModel
Expand Down Expand Up @@ -123,6 +123,10 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig):
"layers.*.mlp.down_proj": "rowwise",
}

attribute_map = {
"num_local_experts": "n_routed_experts",
}

hidden_size: int = 6144
intermediate_size: int = 12288
moe_intermediate_size: int = 2048
Expand All @@ -136,34 +140,29 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig):
index_topk: int = 2048
index_head_dim: int = 128
index_n_heads: int = 32
indexer_types: list[str] | None = None
pretraining_tp = AttributeError()
rope_interleave = AttributeError()
indexer_types: list[str] | None = None
attribute_map = {
"num_local_experts": "n_routed_experts",
}

def __post_init__(self, **kwargs):
self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

# MLP layer types: first 3 dense, rest sparse
if self.mlp_layer_types is None:
self.mlp_layer_types = ["dense"] * min(3, self.num_hidden_layers) + ["sparse"] * (
self.num_hidden_layers - 3
)
moe_layer_freq = kwargs.get("moe_layer_freq", 1)
self.mlp_layer_types = [
"sparse" if i >= 3 and i % moe_layer_freq == 0 else "dense" for i in range(self.num_hidden_layers)
]

# Indexer layer types
if self.indexer_types is None:
pattern = kwargs.pop("index_topk_pattern", None)
freq = kwargs.pop("index_topk_freq", 1)
pattern = kwargs.get("index_topk_pattern")
if pattern is not None:
self.indexer_types = (
[{"F": "full", "S": "shared"}[c] for c in pattern] if isinstance(pattern, str) else list(pattern)
)
else:
# First layer full, then every freq-th layer full, rest shared
freq = max(kwargs.get("index_topk_freq", 1), 1)
offset = kwargs.get("index_skip_topk_offset", 2)
self.indexer_types = [
"full" if (max(i - 1, 0) % freq) == 0 else "shared" for i in range(self.num_hidden_layers)
"full" if (max(i - offset + 1, 0) % freq) == 0 else "shared" for i in range(self.num_hidden_layers)
]
PreTrainedConfig.__post_init__(self, **kwargs)

Expand All @@ -177,8 +176,8 @@ class GlmMoeDsaIndexer(nn.Module):
DeepSeek Sparse Attention (DSA) indexer for selecting top-k tokens.

The Indexer has its own lightweight projections (wq_b, wk) separate from the
main MLA attention. It uses non-interleaved (NeoX/Llama) RoPE, unlike the main attention
which uses interleaved RoPE.
main MLA attention. RoPE uses the same interleaved pair layout as main MLA
attention.

**Cache strategy**: The Indexer manages its own key cache (`_cached_keys`) separately
from the DynamicCache used by MLA attention, since DynamicCache is sized for exactly
Expand Down Expand Up @@ -315,6 +314,9 @@ class GlmMoeDsaAttention(nn.Module):
and SDPA backends. The reference's compressed-cache decode path (which avoids the kv_b_proj
expansion at decode time) is a future optimization that would require a dedicated MLA cache class.

**DSA layer sharing**: full layers run the indexer; shared layers reuse the previous full
layer's top-k indices (`skip_topk` / `next_skip_topk`, derived from `config.indexer_types`).

**FP8 compatibility**: all weight accesses use standard nn.Linear forward calls (never
raw `.weight` access), so FP8-quantized checkpoints work transparently.
"""
Expand Down Expand Up @@ -366,15 +368,15 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int):

self.scaling = self.qk_head_dim ** (-0.5)

self.indexer = GlmMoeDsaIndexer(config, layer_idx)

# Refer: https://arxiv.org/abs/2603.12201 for more details.
# skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices.
# next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices.
# skip_topk: when True, this layer reuses the previous full indexer layer's top-k indices.
# next_skip_topk: when True, the next layer reuses this layer's top-k indices.
# Shared layers have no indexer of their own.
self.skip_topk = config.indexer_types[layer_idx] == "shared"
self.next_skip_topk = (
config.indexer_types[layer_idx + 1] == "shared" if layer_idx < len(config.indexer_types) - 1 else False
)
self.indexer = None if self.skip_topk else GlmMoeDsaIndexer(config, layer_idx)

def forward(
self,
Expand All @@ -398,7 +400,6 @@ def forward(
query_states = query_states.view(batch_size, seq_length, -1, self.qk_head_dim).transpose(1, 2)
# Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D]
q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format

# ===== KV path =====
compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D]
Expand All @@ -412,9 +413,10 @@ def forward(
k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D]
value_states = value_states.transpose(1, 2) # [B, H, S, v_D]

# RoPE on k_pe (single-head rope stream)
# RoPE on q_pe / k_pe (single-head rope stream for k), using interleaved pair layout.
k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D]
k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format
q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1)
k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1)
k_pe = k_pe.expand(-1, k_nope.shape[1], -1, -1) # [B, H, S, rope_D]

# Assemble full Q and K
Expand All @@ -428,6 +430,8 @@ def forward(
# ===== Indexer (DSA sparse mask) =====
# attention_mask is [B, 1, S, T] (4D) for eager and (2D) otherwise but indexer works with [B, S, T] (3D)
if not self.skip_topk or prev_topk_indices is None:
if self.indexer is None:
raise ValueError("Shared DSA layers require top-k indices from a previous full indexer layer.")
indexer_mask = (
attention_mask[:, 0, :, :]
if attention_mask is not None and attention_mask.dim() == 4
Expand Down
7 changes: 7 additions & 0 deletions tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,13 @@ def test_default_mlp_layer_types(self):
config.mlp_layer_types, ["dense", "dense", "dense", "sparse", "sparse", "sparse", "sparse", "sparse"]
)

def test_indexer_types_respect_skip_topk_offset(self):
config = GlmMoeDsaConfig(num_hidden_layers=8, index_topk_freq=4, index_skip_topk_offset=3)
self.assertEqual(
config.indexer_types,
["full", "full", "full", "shared", "shared", "shared", "full", "shared"],
)

@parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION)
@unittest.skip("Won't fix: Blip2 + T5 backbone needs custom input preparation for this test")
def test_eager_matches_sdpa_inference(self, *args):
Expand Down
Loading