【轻量化复刻大模型 / KIMI K3】Gated MLA 梳理

【轻量化复刻大模型 / KIMI K3】Gated MLA 梳理

目录

0x00 概述

知识点：Gated Multi-head Latent Attention（门控多头潜在注意力）— K3 混合架构中的全局注意力组件。

0.1 原始论文

论文 出处
Kimi K3: Open Frontier Intelligence（K3 技术报告，Gated MLA 最终定义 §2.1.2） arXiv:2607.24653（2026-07-27）；Kimi Team；Moonshot AI
DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model（MLA 原始提出者） arXiv:2405.04434（2024-05）；DeepSeek-AI

0.2 直接相关的图/表

K3 架构总览，MLA 与 KDA 层 3:1 交错：

Kimi K3 architecture

Gated MLA 公式（NoPE + 全秩输出门 Eq.7）：

gated mla

0.3 本文说明

本文分为两部分：第一部分是 MLA 原理/代码解读。第二部分是 MLA demo 解读（因为每一部分的文字都太少了）。

其次，关于项目：

  • 项目https://github.com/RossiLH/Model-Viz 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
  • 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
  • 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
  • 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
  • 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
  • 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本（plot_*.py），可生成静态图。
  • 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
  • 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。

0x01. 演进脉络与定位

1.1 脉络

Gated MLA 的谱系：MLA 最早由 DeepSeek-V2 提出，用"压缩-重构"把 KV cache 从 O(T·H·D) 降到 O(T·kv_lora_rank)。K3 继承 MLA 框架，做了两处关键改动使其适配 K3 混合架构。

1.2 同期竞品

方案名 核心机制 优势 劣势
标准 MHA 每 token 存 K/V 精确检索、简单 cache 随 T 线性爆炸
DeepSeek MLA latent 压缩 K/V cache 小、有 RoPE 位置参数扩展需手术
K3 Gated MLA latent 压缩 + NoPE + 输出门 cache 小 + 无位置参数 + 可门控 压缩有损、需 KDA 提供位置
GQA 共享 KV head cache 减半 压缩比有限（~8×）

1.3 定位一句话

不同于 DeepSeek-V2 仍需 RoPE 的 MLA，K3 的 Gated MLA 通过 NoPE（位置交给 KDA 递推） + 全秩输出门 实现"cache 小 + 免位置参数手术 + 逐通道可门控"的全局注意力。

0x02. 原理讲解

2.1 引入 MLA 的必要

KDA 固定大小的状态不可能无损保存无限长历史。

随着上下文不断变长，新 Token 会持续写入、覆写和冲刷旧状态。因此，在需要精确回忆某个细节的任务中，例如“needle in a haystack”任务，线性状态可能逐渐遗忘目标信息，导致精准召回能力下降。遗忘门和 Delta Rule 能改善写入策略，但它们不会让固定容量变成无限容量。

一个自然的折中方案是 Hybrid：

  • KDA 层负责低成本地聚合长历史；
  • Full Attention 或 Sparse Attention 层提供更直接的历史访问路径；
  • 两者结合，在效率和精确回忆之间取平衡。

2.1 直觉类比视角

想象你有一个智能摘要员（latent 压缩）。每来一条新信息（token），它不存全文，只记一句高度压缩的摘要（latent c）放进小本子（latent cache）。要回答问题时，再根据摘要"脑补"细节（up 投影重构 K/V）。

关键两点：

  1. 本子很小 —— 摘要只有 kv_lora_rank（如 512）维，而非原文 H·D（如 96×128）维，1M 条信息只占很小空间。
  2. 它不记位置 —— 摘要员不关心"这是第几条"（NoPE）；位置由搭档 KDA 负责（其递推衰减天然知道先后）。上下文扩到 100 万条也无需"重新编号"手术。

最后，摘要员交出答案前根据当前问题决定每块读多少（输出门 sigmoid(W_g x)）—— 这就是"Gated"。

2.3 准专业水平视角

输入输出

  • 输入：x ∈ R^{B×T×d}
  • 压缩 latent：c = RMSNorm(W_dkv x) ∈ R^{B×T×r}，r = kv_lora_rank
  • 重构：k = W_uk c、v = W_uv c（每 head 展开）

gated-mla-picture

核心公式（单 head）

q = W_upq( RMSNorm(W_dq x) )           # query（q-LoRA 分支）
c = RMSNorm( W_dkv x )                 # 唯一缓存对象
k = W_uk c ;  v = W_uv c               # 注意力时重构
o = softmax( q k^T / sqrt(d_qk) ) v    # 标准因果注意力
y = W_o [ sigmoid(W_g x) ⊙ o ]         # Eq.7 逐通道输出门
  • W_dkv ∈ R^{r×d}：压缩投影（唯一真正"压信息"处）
  • W_uk, W_uv：从 latent 展开 K/V
  • sigmoid(W_g x)：全秩门，逐通道在 (0,1)

与标准做法公式级对比

标准 MHA Gated MLA
记忆 O(T·H·D) O(T·r)
检索 精确 来自压缩 latent，有损
位置 RoPE 无（NoPE）
输出调制 无 sigmoid(W_g x) 逐通道门

2.4 专家视角

  • 为什么压缩 work：K/V 高维表示冗余大，低秩 latent 捕获主要信息是合理低秩近似。代价是逐 token 精确检索被"模糊化"——这正是 K3 用 KDA（线性注意力）+ MLA（全局注意力）混合的原因。
  • NoPE 必要性：位置由 KDA 递推衰减（recency-aware）提供。上下文到 1M 无需位置参数手术（不调 RoPE 频率基、不用 YaRN）——这是 K3 支持 1M 上下文的关键系统设计。
  • 为什么全秩门：KDA 输出门是"head-wise RMSNorm + 低秩门"；MLA 用全秩门让每 token 独立调制每个通道，表达力更强。
  • 论文未解决问题：压缩信息损失在多大程度可被"全局模糊检索"弥补？latent 秩是否需按层自适应，仍开放。

2.5 实现代码

