# xserv — LLM Inference Engine Roadmap

> 从零用 Rust + CUDA 构建一个完整的 LLM 推理引擎，目标是深入理解 LLM Serving 全栈技术。

## 设计决策

| 决策项 | 选择 | 备注 |
|--------|------|------|
| 抽象层级 | Level 0.5 | 自写 CUDA kernel + cuBLAS 可切换，便于 benchmark 对比 |
| 硬件 | 8×RTX 5090 (Blackwell, CC 12.0, 32GB GDDR7) | 纯 PCIe Gen5 x16 互联，无 NVLink (详见下方硬件拓扑) |
| 语言 | Rust + CUDA (C/C++) | Rust FFI 调用 CUDA |
| 起步模型 | GPT-2 124M → Qwen3-8B | 从简单到实用 |
| 精度 | BF16/FP16 | 后期扩展 FP8 |
| Tensor | 自己实现 | 完整学习 tensor 抽象设计 |
| Tokenizer | 自己实现 BPE | 学习分词机制 |
| 权重格式 | safetensors | Rust 友好，零拷贝 mmap |
| Async Runtime | tokio | 成熟稳定，不引入性能问题 |
| API | OpenAI 兼容 | `/v1/chat/completions`，SSE streaming |
| 时间线 | 不限 | 学习为主，每步验证 |

## 硬件拓扑 (dash5, 已确认 2026-05-21)

**GPU**: 8× NVIDIA GeForce RTX 5090, 32607 MiB, Compute Capability 12.0
**CUDA Toolkit**: 12.9 (安装于 `/usr/local/cuda-12.9`，需将 `bin/` 加入 PATH)
**PCIe**: Gen 5 x16 (理论单向 ~64 GB/s，空闲时降频至 Gen 1)

**互联拓扑** (`nvidia-smi topo -m`):
```
        GPU0  GPU1  GPU2  GPU3  GPU4  GPU5  GPU6  GPU7
GPU0     X    PHB   PHB   PHB   NODE  NODE  NODE  NODE
GPU1    PHB    X    PHB   PHB   NODE  NODE  NODE  NODE
GPU2    PHB   PHB    X    PHB   NODE  NODE  NODE  NODE
GPU3    PHB   PHB   PHB    X    NODE  NODE  NODE  NODE
GPU4    NODE  NODE  NODE  NODE   X    PHB   PHB   PHB
GPU5    NODE  NODE  NODE  NODE  PHB    X    PHB   PHB
GPU6    NODE  NODE  NODE  NODE  PHB   PHB    X    PHB
GPU7    NODE  NODE  NODE  NODE  PHB   PHB   PHB    X

PHB  = 同一 PCIe Host Bridge（同组，延迟低）
NODE = 跨 PCIe Host Bridge（跨组，延迟较高）
```

**分组**: GPU 0-3 为一组, GPU 4-7 为一组。组内 PHB 互联，跨组 NODE 互联。

**对设计的影响**:
- **无 NVLink**: AllReduce 带宽受限于 PCIe (~64 GB/s vs NVLink ~450 GB/s)
- **TP 策略**: 当前阶段目标 TP=1/2/4，在同组内 (0-3 或 4-7) 执行，全 PHB 互联
- **跨组并行 (TP=8, PP 等)**: 留待后续扩展
- **CPU Affinity**: GPU 0-3 亲和 CPU 0-127, GPU 4-7 亲和 CPU 0-186（NUMA 0-1）

## 项目结构

```
xserv/
├── Cargo.toml                    # workspace root
├── csrc/                         # CUDA 源文件 (.cu / .cuh)
│   ├── gemm/                     #   GEMM kernels (naive, tiled, tensor core)
│   ├── attention/                #   Attention kernels (naive, flash, paged)
│   ├── normalization/            #   LayerNorm, RMSNorm
│   ├── activation/               #   GELU, SiLU
│   ├── embedding/                #   Embedding lookup, RoPE
│   ├── reduce/                   #   Softmax, argmax, sampling
│   └── quantize/                 #   FP8/INT8 kernels
├── crates/
│   ├── xserv-cuda/               # Phase 1: CUDA FFI, context, stream, allocator
│   ├── xserv-tensor/             # Phase 2: Tensor type, ops dispatch, DType
│   ├── xserv-kernels/            # Phase 3-5: kernel registry (custom + cuBLAS)
│   ├── xserv-tokenizer/          # Phase 7: BPE tokenizer
│   ├── xserv-model/              # Phase 6,8,10: model def + weight loading
│   ├── xserv-runtime/            # Phase 9,11,12: KV cache, paging, scheduler
│   ├── xserv-engine/             # Phase 13: inference engine orchestration
│   ├── xserv-api/                # Phase 13: HTTP server + OpenAI compat
│   ├── xserv-speculative/        # Phase 16: speculative decoding
│   └── xserv-distributed/        # Phase 17: tensor parallelism, NCCL
├── tests/                        # integration tests
├── benches/                      # criterion benchmarks
├── tools/                        # 辅助脚本 (PyTorch reference output 生成等)
└── docs/                         # 每个 phase 的设计文档
```

## Phase 依赖图

```
Phase 0: 项目脚手架 + 环境验证
    │
Phase 1: CUDA FFI 基础设施
    │
Phase 2: Tensor 抽象层
    │
Phase 3: GEMM (naive → tiled → tensor core → cuBLAS)
    │
Phase 4: Transformer Kernels (Norm, Activation, Embedding, RoPE, Softmax)
    │
Phase 5: Attention Kernel (naive MHA)
    │
Phase 6: 模型加载 (safetensors + HF config)
    │              │
    │         Phase 7: BPE Tokenizer (可与 Phase 6 并行)
    │              │
Phase 8: GPT-2 完整推理 ◄──────────── 里程碑 ① CLI 文本生成
    │
Phase 9: KV Cache + Autoregressive Generation
    │
Phase 10: Qwen3-8B 支持 ◄─────────── 里程碑 ② 8B 模型推理
    │
Phase 11: Paged Attention + KV Cache Manager
    │
Phase 12: Continuous Batching + Request Scheduler
    │
Phase 13: HTTP API + SSE Streaming ◄── 里程碑 ③ 端到端 API 可用
    │
Phase 14: Flash Attention (FA2 for SM120)
    │
Phase 15: 性能优化 ◄──────────────── 里程碑 ④ 50% vLLM throughput
    │
Phase 16: Speculative Decoding
    │
Phase 17: Tensor Parallelism (TP=1/2/4) ◄── 里程碑 ⑤ 多卡推理
    │
Phase 18: 量化 (FP8 / INT8)
    │
Phase 19: Multimodal ◄────────────── 里程碑 ⑥ 视觉问答
```

---

## Phase 0: 项目脚手架 + 环境验证

**目标**: 搭建 Cargo workspace，验证 CUDA 工具链，确保开发环境就绪。

**技术要点**:
- Cargo workspace 配置，所有 crate 共享依赖版本
- `build.rs` 中用 `cc` crate 编译 `.cu` 文件的 pipeline
- 验证 CUDA toolkit 版本: **已确认 CUDA 12.9** (`/usr/local/cuda-12.9`)
- 验证 GPU compute capability: **已确认 CC 12.0** (Blackwell)
- 确认 `nvcc` 在 PATH 中 (需要 `export PATH=/usr/local/cuda-12.9/bin:$PATH`)
- ~~运行 `nvidia-smi topo -m` 确认互联拓扑~~ **已确认: 纯 PCIe Gen5, 无 NVLink**

**外部依赖**: `cc` crate（编译 CUDA）

**测试验收**:
- `cargo build` 通过
- 一个最小的 `.cu` kernel（向量加法）能从 Rust 调用并返回正确结果
- 输出 GPU 信息（名称、显存、compute capability）

**设计文档**: `docs/01-cuda-ffi.md`（与 Phase 1 合并）

---

## Phase 1: CUDA FFI 基础设施

**Crate**: `xserv-cuda`

**目标**: 封装 CUDA Runtime API，提供安全的 Rust 抽象层。

### 模块划分

```
xserv-cuda/src/
├── lib.rs
├── error.rs        # CudaError 类型, cudaGetLastError 封装
├── device.rs       # Device 枚举, 设备查询 (属性/数量/当前设备)
├── context.rs      # CUDA context 管理
├── stream.rs       # CudaStream (异步操作流)
├── memory.rs       # GPU 内存分配/释放/拷贝 (H2D, D2H, D2D)
├── allocator.rs    # Caching Allocator (显存池)
└── module.rs       # cuModuleLoad (加载 PTX/cubin, 可选)
```

### 关键技术点

1. **FFI 绑定策略**:
   - 手写 `extern "C"` 绑定核心 CUDA Runtime API（~30 个函数）
   - 不用 bindgen，保持可控和可读
   - 需要绑定的 API: `cudaMalloc`, `cudaFree`, `cudaMemcpy`, `cudaMemcpyAsync`,
     `cudaStreamCreate`, `cudaStreamSynchronize`, `cudaGetDeviceProperties`,
     `cudaSetDevice`, `cudaDeviceSynchronize`, `cudaGetLastError` 等

