Multi-head latent attention

Checkpoint configuration and implementation zai-org/GLM-5 @ c183ef8, huggingface/transformers @ 5e4d630 · parameter count checked against the checkpoint’s tensor headers · How to read

78 transformer blocks › 75 MoE blocks › Sparse latent attention

Query compression

Open inside →×

learnedself.q_a_layernorm(self.q_a_proj(hidden_states))

Compress each position’s 6,144 features into a 2,048-feature query latent.

Why

Queries pass through a low-rank latent as well; the authors of latent attention say this can reduce activation memory during training. DeepSeek-V3 §2.1.1

cQ=xWDQc^{{\htmlClass{vl-stream-q}{Q}}}=xW_{DQ}
c^Q=RMSNorm⁡(cQ)\hat c^{{\htmlClass{vl-stream-q}{Q}}}=\operatorname{RMSNorm}(c^{{\htmlClass{vl-stream-q}{Q}}})
in
xx [B batch, 1 positions, 6144 d]
out
c^Q\hat c^{Q} [B batch, 1 positions, 2048 q latent]
parameters
12,584,960 · × 75 MoE blocks = 943,872,000
Parameters

12,584,960 parameters per instance, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Query down-projection6,144×2,0486{,}144 \times 2{,}04812,582,912
Query latent RMSNorm2,0482{,}0482,048
Per-instance total12,584,960

Across 75 independent MoE blocks: 12,584,960×75=943,872,00012{,}584{,}960 \times 75 = 943{,}872{,}000 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

Two matrices, 6,144 × 2,048 and 2,048 × 16,384, hold 46,137,344 weights; one 6,144 × 16,384 matrix would hold 100,663,296.

Codemodeling_glm_moe_dsa.py 293–396
modeling_glm_moe_dsa.py 293–306
def yarn_get_mscale(scale=1, mscale=1):
    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


def yarn_apply_mscale(rope_parameters, scaling):
    if rope_parameters.get("rope_type", "default") != "default":
        mscale_all_dim = rope_parameters.get("mscale_all_dim", 0)
        scaling_factor = rope_parameters["factor"]
        if mscale_all_dim:
            mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
            scaling = scaling * mscale * mscale
    return scaling
modeling_glm_moe_dsa.py 326–360
        self.q_lora_rank = config.q_lora_rank
        self.qk_rope_head_dim = config.qk_rope_head_dim
        self.kv_lora_rank = config.kv_lora_rank
        self.v_head_dim = config.v_head_dim
        self.qk_nope_head_dim = config.qk_nope_head_dim
        self.qk_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads

        self.is_causal = True
        self.q_a_proj = nn.Linear(self.hidden_size, self.q_lora_rank, bias=config.attention_bias)
        self.q_a_layernorm = GlmMoeDsaRMSNorm(self.q_lora_rank)
        self.q_b_proj = nn.Linear(self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)

        self.kv_a_proj_with_mqa = nn.Linear(
            self.hidden_size,
            config.kv_lora_rank + config.qk_rope_head_dim,
            bias=config.attention_bias,
        )
        self.kv_a_layernorm = GlmMoeDsaRMSNorm(config.kv_lora_rank)
        self.kv_b_proj = nn.Linear(
            config.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
        )

        self.o_proj = nn.Linear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=config.attention_bias,
        )

        self.scaling = yarn_apply_mscale(config.rope_parameters, self.qk_head_dim ** (-0.5))
        # Refer: https://arxiv.org/abs/2603.12201 for more details.
        self.skip_topk = config.indexer_types[layer_idx] == "shared"
        self.indexer = None if self.skip_topk else GlmMoeDsaIndexer(config, layer_idx)
modeling_glm_moe_dsa.py 394–396
        q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states))
        q_states = self.q_b_proj(q_resid).view(query_shape).transpose(1, 2)
        q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
Sources2

Direct evidence

MLA owns a query down-projection, query RMSNorm and up-projection, a joint key–value down-projection that also emits the shared rotary key, a latent RMSNorm, a key–value up-projection and an output projection; the latent RMSNorms are built without ε and take the default 10⁻⁶; the softmax scale is 1/√256, with no YaRN correction under default RoPE; every layer builds its own indexer.

mla-init · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 293–360

Read code excerpt
def yarn_get_mscale(scale=1, mscale=1):
    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


def yarn_apply_mscale(rope_parameters, scaling):
    if rope_parameters.get("rope_type", "default") != "default":
        mscale_all_dim = rope_parameters.get("mscale_all_dim", 0)
        scaling_factor = rope_parameters["factor"]
        if mscale_all_dim:
            mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
            scaling = scaling * mscale * mscale
    return scaling

… lines 307–325 omitted …

        self.q_lora_rank = config.q_lora_rank
        self.qk_rope_head_dim = config.qk_rope_head_dim
        self.kv_lora_rank = config.kv_lora_rank
        self.v_head_dim = config.v_head_dim
        self.qk_nope_head_dim = config.qk_nope_head_dim
        self.qk_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads

        self.is_causal = True
        self.q_a_proj = nn.Linear(self.hidden_size, self.q_lora_rank, bias=config.attention_bias)
        self.q_a_layernorm = GlmMoeDsaRMSNorm(self.q_lora_rank)
        self.q_b_proj = nn.Linear(self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)

        self.kv_a_proj_with_mqa = nn.Linear(
            self.hidden_size,
            config.kv_lora_rank + config.qk_rope_head_dim,
            bias=config.attention_bias,
        )
        self.kv_a_layernorm = GlmMoeDsaRMSNorm(config.kv_lora_rank)
        self.kv_b_proj = nn.Linear(
            config.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
        )

        self.o_proj = nn.Linear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=config.attention_bias,
        )

        self.scaling = yarn_apply_mscale(config.rope_parameters, self.qk_head_dim ** (-0.5))
        # Refer: https://arxiv.org/abs/2603.12201 for more details.
        self.skip_topk = config.indexer_types[layer_idx] == "shared"
        self.indexer = None if self.skip_topk else GlmMoeDsaIndexer(config, layer_idx)

Direct evidence

Queries are produced from the normalized query latent and split into a non-rotated and a rotated part.

mla-queries · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 394–396

Read code excerpt
        q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states))
        q_states = self.q_b_proj(q_resid).view(query_shape).transpose(1, 2)
        q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)