# -*- coding: utf-8 -*-
"""Gated Multi-head Latent Attention（门控多头潜在注意力，K3 报告 §2.1.2）。

解决的问题：标准 softmax 注意力的 KV cache 随序列长度线性增长。在 1M 上下文下，
cache 显存与带宽会成为瓶颈。MLA 的核心思想是**压缩 + 重构**：把每个 token 的
key/value 先压进一个低维潜在向量 ``c = W_dkv x``（``kv_lora_rank`` 维），只在
cache 里存这个压缩后的 latent；真正参与注意力时再用学到的 up 投影把每 head 的
k/v 展开出来。于是 cache 大小由 ``kv_lora_rank`` 决定，而不是
``num_heads * head_dim``。

K3 在继承 DeepSeek-V2 的 MLA 基础上做了两处改动：

1. **NoPE（完全不施加位置编码）**：MLA 层不应用 RoPE，也没有解耦的 rope head
   维。位置信息完全由与之交错的 KDA 层提供 —— KDA 的递推衰减和 delta 更新天然
   位置敏感（recency 由衰减累积编码）。好处是：把上下文扩到 1M **不需要任何
   位置参数手术**（不用调 frequency base、不用 YaRN）。

2. **全秩输出门（Eq. 7）**：``y = W_o [ sigmoid(W_g x) * o ]``，让每个 token 按
   其输入调制它从全局注意力中读取哪些通道。这里 ``W_g`` 是**全秩**的（输出维等于
   拼接后的 v 维），而非低秩门 —— 低秩门在 KDA 上换成了 head-wise RMSNorm。

数值：训练时注意力输出在 FP32 累加（``attn_out_fp32``），用于修正 flash
attention 的偏向舍入误差（biased rounding）。

本项目实现与 ``kimi3-master/kimi3/modules/mla.py`` 逐位对齐（同公式、同投影命名、
同 cache 结构）；SGLang 在独立 ``kimi-k3`` 分支里把输出门融合成单 CUDA kernel
（``mla_output_gate.cuh``：``out = x * sigmoid(gate)``）。
"""

class MLACache:
    """解码缓存：**只存压缩后的 latent**。

    这是 MLA 的意义所在：cache 随 ``kv_lora_rank``（如 512）增长，而不是随
    ``num_heads * head_dim``（如 96*128=12288）增长。相比之下标准 attention 的
    KV cache 每个 token 都要存原始 K/V，MLA 只需存压缩表示，省下约一个数量级。

    ``kv_latent`` 形状 ``(batch, seq, kv_lora_rank)``；``append`` 沿 seq 维追加
    新 latent 并返回完整历史（供当前步展开全部 k/v）。
    """

    __slots__ = ("kv_latent", "length")

    def __init__(self, kv_latent: Tensor, length: int = 0):
        self.kv_latent = kv_latent
        self.length = length

    def append(self, latent: Tensor) -> Tensor:
        """追加 ``(batch, new, rank)``，返回完整 latent 历史。"""
        self.kv_latent = (
            latent if self.kv_latent is None else torch.cat([self.kv_latent, latent], dim=1)
        )
        self.length = self.kv_latent.shape[1]
        return self.kv_latent

    def clone(self) -> "MLACache":
        """深拷贝本 cache（投机解码 / beam search 需要在候选续写间分支状态）。

        与 ``KDACache.clone`` 语义一致：返回独立副本，改动副本不影响原对象。
        MLA 的 cache 只是 ``kv_latent`` 一个压缩 latent 张量，深拷贝即可。
        """
        return MLACache(
            None if self.kv_latent is None else self.kv_latent.clone(),
            length=self.length,
        )

    def truncate(self, length: int) -> None:
        """原地截断 latent 到 ``length`` 个位置（回滚到被接受前缀末尾）。

        投机解码部分接受时，草稿被拒绝的位置对应的 latent 必须丢弃，否则
        下一轮会把“未验证的续写”当成历史。MLA 是逐 token 追加的压缩 latent，
        直接按 seq 维切回 ``length`` 即可精确回滚。
        """
        if self.kv_latent is not None:
            self.kv_latent = self.kv_latent[:, :length]
        self.length = length

