xtrain/docs/03-autograd-engine.md

Phase: Autograd Engine + Op Backward — Design Document

Goal

在 T3 的 Tensor（matmul/transpose/finite-diff harness）之上，交付 tape-based 动态 autograd 引擎 + 一个 tiny 现代 transformer 所需算子的前向 kernel + 解析 backward，每个 backward 都用 T3 的有限差分 harness 对拍通过。

具体三件事：

1. autograd 引擎：define-by-run 的反向自动微分。Var 包一个 Tensor + 可选 grad；每个 op 在 tape 上记一个节点（父节点 + backward 闭包）；backward() 按逆拓扑序遍历，把梯度推给父节点。关键正确性点：梯度累加——一个张量被多个 op 消费（扇出）时，各路梯度必须求和（T3 没有累加路径，在这里实现）。
2. 算子节点：matmul / add / mul / add_bias(broadcast) / scale / rms_norm / silu / swiglu / rope / softmax / cross_entropy，各带前向 CUDA kernel（需要时）+ 解析 backward。
3. Attention 用组合：attn = matmul(softmax(matmul(Q,Kᵀ)·scale), V)。一旦 matmul/softmax/scale 是 autograd 节点，attention 的 backward 自动成立——不写 fused attention backward kernel，只加一个端到端 grad-check 测试。

明确不做（留给 T5/T6）：组装 transformer / 训练 loop / 优化器 / embedding / KV-cache / GQA 重复。本 Phase 只到「算子 backward 逐个对拍通过」。

Module Layout

csrc/ops/nn.cu                         # 所有 T4 算子的 fwd+bwd kernel + launch_*（含 inlined warp/block reduce）

crates/xtrain-cuda/
├── build.rs                           # 新增 nn.cu
└── src/ffi.rs                         # 新增 launch_* 声明（no_cuda 门控）

crates/xtrain-tensor/
├── src/dtype.rs                       # 新增 I32（cross-entropy target 用）
└── src/tensor.rs                      # add/mul/add_bias/sum_rows/rms_norm(+bwd)/silu(+bwd)/
                                       #   rope(+bwd)/softmax(+bwd)/cross_entropy(+bwd)（no_cuda 门控）

crates/xtrain-autodiff/               # 引擎落在这里（已含 grad_check harness，自然归宿）
├── build.rs                           # 新增：检测 nvcc → no_cuda cfg（cfg 不跨 crate 传播）
├── src/
│   ├── lib.rs                         # 导出 tape::Var + ops（no_cuda 门控）
│   ├── finite_diff.rs                 # T3 既有 harness（不动）
│   ├── tape.rs                        # Var / VarNode / backward / 梯度累加
│   └── ops.rs                         # 各算子的 Var 节点构造器
└── tests/autograd.rs                  # 每算子 grad-check + 扇出累加 + 组合 attention（#![cfg(not(no_cuda))]）

为什么引擎放 xtrain-autodiff 而不是新 crate：该 crate 本就是「自动微分」语义的归宿，且已持有 grad_check。前向 kernel/Tensor 方法仍按 T2/T3 约定落在 xtrain-tensor（与 scale/matmul 一致），引擎只是在其上叠 tape。

Key Design Decisions

Tape 设计：Rc<RefCell<VarNode>> + 逆拓扑遍历

pub struct VarNode {
    value: Tensor,            // 前向输出
    grad: Option<Tensor>,     // 反向累加的梯度
    parents: Vec<Var>,        // 计算来源
    backward: Option<BackwardFn>,  // None=叶子
}
pub struct Var(Rc<RefCell<VarNode>>);
type BackwardFn = Box<dyn Fn(&Tensor, &[Var])>;

• Var clone 只是 bump Rc，clone 共享同一节点——这正是「扇出」的识别方式（同一 Rc::as_ptr 在多处出现）。
• backward()：① post-order DFS 建拓扑序（按指针去重）；② 把 loss（必须是标量）的 grad 种子设为 1；③ 逆序遍历，每个节点把自己的 grad 传给父节点的 backward 闭包。
• 闭包签名 Fn(&grad, &parents)：给本节点已累加的 grad 和父节点列表，闭包算出各父的梯度贡献并 push_grad 回去。前向需要 cache 的中间量（softmax 的 y、rms 的 inv_rms、ce 的 probs）用 move 闭包捕获。

梯度累加（扇出求和）——本 Phase 的正确性核心

push_grad(parent, g) 一律走 accumulate：

fn accumulate(&self, g: Tensor) {
    match self.grad.take() {
        None => self.grad = Some(g),                 // 首次
        Some(prev) => self.grad = Some(prev.add(&g)),// 扇出：SUM
    }
}

