Initial K3 snapshot: 0.5B KDA/MLA/MoE train path
Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
This commit is contained in:
dela
@@ -0,0 +1,187 @@
|
||||
schema: superpaper.ledger/v1
|
||||
retired_ids: []
|
||||
paper:
|
||||
id: "kda-project"
|
||||
title: "KDA 训练→推理 手写实现 — 完整笔记"
|
||||
authors: ["dela"]
|
||||
notes_language: zh
|
||||
source:
|
||||
kind: markdown
|
||||
|
||||
coverage:
|
||||
mode: full
|
||||
sections_in:
|
||||
- "KDA 递归核心"
|
||||
- "Gate 激活"
|
||||
- "分块并行计算"
|
||||
- "GVA 分组值注意力"
|
||||
- "KDAAttention 层"
|
||||
- "Gated MLA 矩阵吸收版"
|
||||
- "SiTU-GLU 与 Stable LatentMoE"
|
||||
- "K3 混合架构"
|
||||
- "Attention Residual 深度残差"
|
||||
- "反向传播推导"
|
||||
sections_skipped:
|
||||
- "Triton kernel 细节"
|
||||
- "Docker 部署"
|
||||
- "AttnRes 论文的 kernel 级调度与 pipeline 重叠"
|
||||
|
||||
questions:
|
||||
- id: Q1
|
||||
text: "KDA 的状态更新如何避免 softmax、实现线性复杂度？"
|
||||
- id: Q2
|
||||
text: "safe gate 与 standard gate 的区别是什么？"
|
||||
- id: Q3
|
||||
text: "分块并行如何在保持递归等价的同时利用 GPU 并行？"
|
||||
- id: Q4
|
||||
text: "GVA 的 repeat_interleave + sum 反向是怎么回事？"
|
||||
- id: Q5
|
||||
text: "MLA 矩阵吸收如何避免解压 K/V？"
|
||||
- id: Q6
|
||||
text: "SiTU-GLU 为什么比 SwiGLU 更稳定？"
|
||||
- id: Q7
|
||||
text: "AttnRes 如何把残差流从等权累加换成按内容选择？"
|
||||
- id: Q8
|
||||
text: "Block AttnRes 的两阶段算法为什么和 naive 逐层实现数值等价？"
|
||||
- id: Q9
|
||||
text: "深度残差接入 CausalLM 时怎样避免参数被重复注册？"
|
||||
|
||||
claims:
|
||||
- id: C1
|
||||
text: "KDA 用 delta rule 更新 KV 状态矩阵，不需要 softmax，复杂度 O(T·K·V)"
|
||||
kind: methodological
|
||||
status: core
|
||||
- id: C2
|
||||
text: "safe gate = lower_bound · σ(rate · input)，保证 gate 值在 [lower_bound, 0] 范围内"
|
||||
kind: methodological
|
||||
status: core
|
||||
- id: C3
|
||||
text: "分块计算：chunk 内用下三角解，chunk 间用状态递推，数值等价于 naive recurrent"
|
||||
kind: methodological
|
||||
status: core
|
||||
- id: C4
|
||||
text: "MLA 矩阵吸收：q 吸收 W_UK 后直接与 latent c 内积，永不解压 K/V"
|
||||
kind: methodological
|
||||
status: core
|
||||
- id: C5
|
||||
text: "LatentMoE 通过 latent 接口把 routed 专家限制在半宽空间 ℓ=d/2"
|
||||
kind: methodological
|
||||
status: core
|
||||
- id: C6
|
||||
text: "AttnRes 用逐 token 的深度维 softmax 代替等权残差累加：打分在 RMS 归一化后做，加权和在原始张量上做"
|
||||
kind: methodological
|
||||
status: core
|
||||
- id: C7
|
||||
text: "Block AttnRes 块内退化为普通求和、只让块输出进入源列表，源数从 O(N) 降到 O(N/S)"
|
||||
kind: methodological
|
||||
status: core
|
||||
- id: C8
|
||||
text: "两阶段算法 = inter 块间批量 einsum + intra online-softmax 增量合并，与 naive 逐层实现数值等价 (atol 1e-5)"
|
||||
kind: methodological
|
||||
status: core
|
||||
- id: C9
|
||||
text: "BorrowedSubLayer 用普通 tuple 持有 norm/fn，不注册为子模块，保证参数与 state_dict 键不重复"
|
||||
kind: methodological
|
||||
status: supporting
|
||||
|
||||
symbols:
|
||||
- {name: B, latex: "B", meaning: "batch size", kind: "shape parameter"}
|
||||
- {name: T, latex: "T", meaning: "序列长度", kind: "shape parameter"}
|
||||
- {name: H, latex: "H", meaning: "query/key 头数", kind: "shape parameter"}
|
||||
- {name: HV, latex: "H_V", meaning: "value 头数 (GVA)", kind: "shape parameter"}
|
||||
- {name: G, latex: "G", meaning: "GVA 组数 = HV/H", kind: "shape parameter"}
|
||||
- {name: K, latex: "K", meaning: "key/query 头维度", kind: "shape parameter"}
|
||||
- {name: V, latex: "V", meaning: "value 头维度 (= K)", kind: "shape parameter"}
|
||||
- {name: D, latex: "D", meaning: "hidden_size", kind: "shape parameter"}
|
||||
- {name: C, latex: "C", meaning: "chunk_size", kind: "shape parameter"}
|
||||
- {name: r, latex: "r", meaning: "KV latent rank (kv_lora_rank)", kind: "shape parameter"}
|
||||
- {name: ell, latex: "\\ell", meaning: "MoE latent 宽度 = d/2", kind: "shape parameter"}
|
||||
- {name: S, latex: "S", meaning: "KV 状态矩阵", domain: "[B, HV, K, V]", kind: value}
|
||||
- {name: q, latex: "q", meaning: "query", domain: "[B, T, H, K]", kind: value}
|
||||
- {name: k, latex: "k", meaning: "key", domain: "[B, T, H, K]", kind: value}
|
||||
- {name: v, latex: "v", meaning: "value", domain: "[B, T, HV, V]", kind: value}
|
||||
- {name: g, latex: "g", meaning: "gate (log-space decay)", domain: "[B, T, HV, K]", kind: value}
|
||||
- {name: beta, latex: "\\beta", meaning: "学习率/写入强度", domain: "[B, T, HV]", kind: value}
|
||||
- {name: A_log, latex: "A_{\\log}", meaning: "head-wise 衰减参数 (log-space)", domain: "[HV]", kind: value}
|
||||
- {name: dt_bias, latex: "\\Delta_b", meaning: "per-dim gate bias", domain: "[HV, K]", kind: value}
|
||||
- {name: c, latex: "c", meaning: "KV latent 向量", domain: "[B, T, r]", kind: value}
|
||||
- {name: W_UK, latex: "W_{UK}", meaning: "Key 解压矩阵 (MLA)", domain: "[H, d_q, r]", kind: value}
|
||||
- {name: W_UV, latex: "W_{UV}", meaning: "Value 解压矩阵 (MLA)", domain: "[H, d_v, r]", kind: value}
|
||||
- {name: N, latex: "N", meaning: "AttnRes 原子层数 = 2L", kind: "shape parameter"}
|
||||
- {name: S, latex: "S", meaning: "AttnRes 块大小（原子层）", kind: "shape parameter"}
|
||||
- {name: v_i, latex: "v_i", meaning: "AttnRes 第 i 个源（v_0 = embedding 输出）", domain: "[B, T, D]", kind: value}
|
||||
- {name: w_l, latex: "w_l", meaning: "第 l 层 depth query（零初始化）", domain: "[D]", kind: value}
|
||||
- {name: alpha, latex: "\\alpha_{l,i}", meaning: "深度维 softmax 权重", domain: "[n, B, T]", kind: value}
|
||||
- {name: h_l, latex: "h_l", meaning: "深度注意力聚合出的层输入", domain: "[B, T, D]", kind: value}
|
||||
- {name: b_j, latex: "b_j", meaning: "Block AttnRes 第 j 块的输出", domain: "[B, T, D]", kind: value}
|
||||
- {name: p, latex: "p", meaning: "块内 running partial", domain: "[B, T, D]", kind: value}
|
||||
|
||||
terms:
|
||||
- {canonical: "KDA", aliases: ["Key-Decayed Attention", "键衰减注意力"]}
|
||||
- {canonical: "GVA", aliases: ["Grouped Value Attention", "分组值注意力"]}
|
||||
- {canonical: "MLA", aliases: ["Multi-head Latent Attention", "多头隐变量注意力"]}
|
||||
- {canonical: "MoE", aliases: ["Mixture of Experts", "混合专家"]}
|
||||
- {canonical: "SiTU-GLU", aliases: ["Sigmoid Tanh Unit GLU"]}
|
||||
- {canonical: "delta rule", aliases: ["δ 规则"]}
|
||||
- {canonical: "safe gate", aliases: ["安全门控"]}
|
||||
- {canonical: "matrix absorption", aliases: ["矩阵吸收"]}
|
||||
- {canonical: "AttnRes", aliases: ["Attention Residual", "注意力残差", "深度残差"]}
|
||||
- {canonical: "depth residual", aliases: ["DepthResidual", "深度维残差"]}
|
||||
- {canonical: "online softmax", aliases: ["在线 softmax", "增量 softmax"]}
|
||||
- {canonical: "atomic layer", aliases: ["原子层", "atomic sublayer"]}
|
||||
|
||||
derivations:
|
||||
- id: DER1
|
||||
claim: C1
|
||||
title: "KDA 递归状态更新推导"
|
||||
expand: true
|
||||
figure: null
|
||||
steps:
|
||||
- {id: "1", from: "S_{t-1}", to: "S_{\\mathrm{dec}} = \\exp(g_t) \\odot S_{t-1}", rule: scale}
|
||||
- {id: "2", from: "S_{\\mathrm{dec}}", to: "r_t = v_t - k_t \\cdot S_{\\mathrm{dec}}", rule: definition}
|
||||
- {id: "3", from: "r_t", to: "S_t = S_{\\mathrm{dec}} + (\\beta_t k_t) \\otimes r_t", rule: definition}
|
||||
- {id: "4", from: "S_t", to: "o_t = (q_t \\cdot \\text{scale}) \\cdot S_t", rule: definition}
|
||||
- id: DER2
|
||||
claim: C4
|
||||
title: "MLA 矩阵吸收推导"
|
||||
expand: true
|
||||
figure: null
|
||||
steps:
|
||||
- {id: "1", from: "q \\in [B,T,H,d_q]", to: "q_{\\mathrm{abs}} = q \\cdot W_{UK} \\in [B,T,H,r]", rule: substitute}
|
||||
- {id: "2", from: "q_{\\mathrm{abs}}, c", to: "\\text{score} = q_{\\mathrm{abs}} \\cdot c^T \\in [B,H,T,T]", rule: definition}
|
||||
- {id: "3", from: "\\text{attn}, c", to: "\\tilde{o}_{\\mathrm{lat}} = \\text{attn} \\cdot c \\in [B,H,T,r]", rule: definition}
|
||||
- {id: "4", from: "\\tilde{o}_{\\mathrm{lat}}", to: "\\tilde{o} = \\tilde{o}_{\\mathrm{lat}} \\cdot W_{UV}^T \\in [B,H,T,d_v]", rule: substitute}
|
||||
- id: DER3
|
||||
claim: C8
|
||||
title: "AttnRes 两阶段 online softmax 合并推导"
|
||||
expand: true
|
||||
figure: null
|
||||
steps:
|
||||
- {id: "1", from: "s_{l,i} = \\tilde{w}_l^T \\mathrm{RMS}(v_i)", to: "(m, n, d) = (\\max_i s_i, \\sum_i e^{s_i - m} v_i, \\sum_i e^{s_i - m})", rule: definition}
|
||||
- {id: "2", from: "inter sources b_0..b_{j-1} 固定", to: "一次批量 einsum 'q d, n b t d -> q n b t' 得块内全部 query 的 (m,n,d)", rule: substitute}
|
||||
- {id: "3", from: "单源 partial p", to: "(m, n, d) = (s_p, p, 1)，因为 e^{s_p - m} = 1", rule: definition}
|
||||
- {id: "4", from: "(m_a,n_a,d_a), (m_b,n_b,d_b)", to: "m = \\max(m_a,m_b);\\ n = e^{m_a-m} n_a + e^{m_b-m} n_b;\\ d = e^{m_a-m} d_a + e^{m_b-m} d_b", rule: scale}
|
||||
- {id: "5", from: "(m, n, d)", to: "h_l = n / d，与 forward_naive 逐位一致", rule: definition}
|
||||
|
||||
figures:
|
||||
- id: F1
|
||||
claim: C1
|
||||
title: "KDA 递归状态更新张量图"
|
||||
grammar: tensor-face
|
||||
toolkit: supertensor
|
||||
signals: [shape, contraction, broadcast]
|
||||
status: planned
|
||||
- id: F2
|
||||
claim: C4
|
||||
title: "MLA 矩阵吸收计算流"
|
||||
grammar: tensor-face
|
||||
toolkit: supertensor
|
||||
signals: [shape, contraction, transpose]
|
||||
status: planned
|
||||
- id: F3
|
||||
claim: C7
|
||||
title: "Full vs Block AttnRes 的源列表增长"
|
||||
grammar: tensor-face
|
||||
toolkit: supertensor
|
||||
signals: [shape, contraction]
|
||||
status: planned
|
||||
Reference in New Issue
Block a user