【轻量化复刻大模型 / KIMI K3】KDA 模块交互式演示

罗西的思考2026-11-04 21:25

【轻量化复刻大模型 / KIMI K3】KDA 模块交互式演示

0x00 文件信息

项目 值
文件 gradio_kda_demo.py
依赖 Gradio 6, PyTorch, Matplotlib, NumPy
运行 python experiments\models\llm\kimi_k3\modules\kda\gradio_kda_demo.py
端口 http://localhost:7861

注：

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

0x01 KDA核心数学

KDA (Kimi Delta Attention) 是 K3 的线性注意力机制，核心思想是用固定大小状态矩阵 + delta‑rule 纠错 + 有界衰减实现O(1) 状态、O(T) 并行的序列建模。

1.1 ShortConv: 因果 depthwise 短卷积

K3 的 MLA 层是 NoPE (无位置编码)，ShortConv 给每个通道注入局部位置信息：

css 复制代码 
out[t, c] = Σ_{lag=0}^{K‑1} w[c, θ, lag] · x[t − lag, c] + bias[c]
  • 因果性：只看当前 + 过去 K‑1 个 token (左侧补零)
  • 恒等初始化：weight[:, θ, ‑1] = 1, 其余为 0 → 初始为 no‑op, 不扰动 KDA 衰减动力学
  • 单步 decode: 维护 K‑1 长度的滑动窗口, O(K) 逐 token 更新

1.2 SiTU‑GLU: 有界激活函数 (Eq. 12)