2. **GpuBuffer 抽象**:
   ```rust
   pub struct GpuBuffer {
       ptr: *mut c_void,
       size_bytes: usize,
       device: usize,
   }

   impl Drop for GpuBuffer {
       fn drop(&mut self) { /* cudaFree or return to allocator */ }
   }
   ```
   - `Drop` trait 自动释放，防止 GPU 内存泄漏
   - 不实现 `Clone`（显式 `copy_from` 代替）

3. **Caching Allocator**:
   - 维护 free list（按大小分桶，桶边界: 512B, 1KB, 2KB, ..., 1GB）
   - `alloc(size)`: 在对应桶中找 >= size 的 free block，miss 时 `cudaMalloc`
   - `free(ptr, size)`: 不调 `cudaFree`，放回 free list
   - `trim()`: 真正释放所有 free blocks（OOM 恢复时用）
   - 这是性能关键组件——频繁 `cudaMalloc/cudaFree` 会严重影响 throughput
   - 参考: PyTorch 的 `CUDACachingAllocator` 设计

4. **Stream 管理**:
   - 每个 stream 是独立的 GPU 执行队列
   - Kernel launch 和 memcpy 是异步的（提交到 stream 后立即返回）
   - `stream.synchronize()` 等待该 stream 上所有操作完成
   - 后续用于 overlap compute 和 memory transfer

5. **Error Handling**:
   ```rust
   #[derive(Debug)]
   pub enum CudaError {
       OutOfMemory,
       InvalidDevice,
       LaunchFailure,
       Raw { code: i32, message: String },
   }

   // 所有 CUDA 调用包装为 Result
   pub(crate) fn check(code: cudaError_t) -> Result<(), CudaError>;
   ```

### 测试验收

- [ ] 分配 1GB GPU 内存，H2D 拷贝一个大数组，D2H 拷回，验证数据一致
- [ ] Caching allocator: alloc → free → re-alloc same size，第二次不触发 `cudaMalloc`（通过内部计数验证）
- [ ] 多 stream 并发拷贝两个数组，验证结果正确
- [ ] 设备查询: 打印 GPU name, total memory, compute capability, SM count
- [ ] Benchmark: caching allocator vs 裸 `cudaMalloc` 的分配延迟对比（100 次 alloc/free 循环）

---

## Phase 2: Tensor 抽象层

**Crate**: `xserv-tensor`

**目标**: 实现核心 Tensor 类型，支持 CPU/GPU 存储、多种数据类型、视图操作。

### 核心数据结构

```rust
// --- 数据类型 ---
#[derive(Clone, Copy, PartialEq)]
pub enum DType {
    F32,
    F16,
    BF16,
    // 后期: U8, I8, F8E4M3, F8E5M2
}

impl DType {
    pub fn size_bytes(&self) -> usize; // F32=4, F16=2, BF16=2
}

// --- 设备 ---
#[derive(Clone, Copy, PartialEq)]
pub enum Device {
    Cpu,
    Cuda(usize),  // device ordinal
}

// --- 存储 ---
// 引用计数，支持 view（多个 Tensor 共享同一 Storage）
pub struct Storage(Arc<StorageInner>);

enum StorageInner {
    Cpu { data: Vec<u8> },
    Cuda { buffer: GpuBuffer, device: usize },
}

// --- Tensor ---
pub struct Tensor {
    storage: Storage,
    shape: SmallVec<[usize; 4]>,    // 维度（大多数 tensor <= 4D）
    strides: SmallVec<[usize; 4]>,  // 步长（以元素为单位）
    offset: usize,                   // storage 中的起始偏移（元素数）
    dtype: DType,
    device: Device,
}
```

### 关键技术点

1. **Strided Layout**:
   - 支持 `transpose`, `slice`, `permute` 等操作不拷贝数据，只改 strides/offset
   - `is_contiguous()`: strides 从右到左依次为 1, shape[-1], shape[-1]*shape[-2], ...
   - 非 contiguous tensor 在送入 CUDA kernel 前需要 `contiguous()` 拷贝为连续布局
   - 例: `[3,4]` tensor 的 strides = `[4, 1]`；transpose 后 shape=`[4,3]`, strides=`[1, 4]`

2. **BF16/F16 在 Rust 中的表示**:
   - 使用 `half` crate 的 `bf16` 和 `f16` 类型
   - GPU kernel 中使用 `__nv_bfloat16` / `__half`
   - Tensor 内部存储为 raw bytes，通过 DType dispatch 解释

3. **设备间拷贝**:
   ```rust
   impl Tensor {
       pub fn to(&self, device: Device) -> Tensor; // CPU↔GPU 拷贝
       pub fn to_dtype(&self, dtype: DType) -> Tensor; // 类型转换
   }
   ```

4. **基础操作**（此阶段实现）:
   - **创建**: `zeros`, `ones`, `from_slice`, `rand`, `full`, `arange`
   - **形状**: `reshape`, `view`, `transpose`, `squeeze`, `unsqueeze`, `contiguous`
   - **逐元素** (CPU + GPU kernel): `add`, `mul`, `sub`, `div`
   - **广播 (Broadcasting)**: NumPy 语义，维度从尾部对齐
   - **归约**: `sum`, `max`, `mean`（沿指定轴）

5. **Op Dispatch 机制**:
   ```rust
   // 根据 device 和 dtype dispatch 到不同实现
   pub fn add(a: &Tensor, b: &Tensor) -> Tensor {
       match (a.device(), b.device()) {
           (Device::Cpu, Device::Cpu) => cpu_ops::add(a, b),
           (Device::Cuda(_), Device::Cuda(_)) => cuda_ops::add(a, b),
           _ => panic!("device mismatch"),
       }
   }
   ```

### 测试验收

- [ ] 创建 tensor, reshape, transpose, slice，验证 shape/strides 计算正确
- [ ] 广播加法: `[3,1] + [1,4]` → `[3,4]`，与 numpy 结果对比
- [ ] CPU ↔ GPU 拷贝往返，数据一致
- [ ] BF16 tensor 的基础运算精度验证（与 FP32 结果对比 relative error）
- [ ] View 共享存储: 修改 view 的数据，原 tensor 也应变化
- [ ] Benchmark: GPU 逐元素 kernel vs CPU 的加速比（大数组）

---

## Phase 3: GEMM — 矩阵乘法

**Crate**: `xserv-kernels`
**CUDA 源码**: `csrc/gemm/`

**目标**: 实现 GEMM 的多个版本，从 naive 到 tensor core，同时封装 cuBLAS，建立 benchmark 对比框架。

这是 CUDA kernel 编程的第一个"修罗场"，会深刻理解 GPU 编程的核心概念。

### 实现路线（4 个递进版本）

#### Version 1: Naive GEMM
- 每个 thread 计算输出矩阵 C 的一个元素: `C[i][j] = sum(A[i][k] * B[k][j])`
- grid 维度: `(M/BLOCK, N/BLOCK)`, block 维度: `(BLOCK, BLOCK)`
- **学到**: grid/block 维度规划, global memory access pattern
- **问题**: global memory 访问完全没有局部性，bandwidth 利用率极低
- **预期性能**: ~1-2% cuBLAS

#### Version 2: Tiled GEMM (shared memory)
- 将 A, B 分成 TILE×TILE 的小块，加载到 shared memory
- 每个 thread block 计算 C 的一个 TILE×TILE 输出块
- 内层循环沿 K 维度滑动 tile
- **学到**: shared memory 使用, `__syncthreads()`, bank conflict, memory coalescing
- **关键**: A 的 tile 要按行加载（coalesced），B 的 tile 按列访问需要注意 bank conflict
- **预期性能**: ~10-20% cuBLAS

#### Version 3: Register Tiling + 向量化
- 每个 thread 计算多个输出元素（如 4×4 或 8×8）
- 使用寄存器存储中间结果，减少 shared memory 访问
- 向量化加载: `float4` 一次读 128 bit
- **学到**: register pressure, ILP (Instruction-Level Parallelism), occupancy vs. ILP tradeoff
- **预期性能**: ~30-50% cuBLAS

#### Version 4: Tensor Core GEMM (WMMA)
- 使用 CUDA WMMA API 调用 Tensor Core
- BF16 输入, FP32 累加
- 每次 wmma::mma_sync 计算 16×16×16 矩阵乘
- **学到**: WMMA fragment layout, Tensor Core 编程模型, warp-level 协作
- **关键**: 5090 Blackwell (CC 12.0) 的 Tensor Core 支持 BF16 和 FP8
- **预期性能**: ~60-80% cuBLAS

### cuBLAS 封装

```rust
// 需要封装的 cuBLAS API
extern "C" {
    fn cublasCreate_v2(handle: *mut cublasHandle_t) -> cublasStatus_t;
    fn cublasSetStream_v2(handle: cublasHandle_t, stream: cudaStream_t) -> cublasStatus_t;
    fn cublasGemmEx(
        handle: cublasHandle_t,
        transa: cublasOperation_t, transb: cublasOperation_t,
        m: i32, n: i32, k: i32,
        alpha: *const c_void,
        A: *const c_void, Atype: cudaDataType, lda: i32,
        B: *const c_void, Btype: cudaDataType, ldb: i32,
        beta: *const c_void,
        C: *mut c_void, Ctype: cudaDataType, ldc: i32,
        computeType: cublasComputeType_t,
        algo: cublasGemmAlgo_t,
    ) -> cublasStatus_t;
}
```

支持: BF16×BF16→BF16 (compute=FP32), FP16×FP16→FP16, FP32×FP32→FP32

