Weight-only 4-bit for the gpt-oss MoE experts: weights stored MXFP4 (E2M1 + per-32-element UE8M0 block scale, tools/quantize_mxfp4.py), a fused kernel reads the 4-bit weights and dequantizes on-chip to BF16. Decode (M=1) uses a fused dequant-GEMV (batched_gemv_mxfp4) with shared-memory activation tiling; prefill (M>1) dequantizes to BF16 then reuses the BF16 batched GEMM. MXFP4 is detected by the scale tensor's rank (3-D [E,N,K/32]) vs FP8's 1-D [E]. Verified on dash5 (gpt-oss-20b, TP=2, 5090): byte-identical greedy tokens to FP8/BF16, smallest footprint (13 GB vs 22 GB FP8, 39 GB BF16) — fits one 32 GB 5090 with room for KV cache. NOT a decode speedup: the hand-written W4A16 GEMV (no tensor cores) is less efficient than cuBLASLt's FP8 tensor-core GEMM, so even at half the weight bytes decode is 17.0 ms vs FP8 13.5 ms (faster than BF16 18.8 ms); prefill regresses (350 vs 134 ms, dequant fallback). Committed as a correct memory-optimization foundation. Beating FP8 on speed needs FP4 tensor cores (W4A4, cuBLASLt block-scaled MXFP4) or a Marlin-class kernel; see docs/benchmarks/mxfp4-and-llama-decode.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
xserv
从零用 Rust + CUDA 构建的 LLM 推理引擎,目标是吃透 LLM Serving 全栈技术。
xserv 不依赖 PyTorch / vLLM / TensorRT 等现成框架,自己实现了张量抽象、CUDA kernel、 分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。当前在单张 RTX 5090 上可以 跑通 Qwen3-8B(BF16),并提供一套与 llama.cpp 对比正确性和性能的标准 benchmark。
现状一览
- 模型:GPT-2(124M)、Qwen3-8B(BF16)
- 性能(RTX 5090,Qwen3-8B BF16,贪心解码,单流):约 56 tok/s,约为 HF transformers 的 1.4×、llama.cpp 的 ~0.6×
- 精度:在 AIME 2025 / GSM8K 上与 llama.cpp 同权重对比基本持平(数值保真度验证通过)
- 服务:OpenAI 兼容
/v1/chat/completions,支持 SSE 流式输出 - 关键能力:自写 GEMM / Flash-Attention 2(SM120) / Paged-Attention kernel、 分页 KV cache(含 CPU 换出/换入 弹性显存)、连续批处理(continuous batching)、 CUDA Graph 解码、按显存自适应的 KV 池
这是一个以学习为主的项目,逐 Phase 推进,每步都做数值/端到端验证。
架构
xserv/
├── csrc/ # CUDA 源码 (.cu/.cuh)
│ ├── gemm/ # GEMM (naive / tiled / gemv)
│ ├── attention/ # Flash-Attention 2 (SM120)、Paged-Attention、causal mask
│ ├── normalization/ # LayerNorm / RMSNorm
│ ├── activation/ # GELU / SiLU
│ ├── embedding/ # embedding lookup / RoPE / transpose
│ └── reduce/ # softmax
├── crates/
│ ├── xserv-cuda/ # CUDA FFI、Stream、显存分配器、Pinned 内存、CUDA Graph
│ ├── xserv-tensor/ # Tensor 类型(strided 布局、BF16/F16/F32、CPU↔GPU)
│ ├── xserv-kernels/ # kernel registry(自写 kernel + cuBLAS 可切换)
│ ├── xserv-tokenizer/ # BPE 分词器
│ ├── xserv-model/ # 模型定义(GPT-2 / Qwen3)、权重加载、KV cache、采样
│ └── xserv-server/ # tokio + axum HTTP 服务、调度器
├── tools/ # 辅助脚本 + benchmark 套件(见下)
└── docs/ # 每个 Phase 的设计文档 + benchmark 报告
环境要求
- GPU:NVIDIA,计算能力 SM120(RTX 5090 / Blackwell)。其它架构需调整
CUDA_ARCH。 - CUDA Toolkit:12.9(
nvcc需在PATH,构建.cu依赖它) - Rust:edition 2024(建议较新的 stable 工具链)
- 模型:HuggingFace 目录格式(含
config.json、tokenizer.json、*.safetensors)
构建
export CUDA_HOME=/usr/local/cuda-12.9
export PATH=$CUDA_HOME/bin:$PATH
cargo build --release
如果本地没有 GPU/CUDA,可用远端构建脚本把代码同步到带卡的机器上构建/运行/测试:
./tools/sync-and-build.sh build # 远端 cargo build --release
./tools/sync-and-build.sh test # 远端 cargo test
(远端主机、目录、模型路径在 tools/sync-and-build.sh 顶部配置。)
基本用法
1. 启动 HTTP 服务(OpenAI 兼容)
./target/release/xserv-server /path/to/qwen3-8b \
--port 8080 \
--max-batch 4 \
--max-seq-len 8192 \
--swap-space-gb 8
参数说明:
| 参数 | 含义 | 默认 |
|---|---|---|
--port |
监听端口 | 8080 |
--max-batch |
解码批大小(并发上限) | 4 |
--max-seq-len |
单序列最大长度 | 2048 |
--swap-space-gb |
KV 换出到 CPU 的 pinned 内存大小(0 关闭) | 8 |
请求示例(流式):
curl http://localhost:8080/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "qwen3-8b",
"messages": [{"role": "user", "content": "用一句话解释什么是注意力机制"}],
"max_tokens": 256,
"temperature": 0,
"stream": true
}'
其它端点:GET /health、GET /v1/models。
2. 命令行推理
# 单轮生成
cargo run --release --bin xserv-cli -- /path/to/qwen3-8b --max-tokens 256
# 交互式多轮对话
cargo run --release --bin xserv-chat -- /path/to/qwen3-8b
3. 单机性能基准
# 输出每个 prompt 的 TTFT / TBT / TPOT(JSON)
cargo run --release --bin bench-qwen3 -- /path/to/qwen3-8b --gen-tokens 64 [--cuda-graph]
与 llama.cpp 对比 benchmark
tools/bench/ 提供一套一键对比套件,把 xserv 和 llama.cpp(同一份 BF16 权重)放在
相同负载下,黑盒通过 OpenAI API 对比:
- 性能:TTFT、TPOT、吞吐(单流 + 不同并发)
- 精度:AIME 2025、GSM8K(标准数据集,exact-match 评分)
# 一次性准备(需联网的机器):拉取 llama.cpp 子模块 + 下载数据集
git submodule update --init third_party/llama.cpp # 固定在 tag b9371
HF_ENDPOINT=https://hf-mirror.com python3 -m tools.bench.fetch_datasets
# 一键对比(构建 llama.cpp + 转 GGUF + 构建 xserv + 跑两套 + 出报告)
./tools/sync-and-build.sh bench -- --max-seq-len 8192 --quality-limit 50
./tools/sync-and-build.sh fetch-bench-out
# 报告产物:bench-out/comparison-<时间戳>.{md,json}
设计细节见 docs/16-llama-cpp-comparison.md,结果报告见 docs/benchmarks/llama-cpp-comparison.md。
文档
docs/00-roadmap.md:总体路线图与各 Phase 设计docs/01..15-*.md:CUDA FFI / Tensor / GEMM / Attention / KV cache / 性能优化等每个 Phase 的设计文档docs/16-llama-cpp-comparison.md:llama.cpp 对比基准的设计docs/17-tensor-parallelism.md:张量并行(TP)设计docs/18-pipeline-parallelism.md:流水线并行(PP)设计docs/benchmarks/:各阶段的 benchmark 报告(含pp-sweep.md)
多卡并行(TP / PP)
单机多卡,复用 NCCL(crate xserv-distributed)。两种切法正交、二选一:
- 张量并行
--tp N:按 head / 中间维切每一层,层内用 AllReduce 聚合(每 token2·层数次)。 - 流水线并行
--pp N:按层切成 N 段,相邻段间用 NCCL P2P 传 hidden state(每 token 仅N-1次), 通信量远小于 AllReduce,对无 NVLink 的 PCIe 更友好。
# 组内 GPU 0-3:4 卡张量并行 / 4 卡流水线并行
CUDA_VISIBLE_DEVICES=0,1,2,3 ./target/release/xserv-server /path/to/qwen3-8b --tp 4
CUDA_VISIBLE_DEVICES=0,1,2,3 ./target/release/xserv-server /path/to/qwen3-8b --pp 4
PP 实测(dash5,Qwen3-8B BF16,单流贪心;每卡显存为权重+最小 KV 池):
| 配置 | TTFT | TPOT | tok/s | 每卡显存 |
|---|---|---|---|---|
| 单卡 | 33ms | 17.4ms | 57.5 | 24.0 GB |
| PP=2 | 36ms | 18.1ms | 55.3 | 11.6 / 13.6 GB |
| PP=4 | 36ms | 17.9ms | 55.8 | 7.3 / 5.3 / 5.3 / 9.4 GB |
质量对比(AIME 2025 30 题 + GSM8K 30 题,贪心,xserv 在 GPU 0-3、llama.cpp 在 GPU 4-7 并行):
| 引擎 | PP | AIME | GSM8K |
|---|---|---|---|
| xserv | 1/2/4 | 8 / 7 / 7 (/30) | 29/30 (96.7%) 全部一致 |
| llama | 1/2/4 | 7 / 7 / 7 (/30) | 29/30 (96.7%) 全部一致 |
正确性:hidden state 跨段是 bit-exact BF16 P2P 拷贝,PP=4 输出与单卡逐字节一致(用「单卡×2 vs
PP=4×2」对照确认——单卡自身因 cuBLAS 非确定性 run-to-run 会变,而 PP=4 可复现且落在某次单卡轨迹上)。
GSM8K 12 个格子全是 29/30,xserv 与 llama.cpp 完全一致;AIME 的 ±1 是长生成下贪心对 GEMM 抖动的敏感,
非 PP 或引擎效应。收益在显存(每卡权重+KV ≈ 1/N);v1 为串行流水线,单流 TPOT 基本持平、不优于单卡,
真正的吞吐提升需后续做 microbatch / 1F1B 重叠。完整数据见 docs/benchmarks/pp-sweep.md。
路线图(节选)
已完成 Phase 0–18:CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention → 模型加载 → 分词器 → GPT-2 → KV cache → Qwen3-8B → Paged Attention → 连续批处理 → HTTP API → Flash Attention 2 → 性能优化 → 张量并行(TP) → 流水线并行(PP); 并加入了 llama.cpp 对比基准 与 KV CPU 换出 等基础设施。
后续方向:PP microbatch/1F1B 流水线重叠(吞吐收益)、2D TP×PP、投机解码、量化(FP8 / INT8)、多模态。
许可
MIT