class GatedMLA(nn.Module):
    """完整的门控 MLA 层。

    参数结构（§2.1.2）：
    - ``q_down / q_norm / q_up``（q-LoRA 分支）：query 也走低秩压缩
      ``dim -> q_lora_rank -> qk_dim``；若 ``q_lora_rank is None`` 则直接用
      ``q_proj: dim -> qk_dim`` 全秩投影。
    - ``kv_down / kv_norm / k_up / v_up``：key/value 的压缩-重构。``kv_down``
      是唯一的压缩投影（``dim -> kv_lora_rank``），``k_up / v_up`` 在注意力时
      把 latent 展开成每 head 的 qk 维 / v 维。
    - ``g_proj``：**全秩输出门** ``sigmoid(W_g x)``（Eq. 7）。
    - ``o_proj``：最终输出投影 ``v_dim -> dim``。

    数据流（forward）：
        q = W_upq(RMSNorm(W_dq x))                  (b,h,t,dqk)
        c = RMSNorm(W_dkv x)                        (b,t,kv_lora)  ← 唯一 cache
        k = W_uk(c)  (b,h,t,dqk)  ;  v = W_uv(c)  (b,h,t,dv)
        o = softmax(q k^T / sqrt(dqk)) v            (b,h,t,dv)  [fp32 累加]
        y = W_o [ sigmoid(W_g x) * o ]              (b,t,dim)
    """

    def __init__(self, dim: int, cfg: MLAConfig, *, layer_idx: int = 0, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.cfg = cfg
        self.layer_idx = layer_idx
        h = cfg.num_heads
        self.num_heads = h
        self.qk_head_dim = cfg.qk_head_dim
        self.v_head_dim = cfg.v_head_dim
        # 注意力缩放因子 1/sqrt(qk_head_dim)。
        self.scale = cfg.qk_head_dim**-0.5

        # query 投影：可选 q-LoRA 低秩分支，或全秩直连。
        if cfg.q_lora_rank is None:
            self.q_proj = nn.Linear(dim, cfg.qk_dim, bias=False)
            self.q_down = self.q_up = self.q_norm = None
        else:
            self.q_proj = None
            self.q_down = nn.Linear(dim, cfg.q_lora_rank, bias=False)
            self.q_norm = RMSNorm(cfg.q_lora_rank, eps)
            self.q_up = nn.Linear(cfg.q_lora_rank, cfg.qk_dim, bias=False)

        # key/value 的压缩与重构。
        self.kv_down = nn.Linear(dim, cfg.kv_lora_rank, bias=False)  # W_dkv：压缩
        self.kv_norm = RMSNorm(cfg.kv_lora_rank, eps)
        self.k_up = nn.Linear(cfg.kv_lora_rank, cfg.qk_dim, bias=False)  # W_uk
        self.v_up = nn.Linear(cfg.kv_lora_rank, cfg.v_dim, bias=False)   # W_uv

        self.g_proj = nn.Linear(dim, cfg.v_dim, bias=False)   # 全秩门，Eq. 7
        self.o_proj = nn.Linear(cfg.v_dim, dim, bias=False)

    def _queries(self, x: Tensor) -> Tensor:
        """计算 query。q-LoRA 分支：``W_upq(RMSNorm(W_dq x))``；否则 ``W_q x``。"""
        if self.q_proj is not None:
            return self.q_proj(x)
        assert self.q_down is not None and self.q_norm is not None and self.q_up is not None
        return self.q_up(self.q_norm(self.q_down(x)))

    def forward(
        self,
        x: Tensor,
        *,
        attn_mask: Tensor | None = None,
        cache: MLACache | None = None,
    ) -> Tensor:
        b, t, _ = x.shape
        h, dqk, dv = self.num_heads, self.qk_head_dim, self.v_head_dim

        # query -> (b, h, t, dqk)。
        q = self._queries(x).view(b, t, h, dqk).transpose(1, 2)

        # 压缩 latent；有 cache 则追加历史（decode 时只传当前 token，靠 append 累积）。
        latent = self.kv_norm(self.kv_down(x))
        if cache is not None:
            latent = cache.append(latent)
        kv_len = latent.shape[1]

        # 从 latent 展开每 head 的 k/v。
        k = self.k_up(latent).view(b, kv_len, h, dqk).transpose(1, 2)
        v = self.v_up(latent).view(b, kv_len, h, dv).transpose(1, 2)

        # 训练时在 FP32 累加注意力，修正低精度舍入偏置。
        if self.cfg.attn_out_fp32 and self.training:
            q, k, v = q.float(), k.float(), v.float()

        is_causal = attn_mask is None and kv_len == t
        if hasattr(F, 'scaled_dot_product_attention'):
            out = F.scaled_dot_product_attention(
                q, k, v, attn_mask=attn_mask, is_causal=is_causal, scale=self.scale
            )
        else:
            # torch < 2.0（本仓库 CPU 环境）：手动因果 softmax 注意力。
            attn = torch.matmul(q * self.scale, k.transpose(-2, -1))
            if is_causal:
                mask = torch.triu(torch.full((t, kv_len), float('-inf'), device=x.device), diagonal=1)
                attn = attn + mask
            elif attn_mask is not None:
                attn = attn + attn_mask
            attn = F.softmax(attn, dim=-1)
            out = torch.matmul(attn, v)
        # (b,h,t,dv) -> (b,t,h*dv)，降回输入 dtype。
        out = out.transpose(1, 2).reshape(b, t, h * dv).to(x.dtype)

        # Eq. 7：依赖输入的、逐通道的、全秩输出门。
        # sigmoid(W_g x) 逐通道在 (0,1)，决定该 token 从注意力输出里读取多少。
        out = out * torch.sigmoid(self.g_proj(x))
        return self.o_proj(out)

    def init_cache(self, batch: int, *, device=None, dtype=torch.float32) -> MLACache:
        """为 decode 初始化空 latent cache：``(batch, 0, kv_lora_rank)``。"""
        return MLACache(
            torch.zeros(batch, 0, self.cfg.kv_lora_rank, device=device, dtype=dtype)
        )

    def extra_repr(self) -> str:
        return (
            f"dim={self.dim}, heads={self.num_heads}, qk={self.qk_head_dim}, "
            f"v={self.v_head_dim}, q_lora={self.cfg.q_lora_rank}, "
            f"kv_lora={self.cfg.kv_lora_rank}, nope=True"
        )

2.6 Gated 体现在哪里？

"Gated" 体现在 mla.py 三处，核心是一个全秩输出门（Eq.7）：

1. 门的参数定义（mla.py）

self.g_proj = nn.Linear(dim, cfg.v_dim, bias=False)   # 全秩门，Eq. 7
  • g_proj: dim -> v_dim，v_dim = num_heads * v_head_dim（拼接后的 v 全宽）。
  • "全秩"指它输出维等于 v 全宽（不是低秩），让每个通道都有独立的门值。

2. 门真正作用的计算（mla.py，这是"Gated"的落点）

# Eq. 7：依赖输入的、逐通道的、全秩输出门。
out = out * torch.sigmoid(self.g_proj(x))
return self.o_proj(out)
  • torch.sigmoid(self.g_proj(x)) 产生一个逐通道、值在 (0,1) 的掩码。
  • out * 掩码：逐元素乘，决定该 token 从注意力输出里读多少。这就是"门控" —— 门在"读之前/读之后"逐通道地放行或削弱全局信息。

3. 门的设计意图（模块 docstring，mla.py）

**全秩输出门（Eq. 7）**：``y = W_o [ sigmoid(W_g x) * o ]``

数据流定位：门作用在注意力输出之后、o_proj 之前：

注意力输出 o (b,t,h*dv)
      ↓
sigmoid(g_proj(x))  ← 逐通道 (0,1) 掩码，依赖输入 x
      ↓ 逐元素乘
o * sigmoid(g_proj(x))   ← 门控后的输出
      ↓
o_proj → (b,t,dim)

"Gated" vs 普通 MLA 的区别：标准 MLA（DeepSeek-V2）是 y = W_o o，直接输出；K3 的 Gated MLA 插入了一个依赖输入的逐通道乘法门 sigmoid(W_g x)。这正是类名 GatedMLA 与变量名 g_proj（g = gate）的来源。

对比 KDA：KDA 的输出门（kda.py:254）也是 out * sigmoid(g_proj(x))，但前面多了一步 HeadRMSNorm（kda.py:253）；MLA 的门则直接作用在注意力输出上。两者共用"全秩 sigmoid 门"这个思路。

0x03. 子知识点拆解

Gated MLA
    ├── 3.1 latent 压缩-重构（kv_down/norm/k_up/v_up）
    │   └── 依赖：mla.py:75-78 + mla.py:101-107
    ├── 3.2 q-LoRA 分支（或全秩 q_proj）
    │   └── 依赖：mla.py:66-73 + mla.py:83-87
    ├── 3.3 全秩输出门（Eq.7）
    │   └── 依赖：mla.py:80 + mla.py:128-130
    ├── 3.4 NoPE
    │   └── 依赖：mla.py:141
    ├── 3.5 解码缓存 MLACache（只存 latent）
    │   └── 依赖：mla.py:34-51 + mla.py:132-135
    └── 3.6 FP32 注意力累加
        └── 依赖：mla.py:109-110

代码对应

子知识点 代码位置 变体/切换入口
latent 压缩-重构 mla.py:75-78 kv_lora_rank（config.py:45）
q-LoRA 分支 mla.py:66-73, 83-87 q_lora_rank（config.py:44，None 时全秩 q_proj）
全秩输出门 mla.py:80, 128-130 无
NoPE mla.py:141 恒为 True
解码缓存 mla.py:34-51, 132-135 无
FP32 累加 mla.py:109-110 attn_out_fp32（config.py:46）

0x04 子知识点深度分析

4.0 共享抽象与设计动机

统一符号：r = kv_lora_rank，d_qk/d_v 为每 head 的 qk/v 维，H 为 head 数，d = H·d_v。压缩 latent c 是 MLA 唯一需缓存的张量。

为什么只缓存 latent：标准 MHA 缓存 (k_i,v_i) 每 token O(2H·d_qk)；MLA 缓存 c_i 每 token O(r)。以 K3 配置（H=96，d=128，r=512）为例，压缩比约 2·96·128/512 = 48×。

4.1 子知识点：latent 压缩-重构

4.1.1 原理

  • 直觉：摘要员只记摘要，回答时脑补。
  • 准专业：c = RMSNorm(W_dkv x)，k=W_uk c，v=W_uv c。这是秩 r 的低秩近似（W_uk W_dkv 有效秩 ≤ r）。
  • 专家：等价于 K/V 的低秩分解。信息论上 r 决定可保留的"自由信息"上界；r 越大越接近无损但 cache 越大。

4.1.2 对比基线

  • 标准：直接存 W_k x、W_v x（无损）。
  • 本方案：存 W_dkv x，注意力时展开。
  • 收益：cache 压缩；代价：检索有损。

4.1.3 反方辩手

  • a) 最强反对：压缩丢失逐 token 细节，长程精确检索受损。
  • b) 若失败，最可能错在：假设 K/V 冗余度足够高。
  • c) 胜出：长上下文、cache 内存敏感、容忍模糊检索；失败：需精确逐 token recall 的任务。
  • d) 更稳健：提高 r，或按层混合 r。
  • e) 受影响方：需要精确检索的任务方。

