昇腾论坛

起因

前段时间看到 antirez（Redis 的作者，Salvatore Sanfilippo）在 GitHub 上放出了一个项目 ds4，完整名字叫 DwarfStar 4。

这是一个专门为 DeepSeek V4 Flash 写的本地推理引擎，不依赖 llama.cpp，不是通用 GGUF 运行时，完全从头实现，目标是 Metal（Mac）和 CUDA（NVIDIA）两个后端。

代码量不大，结构清晰，C 写的，有点 Redis 的风格——能用最简单的方式解决问题就不搞复杂。作为一个平时也在做系统优化相关工作的人，看完觉得是个很好的学习材料。

恰好手边有一块 Atlas 300i Duo（华为昇腾 310P 系列的 PCIe 推理卡，48GB × 2），就想着把 ds4 改写一个昇腾版本，自己走一遍推理流程，加深理解。项目放在了这里：ds4-ascend。

效果先说：非常慢，但流程跑通了。这篇主要记录 ds4 的推理原理和适配昇腾过程中遇到的几个关键问题。

ds4 做了什么

专用引擎，不通用

ds4 和 llama.cpp 最大的区别是范围收窄：它只跑 DeepSeek V4 Flash 这一个模型，只支持作者专门打包的 GGUF 文件，不是通用 GGUF 加载器。

这让代码变得非常精简。不需要考虑各种模型的 tensor 布局差异，不需要抽象成通用图，就是一条针对这个模型写死的推理路径。

MLA：DeepSeek 的 KV Cache 压缩

DeepSeek V4 Flash 用的注意力机制叫 MLA（Multi-head Latent Attention），这是它最有意思的地方。

传统 Transformer 的 KV Cache 每层要存 n_heads × head_dim × seq_len 的数据。对于长上下文（比如 100k token），这个量非常大。

MLA 的做法是把 K 和 V 一起压缩成一个低维的"潜在向量"（latent vector）存起来，只有在实际计算注意力时才解压出来。压缩比大约是 8:1 到 10:1。

ds4 把这个压缩向量叫做 compressed KV，存储和加载都是围绕这个设计的。磁盘 KV cache 存的也是这个压缩后的格式，所以一个百万 token 的上下文在 M3 Max 上大约只需要 26GB 内存，这在之前是不可想象的。

MoE：只激活一部分专家

DeepSeek V4 Flash 是 MoE（Mixture of Experts）架构，总参数 284B，但每个 token 推理时只激活大约 37B 的参数（几个 routed expert + shared expert）。

ds4 的量化策略也专门针对这个设计：routed MoE expert 用 2-bit（IQ2_XXS / Q2_K），其他部分（shared expert、projection、routing）保持高精度。这样总模型大约 81GB，在 128GB 内存的 MacBook 上能跑。

推理流程

一次推理大致是这样的：

输入 tokens
  → Embedding 查表
  → N 层 Transformer block：
       - RMSNorm
       - MLA 注意力（Q/K/V 投影 + 压缩 KV + RoPE + softmax + 输出投影）
       - MoE FFN（routing → top-k experts → 加权求和）
  → 最后一层 RMSNorm
  → LM head（线性投影到词表）
  → 采样（greedy / temperature / top-p / min-p）
输出下一个 token

ds4 把这个流程在 Metal/CUDA 上用 kernel 实现，prefill 和 decode 分开处理（prefill 并行处理所有输入 token，decode 每次只处理一个）。

KV Cache 持久化

ds4 的另一个亮点是 on-disk KV cache。服务端把 KV 状态序列化到磁盘，下次相同前缀的请求直接从磁盘加载，跳过 prefill。

磁盘格式是自定义的二进制格式，头部 48 字节，然后是 rendered text（用于前缀匹配的原始 token 解码文本），再是压缩 KV 的实际 tensor 数据。

工具调用的精确 replay 也是存在这个文件里的，确保 DSML tool call 文本和模型采样时完全一致，KV Cache 不会因为重新渲染而失效。

移植到昇腾

昇腾的编程模型

华为昇腾的编程接口叫 CANN（Compute Architecture for Neural Networks），类似 CUDA 的角色。核心 API 是 AscendCL，写 kernel 用 Ascend C（基于 C++ 的 DSL）。

和 CUDA 的主要区别：

没有 CUDA 那么多第三方生态，很多东西要自己写
kernel 通过 算子编译（AOT 或 JIT）生成，不像 CUDA 直接 nvcc 编译
内存模型分 Host（CPU 内存）和 Device（NPU 内存），需要显式搬运

适配思路

ds4 的后端抽象在 ds4_gpu.h 里，Metal 实现在 ds4_metal.m，CUDA 在 ds4_cuda.cu。我的做法是新增一个 ds4_ascend.cpp，用 AscendCL 实现相同的接口。

主要需要实现的算子：

矩阵乘法（GEMM）：投影层、FFN 都需要，昇腾有 aclblasGemmEx 可以直接用
RMSNorm：需要自己写 kernel，或者用 AIPP 算子拼
RoPE（旋转位置编码）：按位旋转，不复杂，但要小心 head dim 的排布
Softmax：attention 权重归一化
量化矩阵乘法（IQ2_XXS / Q2_K）：这是最麻烦的部分

最麻烦的：2-bit 量化 kernel

ds4 里 routed expert 用 IQ2_XXS 格式量化，这是 llama.cpp 的一种特殊 2-bit 量化，带 lookup table 的非均匀量化。在 CUDA 上有针对 Tensor Core 优化的 kernel，在昇腾上没有现成的。

目前 ds4-ascend 的做法是把 2-bit 量化矩阵先在 CPU 侧反量化成 fp16，再搬到 NPU 做普通 fp16 GEMM。这是性能差的根本原因：反量化本身不慢，但每次推理都做一次 Host↔Device 数据搬运，带宽成为瓶颈。

正确的做法是把反量化也写成 NPU 上的 kernel，直接在 NPU 片上展开，避免来回搬运。这部分还没完成。

Atlas 300i Duo 的特殊性

300i Duo 是一块 PCIe 推理卡，物理上两个 Die（每个 48GB），在 CANN 里会枚举成两个独立的 device（device 0 和 device 1）。

ds4 目前是单 device 的，适配时直接用 device 0，另外 48GB 没用上。单 Die 48GB 跑 2-bit 量化的 81GB 模型，只能做层间切分（每次推理一层一层过，不驻留全量权重），速度更慢。

实际运行速度大概是 0.1~0.3 token/s，和 Mac 上几十 token/s 差距巨大。

跑通的意义

慢归慢，流程跑通的意义在于：

验证了推理正确性。输出和 CPU 路径对齐，说明 MLA attention、MoE routing、量化反量化这几个关键步骤都是对的。
搞清楚了 ds4 的内部结构。antirez 的代码写得很直接，但 MLA 的 KV 压缩和 DSML tool call 的精确 replay 不读代码很难理解，跑一遍加深了印象。
暴露了昇腾适配的难点。量化 kernel 是核心瓶颈，后续如果想做到实用，必须把 2-bit 反量化搬到 NPU 上，同时考虑双 Die 的权重切分。

后续

在 NPU 上实现 IQ2_XXS 反量化 kernel，消除 Host↔Device 搬运
双 Die 权重切分，充分利用 96GB 显存
prefill 的 batched attention kernel

代码在 github.com/donge/ds4-ascend，欢迎有昇腾硬件的同学一起搞。

原始项目：antirez/ds4，MIT License

起因