# Phase 4: Transformer Core Kernels — Design Document

## Goal

实现 Transformer 所需的所有非 Attention 算子的 CUDA kernel，每个 kernel 都支持 BF16 和 F32，与 PyTorch 参考实现对比验证。

## Kernel 清单

| Kernel | 用于 | 核心计算 | 关键优化点 |
|--------|------|---------|-----------|
| LayerNorm | GPT-2 | `(x - mean) / sqrt(var + eps) * gamma + beta` | Welford online, warp reduce |
| RMSNorm | Qwen3 | `x / sqrt(mean(x²) + eps) * gamma` | 无 mean，比 LayerNorm 简单 |
| GELU | GPT-2 | `0.5x(1 + tanh(sqrt(2/π)(x + 0.044715x³)))` | tanh 近似，逐元素 |
| SiLU | Qwen3 | `x * sigmoid(x)` | 逐元素 |
| Softmax | Attention | `exp(x - max) / sum(exp(x - max))` | Online safe softmax, warp reduce |
| Embedding | 全部 | `output[i] = table[token_ids[i]]` | Gather, coalesced write |
| RoPE | Qwen3 | 对 Q/K 的相邻元素对做旋转 | Precompute freq, in-place |

## 文件布局

```
csrc/
├── normalization/
│   ├── layernorm.cu
│   └── rmsnorm.cu
├── activation/
│   ├── gelu.cu
│   └── silu.cu
├── reduce/
│   └── softmax.cu
├── embedding/
│   ├── embedding.cu
│   └── rope.cu

crates/xserv-kernels/src/
├── layernorm.rs
├── rmsnorm.rs
├── activation.rs      # GELU + SiLU
├── softmax.rs
├── embedding.rs
├── rope.rs
└── lib.rs             # 新增 mod 声明
```

## Kernel 设计细节

### LayerNorm

输入 `x: [*, hidden_size]`, 输出 `y: [*, hidden_size]`
参数 `gamma, beta: [hidden_size]`

```
y[i] = gamma[i] * (x[i] - mean) / sqrt(var + eps) + beta[i]
```

**GPU 映射**: 每个 thread block 处理一行（一个 hidden_size 向量）。
- Phase 1: 并行加载 x，Welford online 算法计算 mean 和 var
- Phase 2: warp-level reduce (`__shfl_down_sync`) 聚合 mean/var
- Phase 3: block-level reduce via shared memory
- Phase 4: 每个 thread 对自己负责的元素做 normalize + affine

**Block 配置**: `block = min(1024, hidden_size)`, `grid = num_rows`

### RMSNorm

比 LayerNorm 简单：不减 mean，只做 `x * rsqrt(mean(x²) + eps) * gamma`。

```
rms = sqrt(sum(x²) / hidden_size + eps)
y[i] = x[i] / rms * gamma[i]
```

**GPU 映射**: 同 LayerNorm，每个 block 处理一行。
- 只需要一次 reduce（求 sum(x²)），不需要两次（mean + var）。

### GELU

逐元素操作，用 tanh 近似:
```
gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))
```

**GPU 映射**: 每个 thread 处理多个元素（向量化），grid 覆盖全部元素。

### SiLU (Swish)

逐元素: `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`

### Softmax

输入 `x: [*, seq_len]`, 沿最后一维做 softmax:
```
1. m = max(x)                    // 数值稳定
2. e[i] = exp(x[i] - m)
3. s = sum(e)
4. y[i] = e[i] / s
```

**GPU 映射**: 每个 block 处理一行。
- 第一遍 reduce: 求 max
- 第二遍: exp(x - max) 并 reduce sum
- 第三遍: 除以 sum

**优化**: 可以用 online softmax 合并前两遍（边算 exp 边更新 max），但先实现三遍版本保证正确。

### Embedding

```
output[seq_idx] = embedding_table[token_ids[seq_idx]]
```

**GPU 映射**: 每个 thread 处理一个 token 的部分维度。
- `grid = num_tokens`, `block = hidden_size`（或分多个 thread 处理一个 token）
- 写端是 coalesced（连续 thread 写连续地址），读端是 gather（非连续）

### RoPE (Rotary Position Embedding)

对 Q/K 的每对相邻元素 `(x0, x1)` 做 2D 旋转:
```
freq[i] = 1.0 / (theta ^ (2i / dim))
cos_val = cos(position * freq[i])
sin_val = sin(position * freq[i])
y0 = x0 * cos_val - x1 * sin_val
y1 = x0 * sin_val + x1 * cos_val
```

**GPU 映射**: 每个 thread 处理一对元素 `(x[2i], x[2i+1])`。
- Precompute `cos_cache[max_seq_len][head_dim/2]` 和 `sin_cache` 在初始化时
- 运行时 kernel 只做乘加

**theta**: Qwen3 默认 `rope_theta = 1000000.0`