### Kernel Registry（运行时可切换 backend）

```rust
#[derive(Clone, Copy)]
pub enum GemmBackend {
    Naive,
    Tiled,
    RegisterTiled,
    TensorCore,
    CuBlas,
}

pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor;

// 全局默认 backend（可配置）
pub fn set_default_gemm_backend(backend: GemmBackend);
```

### 测试验收

- [ ] 正确性: 所有 5 个 backend 的输出与 cuBLAS 对比，max absolute error < 1e-3 (BF16)
- [ ] Benchmark 表格（用 `criterion` crate）:

| Backend | M=N=K=1024 | M=N=K=4096 | % of cuBLAS |
|---------|------------|------------|-------------|
| Naive | ms | ms | % |
| Tiled | ms | ms | % |
| RegisterTiled | ms | ms | % |
| TensorCore | ms | ms | % |
| cuBLAS | ms | ms | 100% |

- [ ] Profile: 用 `nsys`/`ncu` 分析 naive vs tiled 的 memory throughput 差异
- [ ] 非方阵测试: M=1, N=4096, K=4096 (decode 阶段的典型 shape)

---

## Phase 4: Transformer 核心 Kernels

**Crate**: `xserv-kernels`
**CUDA 源码**: `csrc/normalization/`, `csrc/activation/`, `csrc/embedding/`, `csrc/reduce/`

**目标**: 实现 Transformer 所需的所有非 Attention 算子，每个都有自定义 CUDA kernel。

### Kernel 清单

| Kernel | 用途 | CUDA 文件 | 核心优化点 |
|--------|------|-----------|-----------|
| LayerNorm | GPT-2 | `normalization/layernorm.cu` | Online Welford 算法, warp reduce, 向量化加载 |
| RMSNorm | Qwen3/LLaMA | `normalization/rmsnorm.cu` | 比 LayerNorm 简单（无 mean）, rsqrt |
| GELU | GPT-2 激活 | `activation/gelu.cu` | tanh 近似 vs 精确, 向量化 |
| SiLU (Swish) | Qwen3 激活 | `activation/silu.cu` | `x * sigmoid(x)`, 逐元素 |
| SwiGLU | Qwen3 FFN | `activation/swiglu.cu` | `SiLU(gate) * up`, fused 逐元素 |
| Embedding | token→vector | `embedding/embedding.cu` | Gather 操作, coalesced access |
| RoPE | Qwen3 位置编码 | `embedding/rope.cu` | 复数旋转, precompute freq |
| Softmax | Attention 内 | `reduce/softmax.cu` | Online safe softmax, 数值稳定 |
| Argmax | Greedy sampling | `reduce/argmax.cu` | Parallel reduction |
| TopK | TopK sampling | `reduce/topk.cu` | Bitonic sort 或 radix select |

### 关键学习主题

**Reduction Pattern（核心中的核心）**:

LayerNorm, RMSNorm, Softmax 都涉及对某个维度求和/求最大值。GPU reduction 是分层的:

```
Thread-level:   每个 thread 处理多个元素，本地累加
    ↓
Warp-level:     __shfl_down_sync() 在 warp (32 threads) 内规约
    ↓
Block-level:    shared memory 存各 warp 的结果，再规约
    ↓
Grid-level:     (如果需要) atomic 或两遍 kernel
```

对于 Norm/Softmax，通常 hidden_dim <= 8192，一个 block 就够，不需要 grid-level reduction。

**向量化内存访问**:
- `float4` (128-bit) 一次加载 4 个 float 或 8 个 bf16
- `__nv_bfloat162` 一次处理 2 个 bf16
- 提升 memory throughput，减少 load/store 指令数

**每个 kernel 都实现两个版本**:
1. Custom CUDA kernel（自己写，深入理解）
2. Reference 实现（简单的 Python/numpy，生成 reference output 用于验证）

### 测试验收

- [ ] 每个 kernel 的输出与 PyTorch 参考实现对比
  - 写一个 `tools/generate_reference.py` 脚本，为每个 op 生成 reference input/output，保存为 `.npy`
  - Rust 测试中加载 `.npy` 对比
- [ ] 数值精度: BF16 下 max relative error < 1e-2, FP32 下 < 1e-5
- [ ] RoPE: 验证旋转后的向量与 HF transformers 的 `apply_rotary_pos_emb` 结果一致
- [ ] Softmax: 验证 `sum(output, dim=-1) == 1.0`，验证 numerical stability（大值输入不 overflow）
- [ ] Benchmark: 每个 kernel 与 PyTorch 对应操作的延迟对比

---

## Phase 5: Attention Kernel (Naive 版)

**Crate**: `xserv-kernels`
**CUDA 源码**: `csrc/attention/naive_attention.cu`

**目标**: 实现标准 Multi-Head Attention，不做 Flash/Paged 优化。理解 attention 机制的计算基础。

### 计算流程

```
Input: Q [B, H, S, D], K [B, H, S, D], V [B, H, S, D]
       其中 B=batch, H=num_heads, S=seq_len, D=head_dim

1. scores = Q @ K^T / sqrt(D)           → [B, H, S, S]
2. scores = scores + causal_mask         → 上三角置为 -inf
3. weights = softmax(scores, dim=-1)     → [B, H, S, S]
4. output = weights @ V                  → [B, H, S, D]
```

### 实现方式

**方式一: 组合式（先跑通）**
- 用 Phase 3 的 GEMM (Q@K^T) + Phase 4 的 Softmax + GEMM (weights@V)
- 简单但 materialize 了 S×S 矩阵，内存 O(S²)

**方式二: Fused kernel（理解 Flash Attention 的前置）**
- 一个 kernel 完成整个 attention
- 仍然 materialize S×S（不做 tiling），但减少 kernel launch 和 global memory 读写次数

### Causal Mask

- 不显式构造 mask 矩阵（浪费内存）
- 在 softmax 前对 `scores[i][j]` where `j > i` 写 `-inf`（`-1e9` for BF16）
- 编译期条件判断: `if (col > row) score = -inf;`

### GQA 预备

- 本阶段实现标准 MHA: `num_kv_heads == num_heads`
- Phase 10 扩展为 GQA: `num_kv_heads < num_heads`
- GQA 时 K/V 需要 repeat: 每个 KV head 服务 `num_heads / num_kv_heads` 个 Q head
- 实际实现: 不真正 repeat 数据，在 kernel 中用 `kv_head_idx = q_head_idx / num_groups` 索引

### 测试验收

- [ ] 随机 Q, K, V，输出与 PyTorch `F.scaled_dot_product_attention(is_causal=True)` 对比
- [ ] 验证 causal mask: attention weight 矩阵的上三角全为 0
- [ ] Benchmark 表（记录为 Flash Attention 的 baseline）:

| Seq Length | Latency (ms) | GPU Memory (MB) |
|------------|-------------|-----------------|
| 128 | | |
| 512 | | |
| 2048 | | |
| 4096 | | |
| 8192 | | OOM? |

- [ ] 输出 attention weights 可视化（小规模，验证 causal pattern）

---

## Phase 6: 模型加载

**Crate**: `xserv-model`

**目标**: 从 HuggingFace safetensors 文件加载模型权重到 GPU Tensor。

### 核心组件

#### 1. safetensors 解析

- 使用 `safetensors` crate 读取文件
- 文件结构: 8 bytes header_size + JSON header + raw tensor data
- 支持 mmap 零拷贝读取
- 支持 sharded 文件: `model-00001-of-00003.safetensors`
  - 通过 `model.safetensors.index.json` 查找 tensor → file 的映射

#### 2. HF Config 解析

```rust
#[derive(Deserialize)]
pub struct ModelConfig {
    pub architectures: Vec<String>,
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub num_attention_heads: usize,
    pub num_key_value_heads: usize,    // GQA: 可能 < num_attention_heads
    pub num_hidden_layers: usize,
    pub vocab_size: usize,
    pub max_position_embeddings: usize,
    pub rms_norm_eps: f64,             // Qwen3 用
    pub rope_theta: f64,               // RoPE base frequency
    pub tie_word_embeddings: bool,
    // ... 其他字段按模型按需添加
}
```

#### 3. 权重映射

HuggingFace 命名规范 (以 Qwen3 为例):
```
model.embed_tokens.weight                          → embedding
model.layers.{i}.self_attn.q_proj.weight           → layer[i].attn.q_proj
model.layers.{i}.self_attn.k_proj.weight           → layer[i].attn.k_proj
model.layers.{i}.self_attn.v_proj.weight           → layer[i].attn.v_proj
model.layers.{i}.self_attn.o_proj.weight           → layer[i].attn.o_proj
model.layers.{i}.mlp.gate_proj.weight              → layer[i].mlp.gate
model.layers.{i}.mlp.up_proj.weight                → layer[i].mlp.up
model.layers.{i}.mlp.down_proj.weight              → layer[i].mlp.down
model.layers.{i}.input_layernorm.weight             → layer[i].attn_norm
model.layers.{i}.post_attention_layernorm.weight    → layer[i].ffn_norm
model.norm.weight                                   → final_norm
lm_head.weight                                      → lm_head
```

#### 4. 加载流程

```
safetensors file (disk)
  → mmap (host memory, 零拷贝)
  → dtype check/cast (如 FP32 → BF16)
  → H2D copy → GPU Tensor
  → 按 layer 组织成模型结构
```

