Gahow Wang 6da0972740 speculative: copy_kv_position primitive for tree drafting KV remap
SGLang-style "write-all, copy-move on acceptance" approach: after tree
verification, physically copy an accepted sibling's K/V from its
physical cache slot to the canonical sequential position.

New CUDA kernel: copy_kv_position_kernel in reshape_and_cache.cu.
For one token (src_pos → dst_pos), copies head_dim × num_kv_heads BF16
elements in both K and V pools. Grid = num_kv_heads, block = head_dim.
Cost for one token across 36 layers: ~5.3 MB D2D copy @ 900 GB/s = <6μs.

Rust FFI: copy_kv_position(k_pool, v_pool, block_ids, src_pos, dst_pos,
num_kv_heads, head_dim, block_size, stream).

PagedKVCache method: copy_kv_position(slot, src_pos, dst_pos) — uploads
block_ids for the sequence, calls the kernel per layer. This is the
primitive needed by tree drafting: when a non-primary sibling at cache
position P+2 is accepted as the "true" token for target position P+1,
call copy_kv_position(slot, P+2, P+1) then truncate to P+2.

Next: wire into bench-eagle3 tree drafting loop with top-2 siblings.
2026-07-01 23:09:35 +08:00

xserv

从零用 Rust + CUDA 构建的 LLM 推理引擎,目标是吃透 LLM Serving 全栈技术。

xserv 不依赖 PyTorch / vLLM / TensorRT 等现成框架自己实现了张量抽象、CUDA kernel、 分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。支持 Qwen3-8BBF16gpt-oss-20bMoEBF16/FP8/MXFP4 量化),多卡 TP/PP并提供一套与 llama.cpp 对比正确性和性能的标准 benchmark。

现状一览

  • 模型GPT-2124M、Qwen3-8BBF16、gpt-oss-20b32 专家 top-4 MoEharmony 格式)
  • 性能RTX 5090贪心单流
    • Qwen3-8B BF16 单卡:约 56 tok/sHF transformers 的 1.4×
    • gpt-oss-20b FP8 稀疏 MoE + CUDA Graph decodeTPOT 5.8ms~172 tok/s TP=1/2 同速);同配置 TP=2 全面快于 llama.cpp1.26-1.47×llama 单卡模式2.8ms)仍领先,差距 2.0×
  • 精度GSM8K 全量与 llama.cpp 同权重持平94.5% vs 94.4%FP8/MXFP4 量化无回归
  • 服务OpenAI 兼容 /v1/chat/completionsSSE 流式gpt-oss 量化后可单卡 32GB 服务
  • 关键能力:自写 GEMM / Flash-Attention 2(SM120含 attention sinks + sliding window) / Paged-Attention kernel、分页 KV cacheCPU 换出/换入)、连续批处理、 CUDA Graph 解码Qwen3 单卡 + gpt-oss 全路径整图回放)、Tensor/Pipeline 并行NCCLTP=1/2/4、PP=2/4FP8 W8A8 / MXFP4 W4A16 量化稀疏 top-k MoE decode(只算被路由的专家)

这是一个以学习为主的项目,逐 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 / gpt-oss GLU
│   ├── embedding/         #   embedding lookup / RoPE / transpose
│   ├── moe/               #   MoE top-k 路由、稀疏专家 GEMV、加权求和
│   ├── quantization/      #   FP8 量化/反量化、cuBLASLt FP8 GEMM、MXFP4 GEMV
│   └── 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-distributed/ # NCCL FFI、TP 上下文AllReduce
│   ├── xserv-model/       # 模型定义GPT-2 / Qwen3 / gpt-oss MoE、权重加载、KV cache、采样
│   └── xserv-server/      # tokio + axum HTTP 服务、调度器、TP/PP 引擎
├── tools/                 # 辅助脚本 + benchmark 套件(见下)
└── docs/                  # 每个 Phase 的设计文档 + benchmark 报告

环境要求

  • GPUNVIDIA计算能力 SM120RTX 5090 / Blackwell。其它架构需调整 CUDA_ARCH
  • CUDA Toolkit12.9nvcc 需在 PATH,构建 .cu 依赖它)
  • Rustedition 2024建议较新的 stable 工具链)
  • 模型HuggingFace 目录格式(含 config.jsontokenizer.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 /healthGET /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 / TPOTJSON
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-*.mdCUDA FFI / Tensor / GEMM / Attention / KV cache / 性能优化等每个 Phase 的设计文档
  • docs/16-llama-cpp-comparison.mdllama.cpp 对比基准的设计
  • docs/17-tensor-parallelism.md张量并行TP设计
  • docs/18-pipeline-parallelism.md流水线并行PP设计
  • docs/benchmarks/:各阶段的 benchmark 报告(含 pp-sweep.md

多卡并行TP / PP

单机多卡,复用 NCCLcrate xserv-distributed)。两种切法正交、二选一:

  • 张量并行 --tp N:按 head / 中间维切每一层,层内用 AllReduce 聚合(每 token 2·层数 次)。
  • 流水线并行 --pp N:按层切成 N 段,相邻段间用 NCCL P2P 传 hidden state每 token 仅 N-1 次), 通信量远小于 AllReduce对无 NVLink 的 PCIe 更友好。
# 组内 GPU 0-34 卡张量并行 / 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 实测dash5Qwen3-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/30xserv 与 llama.cpp 完全一致AIME 的 ±1 是长生成下贪心对 GEMM 抖动的敏感, 非 PP 或引擎效应。收益在显存(每卡权重+KV ≈ 1/Nv1 为串行流水线,单流 TPOT 基本持平、不优于单卡, 真正的吞吐提升需后续做 microbatch / 1F1B 重叠。完整数据见 docs/benchmarks/pp-sweep.md

路线图(节选)

已完成 Phase 021CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention → 模型加载 → 分词器 → GPT-2 → KV cache → Qwen3-8B → Paged Attention → 连续批处理 → HTTP API → Flash Attention 2 → 性能优化 → 张量并行TP流水线并行PPgpt-oss MoE + FP8/MXFP4 量化稀疏 top-k MoE decodedecode CUDA Graph 整图回放 并加入了 llama.cpp 对比基准KV CPU 换出 等基础设施。

后续方向非专家权重量化lm_head/qkv/o、稀疏 prefillgrouped GEMM、server 侧 harmony channel 分离、PP microbatch/1F1B、投机解码、多模态。详见 docs/00-roadmap.md 的实际进展记录。

许可

MIT

Description
No description provided
Readme 2 MiB
Languages
Rust 67.5%
Python 15.1%
Cuda 13.5%
Shell 3.9%