原创 罗西的思考 2026-11-04 21:15 北京

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

• 0x00 文件信息
• 0x01 KDA核心数学
  • 1.1 因果深度卷积：因果逐通道短卷积
  • 1.2 SiTU-GLU：有界激活函数（公式12）
  • 1.3 KDA递推方程（公式1）
  • 1.4 有界衰减（公式5）
  • 1.5 分块并行（公式3-4）
  • 1.6 KCP前缀扫描（章节5.1.2）
  • 1.7 KimiDeltaAttention完整注意力层
  • 1.8 3:1 KDA与MLA调度
• 0x02 KDA实现了哪些功能
  • 2.1 覆盖的源模块（共5个）
  • 2.2 覆盖的测试（共90个测试函数）
  • 2.3 15个标签页功能清单
• 0x03 演示实现思路
  • 3.1 整体架构
  • 3.2 核心设计决策
  • 3.3 测试与标签页映射逻辑
  • 3.4 环境与路径
• 0x04 每个标签页的使用说明
  • 标签页1：因果深度卷积感受野 vs 延迟误差
  • 标签页2：因果深度卷积恒等初始化
  • 标签页3：因果深度卷积因果性
  • 标签页4：SiTU-GLU激活曲线对比
  • 标签页5：SiTU-GLU饱和区梯度
  • 标签页6：SiTU-GLU离群点回归
  • 标签页7：KDA原理（4格展示）
  • 标签页8：递推 vs 分块等价
  • 标签页9：有界 vs 无界衰减
  • 标签页10：v1 vs v2门控
  • 标签页11：KCP前缀扫描
  • 标签页12：KDA模块形状与缓存
  • 标签页13：单步匹配递推
  • 标签页14：诊断指标
  • 标签页15：3:1调度可视化
• 0xFF 参考

本文是对项目Gradio演示的解释说明。

0x00 文件信息

注：

• 项目相关链接：https://github.com/RossiLH/Model-Viz 和系列博客是相关尝试，也是学习过程。不确定是否满足读者需求，会反思改进，也欢迎提出宝贵意见。
• 涉及KIMI K3、DeepSeek V4、PI0.5、PI0.7等模型，后续可能跟踪新模型。
• 全CPU复刻，不涉及GPU，也不涉及大规模训练和硬件。
• 模型组件尽量是原始大模型的缩小版，若组件无开源版本，则用近似组件替代。
• 会解释每个组件解决的问题、难点、解法原理，以及组件间的关联，重点是理解主线而非记忆名词。
• 每个主要组件都有可直接运行的Gradio演示，无需GPU；另有若干可视化脚本（plot_*.py），可生成静态图。
• 所有文件行号基于本地源码快照，源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
• 为保证每个文章尽量独立成文，部分内容会在不同文章中重复，系列层面可能存在冗余，敬请理解。

0x01 KDA核心数学

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

1.1 因果深度卷积：因果逐通道短卷积

K3的MLA层是无位置编码，因果深度卷积给每个通道注入局部位置信息：

输出[t, c] = Σ（lag从0到K-1）权重[c, θ, lag] × 输入[t - lag, c] + 偏置[c]

• 因果性：只看当前及过去K-1个token（左侧补零）
• 恒等初始化：权重最后一个时间步设为1，其余为0，初始为无操作，不扰动KDA衰减动力学
• 单步解码：维护K-1长度的滑动窗口，逐token更新复杂度为O(K)

1.2 SiTU-GLU：有界激活函数（公式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上限448
• 饱和区梯度：softcap用tanh截断，饱和区仍有非零梯度，不会像硬clamp那样切断训练
• 对比SwiGLU：SwiGLU(g, u) = silu(g) × u，无界，仅作消融基线

1.3 KDA递推方程（公式1）

KDA维护固定大小状态矩阵S ∈ R^(Dk×Dv)，逐token更新：

状态递推（Delta规则）：
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]：逐通道衰减因子，保留旧记忆
• β_t ∈ (0,1)：写入强度，sigmoid输出
• v_t - S_{t-1}^T k_t：delta-rule纠错项，只写入当前检索与目标的误差，而非盲目覆盖

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

1.4 有界衰减（公式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）在分块并行中，累积保留因子1/γ = exp(-Σg)可能溢出到NaN
• safe_chunk_size(g_min, dtype) = ⌊log(dtype_max / headroom) / (-g_min)⌋，保证数值安全