### 外部依赖

- `safetensors` crate
- `serde` + `serde_json` (解析 config.json)
- `memmap2` (mmap 支持，safetensors crate 可能内置)

### 测试验收

- [ ] 加载 GPT-2 124M (`openai-community/gpt2`)，打印所有 tensor name, shape, dtype
- [ ] 抽查几个 tensor 的前 10 个值，与 PyTorch `from_pretrained` 对比
- [ ] 加载 Qwen3-8B sharded 权重，验证所有 tensor 都成功加载
- [ ] 性能: 测量 8B 模型权重加载时间 (mmap → GPU 全流程)
- [ ] 错误处理: 缺少 tensor、dtype 不匹配、文件不存在等情况

---

## Phase 7: BPE Tokenizer

**Crate**: `xserv-tokenizer`

**目标**: 从零实现 Byte-Pair Encoding tokenizer，兼容 HuggingFace tokenizer.json 格式。

### BPE 算法核心

#### 编码 (encode)

```
输入: "Hello world"
  → pre-tokenize (regex split): ["Hello", " world"]
  → 每个词转为 byte 序列: [72, 101, 108, 108, 111], [32, 119, 111, ...]
  → 初始 token 序列: 每个 byte 是一个 token
  → 反复合并:
      1. 找当前序列中优先级最高的 byte-pair (从 merges 表查)
      2. 合并该 pair
      3. 重复直到无可合并
  → 输出: token IDs
```

#### 解码 (decode)

```
token IDs → 查 vocab 得到 byte 序列 → 拼接 → UTF-8 decode → 文本
```

### 需要处理的细节

1. **Pre-tokenization**:
   - GPT-2 regex: `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+`
   - Qwen3 可能使用不同的 regex pattern（从 `tokenizer.json` 的 `pre_tokenizer` 字段读取）
   - 用 `regex` crate 实现

2. **tokenizer.json 解析**:
   ```json
   {
     "model": {
       "type": "BPE",
       "vocab": {"Hello": 0, "world": 1, ...},
       "merges": ["H e", "He l", "Hel lo", ...]
     },
     "added_tokens": [...],
     "pre_tokenizer": {...},
     "post_processor": {...}
   }
   ```

3. **Special Tokens**:
   - `<|endoftext|>` (GPT-2 EOS)
   - `<|im_start|>`, `<|im_end|>` (Qwen3 ChatML)
   - `<|endoftext|>` (Qwen3 EOS)
   - Special tokens 不参与 BPE merge，直接映射到 ID

4. **Chat Template** (Qwen3 格式):
   ```
   <|im_start|>system
   You are a helpful assistant.<|im_end|>
   <|im_start|>user
   Hello<|im_end|>
   <|im_start|>assistant
   ```

5. **性能优化**:
   - Merge rules 用 `HashMap<(TokenId, TokenId), MergePriority>` 预索引
   - 对于长文本，考虑 priority queue 加速 pair 查找

### 测试验收

- [ ] 加载 GPT-2 tokenizer，encode + decode 一批测试文本，与 Python `AutoTokenizer` 逐 token 对比
- [ ] 加载 Qwen3 tokenizer，同样逐 token 对比
- [ ] 边界情况: 空字符串、纯 emoji (🎉🔥)、中英混合、超长文本 (1MB)
- [ ] Chat template: 给定 messages 列表，生成与 HF `apply_chat_template` 一致的 token 序列
- [ ] Benchmark: encode 1MB 文本的延迟

---

## Phase 8: GPT-2 完整推理 — 里程碑 ①

**Crate**: `xserv-model`

**目标**: 将所有组件串联，实现 GPT-2 的完整推理 pipeline。这是第一次看到模型"说话"。

### 模型结构

```rust
pub struct GPT2 {
    config: GPT2Config,
    wte: Tensor,           // token embedding [vocab_size, hidden_size]
    wpe: Tensor,           // position embedding [max_seq_len, hidden_size]
    layers: Vec<GPT2Block>,
    ln_f: LayerNorm,       // final layer norm
    // lm_head 与 wte 共享权重 (tied embeddings)
}

pub struct GPT2Block {
    ln_1: LayerNorm,
    attn: GPT2Attention,   // MHA: q_proj, k_proj, v_proj, o_proj
    ln_2: LayerNorm,
    mlp: GPT2MLP,          // fc1 (4H) → GELU → fc2 (H)
}
```

### Forward Pass 流程

```
tokens [B, S]
  → wte[tokens] + wpe[0..S]              → hidden [B, S, 768]
  → for each layer:
      → residual = hidden
      → hidden = ln_1(hidden)
      → hidden = attention(hidden)         # Q, K, V 从 hidden 线性变换
      → hidden = hidden + residual         # residual connection
      → residual = hidden
      → hidden = ln_2(hidden)
      → hidden = mlp(hidden)               # Linear→GELU→Linear
      → hidden = hidden + residual
  → hidden = ln_f(hidden)
  → logits = hidden @ wte.T               → [B, S, vocab_size]
  → next_token = sample(logits[:, -1, :]) # 只取最后一个 position
```

### Sampling 策略

```rust
pub struct SamplingParams {
    pub temperature: f32,      // default 1.0
    pub top_k: usize,         // default 50
    pub top_p: f32,            // default 1.0 (disabled)
    pub max_tokens: usize,     // default 256
    pub repetition_penalty: f32, // default 1.0 (disabled)
}
```

实现:
1. **Greedy**: `argmax(logits)`
2. **Temperature**: `logits = logits / temperature` → softmax → sample
3. **Top-K**: 保留 top-k logits，其余置为 -inf → softmax → sample
4. **Top-P (Nucleus)**: 按概率降序排列，累加到概率 >= p → 截断 → 重新 normalize → sample
5. 以上可以组合: temperature → top-k → top-p → sample

### CLI 交互

```
$ cargo run --release --bin xserv-cli -- --model openai-community/gpt2

xserv> The future of AI is
GPT-2> The future of AI is not just about the technology, but about the people
       who are building it. The question is whether we can...

xserv> Once upon a time
GPT-2> Once upon a time, there was a young man who lived in a small village...
```

### 测试验收

- [ ] 加载 `openai-community/gpt2`，prefill "The future of AI is"
- [ ] Prefill logits 与 PyTorch 对比: top-5 token IDs 和对应 logit 值一致
- [ ] Greedy decode 50 tokens，结果应该是连贯英文
- [ ] Temperature/TopK/TopP sampling: 生成多次结果应有变化
- [ ] CLI 交互模式可用

---

## Phase 9: KV Cache + Autoregressive 优化

**Crate**: `xserv-runtime`

**目标**: 实现 KV Cache，将 decode 从 O(S²) 降到 O(S) per step。

### 核心概念: Prefill vs Decode

**Prefill（首次处理 prompt）**:
- 输入: 完整 prompt `[B, S, D]`
- 计算: 所有 token 的 Q, K, V
- Attention: `Q[B,H,S,D] @ K[B,H,S,D]^T` → 完整 S×S 矩阵
- 输出: 缓存 K, V 到 KV cache
- 特点: **Compute-bound**（大矩阵乘法）

**Decode（逐 token 生成）**:
- 输入: 上一步生成的 1 个 token `[B, 1, D]`
- 计算: 只计算新 token 的 Q, K, V
- K_new, V_new append 到 cache
- Attention: `Q[B,H,1,D] @ K_cache[B,H,S+1,D]^T` → `[B,H,1,S+1]`
- 特点: **Memory-bound**（Q 只有 1 行，瓶颈在读 K/V cache）

### KV Cache 设计（简单版，非 paged）

```rust
pub struct KVCache {
    // 每层一对 K/V tensor
    // shape: [batch_size, num_kv_heads, max_seq_len, head_dim]
    k_caches: Vec<Tensor>,   // 索引 = layer_idx
    v_caches: Vec<Tensor>,
    seq_len: usize,          // 当前已填充的长度
    max_seq_len: usize,      // 预分配的最大长度
}

impl KVCache {
    // prefill 时: 写入 [0..prompt_len] 的 K, V
    pub fn fill(&mut self, layer: usize, k: &Tensor, v: &Tensor);

    // decode 时: 在 seq_len 位置写入新的 K, V，返回完整 cache
    pub fn append(&mut self, layer: usize, k: &Tensor, v: &Tensor) -> (&Tensor, &Tensor);
}
```

### Decode Attention Kernel

与 prefill attention 不同，decode 时 Q 只有 1 行:
```
Q [B, H, 1, D] × K_cache [B, H, S, D]^T → scores [B, H, 1, S]
→ softmax → weights [B, H, 1, S]
weights × V_cache [B, H, S, D] → output [B, H, 1, D]
```

优化方向:
- 每个 warp 处理一个 head
- 沿 S 维度做 parallel reduction (dot product + online softmax)
- 重点优化 memory bandwidth（K/V cache 的读取是瓶颈）

### 测试验收

- [ ] 对比有/无 KV cache 的生成结果 → **必须完全一致**（bit-exact for greedy）
- [ ] Benchmark decode 延迟:

| Seq Length | Without Cache (ms/token) | With Cache (ms/token) | Speedup |
|------------|--------------------------|----------------------|---------|
| 128 | | | |
| 512 | | | |
| 2048 | | | |

