diff --git a/README.md b/README.md index 703e17a..f35a2e2 100644 --- a/README.md +++ b/README.md @@ -3,18 +3,23 @@ > 从零用 **Rust + CUDA** 构建的 LLM 推理引擎,目标是吃透 LLM Serving 全栈技术。 xserv 不依赖 PyTorch / vLLM / TensorRT 等现成框架,自己实现了张量抽象、CUDA kernel、 -分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。当前在单张 RTX 5090 上可以 -跑通 **Qwen3-8B**(BF16),并提供一套与 **llama.cpp** 对比正确性和性能的标准 benchmark。 +分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。支持 **Qwen3-8B**(BF16) +和 **gpt-oss-20b**(MoE,BF16/FP8/MXFP4 量化),多卡 TP/PP,并提供一套与 **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 池 +- **模型**:GPT-2(124M)、Qwen3-8B(BF16)、gpt-oss-20b(32 专家 top-4 MoE,harmony 格式) +- **性能**(RTX 5090,贪心,单流): + - Qwen3-8B BF16 单卡:约 56 tok/s(HF transformers 的 1.4×) + - gpt-oss-20b FP8 稀疏 MoE TP=2:**132 tok/s(TPOT 7.2ms),decode 快于 + llama.cpp 同配置**(7.5-8.4ms);llama 单卡模式(2.9ms)仍领先,是下一阶段目标 +- **精度**:GSM8K 全量与 llama.cpp 同权重持平(94.5% vs 94.4%);FP8/MXFP4 量化无回归 +- **服务**:OpenAI 兼容 `/v1/chat/completions`,SSE 流式;gpt-oss 量化后可**单卡 32GB 服务** +- **关键能力**:自写 GEMM / Flash-Attention 2(SM120,含 attention sinks + sliding window) / + Paged-Attention kernel、分页 KV cache(含 **CPU 换出/换入**)、连续批处理、 + CUDA Graph 解码(Qwen3 单卡路径)、**Tensor/Pipeline 并行**(NCCL,TP=1/2/4、PP=2/4)、 + **FP8 W8A8 / MXFP4 W4A16 量化**、**稀疏 top-k MoE decode**(只算被路由的专家) > 这是一个以学习为主的项目,逐 Phase 推进,每步都做数值/端到端验证。 @@ -26,16 +31,19 @@ xserv/ │ ├── gemm/ # GEMM (naive / tiled / gemv) │ ├── attention/ # Flash-Attention 2 (SM120)、Paged-Attention、causal mask │ ├── normalization/ # LayerNorm / RMSNorm -│ ├── activation/ # GELU / SiLU +│ ├── 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-model/ # 模型定义(GPT-2 / Qwen3)、权重加载、KV cache、采样 -│ └── xserv-server/ # tokio + axum HTTP 服务、调度器 +│ ├── 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 报告 ``` diff --git a/docs/00-roadmap.md b/docs/00-roadmap.md index e8c2003..28cb6d4 100644 --- a/docs/00-roadmap.md +++ b/docs/00-roadmap.md @@ -1748,6 +1748,27 @@ Text → Tokenizer → Text Tokens ────────────→ --- +## 实际进展记录(与原计划的分叉,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` | + +**下一步候选(按预期收益排序):** + +| 候选 Phase | 内容 | 预期 | +|---|---|---| +| 21 | **gpt-oss decode CUDA Graph**:把 Phase 15 的 split-graph 方案(`decode_graph.rs`,目前只用于 Qwen3 单卡)推广到 MoE/TP 路径,消除 ~200 launch/token | TPOT 7.0 → ~4-5ms,逼近 llama 单卡 2.9ms | +| 22 | **非专家权重量化**:qkv/o + lm_head(1.16GB/token)仍是 BF16 | TPOT 再省 ~1-1.5ms | +| 23 | **稀疏 prefill**(按专家 permute + grouped GEMM) | 长 prompt TTFT 79 → ~40ms | +| 24 | server 侧 harmony channel 分离(`reasoning_content` 流式输出,对齐 llama-server 行为) | API 易用性 | +| — | Speculative Decoding、多模态(原 16/19) | 推迟 | + ## 里程碑总结 | 里程碑 | Phase | 验收标准 | @@ -1757,7 +1778,9 @@ Text → Tokenizer → Text Tokens ────────────→ | ③ E2E API | 13 | HTTP streaming API, Python OpenAI SDK 可调用, 10 并发正确 | | ④ 性能达标 | 15 | throughput >= 50% vLLM, profiling 报告完成 | | ⑤ 多卡推理 | 17 | TP=2/4 同组 GPU 推理正确, scaling benchmark 完成 | -| ⑥ 多模态 | 19 | 图片输入 → 文字回答, API 端到端 | +| ⑥ MoE 模型(实际) | 19 | gpt-oss-20b 端到端正确, GSM8K 与 llama.cpp 持平 ✅ | +| ⑦ 性能反超(实际) | 20 | 同配置 decode 快于 llama.cpp(TP=2 达成;单卡是 Phase 21+ 目标) ✅ | +| ⑧ 多模态 | 推迟 | 图片输入 → 文字回答, API 端到端 | ## 外部依赖清单 diff --git a/docs/19-gpt-oss-moe.md b/docs/19-gpt-oss-moe.md new file mode 100644 index 0000000..7e74bce --- /dev/null +++ b/docs/19-gpt-oss-moe.md @@ -0,0 +1,118 @@ +# Phase 19: gpt-oss-20b — MoE 模型支持与两次 CUDA 调试实战 + +> 目标:支持 OpenAI gpt-oss-20b(32 专家 top-4 MoE),GSM8K 精度对齐 llama.cpp, +> 并以此为载体做 FP8 / MXFP4 量化。本文档事后整理,重点放在**踩过的坑**: +> 两个教科书级的 CUDA bug 排查过程比结论本身更有学习价值。 +> +> 后续:`docs/20-sparse-moe.md`(稀疏化),benchmark 数据见 +> `docs/benchmarks/{fp8-quantization,mxfp4-and-llama-decode,sparse-moe}.md`。 + +## 1. 模型架构(与 Qwen3 的差异点) + +gpt-oss-20b(`config.json`,已在 dash5 验证): + +| 项 | 值 | 说明 | +|---|---|---| +| layers / hidden | 24 / 2880 | hidden **不是** 128 的倍数的来源(2880 = 22.5×128) | +| heads | 64 Q / 8 KV,head_dim **64** | head_dim ≠ hidden/heads(64×64=4096>2880),GQA n_rep=8 | +| MoE | 32 experts,top-4,expert inter 2880 | router 是普通 Linear [2880→32] + bias | +| attention | **交替 sliding(128)/full**,layer 0 是 sliding | 每层带 **attention sinks**(每 head 一个可学习标量) | +| RoPE | YaRN(theta 150000, factor 32, orig 4096) | attn_factor = 0.1·ln(32)+1 乘在 cos/sin 上 | +| 激活 | clamp 后的 GLU | gate=gu[::2], up=gu[1::2](**交错**), gate≤7, up∈[-7,7], glu=gate·σ(1.702·gate), h=(up+1)·glu | +| 词表 | 201088 | EOS 是**列表** [200002,199999,200012] = `<|return|>`/`<|endoftext|>`/`<|call|>` | +| 其它 | attention_bias=true | q/k/v/o 全部带 bias(Qwen3 没有) | + +**Harmony 对话格式**:gpt-oss 不是普通 chat template,输出分 channel +(`analysis`=思维链,`final`=正式回答),控制 token `<|start|>/<|channel|>/<|message|>/<|end|>`。 +三个坑:(1) system 消息必须含 `Reasoning:` 等 canonical 行,缺了模型 OOD、 +channel 选择不稳定;(2) repetition penalty 会惩罚必须重复出现的控制 token, +导致模型只输出 analysis 不出 final(MoE 默认关掉);(3) 服务端要用多 EOS 判停。 + +## 2. MoE 前向(dense 版,Phase 20 之前) + +```text +router GEMV → topk_softmax(GPU)→ moe_replicate(复制到全部本地专家) +→ batched GEMM gate_up → bias → GLU → batched GEMM down → bias +→ weighted_sum(只取 top-4)→ all-reduce +``` + +要点:top-k 的专家编号始终留在 GPU(`topk_ids`),host 不同步; +dense 的代价(每 token 读全部专家权重)在 Phase 20 用 sparse GEMV 解决。 +TP 用 **expert parallelism**:rank r 拥有专家 [r·E/world, (r+1)·E/world), +weighted_sum 里按 `expert_start + local_experts` 过滤非本地命中, +all-reduce 把各 rank 的部分和加起来——这要求"跳过"语义而不是"乘 0"。 + +## 3. CUDA 调试实战 ①:prefill NaN(flash-attention sinks) + +**症状**:长 prompt(≳192 token)prefill 后输出全 NaN → argmax 落在 +token 201087(`max_by` 平局取最后)或 token 0(`!`)。短 prompt 完全正常。 + +**定位手法**:给每个 stage 加 NaN 检查(环境变量开关,事后移除), +二分出第一个出 NaN 的位置:layer-0 的 `flash_attention_sinks` 输出, +而它的 q/k 输入是干净的 → bug 在 kernel 内部。 + +**根因**:causal 跳过逻辑只剔除"完全在未来"的 kv tile;一个完全滑出 +sliding window(128)的**过去** tile 仍被处理,所有 key 都被 mask 成 -inf +→ `row_max = -inf` → online softmax 里 `expf(-inf-(-inf)) = NaN`, +下一个有效 tile 的修正项 `0·NaN` 把整行毒掉。 + +**修复**:`row_max == -INFINITY` 的 tile 直接跳过(贡献为零)。 +**教训**:online softmax 的"空 tile"是边界条件标配;decode kernel 早就 +防了这个(`local_max==-INFINITY` guard),prefill kernel 漏了—— +**同一逻辑的两份实现要做同样的边界测试**。触发阈值 ~192 token 解释了 +"短测试全过、长对话必炸"的诡异表象。 + +## 4. CUDA 调试实战 ②:decode 间歇性乱码(GEMV 未初始化共享内存) + +**症状**:同一 prompt ~70% 的运行在第二轮对话或长生成中突然输出 +`!!!!`/token 201087/NaN logits,**间歇性** → 不是确定性逻辑错误, +是竞态或未初始化读。只有 gpt-oss 出问题,Qwen3 从不复现。 + +**定位**:逐 stage 检查,第一个出问题的是 decode 的 o_proj 输出 +(maxabs≈1e33),输入干净 → M=1 的 GEMV kernel。 + +**根因**(`gemv.cu`): + +```cuda +if (col >= N) return; // ← 在协作加载 x_shared 和 __syncthreads 之前! +...cooperative load + __syncthreads()... +``` + +当 `N % 128 != 0` 时,最后一个 block 的越界线程提前退出,**没参与** +共享内存装载;在界线程读到未初始化的 smem(且 `__syncthreads` 在有线程 +已退出时是 UB)。命中条件:n=2880 的矩阵(o_proj、MoE gate_up/down)—— +2880 % 128 ≠ 0;而 Qwen3 所有维度都是 128 对齐的,**所以"只有 gpt-oss +不稳定"**。q/k/v(4096)、lm_head(201088)对齐,幸免。 + +**修复**:所有线程先完成装载 + barrier,`col >= N` 检查移到 syncthreads +**之后**。 + +**教训**:`__syncthreads()` 之前的任何 early-return 必须是 **block-uniform** +的。Phase 20 的 sparse GEMV 专门遵守了这条(整个 block 基于同一个 +`topk_ids` 值统一退出,发生在装载之前)。 + +**修复后的验证**:GSM8K 全量 1319 题,xserv 94.5% vs llama.cpp 94.4% +——统计上同一水平,证明两个 kernel bug 就是之前 55% vs 95% 差距的全部原因。 + +## 5. 量化(详见 benchmark 文档) + +- **FP8 W8A8**(`tools/quantize_fp8.py`):per-expert 标量 scale,权重转置 + 存 [E,N,K] 喂 cuBLASLt(Blackwell 要求 transA=T)。两个性能坑: + (1) 每次调用重建 plan + 跑 heuristic → 比 BF16 还慢,修复 = per-shape + plan cache;(2) 逐专家发射 ~768 个小 GEMM,修复 = 单条 strided-batched + 调用 + 把 scale 移到融合的 post-scale kernel。最终 1.41× vs BF16。 +- **MXFP4 W4A16**(`tools/quantize_mxfp4.py`):E2M1 + per-32 UE8M0 块 scale, + 13GB 模型,贪心输出与 BF16 逐字一致,但手写 dequant-GEMV 打不过 + cuBLASLt FP8(带宽效率差),定位为省显存方案。 +- 检测方式:safetensors 的 dtype/scale 秩自动识别,loader 无需配置。 + +## 6. 本阶段的工具沉淀 + +- `bench-gpt-oss`:in-process 推理 + `--forced`(teacher-forced prefill + top-1)/`--forced-decode`(沿参考轨迹逐位置 top-1)——分离"前向算错" + 和"贪心轨迹分叉"的利器。 +- `tools/eval_gsm8k_fast.py`(持久 xserv-chat 管道)、 + `tools/xserv_vs_llama.py`(warm-server 同机对打,计入 llama 的 + reasoning_content)。 +- 经验:**贪心解码不是逐位可复现的**(cuBLAS 非确定性会翻转后段 argmax), + 多卡正确性要用"单卡×2 + 多卡×2 互相比",精度要用基准集而不是逐字 diff。 diff --git a/docs/benchmarks/sparse-moe.md b/docs/benchmarks/sparse-moe.md index 39fb707..318d7ce 100644 --- a/docs/benchmarks/sparse-moe.md +++ b/docs/benchmarks/sparse-moe.md @@ -40,6 +40,12 @@ MXFP4 runs W4A16. Dense path retained for prefill / `num_tokens > 8` and via 13.1 vs 6.6 ms before sparse). Remaining loss: long-prompt TTFT — prefill is still the dense all-expert GEMM; sparse/grouped prefill is the next phase. +**Post-review fixes** (same harness, rerun): removing three leftover +`cudaDeviceSynchronize` from the decode hot path and replacing the CPU-tiled +prefill bias-add (96 D2H/H2D round-trips per prefill) with a GPU broadcast +kernel improved both axes — TPOT 7.19-7.32 → **6.99-7.21 ms**, TTFT +short/medium/long 35/49/94 → **29/42/79 ms**. GSM8K-50: 94% (unchanged). + ## TP=1 head-to-head (single 5090; server now routes gpt-oss tp=1 to the TP engine) | prompt | metric | xserv sparse FP8 | llama MXFP4 |