docs: add design docs + takeaways for Phase 2 and Phase 3

- docs/01-cuda-ffi.md: added takeaways (struct layout pitfall,
  Rust 2024 unsafe changes, caching allocator strategy, etc.)
- docs/02-tensor.md: design doc + takeaways for tensor abstraction
- docs/03-gemm.md: design doc + takeaways for GEMM kernels

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-21 20:59:45 +08:00
parent d77f921a12
commit 51a0f2eb14
3 changed files with 227 additions and 6 deletions

97
docs/02-tensor.md Normal file
View File

@@ -0,0 +1,97 @@
# Phase 2: Tensor Abstraction Layer — Design Document
## Goal
实现核心 Tensor 类型,支持 CPU/GPU 存储、多种数据类型、strided view 操作,作为后续所有算子和模型的数据基础。
## Module Layout
```
crates/xserv-tensor/
├── Cargo.toml
└── src/
├── lib.rs # re-exports
├── dtype.rs # DType enum, TensorDType trait
├── shape.rs # strides 计算, broadcast 规则
├── storage.rs # Storage (Arc引用计数), Device enum
└── tensor.rs # Tensor 主体: 创建, 形状操作, 设备迁移
```
## Key Design Decisions
### DType + TensorDType Trait
```rust
pub enum DType { F32, F16, BF16 }
pub trait TensorDType: Copy + Send + Sync + 'static {
const DTYPE: DType;
fn to_f64(self) -> f64;
fn from_f64(v: f64) -> Self;
}
```
-`half` crate 的 `bf16`/`f16` 表示半精度类型
- `TensorDType` trait 让 `from_slice<T>``as_slice<T>` 有类型安全
- GPU kernel 中通过 `DType` dispatch 到对应的 CUDA 类型 (`__nv_bfloat16` / `float`)
### Storage 引用计数
```rust
pub struct Storage(Arc<StorageInner>);
enum StorageInner {
Cpu { data: Vec<u8> },
Cuda { buffer: GpuBuffer },
}
```
- `Arc` 引用计数让 transpose/slice/reshape 能共享底层数据view 语义)
- 不实现 CoWcopy-on-writeview 只能读不能写
- `to_device()` 总是创建新的 Storage
### Strided Tensor
```rust
pub struct Tensor {
storage: Storage,
shape: SmallVec<[usize; 4]>,
strides: SmallVec<[usize; 4]>,
offset: usize,
dtype: DType,
}
```
- `SmallVec<[usize; 4]>` 避免大多数 tensor (≤4D) 的堆分配
- `strides` 以元素为单位(不是字节)
- `offset` 支持 slice 操作view 到 storage 的中间位置)
- `is_contiguous()` 检查 strides 是否与 shape 匹配
- 非 contiguous 的 tensor 调 `contiguous()` 才能送入 CUDA kernel
### Broadcast 规则
实现了 NumPy-style broadcasting:
- 维度从尾部对齐
- 大小为 1 的维度可以广播到任意大小
- `broadcast_strides()` 将 size=1 维度的 stride 置为 0虚拟广播不复制数据
## Test Plan
- [x] from_slice → shape/strides 正确
- [x] reshape, transpose, squeeze, unsqueeze
- [x] transpose 后 contiguous() 重排数据
- [x] BF16 tensor 的精度验证
- [x] CPU↔GPU roundtrip
- [x] zeros on GPU → 拷回 CPU 验证全 0
- [x] broadcast_shape 单元测试
## Takeaways
1. **`SmallVec` 是正确选择**:绝大多数 tensor ≤ 4D避免了频繁堆分配。LLM 推理中常见的维度是 `[B, S, H]` (3D) 和 `[B, H, S, D]` (4D)。
2. **View 语义的取舍**Arc 共享 storage 实现了零拷贝 transpose/reshape但代价是无法原地修改 view 后的 tensor。对于推理引擎这是可以接受的——推理路径上大部分操作是只读的。
3. **contiguous() 的隐性开销**:非 contiguous tensor 在送入 kernel 前需要 `contiguous()` 拷贝。这意味着 `transpose → matmul` 会产生一次额外拷贝。后续优化方向:在 kernel 中直接支持 strided input。
4. **Rust 2024 edition 变化**`unsafe fn` 内部的 unsafe 调用也需要显式 `unsafe {}` 块,`extern "C"` 块必须加 `unsafe` 前缀。这个 edition 对安全性更严格。
5. **CPU 实现先行**:先在 CPU 上验证逻辑正确性(如 contiguous 重排),再扩展到 GPU。这个策略在后续 phase 中应该继续沿用。