- [ ] 显存占用: KV cache 的实际显存与理论值 (`2 * num_layers * num_kv_heads * seq_len * head_dim * sizeof(bf16)`) 对比
- [ ] GPT-2 decode throughput (tokens/s) 记录为 baseline

---

## Phase 10: Qwen3-8B 支持 — 里程碑 ②

**Crate**: `xserv-model`

**目标**: 扩展模型定义以支持 Qwen3-8B，验证输出正确性。

### 架构对比

| 特性 | GPT-2 (124M) | Qwen3-8B |
|------|-------------|----------|
| Normalization | LayerNorm (pre-LN) | RMSNorm (pre-LN) |
| Position Encoding | Learned absolute (wpe) | RoPE (无单独参数) |
| Attention | MHA (12 heads, 12 KV heads) | GQA (如 32 Q heads, 8 KV heads) |
| Activation | GELU | SwiGLU (SiLU gate) |
| FFN | Linear(H→4H) → GELU → Linear(4H→H) | gate_proj + up_proj → SiLU gate → down_proj |
| Vocab Size | 50,257 | ~152,000 |
| Hidden Size | 768 | 4,096 (8B) |
| Layers | 12 | 36 |
| Tied Embeddings | Yes | No |

### 需要新增/修改的组件

#### 1. GQA (Grouped Query Attention)
```
num_heads = 32, num_kv_heads = 8
每个 KV head 服务 32/8 = 4 个 Q head

Q: [B, 32, S, 128]
K: [B, 8, S, 128]   ← 只有 8 个 KV heads
V: [B, 8, S, 128]

Attention 时:
  kv_head_idx = q_head_idx / (num_heads / num_kv_heads)
  不需要真正 repeat K/V 数据，kernel 中做索引映射
```

#### 2. RoPE (Rotary Position Embedding)
```
对 Q, K 的每对相邻元素做旋转:
  Q' = Q * cos(θ) + rotate_half(Q) * sin(θ)
  K' = K * cos(θ) + rotate_half(K) * sin(θ)
其中 θ_i = pos / (rope_theta^(2i/d))

预计算: freqs = 1.0 / (rope_theta^(2i/d)) for i in 0..d/2
运行时: cos_cache[pos][i] = cos(pos * freqs[i])
```

#### 3. SwiGLU FFN
```
x → gate_proj(x) → SiLU → ⊙ up_proj(x) → down_proj → output

三个 Linear:
  gate_proj: [hidden_size, intermediate_size]
  up_proj:   [hidden_size, intermediate_size]
  down_proj: [intermediate_size, hidden_size]
```

### 模型结构

```rust
pub struct Qwen3Model {
    config: Qwen3Config,
    embed_tokens: Tensor,           // [vocab_size, hidden_size]
    layers: Vec<Qwen3DecoderLayer>,
    norm: RMSNorm,                  // final RMSNorm
    lm_head: Tensor,               // [vocab_size, hidden_size] (not tied)
}

pub struct Qwen3DecoderLayer {
    input_layernorm: RMSNorm,
    self_attn: Qwen3Attention,     // GQA with RoPE
    post_attention_layernorm: RMSNorm,
    mlp: Qwen3MLP,                 // SwiGLU
}
```

### 显存预算 (BF16, 单卡 5090 32GB)

```
模型权重:  8B × 2B = ~16 GB
KV cache:  36 layers × 2(KV) × 8 heads × 4096 tokens × 128 dim × 2B ≈ 5.6 GB
Activation (单请求): ~1 GB
────────────────────────
总计: ~22.6 GB (单请求)，剩余 ~10 GB 可用于更多并发
```

### 测试验收

- [ ] 加载 Qwen3-8B 权重到单张 5090，打印模型结构和参数量
- [ ] Prefill logits 与 HF transformers 对比: 输入 "你好" → top-5 logits 一致
- [ ] 英文生成: "What is the capital of France?" → 生成合理回答
- [ ] 中文生成: "请介绍一下量子计算" → 生成通顺中文
- [ ] 多轮对话:
  ```
  <|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n
  ```
  验证 chat template 格式正确
- [ ] 单请求性能 baseline: prefill latency (ms), decode throughput (tokens/s)

---

## Phase 11: Paged Attention + KV Cache Manager

**Crate**: `xserv-runtime`
**CUDA 源码**: `csrc/attention/paged_attention.cu`

**目标**: 实现 vLLM 的核心创新 — PagedAttention，解决 KV cache 内存碎片化问题。

### 问题

Phase 9 的简单 KV cache 为每个请求预分配 `max_seq_len` 的连续内存:
- 请求只用了 100 tokens 但占了 4096 tokens 的空间 → 内存利用率低
- 连续分配导致碎片化 → 并发请求数受限

### PagedAttention 设计

**核心思想**: 像操作系统的虚拟内存一样，将 KV cache 分成固定大小的 page (block)。

```rust
pub const BLOCK_SIZE: usize = 16;  // 每个 block 存 16 个 token 的 KV

// 物理 KV cache: 预分配的大块 GPU 内存
// k_cache shape: [num_physical_blocks, num_kv_heads, block_size, head_dim]
// v_cache shape: [num_physical_blocks, num_kv_heads, block_size, head_dim]

pub struct BlockAllocator {
    free_blocks: Vec<usize>,          // 空闲物理 block ID 列表
    num_total_blocks: usize,
    ref_counts: Vec<usize>,           // 每个 block 的引用计数 (CoW 用)
}

pub struct BlockTable {
    // logical_block_idx → physical_block_idx
    // 例: 一个 seq_len=50 的请求有 4 个 block (50/16=3.125, 向上取整)
    blocks: Vec<usize>,
}

pub struct PagedKVCacheManager {
    k_cache: Tensor,                  // 所有物理 blocks
    v_cache: Tensor,
    allocator: BlockAllocator,
    block_tables: HashMap<SeqId, BlockTable>,
}
```

### Paged Attention Kernel

与普通 attention 的区别: K/V 不是连续存储，需要通过 block table 间接寻址。

```
输入:
  Q:           [num_seqs, num_heads, head_dim]        (decode 时每个 seq 只有 1 个 query)
  k_cache:     [num_blocks, num_kv_heads, block_size, head_dim]  (物理存储)
  v_cache:     [num_blocks, num_kv_heads, block_size, head_dim]
  block_tables: [num_seqs, max_num_blocks]             (间接寻址表)
  seq_lens:    [num_seqs]                              (每个 seq 的实际长度)

每个 thread block 处理:
  1 个 seq 的 1 个 attention head
  遍历该 seq 的所有 logical blocks
  对每个 block: 查 block_table 得到 physical_block_id → 读取 K/V
  online softmax 累加
  输出: [num_seqs, num_heads, head_dim]
```

### Copy-on-Write (高级，可选)

- 多个 sequence 共享相同 prefix 的 KV blocks（beam search, prompt caching）
- 写入时: 如果 ref_count > 1，先复制该 block 再修改
- 这阶段先不实现，标记为后续优化

### 测试验收

- [ ] 正确性: paged attention 输出与 Phase 9 简单 KV cache 完全一致
- [ ] 内存效率对比:

| 场景 | Naive KV Cache | Paged KV Cache |
|------|---------------|----------------|
| 1 req, seq=100 | 分配 4096 tokens | 分配 7 blocks (112 tokens) |
| 10 req, seq=100-500 | 10×4096 | 按需分配 |
| 最大并发数 (32GB) | N 个 | M 个 (M >> N) |

- [ ] Block allocator: alloc/free 循环，无内存泄漏
- [ ] Benchmark: paged attention kernel vs naive decode attention 延迟对比

---

## Phase 12: Continuous Batching + Request Scheduler

**Crate**: `xserv-runtime`

**目标**: 实现 iteration-level 调度，支持请求的动态加入和退出。

### Static Batching vs Continuous Batching

**Static (朴素)**:
```
Batch 1: [req1, req2, req3] → 等 req1, req2, req3 全部完成
Batch 2: [req4, req5, req6] → ...
问题: req1 完成了但 req3 还在生成 → GPU 空转
```

**Continuous (Orca 论文)**:
```
Iteration 1: [req1, req2, req3] → req1 完成!
Iteration 2: [req2, req3, req4] → req4 动态加入
Iteration 3: [req2, req3, req4] → req3 完成!
Iteration 4: [req2, req4, req5] → req5 动态加入
```

### 核心数据结构

```rust
#[derive(Clone, Copy, PartialEq)]
pub enum SequenceStatus {
    Waiting,      // 在等待队列中
    Prefilling,   // 正在做 prefill
    Decoding,     // 正在做 decode
    Finished,     // 已完成 (EOS / max_len / stop string)
    Preempted,    // 被抢占（显存不够，KV cache 被换出）
}

pub struct Sequence {
    pub id: SeqId,
    pub prompt_tokens: Vec<u32>,
    pub generated_tokens: Vec<u32>,
    pub status: SequenceStatus,
    pub sampling_params: SamplingParams,
    pub block_table: BlockTable,
    pub arrival_time: Instant,
    // 用于 streaming 输出
    pub output_sender: tokio::sync::mpsc::Sender<GeneratedToken>,
}

pub struct Scheduler {
    waiting: VecDeque<Sequence>,
    running: Vec<Sequence>,
    max_num_seqs: usize,           // 最大并发 batch size
    max_num_tokens: usize,         // 单次 iteration 最大 token 数
    block_manager: PagedKVCacheManager,
}
```