4.1.4 计算复杂度

  • FLOPs：注意力本身 O(T²·d_qk) 不变；K/V 投影从 O(T·d·H·d_qk) 降到 O(T·d·r)。
  • 内存：标准 cache O(T·2H·d_qk)，MLA O(T·r)。
  • 参数：新增 W_dkv（r×d）+ W_uk/W_uv（H·d_qk×r）。

4.1.5 参数敏感性

  • 关键超参：kv_lora_rank（r）。
  • r 增大：cache 大、检索准；r 减小：cache 小、表达被压平（实验见 §4.7 E2）。
  • 论文经验：K3 取 r=512（相对 2H·d_qk 是 48× 压缩）。

4.1.6 工程实现

伪代码（对齐 mla.py）：

latent = self.kv_norm(self.kv_down(x))        # (B,T,d)->(B,T,r) 压缩+归一化
if cache is not None: latent = cache.append(latent)  # 追加历史
k = self.k_up(latent).view(b,kv_len,h,dqk).transpose(1,2)  # 展开 K
v = self.v_up(latent).view(b,kv_len,h,dv).transpose(1,2)   # 展开 V

真实代码：mla.py:。

与 kimi3-master 核对：本项目 mla.py 与 kimi3-master/kimi3/modules/mla.py 逐位对齐（同投影命名、同 forward 结构、同 cache）。test_mla.py（增量解码 vs prefill）验证一致。

与 kimi-k3-pytorch-main 差异：kimi-k3-pytorch-main/src/mla/ 的 MLA 无 q-LoRA（query 用单一全秩 W_q），本项目 K3 有 q_down/q_up；二者 KV 压缩-重构一致。

与 SGLang 核对：SGLang K3 分支用 DeepSeek-V2 的 DeepseekV2AttentionMLA + skip_rope=True（NoPE），输出门融合成单 CUDA kernel mla_output_gate.cuh（out = x·sigmoid(gate)，bf16 双舍入位对齐）。本项目 CPU 版用 torch.sigmoid(g_proj(x))，语义一致。

4.1.7 探讨代码

见 experiments/models/llm/kimi_k3/modules/gated_mla/discussion.py：E1 验证门可关断（gate→0 时输出被抹掉到 0.0e+00）；E2 扫描 kv_lora_rank，观察输出离散度（检索丰富度）随压缩下降（r=8 时 std=0.0053，r=48 时 std=0.0215）。

4.1.8 验证 demo

见 experiments/models/llm/kimi_k3/modules/gated_mla/plot_gated_mla.py：plot_gate_suppression.png、plot_rank_scan.png、plot_cache_size.png。

可视化脚本

experiments/models/llm/kimi_k3/modules/gated_mla/plot_gated_mla.py 运行后生成三图：

  • plot_gate_suppression.png（门强度 vs 输出幅度）

    plot_gate_suppression

  • plot_rank_scan.png（kv_lora_rank vs 检索丰富度）

    plot_rank_scan

  • plot_cache_size.png（cache 大小 vs 序列长）

    plot_cache_size

截图说明

  • 门抑制图：横轴 g_proj 偏置，纵轴 MLA 输出 max|·|，S 形曲线 —— 展示门从"全开"到"全关"的抑制能力，支撑"门可关断"。

    • 布局：1×1 单图（折线图，symlog x 轴）

      X 轴：g_proj 权重偏置（-1e4 ~ 50，symlog 刻度，中心 = 0）

      Y 轴：MLA 输出 max|·|

      曲线：

      • 蓝色圆点线：偏置从 -1e4 到 50 时的输出幅度
      • 当偏置 < -10 时，输出迅速趋近 0（门完全关闭）
      • 当偏置 > 0 时，输出稳定（门完全打开）

      你该看什么：曲线的 S 形特征——左侧平台（输出≈0）到右侧平台（输出最大）的过渡。

      理想情况：偏置=-1e4 时输出 < 1e-3，偏置=50 时输出 > 0.1，中间有平滑过渡。

      异常指示：如果偏置=-1e4 时输出仍然很大，说明门没有正确乘到输出上（可能 W_g 没连接到前向路径）。

  • rank 扫描图：横轴 latent 秩，纵轴输出 std，上升曲线 —— 展示 latent 越窄信息越被压平，支撑 4.1.3 的"压缩有损"。

    • 布局：1×1 单图（折线图）

      X 轴：kv_lora_rank（8 ~ 48，线性刻度）

      Y 轴：重构输出 std（检索丰富度代理）

      曲线：

      • 紫色方点线：不同 rank 下的输出 std
      • 随 rank 增大，std 单调递增（信息越丰富）
      • 低 rank 区间（8→16）斜率陡峭（信息损失严重）
      • 高 rank 区间（24→48）斜率趋缓（边际收益递减）

      你该看什么：曲线的单调递增趋势。如果曲线不单调或有异常跳变，说明实现有 bug。

      理想情况：rank=8 时 std 最低（信息被严重压平），rank=48 时 std 最高（接近无压缩）。

      异常指示：如果 rank=8 时 std 很高（与 rank=48 相近），说明压缩没有生效（kv_down 权重可能全零）。

  • cache 大小图：双 log 轴，MLA latent cache 远低于标准 KV cache —— 展示 MLA 核心卖点。

    • 布局：1×1 单图（双对数折线图）

      X 轴：序列长度 T（512 ~ 1M，对数刻度）

      Y 轴：cache 元素数（百万，对数刻度）

      曲线：

      • 红色圆点线：标准 KV cache，斜率 = 1（线性增长）
      • 绿色方点线：MLA latent cache，斜率 = 1（线性增长但截距低约 48×）

      你该看什么：两条线的垂直间距即压缩比。在 T=1M 处，标准 KV 约 24,576M 元素，MLA 约 512M 元素。

      理想情况：两条线平行（都是 O(T)），但绿色线始终在红色线下方约 1.5 个对数单位。

      异常指示：如果两条线交叉或绿色线更高，说明实现有 bug（cache 未正确压缩）。

