【轻量化复刻大模型 / KIMI K3】KDA 模块交互式演示
【轻量化复刻大模型 / KIMI K3】KDA 模块交互式演示
本文是对项目 Gradio Demo 的解释说明。
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 |
注：
- 项目https://github.com/RossiLH/Model-Viz 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
- 涉及 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 给每个通道注入局部位置信息：
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)
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 更新：
状态递推 (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)
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²) 并行：
块内: 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}), 必须用前缀扫描重组：
S_running = M_seg · S_running + S̃_seg
通信量固定 O(Dk·(Dk+Dv)), 与序列长度无关 (对比 softmax KV 的 O(T·Dk·Dv))。
1.7 KimiDeltaAttention 完整层
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):
层 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 整体架构
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：
映射原则：
- 形状/缓存测试 → Tab 10（模块形状与缓存）
- 数值等价性测试 → Tab 6（递推 vs 分块）或 Tab 11（单步匹配）
- 机制验证测试 → 对应原理 Tab（如 KCP→Tab 9，衰减→Tab 7）
- 模型集成测试 → Tab 13（3:1 调度）
- 梯度测试 → 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 环境与路径
# 向上 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 ≈ 信号方差 (任务失败)。
实现：
- 生成"延迟线"任务：输入是随机正弦信号，目标是延迟 D 步后的副本
- 训练不同 kernel 宽度（2, 4, 8, 12, 20）的 ShortConv 拟合该任务
- 绘制 kernel 宽度 vs 最终 MSE 的折线图
- 标注 delay=8 的垂直线——kernel <= delay 时 MSE ≈ 信号方差（失败）
关键代码路径：
_make_delayed_task() → 生成 (x, y)
_fit_conv() → Adam 训练 3000 步，返回最终 MSE
tab_receptive_field() → 绘图 + 返回结论
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| kernel_size | 2‑20 | 4 | 当前选择的 kernel 宽度, 绿色方块标记在折线图上 |
输出图解读
- 红色折线图: 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 衰减动力学, 训练起步更平稳。
实现：
- 恒等初始化：
reset_parameters将 weight 最后一个 tap 设为 1，其余为 0 - 随机初始化：
weight.normal_(std=0.5) - 两种初始化各训练 1500 步，记录每 100 步的 MSE
- 绘制训练曲线对比
输出图解读
- 绿色圆点线: 恒等初始化训练曲线 (初始为 no‑op, 起步 MSE 低)
- 红色方块线: 随机初始化训练曲线 (初始扰动大, 起步更抖)
- Y 轴 log scale: 让训练后期的细微差异可见
为什么重要
KDA 的衰减动力学依赖 ShortConv 不在初始阶段引入大幅扰动。恒等初始化保证 conv(x) ≈ x, 不干扰 KDA 的状态矩阵 S 的收敛轨迹。
Tab 3: ShortConv 因果性
目的
验证 ShortConv 只看当前+过去 token, 修改未来输入不影响过去输出 (论文 §2.1.1)。
输出图解读
- 蓝色圆点线: 原始输出
- 橙色方块线: 未来被扰动后的输出 (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 的无界增长。
实现：
- 在 [x_min, x_max] 范围内计算 SiTU-GLU 和 SwiGLU 的激活值
- 绘制两条曲线 + 上界水平线
- 输出最大绝对值对比
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| x_min | ‑50‑0 | ‑20 | 展示范围左端。减小 → 看到更极端负值的 SwiGLU 无界增长 |
| x_max | 0‑50 | 20 | 展示范围右端。增大 → 看到更极端正值的 SwiGLU 无界增长 |
输出图解读
- 红色线: 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), 不会掐断训练。
实现：
- 计算
softcap(z, β)的梯度 d/dz - 计算
clamp(z, max=β)的梯度 d/dz - 对比：softcap 梯度在饱和区 > 0，硬 clamp 梯度 = 0
输出图解读
- 绿色线: 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 无界通过但低精度下溢出。
输出图解读
左图 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 设计的真实优势：
数学事实
- SiTU-GLU 的上界是
beta_g × beta_u = 4 × 25 = 100（不是 25），所以 200 离群点确实超出范围（被截到 ~100） - 但 SwiGLU 无界，会去"硬扛" 200 离群点，梯度被离群点拉爆，连带把 98% 正常样本（值域 ±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) → 衰减更慢, 旧记忆保留更久 |
输出图解读
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 (块内并行) 逐位一致—两种计算路径的数值等价性。
实现：
- 同一组输入分别送入
kda_recurrent和kda_chunkwise - 对比输出和状态的 max abs 误差
- 绘制输出对比曲线、逐位误差柱状图、状态差异热图
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| chunk_size | 1/2/4/8/16/32 | 8 | 分块大小。1 = 退化为逐 token; 增大 → 块内并行度更高 |
输出图解读
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
输出图解读
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 有界) 的门控行为和数值差异。
实现：
- 绘制两条门控曲线 g(z)
- 绘制 alpha 随 A 变化的范围
- 手动实现单步 KDA（因为
kda_step只支持 v2），逐 token 跑 v1/v2 对比 - 输出状态范数和数值差异
为什么需要手写单步：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 差异更明显 |
输出图解读
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)。
实现：
- 将 64-token 序列分成 4 段（各 16 token）
- 错误做法：各段 S=0 独立跑，结果直接相加
- 正确做法：KCP 前缀扫描
S = M·S + S̃ - 对比两者对精确递推的误差
- 绘制通信量对比（KDA 固定 vs softmax KV 随长度增长）
输出图解读
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 传递状态时, 只需传这个矩阵：
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 模块的参数分布和输入输出形状。
实现：
- 用
KDAConfig+Kimi3Config构造一个Kimi3Block - 随机输入 forward，统计参数量
- 绘制参数饼图和形状条形图
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| Batch | 1‑8 | 2 | 批大小 |
| Heads | 1‑8 | 4 | 注意力头数。增大 → 参数数量线性增长 |
| Dk | 8‑64 | 16 | 键维度。增大 → S 矩阵更大, 参数更多 |
| Dv | 8‑64 | 16 | 值维度。增大 → S 矩阵更大, 参数更多 |
输出图解读
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 路径) 逐位一致。
实现：
- 同一组输入送入
kda_recurrent（全序列）和循环kda_step（逐 token） - 注意
kda_step返回 3D(B,H,Dv)，需要unsqueeze(2)再cat成4D - 对比输出和状态的 max abs 误差
输出图解读
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 返回的运行时监控指标, 用于检测训练中的数值异常。
实现：
- 构造
KimiDeltaAttention层，随机输入 forward - 调用
build_kda_diagnostics拿指标字典 - 绘制 4 子图：alpha 统计、状态范数、输出统计、梯度范数
- 注意：需要
requires_grad=True才能计算梯度
输出图解读
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 层间交替调度, 验证调度序列和模型前向正确性。
实现：
- 用
Kimi3Config构造Kimi3Model - 调用
cfg.model.layer_schedule()获取调度序列 - 绘制彩色条形图（蓝=KDA，橙=MLA）
- 执行 forward 验证输出形状
UI 组件与参数影响
| 参数 | 范围 | 默认 | 调节效果 |
|---|---|---|---|
| n_layers | 2‑16 | 8 | 总层数。增大 → 调度条更长, KDA/MLA 交替更明显 |
输出图解读
2×1 布局：
上: 调度条
- 彩色条形图, 蓝色 = KDA (线性注意力), 橙色 = MLA (softmax 注意力)
- 每格标注层号和类型 (如 "L0 KDA")
- 你该看什么: KDA 和 MLA 交替出现, 每 4 层中 3 KDA + 1 MLA
下: 模型前向验证
- 两根水平柱 = 输入 tokens 和模型输出的元素数
- 条旁标注 shape tuple
- 你该看什么: 输出 shape 正确, 前向无报错
为什么重要
K3 的混合架构 (KDA 负责长程线性回忆 + MLA 负责精确近程检索) 依赖正确的层调度。如果调度逻辑有误 (如全 KDA 或全 MLA), 模型能力会退化。此 Tab 可视化调度序列并验证前向通过。
浙公网安备 33010602011771号