【轻量化复刻大模型 / 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
端口

---

注：

• 项目 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
• 涉及 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 (γ^{})·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：

映射原则：

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 环境与路径

# 向上 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 ≈ 信号方差（失败）

关键代码路径：

_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 衰减动力学, 训练起步更平稳。

实现：

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

输出图解读

• 绿色圆点线: 恒等初始化训练曲线 (初始为 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 的无界增长。

实现：

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

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), 不会掐断训练。

实现：

1. 计算 softcap(z, β) 的梯度 d/dz
2. 计算 clamp(z, max=β) 的梯度 d/dz
3. 对比：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 设计的真实优势：

数学事实

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) → 衰减更慢, 旧记忆保留更久

输出图解读

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; 增大 → 块内并行度更高

输出图解读

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 有界) 的门控行为和数值差异。

实现：

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 差异更明显

输出图解读

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 随长度增长）

输出图解读

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 模块的参数分布和输入输出形状。

实现：

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 矩阵更大, 参数更多

输出图解读

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 误差

输出图解读

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 才能计算梯度

输出图解读

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 交替更明显

输出图解读

2×1 布局：

上: 调度条

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

下: 模型前向验证

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

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

---

---

0xFF 参考