实测数值回填
图 关键读数 关键推论
plot_gate_suppression gate_bias=-1e4: out=0.0；0: 2.2e-02；+50: 4.3e-02 门强抑制、输入驱动
plot_rank_scan r=8: std=0.0053；r=48: std=0.0215 秩越低表达越压平
plot_cache_size T=1M: 标准KV=25770M，MLA=537M 元素 压缩约 48×（支撑 MLA 核心卖点）

4.2 子知识点：全秩输出门（Eq.7）

4.2.1 原理

  • 直觉：每个 token 像"读者"，决定从全局信息里读多少。
  • 准专业：y = W_o [ sigmoid(W_g x) ⊙ o ]。sigmoid(W_g x) 逐通道在 (0,1)（g_proj: d -> H·d_v，全秩）。
  • 专家：门是输入依赖的、逐通道的乘法调制。相比加法偏置，乘法门能更直接地"打开/关闭"某条通道。

4.2.2 对比基线

  • 标准 MLA：y = W_o o（无门）。
  • 本方案：y = W_o [ sigmoid(W_g x) ⊙ o ]。
  • 收益：逐通道可门控；
  • 代价：新增参数 W_g（H·d_v × d）与一次额外 matmul。

4.2.3 反方辩手

  • a) 最强反对：门若训练中全关，MLA 退化为直通，全局注意力失效。
  • b) 若失败，最可能错在：初始化让门过早饱和到 0。
  • c) 胜出：需要模型自适应忽略某些全局通道；失败：门饱和（梯度消失）。
  • d) 更稳健：门加偏置初始化（sigmoid 偏置为正），避免过早关断。
  • e) 受影响方：若门退化，全局注意力失效，拖累需要全局信息的任务。

4.2.4 计算复杂度

  • FLOPs：o 是 O(T·H·d_v)，门 sigmoid(W_g x) 是 O(T·d·H·d_v)（与 o_proj 同级）。
  • 内存：门与 o 同形状，无额外 cache。
  • 参数：W_g 为 H·d_v × d，与 o_proj 相同规模。

4.2.5 参数敏感性

  • 关键：W_g 的尺度决定门在 (0,1) 的分布；初始化小权重 → 门 ≈ 0.5（中性）。
  • 实测（discussion.py E1）：全 1 输入、W_g 推向 -1e4 时输出被抹到 0.0e+00；随机输入时门值范围约 [0.13, 0.86]（输入依赖、逐通道）。

4.2.6 工程实现

out = out * torch.sigmoid(self.g_proj(x))   # 全秩门，逐通道
return self.o_proj(out)

与 SGLang 核对：SGLang mla_output_gate.cuh 实现 out = x·sigmoid(gate)（bf16），与本项目 torch.sigmoid(g_proj(x)) 语义逐位一致，仅做 kernel 融合与双舍入对齐优化。

4.2.7 探讨代码

experiments/models/llm/kimi_k3/modules/gated_mla/discussion.py E1 直接操纵 g_proj 验证门能彻底关断输出（0.0223 → 0.0e+00），证明门是强抑制机制；同时展示随机输入下门值随输入变化（全秩、输入依赖）。

4.2.8 验证 demo

plot_gate_suppression.png 展示门偏置 vs 输出幅度曲线。

plot_gate_suppression

4.3 子知识点：NoPE（无位置编码）

4.3.1 原理

  • 直觉：摘要员不记"第几条"，位置交给搭档 KDA。
  • 准专业：MLA 完全不含 RoPE 参数（无 rotary_emb、无 rope head 维）。mla.py 的 extra_repr 恒标记 nope=True。
  • 专家：位置信息由 KDA 层的递推衰减（recency-aware）提供。这让 MLA 层上下文扩展无需位置参数手术，是 K3 支持 1M 上下文的关键系统协同设计。

4.3.2 对比基线

  • 标准 MLA（DeepSeek-V2）：带 RoPE（解耦 rope head 维）。
  • 本方案：NoPE。
  • 收益：免位置参数手术、cache 再省一个 rope 维；代价：MLA 层本身无位置先验，须依赖 KDA。

4.3.3 反方辩手

  • a) 最强反对：若 KDA 的递推无法可靠提供位置/顺序信息，MLA 会混淆远距离 token 的相对顺序。
  • b) 若失败，最可能错在：KDA 衰减的"recency"编码不够强以承载全部位置语义。
  • c) 胜出：长上下文、KDA 比例高（3:1）；失败：KDA 过少、需绝对位置的任务。
  • d) 更稳健：保留少量 rope 维（如 K3 用 NoPE 前几层）。但 K3 选择全 NoPE。
  • e) 受影响方：依赖精确顺序推理的任务。

4.3.4 计算复杂度

  • 相比带 RoPE 的 MLA，NoPE 省去每步 RoPE 旋转计算（O(T·H·d_qk)）与 rope 维 cache。

4.3.5 参数敏感性

  • 无位置超参；test_mla.py:50-56 断言参数/缓冲区不含 rope/freq/pos/alibi/yarn。

4.3.6 工程实现

  • mla.py 无任何 RoPE 相关层；extra_repr 输出 nope=True。
  • 与 SGLang 核对：SGLang kimi_linear.py:464 传入 skip_rope=True 关闭 RoPE，deepseek_v2.py 里 self.rotary_emb = None —— 与本项目 NoPE 完全一致。

4.4 子知识点：解码缓存 MLACache

4.4.1 原理

  • MLACache 只存压缩 latent (B, T, r)。append 沿 seq 维 torch.cat 追加当前 token 的 latent。
  • 增量解码（test_mla.py:78-85）验证：逐 token 解码拼接结果与整段 prefill 逐位一致（max diff ≈ 0）。

4.4.2 对比基线

  • 标准 MHA cache：每 token 存 K、V（O(2H·d_qk)）。
  • 本方案：每 token 只存 latent（O(r)）。
  • 收益：cache 内存 48× 压缩；代价：decode 时仍需 k_up/v_up 展开。

4.4.3 反方辩手

  • a) 最强反对：latent 无法像原始 K/V 那样精确逐 token 检索。
  • b) 若失败，错在假设 latent 含足够检索信息。
  • c) 胜出：长上下文、内存受限；失败：精确记忆任务。
  • d) 稳健做法：保留一个小的"精确 cache"作为 MLA 补充（K3 用 KDA 承担）。
  • e) 受影响方：长程精确检索任务。

4.4.4 计算复杂度

  • cache 内存：O(T·r) 每序列。
  • decode 每步：需对全量 latent 重算 k_up/v_up（无吸收优化，本项目为教学性简化）；SGLang 用 matrix absorption（w_kc/w_vc 折叠）优化。