1.5 分块并行（公式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 = 因果深度卷积(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注意力，压缩潜在）：

层0: KDA 层1: MLA 层2: KDA 层3: MLA ...

KDA层负责长程线性回忆（固定状态），MLA层负责精确近程检索（KV缓存）。

0x02 KDA实现了哪些功能

2.1 覆盖的源模块（5个）

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

2.3 15个标签页功能清单

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个标签页函数（每个返回 Tuple[Figure, str]）
└── build_demo() + __main__（Gradio界面构建 + 启动）

3.2 核心设计决策

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

每个标签页函数的签名统一为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提供组件级验证

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

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

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

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

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

3.3 测试→标签页映射逻辑

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

映射原则：

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

示例映射：

test_shape_and_no_bias → 标签页1（感受野）
test_identity_initialization → 标签页2（恒等初始化）
test_causality → 标签页3（因果性）
test_output_is_bounded → 标签页4（激活曲线）
test_gradients_survive_saturation → 标签页5（饱和区梯度）
test_swiglu_baseline_explodes → 标签页6（离群点回归）
test_chunkwise_equals_recurrence → 标签页8（递推 vs 分块）
test_step_matches_sequence → 标签页13（单步匹配递推）
test_kcp_prefix_scan → 标签页11（KCP前缀扫描）
test_module_forward_shapes → 标签页12（形状与缓存）
test_unbounded_decay_would_overflow → 标签页9（有界 vs 无界衰减）
test_kimi3_model → 标签页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 每个标签页的使用说明

标签页1：因果深度卷积感受野 vs 延迟误差

目的：验证因果深度卷积的局部感受野不足以捕获长程依赖。当kernel_size < delay时，MSE≈信号方差（任务失败）。

实现：

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

关键代码路径：

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

UI组件与参数影响

输出图解读

• 红色折线图：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。

为什么重要：

因果深度卷积是K3中NoPE位置编码的替代方案，kernel_size决定了它能捕获多长的局部模式。此标签页量化展示"局部感受野的局限"——delay=8的任务需要kernel至少9才能重建。

标签页2：因果深度卷积恒等初始化

目的：展示恒等初始化（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的衰减动力学依赖因果深度卷积不在初始阶段引入大幅扰动。恒等初始化保证conv(x)≈x，不干扰KDA的状态矩阵S的收敛轨迹。

标签页3：因果深度卷积因果性

目的：验证因果深度卷积只看当前+过去token，修改未来输入不影响过去输出。

输出图解读

• 蓝色圆点线：原始输出
• 橙色方块线：未来被扰动后的输出（t≥10处输入+100）
• 灰色虚线：扰动起点（第10位）
• 需要关注：两条线在t<10处几乎完全重合（因果性成立），t≥10处分开（扰动生效）

进一步解释：

• 扰动：把输入x在t=10..19这10个时间步上整体+100，前10步保持原样
• 扰动起点：被修改的第一个时间步索引，即t=10

为什么重要：

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

标签页4：SiTU-GLU激活曲线对比

目的：展示SiTU-GLU的有界性，对比SwiGLU的无界增长。

实现：

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

UI组件与参数影响

输出图解读

• 红色线：SwiGLU（无界，|x|大时线性增长）
• 绿色线：SiTU-GLU（有界，被±100水平线截住）
• 灰色水平虚线：上界±100

调节建议

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

标签页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截断保证梯度永远非零，饱和区的参数仍能被微调。

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

数值分解（训练1800步后）：

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以内，牺牲少量拟合精度换取数值安全。

标签页7：KDA原理（4格）

目的：展示KDA的三个核心概念：固定大小记忆（S矩阵）、delta-rule纠错、有界衰减门控。

实现（2×2布局）：

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

UI组件与参数影响

输出图解读

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不会溢出

标签页8：递推 vs 分块等价

目的：验证递推（逐token）与块内并行两种计算路径逐位一致——两种计算路径的数值等价性。

实现：

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

UI组件与参数影响

输出图解读

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)串行)，分块是生产路径 (块内并行)。如果两者不等价，生产推理结果就是错的。此标签页是KDA实现正确性的主锚点。

标签页9：有界 vs 无界衰减

目的：展示无界softplus门控在chunkwise计算中会溢出到NaN，有界sigmoid门控保证数值安全。

实现（2×2布局）：

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

输出图解读

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 训练崩溃。

标签页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组件与参数影响

输出图解读

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 最终选择。此标签页量化展示 v1 的无界性如何导致 alpha 范围不受控, 以及两种门控在相同输入下的数值差异。

标签页11：KCP前缀扫描

目的：验证 KDA 状态跨 segment 不可加，必须用 KCP 前缀扫描重组。

实现：

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 增长。

标签页12：KDA模块形状与缓存

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

实现：

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

UI组件与参数影响

输出图解读

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) 上。

标签页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 布局：

左: 输出逐位对比

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

中: 逐位误差柱状图

• 绿色柱 = 每个时间步的 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) 串行), 分块是生产路径 (块内并行)。如果两者不等价, 生产推理结果就是错的。此标签页是 KDA 实现正确性的主锚点。