78 transformer blocks › 3 dense blocks › Sparse latent attention
learnedself.indexer
Score every position up to its own for each query cheaply and keep the 2,048 best; attention then uses only those.
Why
The indexer scores every pair of positions with 32 heads of 128 features and one key per position shared by them, against the attention’s 64 heads of 256. The authors of DeepSeek Sparse Attention say that an indexer with a small number of heads is computationally efficient. DeepSeek-V3.2-Exp report §1
GLM-5’s authors say the model integrates DeepSeek Sparse Attention, which largely reduces the cost of deployment while preserving long-context capacity. GLM-5 model card, section Introduction
Attending to every earlier position costs O(L²) for a context of L tokens. The authors of DeepSeek Sparse Attention say that choosing k positions for each query brings the main attention to O(Lk). The pinned implementation’s eager and SDPA paths apply the choice as a mask over every position (modeling_glm_moe_dsa.py 431–441). DeepSeek-V3.2-Exp report §3
All 12 steps
- in
- [B batch, T positions, 6144 d] [B batch, T positions, 2048 q latent] [T positions, 32 pairs] [T queries, T keys]
- out
- [B batch, T queries, T keys]
- parameters
- 9,371,904 · × 3 dense blocks = 28,115,712
Parameters
9,371,904 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Indexer queries | 8,388,608 across its components | 8,388,608 |
| Indexer key | 786,688 across its components | 786,688 |
| Head weight projection | 196,608 | |
| Per-instance total | 9,371,904 | |
Across 3 independent dense blocks: parameters.
Derived from the pinned configuration and transformers implementation, and equal to the count of the checkpoint’s tensors read from its safetensors headers, less the multi-token-prediction layer the implementation does not load. The total includes the routers’ 19,200 correction-bias values, which the implementation holds as buffers.
Explanation
This is DeepSeek Sparse Attention, which GLM-5 takes from DeepSeek-V3.2: each query attends to at most 2,048 positions, the ones chosen here. With 32 heads of 128 features and one key shared by all heads, scoring every pair of positions costs far less than attention itself. Every block runs its own indexer.
With 2,048 or fewer positions in the sequence, every permitted position is kept and attention is dense.
Codemodeling_glm_moe_dsa.py 177–649
def __init__(self, config: "GlmMoeDsaConfig", layer_idx: int):
super().__init__()
self.config = config
self.layer_idx = layer_idx
self.hidden_size: int = config.hidden_size
self.n_heads: int = config.index_n_heads
self.head_dim: int = config.index_head_dim
self.qk_rope_head_dim: int = config.qk_rope_head_dim
self.index_topk: int = config.index_topk
self.q_lora_rank: int = config.q_lora_rank
self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)
self.wk = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6)
self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False)
self.softmax_scale = self.head_dim**-0.5 _keep_in_fp32_modules = ["indexer.weights_proj"] @torch.no_grad()
def forward(
self,
hidden_states: torch.Tensor,
q_resid: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor],
attention_mask: torch.Tensor,
position_ids: torch.Tensor, # Kept for BC
past_key_values: Cache | None = None,
) -> torch.Tensor:
"""
Selects the top-k tokens per query for DeepSeek Sparse Attention (DSA).
Same as [`DeepseekV32Indexer.forward`], but the indexer applies **interleaved** RoPE
rather than the non-interleaved half-split RoPE used by DeepSeek-V3.2.
Args:
hidden_states: Input hidden states `[B, S, hidden_size]`.
q_resid: Query residual from `q_a_layernorm(q_a_proj(x))`, shape `[B, S, q_lora_rank]`.
position_embeddings: `(cos, sin)` from RotaryEmbedding.
attention_mask: Causal mask, broadcastable to `[B, S, T]`.
past_key_values: Cache object containing the indexer key cache for this layer.
Returns:
`torch.Tensor`: the `int32` top-k token indices of shape `[B, S, topk]`. The eager / SDPA paths
turn these into an additive sparse mask; the `flash-mla` kernel consumes them directly.
"""
batch_size, seq_len, _ = hidden_states.shape
cos, sin = position_embeddings
q = self.wq_b(q_resid) # [B, S, H*D]
q = q.view(batch_size, seq_len, self.n_heads, self.head_dim) # [B, S, H, D]
q_rot, q_pass = torch.split(q, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
k = self.k_norm(self.wk(hidden_states)).unsqueeze(2) # [B, S, 1, D]
k_rot, k_pass = torch.split(k, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
# GLM-MoE-DSA uses interleaved RoPE in the indexer
q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin, unsqueeze_dim=2)
q = torch.cat([q_rot, q_pass], dim=-1) # [B, S, H, D]
k = torch.cat([k_rot, k_pass], dim=-1).squeeze(2) # [B, S, D]
if past_key_values is not None:
k = past_key_values.update_indexer(k, self.layer_idx)
scores = torch.matmul(q.float(), k.transpose(-1, -2).float().unsqueeze(1)) * self.softmax_scale
scores = F.relu(scores)
# Weight per head and sum across heads: [B, S, 1, H] @ [B, S, H, T] → [B, S, T]
weights = self.weights_proj(hidden_states.to(self.weights_proj.weight.dtype)).float() * (self.n_heads**-0.5)
index_scores = torch.matmul(weights.unsqueeze(-2), scores).squeeze(-2)
# Causality needs to be taken into account when computing scores so padding tokens don't affect computation
if attention_mask.dtype == torch.bool:
index_scores = index_scores.masked_fill(~attention_mask, float("-inf"))
else:
index_scores = index_scores + attention_mask
topk = min(self.index_topk, index_scores.shape[-1])
return index_scores.topk(topk, dim=-1).indices.to(torch.int32) # [B, S, topk] scores = torch.matmul(q.float(), k.transpose(-1, -2).float().unsqueeze(1)) * self.softmax_scale
scores = F.relu(scores)
# Weight per head and sum across heads: [B, S, 1, H] @ [B, S, H, T] → [B, S, T]
weights = self.weights_proj(hidden_states.to(self.weights_proj.weight.dtype)).float() * (self.n_heads**-0.5)
index_scores = torch.matmul(weights.unsqueeze(-2), scores).squeeze(-2)Sources3
Direct evidence
The indexer owns a query projection from the query latent (32 heads × 128 features), one key projection from the block input followed by LayerNorm, and a head-weight projection; its score scale is 1/√128; the head-weight projection is kept in float32 when loading.
indexer-init · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 177–649
Read code excerpt
def __init__(self, config: "GlmMoeDsaConfig", layer_idx: int):
super().__init__()
self.config = config
self.layer_idx = layer_idx
self.hidden_size: int = config.hidden_size
self.n_heads: int = config.index_n_heads
self.head_dim: int = config.index_head_dim
self.qk_rope_head_dim: int = config.qk_rope_head_dim
self.index_topk: int = config.index_topk
self.q_lora_rank: int = config.q_lora_rank
self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)
self.wk = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6)
self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False)
self.softmax_scale = self.head_dim**-0.5… lines 194–648 omitted …
_keep_in_fp32_modules = ["indexer.weights_proj"]Direct evidence
Indexer queries and the shared key rotate their first 64 features with interleaved RoPE; scores are the ReLU of scaled query–key products in float32, weighted per head by the head weights divided by √32 and summed; the causal mask is added and the top 2,048 positions are kept. There is no Hadamard transform, and it runs without gradients.
indexer-forward · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 195–253
Read code excerpt
@torch.no_grad()
def forward(
self,
hidden_states: torch.Tensor,
q_resid: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor],
attention_mask: torch.Tensor,
position_ids: torch.Tensor, # Kept for BC
past_key_values: Cache | None = None,
) -> torch.Tensor:
"""
Selects the top-k tokens per query for DeepSeek Sparse Attention (DSA).
Same as [`DeepseekV32Indexer.forward`], but the indexer applies **interleaved** RoPE
rather than the non-interleaved half-split RoPE used by DeepSeek-V3.2.
Args:
hidden_states: Input hidden states `[B, S, hidden_size]`.
q_resid: Query residual from `q_a_layernorm(q_a_proj(x))`, shape `[B, S, q_lora_rank]`.
position_embeddings: `(cos, sin)` from RotaryEmbedding.
attention_mask: Causal mask, broadcastable to `[B, S, T]`.
past_key_values: Cache object containing the indexer key cache for this layer.
Returns:
`torch.Tensor`: the `int32` top-k token indices of shape `[B, S, topk]`. The eager / SDPA paths
turn these into an additive sparse mask; the `flash-mla` kernel consumes them directly.
"""
batch_size, seq_len, _ = hidden_states.shape
cos, sin = position_embeddings
q = self.wq_b(q_resid) # [B, S, H*D]
q = q.view(batch_size, seq_len, self.n_heads, self.head_dim) # [B, S, H, D]
q_rot, q_pass = torch.split(q, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
k = self.k_norm(self.wk(hidden_states)).unsqueeze(2) # [B, S, 1, D]
k_rot, k_pass = torch.split(k, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
# GLM-MoE-DSA uses interleaved RoPE in the indexer
q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin, unsqueeze_dim=2)
q = torch.cat([q_rot, q_pass], dim=-1) # [B, S, H, D]
k = torch.cat([k_rot, k_pass], dim=-1).squeeze(2) # [B, S, D]
if past_key_values is not None:
k = past_key_values.update_indexer(k, self.layer_idx)
scores = torch.matmul(q.float(), k.transpose(-1, -2).float().unsqueeze(1)) * self.softmax_scale
scores = F.relu(scores)
# Weight per head and sum across heads: [B, S, 1, H] @ [B, S, H, T] → [B, S, T]
weights = self.weights_proj(hidden_states.to(self.weights_proj.weight.dtype)).float() * (self.n_heads**-0.5)
index_scores = torch.matmul(weights.unsqueeze(-2), scores).squeeze(-2)
# Causality needs to be taken into account when computing scores so padding tokens don't affect computation
if attention_mask.dtype == torch.bool:
index_scores = index_scores.masked_fill(~attention_mask, float("-inf"))
else:
index_scores = index_scores + attention_mask
topk = min(self.index_topk, index_scores.shape[-1])
return index_scores.topk(topk, dim=-1).indices.to(torch.int32) # [B, S, topk]Direct evidence
The index score is computed in float32, as the head-weighted sum of ReLU(q·k/√128) with the head weights divided by √32; no FP8 kernel is involved.
index-kernel · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 239–244
Read code excerpt
scores = torch.matmul(q.float(), k.transpose(-1, -2).float().unsqueeze(1)) * self.softmax_scale
scores = F.relu(scores)
# Weight per head and sum across heads: [B, S, 1, H] @ [B, S, H, T] → [B, S, T]
weights = self.weights_proj(hidden_states.to(self.weights_proj.weight.dtype)).float() * (self.n_heads**-0.5)
index_scores = torch.matmul(weights.unsqueeze(-2), scores).squeeze(-2)