4.4.5 参数敏感性

  • 无独立超参；随 kv_lora_rank 变化。

4.4.6 工程实现

def append(self, latent):  # mla.py
    self.kv_latent = latent if self.kv_latent is None else torch.cat([self.kv_latent, latent], dim=1)
    self.length = self.kv_latent.shape[1]
    return self.kv_latent

与 SGLang 核对：SGLang 用 MLATokenToKVPool 分页存储 latent（服务端优化），本项目的 MLACache 是教学性连续存储。语义一致。

4.5 子知识点：FP32 注意力累加

4.5.1 原理

  • attn_out_fp32：训练时把 q/k/v 升到 fp32 再做 SDPA，修正 flash attention 的偏向舍入误差。
  • 本项目 CPU 路径（torch<2.0）手动实现因果 softmax 注意力，同样支持 fp32 累加（mla.py:109-110）。

4.5.2 对比基线

  • 标准：bf16/fp16 直接算注意力。
  • 本方案：训练时 fp32 累加。
  • 收益：数值更准；代价：训练更慢、显存更高。

4.5.3 反方辩手

  • a) 最强反对：fp32 累加拖慢训练。
  • b) 失败点：精度收益 < 速度损失（在不需要高精度的场景）。
  • c) 胜出：训练稳定性敏感；失败：追求极致吞吐。
  • d) 稳健做法：仅对累加部分 fp32（如 SGLang 的做法）。
  • e) 受影响方：训练方（速度）。

4.5.4 计算复杂度

  • 注意力 FLOPs 不变，但 fp32 运算带宽更高。

4.5.5 参数敏感性

  • 超参 attn_out_fp32（config.py:46），默认 True。

4.5.6 工程实现

if self.cfg.attn_out_fp32 and self.training:
    q, k, v = q.float(), k.float(), v.float()   # mla.py

【附录】Gated MLA 模块交互式演示

0x01 实现概述

1.1 模块范围

本演示覆盖 src/models/llm/kimi_k3/modules/mla.py 中的 GatedMLA 和 MLACache 两个类，对应 K3 报告 §2.1.2 的门控多头潜在注意力。

1.2 实现思路

核心卖点是压缩 + 重构：把每个 token 的 key/value 压进低维潜在向量 c = W_dkv x（kv_lora_rank 维），cache 里只存这个 latent；真正参与注意力时再展开成每 head 的 k/v。

重构策略：plot_gated_mla.py 的三个绘图函数已重构为返回 (fig, text) 元组，Gradio demo 直接 import 这些函数，避免维护两份代码。独立实验（incremental decode、NoPE、output gate、FP32）内联定义。

1.3 目录结构

experiments/models/llm/kimi_k3/modules/gated_mla/
├── discussion.py          # 探讨实验（门可关断 + latent 压缩信息损失）
├── plot_gated_mla.py      # 可视化脚本（3 个 plot 函数，返回 fig+text）
├── gradio_mla_demo.py     # Gradio 交互式演示（7 Tabs）

1.4 运行

python experiments\models\llm\kimi_k3\modules\gated_mla\gradio_mla_demo.py

* Running on local URL:  http://0.0.0.0:7860

0x02 核心数学

2.1 总体公式

输入: x ∈ R^{B×T×dim}                    # 隐藏状态
输出: y ∈ R^{B×T×dim}                    # 注意力输出

# Query 路径（可选 q-LoRA）
q = W_upq(RMSNorm(W_dq x))              # (B,T,qk_dim) → reshape → (B,H,T,dqk)

# Key/Value 压缩-重构（MLA 核心）
c = RMSNorm(W_dkv x)                     # (B,T,kv_lora) ← 唯一 cache
k = W_uk(c)                              # (B,H,T,dqk)  从 latent 展开
v = W_uv(c)                              # (B,H,T,dv)   从 latent 展开

# Softmax 注意力
o = softmax(q k^T / √dqk) v             # (B,H,T,dv)   [FP32 累加]

# 全秩输出门 (Eq. 7)
y = W_o [ sigmoid(W_g x) * o ]          # (B,T,dim)

2.2 子机制分解

2.2.1 压缩缓存 (MLACache)

公式：c = RMSNorm(W_dkv x)，其中 W_dkv ∈ R^{kv_lora_rank × dim}

直觉解释：将高维隐藏状态压缩到低维 latent，cache 只需存储 kv_lora_rank 维而非 num_heads × head_dim 维。

关键性质：

  • Cache 大小：O(B × T × kv_lora_rank) vs 标准注意力 O(B × T × num_heads × head_dim)
  • 压缩比：2 × num_heads × head_dim / kv_lora_rank（典型 K3 配置下约 48×）
  • clone() 支持投机解码 / beam search 的状态分支
  • truncate() 支持投机解码拒绝后的精确回滚

2.2.2 全秩输出门 (Eq. 7)

公式：y = W_o [ sigmoid(W_g x) * o ]

直觉解释：每个 token 按其输入调制它从全局注意力中读取哪些通道。W_g 是全秩的（输出维 = v_dim），而非低秩门。

关键性质：

  • 值域：sigmoid(W_g x) ∈ (0, 1)，逐通道调制
  • 输入依赖：不同输入产生不同门值（token 自适应）
  • 可压制：W_g → -∞ 时 sigmoid → 0，输出被完全抹掉
  • 全秩：每个输出通道可被独立调制（vs KDA 的 head-wise RMSNorm 低秩门）

2.2.3 NoPE（无位置编码）

设计：MLA 层不应用 RoPE、ALiBi、YaRN 等位置编码。

直觉解释：位置信息完全由与之交错的 KDA 层提供。KDA 的递推衰减天然位置敏感（recency 由衰减累积编码）。

关键性质：

  • 上下文扩到 1M 不需要位置参数手术
  • MLA 层参数名白名单不含 rope/freq/pos 等关键词
  • qk_head_dim × num_heads == qk_dim 配置一致性

2.2.4 FP32 注意力累加

公式：训练时 q, k, v → float() → 注意力计算 → 输出降回输入 dtype

直觉解释：修正 flash attention 的偏向舍入误差（biased rounding）。训练时在 FP32 累加，推理时降精度。

关键性质：通过 MLAConfig.attn_out_fp32 开关控制

2.3 关键洞察

关键洞察：MLA 的 cache 只存压缩 latent c，不存展开的 k/v。decode 时每步追加一个 latent，重新展开全部历史做注意力——这保证了 incremental decode 与 prefill 的逐位一致性，同时 cache 大小从 O(H·D) 降到 O(kv_lora)。

0x03 设计思路与实现细节

3.1 复用策略

Gradio demo 直接调用 plot_gated_mla.py 的三个绘图函数（plot_gate_suppression, plot_rank_scan, plot_cache_size），避免维护两份代码。独立实验（incremental decode、NoPE、output gate、FP32）因逻辑较复杂，内联定义。