### 调度循环 (Engine 主循环)

```rust
loop {
    // Step 1: 回收已完成的 sequence
    //   - 释放其 KV cache blocks
    //   - 从 running 移除

    // Step 2: 检查能否 admit 新请求
    //   - 条件: running.len() < max_num_seqs
    //           && 有足够的 free blocks 给新请求的 prompt
    //   - FCFS 从 waiting 取

    // Step 3: 划分 prefill / decode
    //   - 新加入的 sequence: prefill (处理完整 prompt)
    //   - 已在 running 的: decode (生成 1 个 token)

    // Step 4: 组装 batch input
    //   - Prefill: 各 seq 的 prompt tokens, 需要 padding 或 ragged batch
    //   - Decode: 各 seq 的最后一个 token

    // Step 5: Forward pass
    //   - prefill 和 decode 可以混合在一个 forward 中
    //   - 或者分开处理 (先 prefill, 再 decode)

    // Step 6: Sampling
    //   - 对每个 seq 的 logits 进行采样

    // Step 7: 更新状态
    //   - 将新 token append 到 sequence
    //   - 检查是否完成 (EOS / max_len)
    //   - 通过 channel 发送新 token 给 API 层
}
```

### Preemption（显存不足时的抢占）

当显存不足以 admit 新请求时:
1. **Swap**: 将低优先级 seq 的 KV cache 从 GPU 换到 CPU（复杂，后续再做）
2. **Recompute**: 丢弃低优先级 seq 的 KV cache，后续重新 prefill（简单，先实现这个）

### 测试验收

- [ ] 模拟 10 个请求在不同时间到达，所有请求都得到正确的生成结果
- [ ] 短请求完成后，新请求立即加入 batch（观察 log）
- [ ] Throughput 对比:

| 方式 | 20 请求总耗时 | Token/s |
|------|-------------|---------|
| 串行 (batch=1) | | |
| Static batch=4 | | |
| Continuous batch | | |

- [ ] 压力测试: 100 个并发请求，全部正确完成，无 hang/crash

---

## Phase 13: HTTP API + SSE Streaming — 里程碑 ③

**Crate**: `xserv-engine`, `xserv-api`

**目标**: 提供 OpenAI 兼容的 HTTP API，支持 SSE streaming。第一个端到端可用的里程碑。

### 技术栈

- **HTTP**: `axum` (Rust async web framework)
- **Async**: `tokio`
- **JSON**: `serde_json`
- **SSE**: `axum` 内置 SSE 支持 (`axum::response::sse`)

### API 端点

```
POST /v1/chat/completions    # 主要端点 (ChatML 格式)
POST /v1/completions         # 纯文本补全
GET  /v1/models              # 列出可用模型
GET  /health                 # 健康检查
```

### 请求/响应格式 (OpenAI 兼容)

**Chat Completion Request**:
```json
{
  "model": "qwen3-8b",
  "messages": [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What is 1+1?"}
  ],
  "stream": true,
  "temperature": 0.7,
  "top_p": 0.9,
  "max_tokens": 256,
  "stop": ["\n\n"]
}
```

**SSE Streaming Response**:
```
data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-8b","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}

data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-8b","choices":[{"index":0,"delta":{"content":"The"},"finish_reason":null}]}

data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-8b","choices":[{"index":0,"delta":{"content":" answer"},"finish_reason":null}]}

data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","created":1234567890,"model":"qwen3-8b","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}

data: [DONE]
```

**Non-streaming Response**:
```json
{
  "id": "chatcmpl-xxx",
  "object": "chat.completion",
  "created": 1234567890,
  "model": "qwen3-8b",
  "choices": [{
    "index": 0,
    "message": {"role": "assistant", "content": "The answer is 2."},
    "finish_reason": "stop"
  }],
  "usage": {
    "prompt_tokens": 25,
    "completion_tokens": 8,
    "total_tokens": 33
  }
}
```

### 架构分层

```
Client (curl / Python OpenAI SDK)
    │
    ▼
┌─────────────────────────────────────┐
│  xserv-api (axum HTTP server)       │
│  - 解析请求, 验证参数               │
│  - apply chat template              │
│  - 将请求提交给 engine              │
│  - 从 channel 接收 token, 编码为 SSE│
└────────────┬────────────────────────┘
             │ InferenceRequest (通过 channel)
             ▼
┌─────────────────────────────────────┐
│  xserv-engine (推理引擎)            │
│  - 独立的 OS thread (非 async)      │
│  - 运行 scheduler 调度循环          │
│  - 管理 model + KV cache            │
│  - 每生成一个 token, 通过 channel   │
│    发送给 API 层                     │
└─────────────────────────────────────┘
```

**关键设计决策**:
- Engine 跑在独立 OS thread（避免 GPU 同步操作 block tokio runtime）
- API ↔ Engine 通过 `tokio::sync::mpsc` channel 通信
- 每个请求有独立的 `mpsc::Sender/Receiver` 用于 token streaming

### 测试验收

- [ ] `curl` 测试:
  ```bash
  curl http://localhost:8080/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{"model":"qwen3-8b","messages":[{"role":"user","content":"Hello"}],"stream":true}'
  ```
  看到 SSE 逐 token 输出

- [ ] Python OpenAI SDK 测试:
  ```python
  from openai import OpenAI
  client = OpenAI(base_url="http://localhost:8080/v1", api_key="unused")
  for chunk in client.chat.completions.create(
      model="qwen3-8b",
      messages=[{"role": "user", "content": "What is 1+1?"}],
      stream=True
  ):
      print(chunk.choices[0].delta.content or "", end="", flush=True)
  ```

- [ ] 非 streaming 模式也能正常工作
- [ ] 并发 10 个请求，全部正确完成
- [ ] 多轮对话: 连续发两轮消息（第二轮包含 history），验证上下文连贯
- [ ] `/v1/models` 返回已加载的模型列表
- [ ] 错误处理: 无效参数返回 400, 模型不存在返回 404

---

## Phase 14: Flash Attention (FA2 for SM120)

**Crate**: `xserv-kernels`
**CUDA 源码**: `csrc/attention/flash_attention.cu`

**目标**: 实现 Flash Attention 的 CUDA kernel，大幅降低 attention 的显存占用并提升速度。

### 硬件适配说明

Flash Attention 已发展到第 4 代 (FA4, arxiv 2603.05451)，但各版本有明确的硬件依赖:

| 版本 | 目标架构 | 关键硬件特性 | RTX 5090 兼容 |
|------|---------|------------|--------------|
| FA2 | 通用 CUDA (SM75+) | 标准 shared memory + HMMA | **是** ✅ |
| FA3 | Hopper SM90 (H100) | TMA + WGMMA + warp specialization | 否 |
| FA4 | Blackwell SM100 (B200/B300) | TMEM + async MMA + 2-CTA mode | 否 |

**RTX 5090 (SM120, CC 12.0) 使用的是消费级 Blackwell 架构 (GB202)，与数据中心 Blackwell (B200, SM100) 是不同的硅片设计。SM120 物理上没有 TMEM (Tensor Memory) 子系统，因此 FA4 的 kernel 无法在 5090 上运行。这不是软件限制，是硬件级差异。**

因此本项目实现 **FA2 算法**，使用标准 CUDA (shared memory + HMMA)。FA2 的核心优化——online softmax tiling、O(1) 显存占用——在任何架构上都有效。

### 核心思想

标准 attention 的问题:
```
S = Q @ K^T         ← 需要 O(S²) 显存存储 S×S 矩阵
P = softmax(S)      ← 需要完整的 S×S 才能做 softmax
O = P @ V
```

Flash Attention 的解法:
- **不 materialize S×S 矩阵**
- 将 Q, K, V 分成 tiles，在 SRAM (shared memory) 中计算
- 使用 **online softmax trick**: 边算边更新 running max 和 running sum

### 算法 (Forward Pass, FA2)

FA2 相比 FA1 的改进: 外层循环遍历 Q tiles (而非 K/V)，减少 HBM 读写次数。

```
Br, Bc = tile sizes for Q and K/V respectively

for each Q tile (q_start..q_start+Br):                    ← 外层: Q tiles
    load Q_tile [Br, D] to shared memory
    initialize: O_tile = 0, l = 0, m = -inf               // running sum and max

    for each K,V tile (kv_start..kv_start+Bc):             ← 内层: K/V tiles
        load K_tile [Bc, D], V_tile [Bc, D] to shared memory

        // Compute attention scores for this tile pair
        S_tile = Q_tile @ K_tile^T              // [Br, Bc], in registers/SRAM

        // Apply causal mask (skip if kv_start > q_start + Br)
        if causal: mask upper triangle of S_tile

        // Online softmax update
        m_new = max(m, rowmax(S_tile))          // new running max
        P_tile = exp(S_tile - m_new)            // safe exp
        l_new = exp(m - m_new) * l + rowsum(P_tile)  // update running sum

        // Rescale and accumulate output
        O_tile = diag(exp(m - m_new)) * O_tile + P_tile @ V_tile
        m = m_new
        l = l_new

    O_tile = O_tile / l                         // final normalization
    write O_tile [Br, D] to global memory (HBM)
```

### 实现要点

