# Phase 3: GEMM — Design Document

## Goal

实现矩阵乘法的多个版本（naive → tiled → cuBLAS），建立 benchmark 对比框架，深入理解 GPU 编程中的内存访问模式和优化手段。

## Module Layout

```
csrc/gemm/
├── naive.cu          # 每个 thread 算一个输出元素
└── tiled.cu          # shared memory tiling, 32x32 tiles

crates/xserv-kernels/
├── build.rs          # 编译 .cu + 链接 cublas
└── src/
    ├── lib.rs
    └── gemm.rs       # FFI 封装, GemmBackend enum, matmul(), CublasContext
```

## Kernel Implementations

### Version 1: Naive GEMM

```
Grid:  (ceil(N/16), ceil(M/16))
Block: (16, 16)
每个 thread: C[row][col] = sum_k(A[row][k] * B[k][col])
```

- 每个 thread 独立遍历 K 维度做点积
- 所有读取走 global memory，无局部性优化
- BF16 版本在 FP32 中累加（`__bfloat162float` → 累加 → `__float2bfloat16`）

### Version 2: Tiled GEMM (Shared Memory)

```
TILE_SIZE = 32
Grid:  (ceil(N/32), ceil(M/32))
Block: (32, 32) = 1024 threads

每个 tile iteration:
  1. 协作加载 A[tile] 和 B[tile] 到 shared memory
  2. __syncthreads()
  3. 在 shared memory 中做 32 次乘加
  4. __syncthreads()
```

- 每个 global memory 读取被 TILE_SIZE 个 thread 复用
- 理论上减少 global memory 访问 TILE_SIZE 倍
- BF16 版本同样在 shared memory 中存 float（FP32 累加）

### Version 3: cuBLAS

- `cublasGemmEx` 支持混合精度
- **Row-major 适配**：cuBLAS 使用 column-major 布局，我们的 tensor 是 row-major
  - 利用恒等式：`C = A @ B` (row-major) ⟺ `C^T = B^T @ A^T` (col-major)
  - 传入 `CUBLAS_OP_N`，让 cuBLAS 把我们的 row-major 数据当作 col-major 的转置
  - 参数：`m=N, n=M, k=K, lda=N (B), ldb=K (A), ldc=N (C)`

### Backend Registry

```rust
pub enum GemmBackend { Naive, Tiled, CuBlas }
pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor;
```

运行时可切换 backend，方便 benchmark 对比和逐步替换。

## CublasContext

RAII 封装 `cublasHandle_t`，Drop 时调 `cublasDestroy_v2`。
目前每次 matmul 创建一个新 handle，后续优化为全局复用。

## Test Plan

- [x] F32: naive/tiled/cuBLAS × small(4)/medium(64-256)/rect(65x33x97) 
- [x] BF16: naive/tiled/cuBLAS × small/medium
- [x] 三种 backend 在相同输入上输出一致（cross-backend consistency）
- [x] 非方阵测试（M≠N≠K）
- [x] 1024x1024 cuBLAS 验证

## Takeaways

1. **Row-major vs Column-major 陷阱**：这是 GEMM 实现中最容易出错的地方。cuBLAS 的 column-major 假设与 C/Rust 的 row-major 冲突。理解 `C=AB` ⟺ `C^T=B^T A^T` 这个恒等式是关键。实际做法：不做任何显式转置，只是交换 A/B 的传入顺序和调整 leading dimension 参数。

2. **BF16 的累加精度**：BF16 只有 ~3 位有效数字（vs FP32 的 ~7 位）。如果在 BF16 中累加 K 次乘法，误差会快速放大。正确做法是**在 FP32 中累加，最后才转回 BF16**。我们的 naive 和 tiled kernel 都遵循了这一点（`float sum = 0.0f`）。cuBLAS 通过 `CUBLAS_COMPUTE_32F` 参数控制。

3. **Shared memory tiling 的核心思想**：global memory 带宽是 GPU 计算的主要瓶颈。通过 shared memory tiling，每个数据从 global memory 读一次，被 TILE_SIZE 个 thread 复用。对于 TILE_SIZE=32，理论上减少 32 倍 global memory 访问。

4. **`__syncthreads()` 的位置关键**：tile 加载后必须同步（确保所有 thread 写完 shared memory），计算后也要同步（防止下一轮加载覆盖还在使用的数据）。漏掉任何一个 sync 都会产生 race condition 导致结果错误。

5. **cuBLAS handle 开销**：每次 matmul 创建/销毁 handle 有~0.1ms 开销。生产环境应全局复用一个 handle。Phase 15（性能优化）时需要修复这个问题。

6. **`error::check` 需要 pub**：Phase 1 中 `check()` 是 `pub(crate)`，Phase 3 需要跨 crate 调用。反思：基础设施 crate 的错误处理函数应该从一开始就设计为 public API。

## 后续优化方向（Phase 15）

- Register tiling（每个 thread 算多个输出元素）
- Tensor Core WMMA（利用 5090 的硬件加速）
- CublasContext 全局复用
- 非 contiguous input 支持（避免 matmul 前的拷贝）