3.2 测试驱动

每个 Tab 对应 test_mla.py 中的一个或多个测试函数：

Tab 对应测试
Cache Size test_latent_cache_is_small
Gate Suppression test_gate_can_suppress_the_attention_output
Rank Scan （discussion.py E2 实验）
Incremental Decode test_incremental_decoding_matches_prefill
NoPE test_no_positional_parameters_exist
Output Gate test_output_gate_is_input_dependent_and_full_rank
FP32 Accumulation test_attention_output_kept_in_fp32_during_training

3.3 数学规律总结

  1. 压缩比：2 × H × D / kv_lora_rank，K3 典型配置下约 48×
  2. 门抑制：sigmoid 的 S 形曲线，偏置 < -10 时输出趋近 0
  3. 信息-压缩权衡：rank 增大，输出 std 单调递增，但边际收益递减
  4. 一致性：incremental decode 与 prefill 的误差 < 1e-5（float32 精度下为 0）

0x04 每个 Tab 的详细说明

Tab 1: Cache Size

目的

展示 MLA 的核心卖点：latent 压缩使 cache 大小从 O(T·num_heads·head_dim) 降到 O(T·kv_lora_rank)。对应 K3 报告中"MLA 使 1M 上下文可行"的论述。

背景概念

  • 标准 KV cache：每个 token 存 K 和 V 两个张量，大小 = T × 2 × num_heads × head_dim
  • MLA latent cache：每个 token 只存压缩 latent c，大小 = T × kv_lora_rank
  • 压缩比：2 × num_heads × head_dim / kv_lora_rank

UI 组件与参数影响

参数 范围 默认 调节效果
无参数 - - 点击按钮直接运行

输出图解读

tab-1

布局：1×1 单图（双对数折线图）

X 轴：序列长度 T（512 ~ 1M，对数刻度）

Y 轴：cache 元素数（百万，对数刻度）

曲线：

  • 红色圆点线：标准 KV cache，斜率 = 1（线性增长）
  • 绿色方点线：MLA latent cache，斜率 = 1（线性增长但截距低约 48×）

你该看什么：两条线的垂直间距即压缩比。在 T=1M 处，标准 KV 约 24,576M 元素，MLA 约 512M 元素。

理想情况：两条线平行（都是 O(T)），但绿色线始终在红色线下方约 1.5 个对数单位。

异常指示：如果两条线交叉或绿色线更高，说明实现有 bug（cache 未正确压缩）。

为什么重要

这个验证展示了 MLA 存在的根本意义——没有压缩，1M 上下文的 KV cache 需要 ~96GB 显存（fp32），而 MLA latent cache 只需 ~2GB。如果验证失败，说明 cache 压缩没有生效，MLA 层退化为标准注意力。

Tab 2: Gate Suppression

目的

验证输出门 sigmoid(W_g x) 的强抑制能力：门越负 → sigmoid→0 → 输出被完全抹掉。对应 K3 报告 Eq. 7 及 test_mla.py 的 test_gate_can_suppress_the_attention_output。

背景概念

  • 输出门：y = W_o [ sigmoid(W_g x) * o ]，W_g 是全秩投影
  • 抑制机制：训练中若门全关（W_g → -∞），MLA 层退化为恒等直通
  • 风险：门过强抑制会削弱全局信息注入

UI 组件与参数影响

参数 范围 默认 调节效果
无参数 - - 点击按钮直接运行

输出图解读

tab-2

布局：1×1 单图（折线图，symlog x 轴）

X 轴：g_proj 权重偏置（-1e4 ~ 50，symlog 刻度，中心 = 0）

Y 轴：MLA 输出 max|·|

曲线：

  • 蓝色圆点线：偏置从 -1e4 到 50 时的输出幅度
  • 当偏置 < -10 时，输出迅速趋近 0（门完全关闭）
  • 当偏置 > 0 时，输出稳定（门完全打开）

你该看什么：曲线的 S 形特征——左侧平台（输出≈0）到右侧平台（输出最大）的过渡。

理想情况：偏置=-1e4 时输出 < 1e-3，偏置=50 时输出 > 0.1，中间有平滑过渡。

异常指示：如果偏置=-1e4 时输出仍然很大，说明门没有正确乘到输出上（可能 W_g 没连接到前向路径）。

为什么重要

输出门是 K3 MLA 的关键创新——它让每个 token 能独立决定"从注意力读多少"。如果门不能有效抑制输出，训练时模型无法关闭不需要的注意力路径，Eq. 7 的调制语义失效。

Tab 3: Rank Scan

目的

扫描不同 kv_lora_rank，观察重构输出的离散度（std，作为检索丰富度的代理指标）。验证 latent 压缩的信息损失：rank 越低 → 信息被压平 → 检索越粗糙。对应 discussion.py 的 E2 实验。

背景概念

  • kv_lora_rank：压缩维度，决定 latent 的容量
  • 检索丰富度：用重构输出的 std 代理——std 越低说明信息越被压平
  • 权衡：低 rank → 省 cache → 信息损失；高 rank → 费 cache → 信息保真

UI 组件与参数影响

参数 范围 默认 调节效果
无参数 - - 点击按钮直接运行

输出图解读

tab-3

布局：1×1 单图（折线图）

X 轴：kv_lora_rank（8 ~ 48，线性刻度）

Y 轴：重构输出 std（检索丰富度代理）

曲线：

  • 紫色方点线：不同 rank 下的输出 std
  • 随 rank 增大，std 单调递增（信息越丰富）
  • 低 rank 区间（8→16）斜率陡峭（信息损失严重）
  • 高 rank 区间（24→48）斜率趋缓（边际收益递减）

你该看什么：曲线的单调递增趋势。如果曲线不单调或有异常跳变，说明实现有 bug。

理想情况：rank=8 时 std 最低（信息被严重压平），rank=48 时 std 最高（接近无压缩）。

异常指示：如果 rank=8 时 std 很高（与 rank=48 相近），说明压缩没有生效（kv_down 权重可能全零）。

为什么重要

这个实验量化了 MLA 压缩的代价：省 cache 不是免费的，信息损失是固有的。K3 选择 kv_lora_rank=512 是在 cache 节省和信息保真之间的平衡点。如果验证失败（std 不随 rank 变化），说明 latent 压缩机制有问题。

Tab 4: Incremental Decode

目的

验证 incremental decode（逐 token 追加 latent 到 cache）与 prefill（一次性整段输入）的输出逐位一致。对应 test_mla.py 的 test_incremental_decoding_matches_prefill。

背景概念

  • Prefill：一次性输入整段序列，得到完整输出
  • Incremental Decode：从空 cache 开始，每步输入一个 token，追加 latent 到 cache，重新展开全部历史做注意力
  • 一致性要求：两种方式的输出必须逐位一致（atol < 1e-5）

UI 组件与参数影响

参数 范围 默认 调节效果
无参数 - - 点击按钮直接运行