1. **Tile 大小选择**:
   - 5090 SM120: shared memory per SM = 100 KB (需实测确认)
   - 需同时存 Q_tile, K_tile, V_tile, S_tile
   - BF16: Q_tile [Br, D] = Br × 128 × 2B; K_tile [Bc, D] = Bc × 128 × 2B
   - S_tile [Br, Bc] 保持 FP32 = Br × Bc × 4B
   - 推荐起步: Br=Bc=64, head_dim=128 → 共需 ~100KB shared memory
   - 优化版: Br=Bc=128 需要更多 shared memory, 可能需要拆分

2. **Causal mask 优化**:
   - 如果 K/V tile 完全在 Q tile 的"未来"（kv_start > q_end）→ 跳过整个 tile
   - 减少约 50% 的计算量

3. **BF16 精度**:
   - S_tile, P_tile 的计算在 FP32 中进行（累加精度）
   - Q, K, V 的加载用 BF16（节省 bandwidth）
   - 最终 O 转回 BF16 写出

4. **GQA 支持**:
   - K/V heads 数量 < Q heads 时，kernel 中做 `kv_head = q_head / num_groups` 索引
   - 不需要 repeat_kv 操作，直接在 kernel 内部解决

5. **Decode attention 特化**:
   - Decode 时 Q 只有 1 行 (Br=1)，退化为 vector-matrix attention
   - 可以写一个专门的 decode attention kernel (类似 FlashDecoding)
   - 沿 KV sequence 维度做 parallel reduction

### 测试验收

- [ ] 正确性: 与 Phase 5 naive attention 对比, max error < 1e-2 (BF16)
- [ ] 显存: Flash Attention 不随 S 平方增长

| Seq Length | Naive VRAM | Flash VRAM | Naive Time | Flash Time |
|------------|-----------|------------|------------|------------|
| 512 | MB | MB | ms | ms |
| 2048 | MB | MB | ms | ms |
| 8192 | OOM? | MB | OOM? | ms |
| 32768 | OOM | MB | OOM | ms |

- [ ] 集成到 Qwen3-8B，端到端 decode latency 对比
- [ ] Profile: `ncu` 分析 compute utilization, memory throughput
- [ ] GQA 支持: 无 repeat_kv 开销

---

## Phase 15: 性能优化 — 里程碑 ④

**目标**: 系统性 profiling + 优化，向 50% vLLM throughput 目标冲刺。

### 优化方向

#### 1. Kernel Fusion
减少 memory-bound kernel 之间的 global memory 读写:
- **Residual Add + RMSNorm**: 一次读 hidden + residual，写出 normed output
- **SiLU + Elementwise Mul** (SwiGLU 内部): 一次读 gate + up, 写出 SiLU(gate)*up
- **Bias Add + Activation**: Linear 的 bias + 激活函数合并

原则: 两个逐元素 kernel 之间如果有 global memory 读写，就值得融合。

#### 2. CUDA Graphs
- Decode 阶段每步的 kernel 序列是固定的（shape 不变）
- 用 `cudaStreamBeginCapture` / `cudaStreamEndCapture` 捕获一次
- 后续用 `cudaGraphLaunch` 重放（消除 kernel launch overhead）
- **注意**: batch size 变化时需要重新 capture

#### 3. Memory 优化
- **权重预加载**: 启动时加载到 GPU，推理路径上零分配
- **Activation reuse**: 同一层的中间结果用完立即释放/复用
- **Pinned memory**: H2D/D2H 用 `cudaMallocHost`（pinned）提升拷贝带宽

#### 4. Compute 优化
- 确保所有 GEMM 走 Tensor Core (BF16)
- Decode attention: 优化 memory bandwidth 利用率
- Prefill: chunked processing（控制峰值显存，允许更大 batch）

#### 5. Scheduling 优化
- Prefill-Decode disaggregation: prefill 和 decode 分开 batch
  - 原因: prefill 是 compute-bound, decode 是 memory-bound, 混合导致两边都不优
- Dynamic batch size: 根据当前 running seqs 的 seq_len 动态调整

### Profiling 工具使用

```bash
# 整体 timeline (哪个 kernel 最耗时)
nsys profile --stats=true ./target/release/xserv-server

# 单个 kernel 分析 (occupancy, memory throughput)
ncu --target-processes all --set full ./target/release/xserv-server

# 自定义 Rust 计时
# 在 engine 循环中记录每个 phase 的耗时
```

### 测试验收

- [ ] 安装 vLLM，同一台机器跑 Qwen3-8B
- [ ] Benchmark 对比:

| Metric | vLLM | xserv | Ratio |
|--------|------|-------|-------|
| Prefill latency (ms, 128 tokens) | | | |
| Decode throughput (tokens/s, batch=1) | | | |
| Decode throughput (tokens/s, batch=16) | | | |
| Max concurrent requests (32GB) | | | |

- [ ] 目标: xserv throughput >= 50% vLLM
- [ ] Profiling 报告: 每个组件的耗时占比 pie chart
- [ ] 无功能回归: 所有之前的集成测试通过

---

## Phase 16: Speculative Decoding

**Crate**: `xserv-speculative`

**目标**: 用小模型（draft model）加速大模型（target model）的 decode 阶段。

### 算法

```
γ = 4 (speculative tokens 数量)

1. Draft model 自回归生成 γ 个 token: t1, t2, t3, t4
2. Target model 一次 forward 处理这 γ+1 个 position
   (等价于一次 prefill: [last_accepted_token, t1, t2, t3, t4])
   得到 γ+1 个 logits
3. Rejection sampling:
   for i in 0..γ:
     p = target_prob[i][t_{i+1}]   // target model 给 draft token 的概率
     q = draft_prob[i][t_{i+1}]    // draft model 给该 token 的概率
     if random() < min(1, p/q):
       accept t_{i+1}
     else:
       reject, resample from adjusted distribution
       break
4. 至少接受 1 个 token, 期望接受 ~γ×acceptance_rate 个
```

### 关键点

- **无损**: rejection sampling 保证输出分布与纯 target model 一致
- **加速条件**: draft model 足够快且与 target 分布接近
- **Draft model 选择**: Qwen3-0.5B / Qwen3-1.5B 作为 Qwen3-8B 的 draft

### KV Cache 处理

- Draft model 有自己的 KV cache
- Target model 验证时，accepted tokens 的 KV 可以复用（不用重算）
- Rejected 位置之后的 KV 需要丢弃

### 测试验收

- [ ] 验证: speculative decode 100 条不同 prompt，输出分布与标准 decode 无统计差异
- [ ] Acceptance rate 统计 (期望 60-80% per token)
- [ ] 端到端加速:

| Method | Tokens/s (batch=1) | Speedup |
|--------|--------------------|---------|
| Standard decode | | 1.0x |
| Speculative (γ=4) | | ~2-3x |

---

## Phase 17: Tensor Parallelism (TP=1/2/4) — 里程碑 ⑤

**Crate**: `xserv-distributed`

**目标**: 实现 Tensor Parallelism，支持 TP=2 和 TP=4（同组 GPU 内），跑通多卡推理。PP 及更大规模并行留待后续扩展。

### 通信后端: NCCL

NVIDIA Collective Communication Library，提供高效的 multi-GPU 通信原语。

需要封装的操作:
```rust
// NCCL FFI
ncclAllReduce(sendbuff, recvbuff, count, datatype, op, comm, stream)
ncclAllGather(sendbuff, recvbuff, count, datatype, comm, stream)
ncclBroadcast(sendbuff, recvbuff, count, datatype, root, comm, stream)
```

### Tensor Parallelism 策略 (Megatron-LM 风格)

每层只需要 **2 次 AllReduce**:

#### Attention 部分
```
Column Parallel: Q/K/V proj 按 head 维度切分
  GPU 0: q_proj[:, :hidden/TP], k_proj[:, :hidden/TP], v_proj[:, :hidden/TP]
  GPU 1: q_proj[:, hidden/TP:], k_proj[:, hidden/TP:], v_proj[:, hidden/TP:]
  → 每卡计算自己那部分 heads 的 attention → 无需通信

Row Parallel: o_proj 按行切分
  每卡计算部分输出 → AllReduce 求和 → 得到完整 output
```

#### FFN (SwiGLU) 部分
```
Column Parallel: gate_proj, up_proj 按列切分
  每卡计算部分 intermediate features → 无需通信

Row Parallel: down_proj 按行切分
  每卡计算部分输出 → AllReduce 求和 → 得到完整 output
```

#### 其他
- **Embedding**: vocab 按行切分，AllGather 拼接
- **RMSNorm**: 每卡独立计算（输入是 AllReduce 后的完整 tensor）
- **lm_head**: 按列切分，AllGather 拼接 logits

### 权重分片

启动时:
1. Rank 0 加载完整权重
2. 按 TP 策略切分
3. AllGather 或 Scatter 分发到各卡

或者:
- 每个 rank 独立加载，只读取属于自己的那部分（更高效）

### 互联拓扑 (已确认)

**纯 PCIe Gen5 x16, 无 NVLink**。GPU 分两组: 0-3 (PHB) 和 4-7 (PHB)，跨组走 NODE。

**TP 部署策略** (当前阶段目标: TP=1/2/4):
- **TP=2**: 同组内任意两卡 (如 GPU0+GPU1)，PCIe PHB 延迟最低
- **TP=4**: 同组 4 卡 (GPU 0-3 或 GPU 4-7)，全 PHB 互联
- **PCIe Gen5 x16 带宽**: 理论 ~64 GB/s 单向，实测 AllReduce 有效带宽约 40-50 GB/s
- **后续扩展**: TP=8 (跨组) 和 Pipeline Parallelism 留待后续阶段