任何节点（叶子或中间）都累加：中间节点需要完整 grad 才能继续链式；叶子的累加结果就是输出。一个张量喂多个消费者时，多路 push_grad 自动求和。mul(&x, &x) 这类「同一 Var 进同一节点两次」也正确：parents=[x,x]（同指针），两次 push_grad 累加，拓扑去重保证 x 只遍历一次但收齐两路。测试 fanout_grad_accumulation 专门验证：y=x∘x + x∘x，dL/dx 须 = 4x（四处 x 全部求和）。

各算子 backward 数学

记上游梯度为 d（=本节点输出的梯度）。

op	forward	backward
matmul	C=A@B	dA=d@Bᵀ, dB=Aᵀ@d（复用 T3 matmul_backward）
add	a+b	da=d, db=d
mul	a∘b	da=d∘b, db=d∘a
add_bias	x[r,c]+bias[c]	dx=d, dbias[c]=Σ_r d[r,c]（沿广播维求和）
scale	x·α	dx=d·α
silu	x·σ(x)	dx=d·(σ + x·σ·(1−σ)), σ=σ(x)
swiglu	silu(g)∘u	由 silu+mul 组合自动得
rope	rotate_half 旋转	RoPE 是正交变换，dx = 用逆（转置）旋转作用于 d（角度 +θ 的转置 ≡ −θ）
softmax	row-wise safe softmax → y	Jacobian：dx[r,c]=y[r,c]·(d[r,c] − Σ_c' d·y)
cross_entropy	mean NLL(softmax(x), tgt)	dx = (probs − onehot)/rows，再乘上游标量 grad

RMSNorm（y[r,c]=x[r,c]·ir·γ[c], ir=rsqrt(mean(x²)+eps)）： 设 g[c]=d[r,c]·γ[c]，n=cols，

dx[r,c]   = ir·g[c] − x[r,c]·ir³/n·Σ_c'(g[c']·x[r,c'])
dγ[c]     = Σ_r d[r,c]·x[r,c]·ir

前向 cache 每行 inv_rms[r]，backward 直接复用，避免重算 reduce。

RoPE 反向推导：前向是 2×2 旋转矩阵 R(θ)，正交 ⇒ Rᵀ = R(−θ)。故

dx[i]      = d[i]·cos + d[i+h]·sin
dx[i+h]    = d[i+h]·cos − d[i]·sin

position=0 时旋转是恒等，backward 也恒等（sanity check）。

Softmax 反向推导：∂y_i/∂x_j = y_i(δ_ij − y_j)，链式后 dx_i = Σ_j d_j·y_i(δ_ij − y_j) = y_i(d_i − Σ_j d_j y_j)，即每行减去 Σ(d∘y) 后乘 y。

Cross-entropy 反向推导：L=−log softmax(x)[t]，softmax+NLL 的经典结果 ∂L/∂x_c = softmax_c − [c=t]；取 batch 平均 ⇒ 除以 rows。kernel 把 scale=upstream/rows 折进去。

Attention 用组合，不写 fused kernel

Kᵀ      = transpose(K)
scores  = scale(matmul(Q, Kᵀ), 1/√d)   # [s,s]
probs   = softmax(scores)
out     = matmul(probs, V)              # [s,d]

每一步都是已有 autograd 节点，backward() 自动沿 matmul→softmax→scale→matmul 链回传，得到 dQ/dK/dV，无需手写 attention backward。测试 attention_composed_bwd 单头单 batch 端到端 grad-check Q/K/V 三者。（transpose 在测试里用一个临时 Var::from_op 节点包，因为引擎暂未把 transpose 列为 op——T5 若需要再补。）

kernel 实现要点

• nn.cu 自带 inlined warp/block_reduce_sum/max（不引外部头文件，与现有 csrc/ 单文件风格一致）；block-reduce 末尾广播到全 block，便于 softmax/rms 的「标量广播」模式。
• 每个 op 各自 cudaDeviceSynchronize()（T3 约定，无 stream）。
• 全 F32、row-major、contiguous；cross-entropy target 用新增的 DType::I32。

验证方法

模板沿用 T3 gemm.rs::run_bwd：标量 loss L = sum(W∘out)，W 固定随机 ⇒ 上游 dOut = W；跑 op 的 backward() 拿 .grad()，对每个输入用 grad_check 与中心差分对拍。

• 每算子一个 grad-check（线性/双线性 op 用大 eps=1e-2、rel_tol=2e-2；非线性 op 用 eps=1e-3、rel_tol=3e-2、atol=1e-3 压住近零梯度）。
• 扇出累加：fanout_grad_accumulation，验证 dL/dx=4x。
• 组合 attention：attention_composed_bwd，端到端 grad-check dQ/dK/dV。
• 全部 #![cfg(not(no_cuda))] 门控；本地只 cargo check/fmt，构建+实跑在 dash5（8× RTX 5090, sm_120），capture 每 op 的 pass + max rel-err。