scss 复制代码 
softcap(x, β) = β · tanh(x / β)
SiTU‑GLU(g, u) = softcap(g, β_g) · σ(g) · softcap(u, β_u)
  • 有界性 : |SiTU‑GLU| ≤ β_g · β_u = 4 × 25 = 100 < FP8_E4M3_MAX = 448
  • 饱和区梯度 : softcap 用 tanh 截断, 饱和区仍有非零梯度 (tanh' > 0), 不会像硬 clamp 那样掐断训练
  • 对比 SwiGLU : SwiGLU(g, u) = silu(g) · u, 无界, 仅作消融基线

1.3 KDA 递推方程 (Eq. 1)

KDA 维护一个固定大小状态矩阵 S ∈ R^{Dk×Dv}, 逐 token 更新：

arduino 复制代码 
状态递推 (Delta Rule) :
 S_t = Diag(a_t) · S_{t‑1} + β_t · k_t · (v_t − S_{t‑1}^T k_t)^T
读出：
 o_t = S_t^T · q_t

其中：

  • q, k ∈ R^{Dk}: 查询和键, L2 归一化
  • v ∈ R^{Dv}: 值
  • a_t = exp(g_t) ∈ (0, 1]: 逐 channel 衰减因子 (保留旧记忆)
  • β_t ∈ (0, 1): 写入强度 (sigmoid 输出)
  • v_t − S_{t‑1}^T k_t: delta‑rule 纠错项---只写入"当前检索与目标的误差", 而非盲目覆盖

delta‑rule 的直觉: 当同一 key 反复出现、value 逐渐改变时, S 中该 key 对应的检索值会收敛到最新 value (而非被覆盖或遗忘)。

1.4 有界衰减 (Eq. 5)

ini 复制代码 
g = g_min · sigmoid(exp(A_h)) · z    z = W_a^up · W_a^down · x + b_a
  • g ∈ (g_min, 0), 严格有下界 → a = exp(g) ∈ (exp(g_min), 1)
  • 为什么必须有界 : 无界门控 (如 softplus) 在 chunkwise 并行中, 累积保留因子 1/γ = exp(‑Σg) 可能溢出到 NaN
  • safe_chunk_size(g_min, dtype) = ⌊log(dtype_max / headroom) / (‑g_min)⌋, 保证数值安全

1.5 分块并行 (Eqs. 3‑4)

块内并行、块间串行, 用 UT 变换 (三角求解) 实现块内 O(T·D²) 并行：

ini 复制代码 
块内: O = (Γ·Q)·S_θ + A·Ṽ    A = Tril[(Q·Γ)(K/Γ)^T]
UT变换: Ṽ = (I+M)^{‑1}·Diag(β)·V    M = tril(Diag(β)·K·K^T, −1)
块末:: S_C = Diag(γ^{<1→C>})·S_θ + Σ_r (γ^{<r+1→C>})·k_r·Ṽ_r^T

1.6 KCP 前缀扫描 (§5.1.2)

KDA 状态跨 segment 不可加 (S_{seg1} + S_{seg2} ≠ S_{full}), 必须用前缀扫描重组：

ini 复制代码 
S_running = M_seg · S_running + S̃_seg

通信量固定 O(Dk·(Dk+Dv)), 与序列长度无关 (对比 softmax KV 的 O(T·Dk·Dv))。

1.7 KimiDeltaAttention 完整层

scss 复制代码 
q, k, v = ShortConv(W_{q/k/v} · x)
q = L2Norm(SiLU(q)),  k = L2Norm(SiLU(k)),  v = SiLU(v)
g = g_min · sigmoid(exp(A_h) · (W_a^up · W_a^down · x + b_a))   # 有界衰减
β = sigmoid(W_β · x)                                            # 写入强度
(o, S) = kda_recurrent / kda_chunkwise / kda_step               # 核心递推
output = W_o · [sigmoid(W_g · x) ⊙ HeadRMSNorm(o)]              # 输出门

1.8 3:1 KDA/MLA 调度

K3 每 4 层中 3 层用 KDA (线性注意力, O(1) 状态) + 1 层用 MLA (softmax 注意力, 压缩 latent):

yaml 复制代码 
层 0: KDA  层 1: MLA  层 2: KDA  层 3: MLA   ...

KDA 层负责长程线性回忆 (固定状态), MLA 层负责精确近程检索 (KV cache)。

0x02 KDA 实现了哪些功能

2.1 覆盖的源模块 (5 个)

源文件 模块职责 导入的关键符号
short_conv.py 因果 depthwise 短卷积, 给 NoPE 架构注入局部位置信息 ShortConv
situ_glu.py SiTU‑GLU 有界激活函数, FP8 数值安全 SiTUGLU, SwiGLU, situ_glu, softcap
kda_ref.py KDA 参考实现: 递推 ground truth、单步 decode、分块并行 kda_recurrent, kda_step, kda_chunkwise, safe_chunk_size, max_reciprocal_decay
kda.py KDA 生产模块: 完整 KimiDeltaAttention (投影+ShortConv+递推+输出门) KimiDeltaAttention, KDACache, build_kda_diagnostics
model.py Kimi3Block / Kimi3Model: 3:1 KDA/MLA 交替调度 Kimi3Block, Kimi3Model

2.2 覆盖的测试（90 个测试函数）

测试文件 测试数 覆盖的 Tab
test_short_conv.py 7 Tab 1, 2, 3
test_situ_glu.py 11 Tab 4, 5, 6
test_kda.py 16 Tab 5--12
test_k3_model.py 4 Tab 13
其他模块测试 ~52 间接覆盖

2.3 15 个 Tab 功能清单

# 区 Tab 名 控件 产出
1 ShortConv 感受野 vs 延迟误差 kernel_size 滑块 折线图：kernel 宽度 vs 延迟任务重建 MSE
2 ShortConv 恒等初始化 按钮 训练曲线：恒等 vs 随机初始化收敛轨迹
3 ShortConv 因果性 按钮 折线图：扰动前后输出对比
4 SiTU-GLU 激活曲线对比 x_min/x_max 滑块 折线图：SiTU-GLU vs SwiGLU 有界性
5 SiTU-GLU 饱和区梯度 按钮 折线图：softcap 梯度 vs 硬 clamp 梯度
6 SiTU-GLU 离群点回归 按钮 训练曲线：SwiGLU vs SiTU-GLU 含离群点 MSE
7 KDA Ref KDA 原理 (4格) seed/T/g_min 滑块 2×2 图：S矩阵热图、单步写入、delta rule 纠错、有界门控
8 KDA Ref 递推 vs 分块等价 chunk_size 下拉 3 子图：输出对比、逐位误差、状态差异热图
9 KDA Ref 有界 vs 无界衰减 按钮 2×2 图：1/gamma 上界、实测溢出、safe_chunk_size、失败 case
10 KDA Ref v1 vs v2 门控 A_log 范围滑块 3 子图：门控曲线、alpha 范围、数值差异
11 KDA Ref KCP 前缀扫描 按钮 2 子图：直接相加 vs KCP 误差、通信量对比
12 KDA Module 形状与缓存 B/H/Dk/Dv 滑块 饼图+条形图：参数分布、输入输出形状
13 KDA Module 单步匹配递推 按钮 3 子图：输出对比、逐位误差、状态差异热图
14 KDA Module 诊断指标 按钮 4 子图：alpha 统计、状态范数、输出统计、梯度流
15 Model 3:1 调度可视化 n_layers 滑块 调度条+前向验证

0x03 演示实现思路

3.1 整体架构

scss 复制代码 
gradio_kda_demo.py
├── sys.path 配置（向上 5 层找到 src/）
├── 源模块导入（kda_ref / kda / short_conv / situ_glu / model）
├── 辅助函数（_make_kda_inputs / _make_delayed_task / _fit_conv）
├── 13 个 Tab 函数（每个返回 Tuple[Figure, str]）
└── build_demo() + __main__（Gradio 界面构建 + 启动）

3.2 核心设计决策

决策 1：所有函数返回 (fig, text) 元组

每个 Tab 函数的签名统一为 func(...) -> Tuple[plt.Figure, str]：

  • fig：Matplotlib Figure 对象，由 Gradio 的 gr.Plot 组件原生渲染
  • text：结论字符串，由 gr.Textbox 组件显示

这样做的好处：

  • Gradio 6 的 gr.Plot.postprocess 自动处理 Figure→base64 转换
  • 不需要手动 plt.savefig 或 fig2arr
  • 函数可以独立测试（不依赖 Gradio 运行时）

决策 2：直接调用源模块，不复制代码

所有可视化函数直接 import 并调用 src/models/llm/kimi_k3/modules/ 下的源模块：

  • kda_recurrent / kda_step / kda_chunkwise 提供真实计算
  • KimiDeltaAttention 提供完整的生产层
  • ShortConv / SiTUGLU 提供组件级验证

唯一"复制"的代码是 Tab 8 (v1 vs v2) 中的手写单步 KDA------因为 kda_ref.py 的 kda_step 只支持 v2，需要手动实现 v1 的 softplus 门控路径来对比。

决策 3：Gradio 独有实验内联定义

13 个 Tab 中有 7 个是 Gradio 独有的（不在现有 plot 脚本中）：

  • Tab 5: KDA 原理 4 格------组合了 S 热图、单步拆解、delta rule、门控
  • Tab 6: 递推 vs 分块------直接对比两种计算路径
  • Tab 8: v1 vs v2 门控------手动实现两种门控做数值对比
  • Tab 10: 模块形状------实时构造 Kimi3Block 并统计参数
  • Tab 11: 单步匹配------逐 token 调用 kda_step 验证与递推一致
  • Tab 12: 诊断指标------调用 build_kda_diagnostics 展示运行时指标
  • Tab 13: 3:1 调度------实时构造 Kimi3Model 并可视化层调度

这些实验没有现成的 plot 脚本，因此直接在 gradio 文件中定义。

3.3 测试→Tab 映射逻辑

每个测试函数都被映射到最能展示其验证内容的 Tab：

映射原则：

  1. 形状/缓存测试 → Tab 10（模块形状与缓存）
  2. 数值等价性测试 → Tab 6（递推 vs 分块）或 Tab 11（单步匹配）
  3. 机制验证测试 → 对应原理 Tab（如 KCP→Tab 9，衰减→Tab 7）
  4. 模型集成测试 → Tab 13（3:1 调度）
  5. 梯度测试 → Tab 12（诊断指标）

示例映射：

复制代码 
test_shape_and_no_bias              → Tab 1（感受野）
test_identity_initialization        → Tab 2（恒等初始化）
test_causality                      → Tab 3（因果性）
test_output_is_bounded              → Tab 4（激活曲线）
test_gradients_survive_saturation   → Tab 5（饱和区梯度）
test_swiglu_baseline_explodes       → Tab 6（离群点回归）
test_chunkwise_equals_recurrence    → Tab 8（递推 vs 分块）
test_step_matches_sequence          → Tab 13（单步匹配递推）
test_kcp_prefix_scan                → Tab 11（KCP 前缀扫描）
test_module_forward_shapes          → Tab 12（形状与缓存）
test_unbounded_decay_would_overflow → Tab 9（有界 vs 无界衰减）
test_kimi3_model                    → Tab 15（3:1 调度）

3.4 环境与路径

ini 复制代码 
# 向上 5 层找到项目根目录（文件在 experiments/models/llm/kimi_k3/modules/）
_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
_p = _THIS_DIR
for _ in range(5):
    _p = os.path.dirname(_p)
_SRC = os.path.join(_p, "src")  # → Model-Viz/src/
sys.path.insert(0, _SRC)

注意：向上层数取决于文件深度。linear_attn 的 demo 只需 4 层，kimi_k3 需要 5 层。

0x04 每个 Tab 的使用说明

Tab 1: ShortConv 感受野 vs 延迟误差

目的 验证 ShortConv 的局部感受野不足以捕获长程依赖 (论文 §2.1.1)。当 kernel_size < delay 时, MSE ≈ 信号方差 (任务失败)。

实现：

  1. 生成"延迟线"任务：输入是随机正弦信号，目标是延迟 D 步后的副本
  2. 训练不同 kernel 宽度（2, 4, 8, 12, 20）的 ShortConv 拟合该任务
  3. 绘制 kernel 宽度 vs 最终 MSE 的折线图
  4. 标注 delay=8 的垂直线------kernel <= delay 时 MSE ≈ 信号方差（失败）

关键代码路径：

scss 复制代码 
_make_delayed_task() → 生成 (x, y)
_fit_conv()          → Adam 训练 3000 步，返回最终 MSE
tab_receptive_field() → 绘图 + 返回结论

UI 组件与参数影响

参数 范围 默认 调节效果
kernel_size 2‑20 4 当前选择的 kernel 宽度, 绿色方块标记在折线图上

输出图解读

tab-1

  • 红色折线图 : X 轴 = kernel_size, Y 轴 = 延迟任务重建 MSE (log scale)
  • 灰色虚线: delay=8 的位置 ---kernel ≤ delay 时 MSE 高 (无法捕获)
  • 绿色方块: 当前滑块选择的 kernel_size 对应的 MSE
  • 灰色水平点线: MSE=1.0 (信号方差基线)
  • 你该看什么: kernel 从 2 增大到 20 时, MSE 应先高后低, 在 kernel ≈ delay 处有拐点

进一步解释：

拐点不是几何拐点，而是对数轴上 MSE 下降速率变化最大处 。实测 MSE 序列 1.5e-2, 9.5e-3, 1.9e-3, 2.5e-4, 8.8e-5（kernel = 2, 4, 8, 12, 20），相邻下降倍数 1.6x, 4.9x, 7.7x, 2.8x → 最大下降发生在 kernel=8→12 区间，故拐点在 kernel=12。

为什么重要

ShortConv 是 K3 中 NoPE 位置编码的替代方案, kernel_size 决定了它能捕获多长的局部模式。此 Tab 量化展示"局部感受野的局限"---delay=8 的任务需要 kernel 至少 9 才能重建。

Tab 2: ShortConv 恒等初始化

目的：展示恒等初始化 (conv 初始为 no‑op) 不扰动 KDA 衰减动力学, 训练起步更平稳。

实现：

  1. 恒等初始化：reset_parameters 将 weight 最后一个 tap 设为 1，其余为 0
  2. 随机初始化：weight.normal_(std=0.5)
  3. 两种初始化各训练 1500 步，记录每 100 步的 MSE
  4. 绘制训练曲线对比

输出图解读

tab-2

  • 绿色圆点线: 恒等初始化训练曲线 (初始为 no‑op, 起步 MSE 低)
  • 红色方块线: 随机初始化训练曲线 (初始扰动大, 起步更抖)
  • Y 轴 log scale: 让训练后期的细微差异可见

为什么重要

KDA 的衰减动力学依赖 ShortConv 不在初始阶段引入大幅扰动。恒等初始化保证 conv(x) ≈ x, 不干扰 KDA 的状态矩阵 S 的收敛轨迹。

Tab 3: ShortConv 因果性

目的

验证 ShortConv 只看当前+过去 token, 修改未来输入不影响过去输出 (论文 §2.1.1)。

tab-3

输出图解读

  • 蓝色圆点线: 原始输出
  • 橙色方块线: 未来被扰动后的输出 (t≥10 处输入 +100)
  • 灰色虚线: 扰动起点 (第 10 位)
  • 你该看什么: 两条线在 t<10 处几乎完全重合 (因果性成立), t≥10 处分开 (扰动生效)

进一步解释：

  • 扰动 = 把输入 x 在 t=10..19 这 10 个时间步上整体 +100（代码 xp[:, 10:] += 100.0），前 10 步保持原样。
  • 扰动起点 = 被修改的第一个时间步索引，即 t=10。图中 axvline(9) 画在索引 9（未扰动段右边界）作为视觉分隔线。文字已重写说明。

为什么重要

因果性是自回归生成的硬性要求---decode 时未来 token 尚未生成, 卷积绝不能偷看。此 Tab 直接可视化"修改未来不影响过去"。

Tab 4: SiTU‑GLU 激活曲线对比

目的

展示 SiTU‑GLU 的有界性 (|f(x)| ≤ β_g·β_u = 100), 对比 SwiGLU 的无界增长。

实现：

  1. 在 x_min, x_max 范围内计算 SiTU-GLU 和 SwiGLU 的激活值
  2. 绘制两条曲线 + 上界水平线
  3. 输出最大绝对值对比

UI 组件与参数影响

参数 范围 默认 调节效果
x_min ‑50‑0 ‑20 展示范围左端。减小 → 看到更极端负值的 SwiGLU 无界增长
x_max 0‑50 20 展示范围右端。增大 → 看到更极端正值的 SwiGLU 无界增长

输出图解读

tab-4

  • 红色线: SwiGLU (无界, |x| 大时线性增长)
  • 绿色线: SiTU‑GLU (有界, 被 ±100 水平线截住)
  • 灰色水平虚线: 上界 ±100

调节建议

  • 想看 SwiGLU 溢出幅度 → 增大 x_max 到 40, SwiGLU 会远超 100
  • 想看 SiTU‑GLU 饱和平滑性 → 增大 x_max 到 50, SiTU‑GLU 平稳趋近 100

Tab 5: SiTU‑GLU 饱和区梯度

目的

展示 softcap 在饱和区保持非零梯度 (优于硬 clamp), 不会掐断训练。

实现：

  1. 计算 softcap(z, β) 的梯度 d/dz
  2. 计算 clamp(z, max=β) 的梯度 d/dz
  3. 对比：softcap 梯度在饱和区 > 0，硬 clamp 梯度 = 0

输出图解读

tab-5

  • 绿色线: softcap 梯度 d/dz (饱和区仍 > 0, 光滑衰减)
  • 红色线: 硬 clamp 梯度 d/dz (饱和区 = 0, 梯度消失)
  • 你该看什么 : |z| > β=4 时, softcap 梯度虽小但非零, clamp 梯度精确为 0

为什么重要

FP8 训练中激活值经常进入饱和区。如果梯度在那里消失 (硬 clamp), 对应参数就停止学习。softcap 的 tanh 截断保证梯度永远非零, 饱和区的参数仍能被微调。

Tab 6: SiTU‑GLU 离群点回归

目的

展示 SiTU‑GLU 对大离群激活的保守行为---有界截断使 loss 略高但数值稳定, SwiGLU 无界通过但低精度下溢出。

输出图解读

tab-6

左图 MSE 曲线，右图离群样本预测散点（直观展示 SwiGLU 过冲、SiTU-GLU 光滑截平）。

右图直观展示：

  • SwiGLU（红圆） ：预测范围 21, 178，震荡剧烈------同一个真实 y≈200，预测从 21 到 178 跨度 157
  • SiTU-GLU（绿方） ：预测范围 75, 194，更集中在 100 附近（截平线的存在让它把"做不到的"压缩到上界）

虽然两者都拟合不到 200，但 SiTU-GLU 的预测方差更小，离 y=x 这条完美线整体偏差反而更近。

右图亮点 ：直观展示 SwiGLU 预测震荡剧烈 （21,178，跨度 157），SiTU-GLU 预测更集中 （75,194，跨度 119），且整体离 y=x 完美线更近。绿色截平线 beta_g*beta_u=100 也画出来了。

你该看什么 ：左图中，绿色（SiTU-GLU）曲线在红色（SwiGLU）下方 ------SiTU-GLU 的 MSE 反而更低。这不是 bug，而是 SiTU-GLU 设计的真实优势：

数学事实

  1. SiTU-GLU 的上界是 beta_g × beta_u = 4 × 25 = 100（不是 25），所以 200 离群点确实超出范围（被截到 ~100）
  2. 但 SwiGLU 无界，会去"硬扛" 200 离群点，梯度被离群点拉爆，连带把 98% 正常样本（值域 ±3）的拟合也带歪
  3. SiTU-GLU 在 100 处光滑截平 （tanh' 在饱和区仍非零），相当于对离群点自动做梯度裁剪，让优化器专注拟合正常样本

数值分解（训练 1800 步后）

正常样本 MSE (98%) 离群样本 MSE (2%) 总 MSE
SwiGLU 120.05 7647.0 236.19
SiTU-GLU 124.08 6168.2 223.43

SiTU-GLU 赢在两头：

  • 离群点拟合误差更小（6168 vs 7647）：因为它只到 100，反而离 200 比 SwiGLU 的"过冲 + 震荡"更近
  • 正常样本拟合误差几乎相同（124 vs 120）：截平没伤害正常区间

物理意义

  • FP8 安全：SiTU-GLU 激活 ≤100 < FP8 E4M3 上限 448，不会溢出
  • 鲁棒性：自动抑制离群点对梯度的主导，等价于内置 gradient clipping
  • 不是"有界的代价" ：在这个任务上 SiTU-GLU 反而更优，因为它防止了 SwiGLU 的"离群点过拟合 + 正常样本被牺牲"

为什么重要

真实训练中 2% 的离群激活可能把 SwiGLU 推到 FP8 溢出 (>448), 导致 NaN。SiTU‑GLU 用 softcap 截断到 100 以内, 牺牲少量拟合精度换取数值安全。

Tab 7: KDA 原理 (4格)

目的

展示 KDA 的三个核心概念: 固定大小记忆 (S 矩阵)、delta‑rule 纠错、有界衰减门控。

实现（2×2 布局）：

  • 左上 ：调用 kda_recurrent 拿最终状态 S，绘制热图
  • 右上 ：手动构造 t=0 状态，展示单步写入量 β·k·(v - Sᵀk)ᵀ
  • 左下：构造"同一 key 反复出现、v 逐渐改变"的场景，展示检索值收敛
  • 右下 ：绘制 g = g_min·sigmoid(z) 的门控曲线

UI 组件与参数影响

参数 范围 默认 调节效果
seed 0‑100 7 随机种子, 改变 q/k/v 的具体取值
T 8‑128 32 序列长度。增大 → S 矩阵更丰富, 单步写入更小
g_min ‑10.0‑‑1.0 ‑5.0 衰减下界。增大 (趋近 0) → 衰减更慢, 旧记忆保留更久

输出图解读

tab-7

2×2 布局：

左上: 状态矩阵 S 热图

  • RdBu_r 色图, 展示 S ∈ R^{Dk×Dv} 的最终状态
  • 你该看什么: S 是固定大小 (不随 T 增长), 这是 KDA O(1) 状态的核心

右上 ------ 单步三件事拆解：每来一个 token，状态只做三件事：

  • 衰减 Diag(alpha) S：旧记忆先按通道乘一个 alpha=exp(g)<=1 的保留因子（旧信息"忘掉一点"）；
  • 检索 S^T k：用当前 key 在记忆里"查一下"已存了什么；
  • 写入 beta * k (v - S^T k)^T：只把误差 v - 检索到的记忆 写进状态（delta rule 的核心）。

左下: delta rule 纠错

  • 同一 key 反复出现、v 逐渐改变时, 检索值的收敛轨迹
  • 你该看什么: 检索值逐步收敛到最新 v, 而非被覆盖或遗忘

右下: 有界 log‑decay 门控

  • g = g_min·sigmoid(z) 的曲线, 永远在 (g_min, 0) 之间
  • 红色虚线: 下界 g_min
  • 你该看什么: g 严格有下界 → a = exp(g) 有正下界 → chunkwise 不会溢出

Tab 8: 递推 vs 分块等价

目的

验证 Eq.1 (逐 token 递推) 与 Eq.3‑4 (块内并行) 逐位一致---两种计算路径的数值等价性。

实现：

  1. 同一组输入分别送入 kda_recurrent 和 kda_chunkwise
  2. 对比输出和状态的 max abs 误差
  3. 绘制输出对比曲线、逐位误差柱状图、状态差异热图

UI 组件与参数影响

参数 范围 默认 调节效果
chunk_size 1/2/4/8/16/32 8 分块大小。1 = 退化为逐 token; 增大 → 块内并行度更高

输出图解读

tab-8

1×3 布局：

左: 输出逐位对比

  • 蓝色圆点 = 递推输出, 橙色方块 = 分块输出
  • 理想情况: 两条线完全重合

中: 逐位误差柱状图

  • 绿色柱 = 每个时间步的 abs 误差 (Y 轴 log scale)
  • 标题显示 max 误差 (应 < 1e‑10)
  • 你该看什么: 误差在机器精度级别 (~1e‑15), 说明两种计算路径逐位一致

右: 状态差异热图

  • hot 色图展示 |S_rec − S_chk| 的逐元素差异
  • 理想情况: 全黑 (差异为 0)

进一步解释：三张图都是同一份误差（量级都在 1e-16，fp64 机器精度）：

  • 左图：直接叠加原始曲线，1e-16 远小于线宽 → 视觉重合（不是没误差，是看不见）
  • 中图 ：把每个时间步的 |递推-分块| 用 log 轴画，能看到 1e-16~1e-18 浮点噪声
  • 右图：把状态矩阵 S (16×16) 每个元素的 abs 差画成热图，看误差的空间分布

为什么重要

递推是 ground truth (O(T) 串行), 分块是生产路径 (块内并行)。如果两者不等价, 生产推理结果就是错的。此 Tab 是 KDA 实现正确性的主锚点。

Tab 9: 有界 vs 无界衰减

目的 展示无界 softplus 门控在 chunkwise 计算中会溢出到 NaN, 有界 sigmoid 门控保证数值安全 (论文 §4.2.7)。

实现（2×2 布局）：

  • 左上：理论 1/gamma 上界 vs tile 大小（对数坐标）
  • 右上：tile=16 实测------有界有限，无界 NaN
  • 左下：safe_chunk_size vs g_min 曲线
  • 右下：chunk=64 失败 case

输出图解读

tab-9

2×2 布局：

左上： 1/γ 上界 vs tile 大小 (log-log)

  • 有界版： 1/γ ≤ exp(-g_min·tile)，数据无关硬上界
  • 红色虚线: BF16 上限 (3.4e38)
  • 灰色点线: 论文 tile=16
  • 关键 : tile=16 时 1/γ≈1e35，刚好低于 BF16 上限；tile=32 时约 1e68，远超 BF16 上限

右上： tile=16 实测 (fp32)

  • 固定 chunk=16，换门控 → 测试"门控是否决定溢出"。蓝色 = 输出有限，红色 = 输出 NaN（柱顶白字标 NaN）。softplus z+6 时 1/γ→inf，红柱。
  • bounded(K3) 1/γ≈6.1e22 有限； softplus(z+6) 1/γ→inf → 输出 NaN
  • 用 fp32 是因为 bf16 值域上限与 fp32 几乎相同（都是 e38 量级），溢出行为一致； fp32 精度更高，能更清楚地看到 1/γ 从有限到 inf 的跃迁

左下： safe_chunk_size vs g_min

  • 圆点实线 = bf16, 方块虚线 = fp32
  • 两条曲线完全重合------因为 bf16.max≈3.39e38 与 fp32.max≈3.40e38 几乎相同（指数都是 38）
  • bf16 和 fp32 的真正差别在精度（尾数 7 位 vs 23 位），不在值域上限； 所以"bf16 更早溢出"是误解

右下： bounded chunk=16 vs 无界 chunk=64 (fp32)

  • bounded 用 chunk=16 vs 无界 softplus 用更大 chunk=64 → 测试"无界门控调大 chunk 能不能救"。即使 chunk=64 也无济于事（红柱）。
  • 无界 softplus 即使 chunk=64: 1/γ→inf → NaN
  • 有界衰减的关键不是"避免 dtype 上限"， 而是让 1/γ 有数据无关的硬上界------这是 chunkwise 能稠密化的前提

为什么重要 这是 K3 选择 sigmoid 门控而非 softplus 的根本原因---数值安全性。无界门控在 chunkwise 并行中, 累积保留因子 1/γ 可能溢出, 导致 NaN 训练崩溃。

Tab 10: v1 vs v2 门控

目的 对比 KDA v1 (softplus 无界) 和 v2 (sigmoid 有界) 的门控行为和数值差异。

实现：

  1. 绘制两条门控曲线 g(z)
  2. 绘制 alpha 随 A 变化的范围
  3. 手动实现单步 KDA （因为 kda_step 只支持 v2），逐 token 跑 v1/v2 对比
  4. 输出状态范数和数值差异

为什么需要手写单步 ：kda_ref.py 的 kda_step 内部硬编码了 v2 的 sigmoid 门控，无法切换到 v1。为了对比两种门控，必须手动实现包含 if version == "v1" 分支的单步函数。

UI 组件与参数影响

参数 范围 默认 调节效果
A_log_min ‑5.0‑0.0 ‑2.0 A 范围下界, 控制门控曲线的缩放
A_log_max 0.0‑5.0 2.0 A 范围上界。增大 → v1 和 v2 差异更明显

输出图解读

tab-10

1×3 布局：

左: 门控曲线 g(z)

  • 红色 = v1:‑softplus(z) (无下界, 可趋向 ‑∞)
  • 绿色 = v2: ‑5·sigmoid(z) (有下界 g_min=‑5)
  • 灰色虚线: g_min=‑5 和 g=0

中: alpha 取值范围 (fill_between)

  • 横轴 = A（控制 log-decay 强度），纵轴 = alpha = exp(g)（衰减系数）
  • 实线 = 下界（z 取最不保守值时 alpha 能跌到多低）
  • 浅色填充带 = alpha 的取值范围（从 min 到 max）
  • 红色实线 (v1 下界): 随 A 增大从 ~0.4 一路跌到 0，没有任何托底
  • 绿色实线 (v2 下界): 不论 A 多大都被 exp(g_min)=0.0067 这条绿色虚线硬托住
  • → v1/v2 上界都贴 1，所以『差异全在下界』------画下界曲线就够了
  • v1 vs v2 状态差异 (16 步后): 0.451, 输出差异: 0.021
  • 结论: v2 的有界衰减保证 chunkwise 数值安全（1/γ 有数据无关硬上界），是 K3 最终选择

右: v1 vs v2 状态范数演进

  • 红色圆点线 = v1 逐 token 状态范数 ||S||
  • 绿色方块线 = v2 逐 token 状态范数 ||S||
  • 你该看什么: v1 的状态范数可能发散 (无界衰减), v2 的状态范数保持有界

为什么重要

v2 是 K3 最终选择。此 Tab 量化展示 v1 的无界性如何导致 alpha 范围不受控, 以及两种门控在相同输入下的数值差异。

Tab 11: KCP 前缀扫描

目的 验证 KDA 状态跨 segment 不可加, 必须用 KCP 前缀扫描重组 (论文 §5.1.2)。

实现：

  1. 将 64-token 序列分成 4 段（各 16 token）
  2. 错误做法：各段 S=0 独立跑，结果直接相加
  3. 正确做法 ：KCP 前缀扫描 S = M·S + S̃
  4. 对比两者对精确递推的误差
  5. 绘制通信量对比（KDA 固定 vs softmax KV 随长度增长）

输出图解读

tab-11

1×2 布局：

左: 跨段状态重组误差

  • 两根柱: 直接相加 vs KCP 前缀扫描, 对精确递推的 max abs 误差
  • 直接相加: 误差巨大 (错误做法)
  • KCP 前缀扫描: 误差 < 1e‑10 (正确做法)

右: 通信量对比 (log‑log)

  • softmax KV 通信量随序列长度线性增长
  • KDA 通信量固定 (与序列长度无关的水平线)

为什么重要 KDA 的固定状态大小意味着上下文并行 (CP) 的通信量与序列长度无关---这是 KDA 相对 softmax attention 的核心工程优势。但前提是必须用正确的 KCP 前缀扫描, 而非简单相加。

为什么通信量与序列长度无关: KDA 的状态 S 是固定大小 Dk×Dv 的矩阵 (不随序列长度 T 增长)。CP 跨 segment 传递状态时, 只需传这个矩阵：

scss 复制代码 
KDA CP 通信量 = O(Dk·(Dk + Dv)) 字节/层/rank    ← 固定, 与 T 无关
softmax CP 通信量 = O(T·Dk·Dv) 字节/层/rank    ← 随 T 线性增长

KDA 的状态是"压缩过的"---无论序列多长, 状态始终是 Dk×Dv 的矩阵。softmax 的 KV cache 则是"未压缩的"---每个 token 的 K/V 都要保留, 所以随 T 增长。

Tab 12: KDA Module 形状与缓存

目的

实时展示 KimiDeltaAttention 模块的参数分布和输入输出形状。

实现：

  1. 用 KDAConfig + Kimi3Config 构造一个 Kimi3Block
  2. 随机输入 forward，统计参数量
  3. 绘制参数饼图和形状条形图

UI 组件与参数影响

参数 范围 默认 调节效果
Batch 1‑8 2 批大小
Heads 1‑8 4 注意力头数。增大 → 参数数量线性增长
Dk 8‑64 16 键维度。增大 → S 矩阵更大, 参数更多
Dv 8‑64 16 值维度。增大 → S 矩阵更大, 参数更多

输出图解读

tab-12

1×2 布局：

左: 参数分布饼图

  • 按模块名分组 (如 q_proj, k_proj, v_proj, o_proj 等)
  • 标题显示总参数量

右: 输入输出形状条形图

  • 蓝色 = 输入 x 元素数, 绿色 = 输出 y 元素数
  • 条旁标注具体 shape tuple

为什么重要

快速验证模块配置的参数效率和形状正确性。KDA 的参数主要在投影矩阵 (W_q, W_k, W_v, W_o) 和低秩衰减路径 (W_a^up, W_a^down) 上。

Tab 13: 单步匹配递推

目的 验证 kda_step (decode 路径) 与 kda_recurrent (prefill 路径) 逐位一致。

实现：

  1. 同一组输入送入 kda_recurrent（全序列）和循环 kda_step（逐 token）
  2. 注意 kda_step 返回 3D (B,H,Dv)，需要 unsqueeze(2) 再 cat 成4D
  3. 对比输出和状态的 max abs 误差

输出图解读

tab-13

1×3 布局：

左: 递推 vs 单步输出对比

  • 蓝色圆点 = 递推输出, 橙色方块 = 单步输出
  • 理想情况: 两条线完全重合

中: 逐位误差柱状图

  • 绿色柱 = 每个时间步的 abs 误差 (Y 轴 log scale)
  • 标题显示 max 误差 (应 < 1e‑10)
  • 你该看什么: 误差在机器精度级别 (~1e‑15), 说明两种计算路径逐位一致

右: 状态差异热图

  • hot 色图展示 |S_rec − S_step| 的逐元素差异
  • 理想情况: 全黑 (差异为 0)

为什么重要 prefill 用 kda_recurrent (全序列一次算完), decode 用 kda_step (逐 token 增量更新)。如果两者不等价, decode 阶段的输出就会偏离 prefill 阶段, 导致生成质量下降。此 Tab 是 prefill/decode 一致性的主锚点。

Tab 14: 诊断指标

目的 展示 build_kda_diagnostics 返回的运行时监控指标, 用于检测训练中的数值异常。

实现：

  1. 构造 KimiDeltaAttention 层，随机输入 forward
  2. 调用 build_kda_diagnostics 拿指标字典
  3. 绘制 4 子图：alpha 统计、状态范数、输出统计、梯度范数
  4. 注意：需要 requires_grad=True 才能计算梯度

输出图解读

tab-14

2×2 布局：

左上: Alpha 统计

  • 水平柱状图, 展示 alpha_min/max/mean 等量
  • 你该看什么 : alpha 应在 (exp(g_min), 1) ≈ (0.0067, 1) 范围内

右上: 逐 head 状态矩阵范数

  • 每根柱 = 一个 head 的 S 矩阵 Frobenius 范数
  • 标题显示总范数
  • 你该看什么: 各 head 范数不应爆炸 (有界衰减保证), 且分布应相对均匀

左下: 输出统计

  • 两根柱 = 输出的 mean 和 std
  • 你该看什么: mean 接近 0, std 不异常大

右下: 梯度范数 (按模块)

  • 水平柱状图, 每根 = 一个子模块的参数梯度范数
  • 每根柱旁标注科学计数法数值 (如 1.2e‑03)
  • 你该看什么: 所有柱 > 0 → 梯度流通, 无死区
  • 注意: alpha_up/alpha_down 数值极小是正常的---KDA 初始化时低秩衰减路径权重接近零

为什么重要 KDA 训练中最常见的故障模式: alpha 范围异常、状态范数爆炸、梯度消失。此诊断指标提供实时监控窗口, 无需手写 print 调试。

Tab 15: 3:1 调度可视化

目的 展示 K3 的 KDA/MLA 层间交替调度, 验证调度序列和模型前向正确性。

实现：

  1. 用 Kimi3Config 构造 Kimi3Model
  2. 调用 cfg.model.layer_schedule() 获取调度序列
  3. 绘制彩色条形图（蓝=KDA，橙=MLA）
  4. 执行 forward 验证输出形状

UI 组件与参数影响

参数 范围 默认 调节效果
n_layers 2‑16 8 总层数。增大 → 调度条更长, KDA/MLA 交替更明显

输出图解读

tab-15

2×1 布局：

上: 调度条

  • 彩色条形图, 蓝色 = KDA (线性注意力), 橙色 = MLA (softmax 注意力)
  • 每格标注层号和类型 (如 "L0 KDA")
  • 你该看什么: KDA 和 MLA 交替出现, 每 4 层中 3 KDA + 1 MLA

下: 模型前向验证

  • 两根水平柱 = 输入 tokens 和模型输出的元素数
  • 条旁标注 shape tuple
  • 你该看什么: 输出 shape 正确, 前向无报错

为什么重要 K3 的混合架构 (KDA 负责长程线性回忆 + MLA 负责精确近程检索) 依赖正确的层调度。如果调度逻辑有误 (如全 KDA 或全 MLA), 模型能力会退化。此 Tab 可视化调度序列并验证前向通过。

TransFormer-封面

0xFF 参考

本文使用 markdown.com.cn 排版

相关推荐
拓恒技术34 分钟前
从“可见光漏检”到“热异常锁定”：无人机河道排污口识别与溯源的技术链路
人工智能·视觉检测·无人机
2601_968900771 小时前
按量计费接口实战：从接入、预算到对账的完整链路
人工智能
七牛云行业应用1 小时前
MIT “权重空间编程”：GPT-6 Astra 修改 0.016% 的 Qwen 权重，MATH-500 从 7.9% 升至 48.1%
人工智能·gpt·ai
故七月1 小时前
AI 答案里的竞品共现：怎么看出赛道格局
大数据·人工智能
正在走向自律1 小时前
【3/5】人工智能训练师 三级/高级工 考证学习方向：从“独立执行”到“设计、归因与全链条负责”
人工智能·学习·人工智能训练师·注方案设计·数据质控负责人·模型评测与调优辅助·ai训练项目交付
Fatbobman(东坡肘子)1 小时前
快速消失的硬盘容量 -- 肘子的 Swift 周报 #160
开发语言·人工智能·redis·ios·swiftui·iphone·swift
Flynn Liang2 小时前
AI 应用的安全红线：Spring AI 下的 Prompt 注入防护、PII 脱敏、输出审计与密钥管理
人工智能·安全·spring·数据泄露·springai·prompt注入
朗迪锋2 小时前
SoErgo重塑“四特”场景下的人机协同与安全边界
人工智能·安全·人机工程
stormzhangV2 小时前
未来十年的建议
人工智能·aigc·agent
热门推荐
01国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）02历年考研数学一、数学二、数学三真题试卷及答案PDF03GitHub 镜像站点04几个好用的ip纯净度检测网站05昇腾910B系列芯片完全指南：四款型号对比、服务器选型与大模型部署实战06苹果iPA游戏软件资源下载网站07蜘蛛磁力 搜索引擎大全，如何使用蜘蛛磁力查找磁力链接08如何新建文件夹？ 电脑新建文件夹的4种方法09Dify：一站式AI应用开发平台全解析与Windows部署指南10Android 屏幕适配系列开篇：核心概念、官方机制与主流方案