### 测试验收

- [ ] TP=2: Qwen3-8B 输出与单卡 (TP=1) 完全一致
- [ ] TP=4: 每卡权重显存占用约 1/4
- [ ] Scaling benchmark (同组 GPU 0-3):

| TP Size | Prefill (tokens/s) | Decode (tokens/s) | Scaling Efficiency |
|---------|--------------------|--------------------|-------------------|
| 1 | | | 100% |
| 2 | | | % |
| 4 | | | % |

- [ ] AllReduce latency 测量 (不同消息大小，同组 PHB 互联)

---

## Phase 18: 量化 (FP8 / INT8)

**Crate**: `xserv-kernels` (新增量化 kernel), `xserv-model` (量化加载)
**CUDA 源码**: `csrc/quantize/`

**目标**: 降低模型显存占用和提升计算吞吐。

### 量化方式

#### 1. Weight-Only INT8
- 只量化权重，activation 保持 BF16
- Per-channel scale: 每个输出 channel 一个 `scale` 和 `zero_point`
- GEMM: INT8 × BF16 (dequantize on-the-fly)
- 适用于 memory-bound 场景（decode）

#### 2. FP8 (E4M3 / E5M2)
- 5090 Blackwell (CC 12.0) 原生支持 FP8 Tensor Core
- 权重和 activation 都量化为 FP8
- Dynamic scaling: 每个 tensor 运行时计算 `amax`，确定 scale factor
- GEMM: FP8 × FP8 → BF16/FP32 accumulate
- 适用于 compute-bound 场景（prefill）

#### 3. GPTQ / AWQ (高级, 可选)
- INT4 weight quantization
- 需要 calibration data
- 更复杂但压缩率更高

### FP8 实现要点

```rust
pub enum DType {
    // ... existing
    F8E4M3,    // 1 sign + 4 exponent + 3 mantissa (范围小，精度高)
    F8E5M2,    // 1 sign + 5 exponent + 2 mantissa (范围大，精度低)
}
```

Dynamic scaling:
```
scale = amax(tensor) / fp8_max   // fp8_max = 448.0 for E4M3
tensor_fp8 = cast_to_fp8(tensor / scale)
// GEMM 后: output = output * scale_A * scale_B
```

### 测试验收

- [ ] 精度:

| Quantization | Perplexity (WikiText-2) | vs BF16 |
|-------------|------------------------|---------|
| BF16 (baseline) | X.XX | — |
| FP8 E4M3 | X.XX | +0.XX |
| INT8 weight-only | X.XX | +0.XX |

- [ ] 显存: FP8 权重占用约 BF16 的一半 (~8 GB for 8B model)
- [ ] 性能: FP8 GEMM throughput vs BF16 GEMM

---

## Phase 19: Multimodal — 里程碑 ⑥

**Crate**: `xserv-model` (新增 vision encoder)

**目标**: 支持 vision-language 模型，接受图片+文本输入。

### 目标模型

Qwen-VL 系列（或类似 architecture）。典型结构:
```
Image → ViT Encoder → Visual Tokens → Projector → LLM Input
Text  → Tokenizer   → Text Tokens   ────────────→ LLM Input
```

### 需要新增的组件

#### 1. ViT (Vision Transformer) Encoder
- **Patch Embedding**: 将图片切成 14×14 patches, 每个 patch 线性投影
  - 输入: [B, 3, H, W] → 输出: [B, num_patches, hidden_dim]
- **ViT Blocks**: 标准 Transformer Encoder blocks
  - Multi-Head Self-Attention (无 causal mask)
  - FFN
  - LayerNorm
- 参数量: 通常 300M-600M

#### 2. Visual-Language Projector
- MLP: 将 ViT 输出维度映射到 LLM embedding 维度
- 可能包含 pooling / resampling (减少 visual token 数量)

#### 3. 图片预处理
- Resize to target resolution (e.g., 448×448)
- Normalize: ImageNet mean/std
- 使用 `image` crate 在 CPU 上完成

#### 4. 输入拼接
```
[<image_token>] × num_visual_tokens + [text tokens]
```
- 在 embedding 层面拼接
- LLM 处理混合 sequence

### API 扩展

```json
{
  "model": "qwen-vl",
  "messages": [{
    "role": "user",
    "content": [
      {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}},
      {"type": "text", "text": "What is in this image?"}
    ]
  }]
}
```

### 外部依赖

- `image` crate (图片加载和预处理)
- `base64` crate (解码 base64 图片)

### 测试验收

- [ ] 加载 Qwen-VL 模型，输入一张猫的图片 + "What is in this image?"
- [ ] 生成合理的图片描述
- [ ] 与 HF transformers 输出对比
- [ ] API 端到端: HTTP POST 含 base64 图片 → streaming 文字回答
- [ ] 纯文本请求不受影响（向后兼容）

---

## 实际进展记录(与原计划的分叉,2026-06 更新)

Phase 0–17 按计划完成。Phase 18 起实际路线偏离了上面的原计划
(speculative decoding 与多模态推迟),实际走向是 MoE + 量化 + 稀疏化:

| 实际 Phase | 内容 | 文档 |
|---|---|---|
| 18 | Pipeline Parallelism(PP=2/4) | `18-pipeline-parallelism.md`、`benchmarks/pp-sweep.md` |
| 19 | **gpt-oss-20b MoE**:harmony 格式、attention sinks + sliding window、YaRN;两个 CUDA bug 实战(prefill sinks NaN、GEMV 未初始化 smem);GSM8K 94.5% 对齐 llama.cpp;FP8 W8A8 / MXFP4 W4A16 量化 | `19-gpt-oss-moe.md`、`benchmarks/{fp8-quantization,mxfp4-and-llama-decode}.md` |
| 20 | **稀疏 top-k MoE decode**:只算被路由的专家,decode 13.9→7.0ms,TP=2 下 decode/TTFT 全面快于 llama.cpp 同配置;gpt-oss 单卡 serving | `20-sparse-moe.md`、`benchmarks/sparse-moe.md` |
| 21 | **decode CUDA Graph + GPU argmax**:整个 decode step 录成一个图回放(thread-local launch stream、retained-warmup 分配策略、NCCL capture);greedy 采样换 GPU argmax。TPOT 7.5→5.9ms(TP=1)/ 5.8ms(TP=2);TP=2 全面领先 llama(1.26-1.47×),TP=1 差距 2.5×→2.0× | `21-cuda-graph-decode.md` |

**下一步候选(按预期收益排序):**

| 候选 Phase | 内容 | 预期 |
|---|---|---|
| 22 | **非专家权重量化**:qkv/o + lm_head(1.16GB/token)仍是 BF16 | TPOT 再省 ~1.5ms |
| 23 | **稀疏 prefill**(按专家 permute + grouped GEMM) | 长 prompt TTFT 51-75 → ~30ms |
| 24 | server 侧 harmony channel 分离(`reasoning_content` 流式输出,对齐 llama-server 行为) | API 易用性 |
| — | Speculative Decoding、多模态(原 16/19) | 推迟 |

## 里程碑总结

| 里程碑 | Phase | 验收标准 |
|--------|-------|---------|
| ① GPT-2 推理 | 8 | CLI 输入 prompt, GPT-2 生成连贯文本, logits 与 PyTorch 一致 |
| ② Qwen3-8B 推理 | 10 | 8B 模型中英文对话, 多轮 chat template 正确 |
| ③ E2E API | 13 | HTTP streaming API, Python OpenAI SDK 可调用, 10 并发正确 |
| ④ 性能达标 | 15 | throughput >= 50% vLLM, profiling 报告完成 |
| ⑤ 多卡推理 | 17 | TP=2/4 同组 GPU 推理正确, scaling benchmark 完成 |
| ⑥ MoE 模型(实际) | 19 | gpt-oss-20b 端到端正确, GSM8K 与 llama.cpp 持平 ✅ |
| ⑦ 性能反超(实际) | 20 | 同配置 decode 快于 llama.cpp(TP=2 达成;单卡是 Phase 21+ 目标) ✅ |
| ⑧ 多模态 | 推迟 | 图片输入 → 文字回答, API 端到端 |

## 外部依赖清单

| Crate | 用途 | 引入 Phase |
|-------|------|-----------|
| `cc` | build.rs 编译 .cu 文件 | 0 |
| `half` | f16 / bf16 Rust 类型 | 2 |
| `smallvec` | Tensor shape / strides (栈分配) | 2 |
| `safetensors` | 权重文件解析 | 6 |
| `serde` + `serde_json` | JSON 序列化 | 6 |
| `memmap2` | 文件 mmap (safetensors 可能内置) | 6 |
| `regex` | BPE pre-tokenization | 7 |
| `rand` | Sampling (随机数) | 8 |
| `tokio` | Async runtime | 13 |
| `axum` | HTTP server | 13 |
| `criterion` | Benchmark framework | 3+ |
| `image` | 图片加载 (multimodal) | 19 |
| `base64` | Base64 decode (multimodal API) | 19 |

**不使用**: `candle`, `burn`, `tch`, `tokenizers`, `cudarc` — 核心组件全部自己实现。