输出图解读

tab-4

布局：1×2 双联图

左图：逐 token 误差柱状图

  • X 轴：Token 位置（0 ~ 9）
  • Y 轴：abs 误差（head 0, dim 0）
  • 所有柱应接近 0（误差 < 1e-5）

右图：误差热图

  • 颜色越深（红/黑）= 误差越大
  • 理想情况：全图为均匀深色（误差接近 0）

你该看什么：所有 token 位置的误差都应该 < 1e-5。如果某个位置误差突然增大，说明 cache 追加顺序有误。

理想情况：max abs 误差 < 1e-5，所有柱几乎不可见。

异常指示：如果 max 误差 > 1e-3，说明 incremental decode 与 prefill 不一致——可能是 cache.append 逻辑有误，或 latent 展开顺序错了。

为什么重要

这是 MLA 缓存机制正确性的关键测试。自回归推理必须保证"每一步看到的注意力视野"与"prefill 一次性算完对应前缀"完全一致。若不一致，decode 阶段的输出分布偏离训练时分布，模型质量严重下降。

Tab 5: NoPE

目的

验证 MLA 层不含任何位置编码参数（NoPE 设计）。对应 test_mla.py 的 test_no_positional_parameters_exist。

背景概念

  • RoPE：旋转位置编码（Rotary Position Embedding）
  • ALiBi：线性偏置注意力
  • YaRN：扩展 RoPE 的方法
  • NoPE：MLA 层完全不使用位置编码，位置信息由交错的 KDA 层提供

UI 组件与参数影响

参数 范围 默认 调节效果
无参数 - - 点击按钮直接运行

输出图解读

tab-5

布局：1×2 双联图

左图：参数组分布饼图

  • 展示 MLA 层各投影组（q_down/q_up/kv_down/k_up/v_up/g_proj/o_proj）的参数占比
  • 没有 rope/freq/pos 等位置相关组

右图：白名单检查柱状图

  • X 轴：禁用词（rope/freq/inv_freq/pos/alibi/yarn）
  • Y 轴：匹配参数数
  • 所有柱应为 0

你该看什么：所有白名单检查柱都为 0。如果有任何非零，说明误加了位置编码参数。

理想情况：6 个检查项全部为 0，参数分布合理。

异常指示：如果 "pos" 或 "rope" 柱非零，说明有人在 MLA 层加了位置编码——这违背了 K3 的 NoPE 设计。

为什么重要

NoPE 是 K3 与 DeepSeek 路线的核心分歧点。K3 的 MLA 不做位置编码，位置信息全交给 KDA。若有人误加了 RoPE，会导致：(1) 位置信息冗余（KDA 和 MLA 都编码位置）；(2) 1M 上下文扩展需要位置参数手术（调 frequency base），违背 NoPE 的设计初衷。

Tab 6: Output Gate

目的

验证输出门（Eq. 7）的三个关键性质：输入依赖、全秩、可压制输出。综合 test_mla.py 的 test_output_gate_is_input_dependent_and_full_rank 和 test_gate_can_suppress_the_attention_output。

背景概念

  • 输入依赖：门值随输入变化（token 自适应），而非常数
  • 全秩：g_proj.weight.shape = (v_dim, dim)，输出维 = num_heads × v_head_dim
  • 可压制：W_g → -∞ 时 sigmoid → 0，输出完全归零

UI 组件与参数影响

参数 范围 默认 调节效果
无参数 - - 点击按钮直接运行

输出图解读

tab-6

布局：1×3 三联图

左图：输入依赖性

  • 两个不同输入（x1, x2）的门值分布直方图
  • 蓝色 = x1，橙色 = x2
  • 两个分布应有明显差异（证明输入依赖）

中图：门值范围

  • 大量随机输入的 sigmoid(W_g x) 分布
  • 值域 (0, 1)，红线 = 0（完全抑制），蓝线 = 1（完全通过）
  • 分布应集中在 (0, 1) 内

右图：压制能力

  • 与 Gate Suppression Tab 类似，但更聚焦于门的极端行为
  • 偏置→-∞ 时输出→0，偏置→+∞ 时输出稳定

你该看什么：

  • 左图：两个分布不重叠（输入依赖）
  • 中图：分布完全在 (0, 1) 内（sigmoid 值域正确）
  • 右图：偏置=-1e4 时输出 < 1e-3（可压制）

理想情况：三个子图都符合预期，文本输出确认三个性质全部通过。

异常指示：如果左图两个分布完全重叠，说明门不依赖输入（可能 W_g 全零）；如果中图分布超出 (0, 1)，说明 sigmoid 没有正确应用。

为什么重要

全秩输出门是 K3 MLA 的关键创新——它让每个 token 能在所有输出通道上独立决定"从注意力读多少"。若门是低秩的，某些通道永远无法被独立调制，表达能力受限。若门不能压制输出，训练时模型无法关闭不需要的注意力路径。

Tab 7: FP32 Accumulation

目的

验证 MLAConfig.attn_out_fp32 配置：训练时注意力输出在 FP32 累加，修正 flash attention 的偏向舍入误差。对应 test_mla.py 的 test_attention_output_kept_in_fp32_during_training。

背景概念

  • 偏向舍入误差：flash attention 在 FP16/BF16 下累加时存在系统性偏向
  • FP32 累加：训练时将 q/k/v 升级为 FP32 计算注意力，输出降回输入 dtype
  • 配置开关：MLAConfig.attn_out_fp32=True 启用，推理时通常关闭

UI 组件与参数影响

参数 范围 默认 调节效果
无参数 - - 点击按钮直接运行

输出图解读

tab-7

布局：1×2 双联图

左图：输出统计对比

  • 两根柱子：FP32 累加 vs 默认
  • 展示输出 mean 值
  • dtype 标注（FP32 累加应为 float32，默认也应为 float32，因为输入是 float32）

右图：输出分布对比

  • 两个直方图叠加：FP32 累加（绿色）vs 默认（红色）
  • 两者分布应非常接近（因为输入是 FP32，配置差异主要影响内部计算精度）

你该看什么：

  • 两个分布的形状应相似（主要差异在内部累加精度，不在最终输出 dtype）
  • FP32 累加的输出可能有更小的数值偏差

理想情况：两个分布重叠度高，max abs 差异 < 1e-5。

异常指示：如果两个分布差异很大，说明 FP32 累加开关没有正确影响计算路径（可能 q/k/v 的 float() 转换没有生效）。

为什么重要

FP32 累加是训练稳定性的重要保障。如果训练时悄悄降到 FP16/BF16，误差特性改变，可能导致训练不稳定或收敛变慢。该测试验证配置开关确实通过前向路径生效，而不是写死在代码里。

TransFormer-封面

0xFF 参考

posted @ 2026-11-09 20:28  罗西的思考  阅读(15)  评论(0)    收藏  举报