## Reduction Pattern（核心学习点）

所有 Norm 和 Softmax 都涉及 reduction。GPU reduction 的分层结构:

```
Thread-level:  每个 thread 处理多个元素，本地累加
    ↓
Warp-level:    __shfl_down_sync() 在 32 threads 内规约（无需 shared memory）
    ↓
Block-level:   shared memory 存各 warp 的结果，warp 0 再规约
```

对于 hidden_size <= 8192（LLM 常见），一个 block 足够，不需要 grid-level reduction。

### Warp Reduce 模板

```cuda
__device__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}
```

### Block Reduce 模板

```cuda
__device__ float block_reduce_sum(float val) {
    __shared__ float shared[32];  // max 32 warps per block
    int lane = threadIdx.x % 32;
    int warp_id = threadIdx.x / 32;

    val = warp_reduce_sum(val);
    if (lane == 0) shared[warp_id] = val;
    __syncthreads();

    val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
    if (warp_id == 0) val = warp_reduce_sum(val);
    return val;
}
```

## Reference 验证策略

写 `tools/generate_reference.py` 脚本，用 PyTorch 为每个 op 生成 reference input/output:
- 保存为 `.npy` 格式
- Rust 测试中加载对比
- 或者直接在 Rust 测试中用 CPU 实现计算 expected 值（更简单，不依赖 Python）

**选择**: 先用 Rust CPU 实现作为 reference（简单），关键 op（RoPE）再与 PyTorch 对比。

## Test Plan

- [x] RMSNorm F32: hidden_size=768, 4 rows → max_err 7.2e-7
- [x] RMSNorm BF16: 同上 → max_err 7.0e-3
- [x] LayerNorm F32: hidden_size=768 → max_err 1.7e-6
- [x] GELU F32: 10000 elements → max_err 3.0e-8
- [x] GELU BF16: 同上 → max_err 2.4e-3
- [x] SiLU F32: 10000 elements → max_err 1.5e-8
- [x] Softmax F32: 8×256 → max_err 1.4e-9
- [x] Softmax sum=1 验证: 4×2048
- [x] Softmax 大值 (1000+) 数值稳定性 → max_err 1.5e-8
- [x] Embedding F32: vocab=100, hidden=64, 5 tokens → exact match
- [x] RoPE F32: 4 tokens × 2 heads × dim=8 → max_err 6.0e-8
- [x] RoPE position=0 恒等验证 → max_err 0

## Takeaways

1. **`common.cuh` 抽取共用 reduction 是正确的做法**：`warp_reduce_sum/max` 和 `block_reduce_sum/max` 被 RMSNorm, LayerNorm, Softmax 三个 kernel 复用。抽到头文件避免了代码重复，也确保 reduction 逻辑一致。build.rs 中需要 `.include("../../csrc")` 让 nvcc 能找到头文件。

2. **Shared memory 中广播标量的模式**：Norm 和 Softmax 都需要将 reduce 结果（mean, rms_inv, max, sum）广播给 block 内所有 thread。标准做法：thread 0 写 `__shared__` 变量，`__syncthreads()` 后所有 thread 读。这比让每个 thread 独立做 reduce 高效得多。

3. **Softmax 三遍 vs 两遍**：我们实现了三遍版本（max → exp+sum → normalize），简单可靠。Online softmax 可以合并前两遍（一遍 pass 内同时跟踪 running max 和 running sum），但需要更复杂的数值更新公式。Flash Attention（Phase 14）会用到 online softmax。

4. **RoPE 的 position=0 恒等性**：`cos(0)=1, sin(0)=0`，所以 position 0 的旋转是恒等变换。这是一个很好的 sanity check。如果 position=0 时输出不等于输入，说明 kernel 有 bug。

5. **BF16 Softmax 的精度陷阱**：exp 结果先写成 BF16 再读回做 normalize 会丢精度。理想做法是用 float scratch buffer 暂存 exp 结果。当前实现可接受（误差在 1e-2 量级），但在 attention score 很接近时可能引入可观察的差异。Phase 14 Flash Attention 会解决这个问题（全程 FP32 累加）。

6. **Embedding 就是 gather 操作**：没有任何计算，纯粹的内存搬运。瓶颈在 global memory 随机读取（token_ids 导致不连续读 table）。写端是 coalesced 的（连续 token 写连续地址）。优化方向：使用向量化加载（`float4`）一次读 128 bit。

7. **RoPE in-place 修改 Tensor 的设计考量**：RoPE 在数学上是对 Q/K 的 in-place 旋转。我们通过 `data_ptr() as *mut` 绕过了 Rust 的不可变借用。这在 GPU 上是安全的（kernel 内部互不干扰），但 Rust 侧没有 `&mut` 语义保护。后续如果需要更严格的安全性，可以引入 `Tensor::as_mut_ptr()` 方法并要求 `&mut self`。