【DGX Spark 实战】面向 128GB 统一内存的DeepSeek-V4 流式推理引擎设计及实现
本文详细介绍了如何在 128GB 统一内存的 DGX Spark（GB10） 单卡上运行 284B 参数、158GB 权重的 DeepSeek-V4-Flash模型，核心思路是利用 MoE 架构稀疏激活特性，仅按需加载激活的专家权重。
摘要：这是一个把大象装进冰箱的故事。感谢Deepseek v4 flash( API )+ Claude Code大力协助。源代码 、转换后的模型文件。
本文记录技术实现过程，初步效果冷启动大约0.5词元/秒，预热后1~3词元/秒。没有实际应用的意义,但是可以打破固有思维，MoE的大语言模型可以在小于模型体积的显存机器上跑一跑。
参考 如何本地运行DeepSeekV4
一、问题背景
DeepSeek-V4-Flash，参数量 284B，模型权重 158GB。目标硬件为 NVIDIA DGX Spark（GB10 Blackwell），配备 CPU-GPU 统一内存架构，物理容量 128GB。
模型权重超出硬件内存约 30GB，导致标准推理方案难以实施。传统方案需要至少两台 DGX Spark。CPU-GPU 显存卸载在统一内存架构下无效——CPU 与 GPU 共享同一 128GB 内存池。模型量化方案已达 FP4 下限，进一步量化不可行。
DeepSeek-V4 采用 MoE（Mixture of Experts）架构，每层包含 256 个 routed experts，但每个 token 仅激活其中 6 个。这意味着无需令全部 256 个专家同时驻留于 GPU 内存，仅需确保当前激活的专家可用即可。
基于以上观察，本文实现了一种流式推理引擎：非专家权重（embedding, attention, norm, gate, shared expert）永久常驻统一内存，专家权重按需从 NVMe 加载并经由 LFU（Least Frequently Used）缓存管理。158GB 模型的推理总统一内存占用约 110+GB（配置 90GB 专家缓存），接近 128GB 硬件上限。在此基础上引入 CUDA Graph 加速（MoE-only Graph 模式），热缓存稳态 decode 速度 2.31 tok/s，缓存命中率 96.1%。
二、设计原理
2.1 非对称加载策略
模型的 69,143 个张量依据访问模式分为两类：
| 类别 | 大小 | 策略 |
|---|---|---|
| Attention 投影权重 | 6.505 GB | 永久常驻 |
| 共享专家（shared expert） | 1.031 GB | 永久常驻 |
| Embedding | 0.986 GB | 永久常驻 |
| Head 投影 | 0.986 GB | 永久常驻 |
| Hyper-Connection 参数 | 0.129 GB | 永久常驻 |
| MoE Gate + Layer Norm | 0.103 GB | 永久常驻 |
| MTP 头 | 0.031 GB | 永久常驻 |
| 非专家权重合计 | 9.77 GB | 永久常驻 GPU |
| Routed 专家权重（FP4） | 140.25 GB | 按需加载 + LFU 缓存 |
非专家权重大小为 9.77 GB，仅占总模型权重的 6.5%，但位于每步推理的关键路径上——将其常驻内存是代价收益比最优的决策。
2.2 专家权重访问模式与缓存策略
每个专家权重约 13 MB（FP4 量化，6 个张量：w1/w2/w3 各包含 weight 和 scale）。从 NVMe 缓存热读取耗时约 0.77 ms（page cache 命中），冷读取约 4 ms。43 层 × 每层 6 个激活专家 = 每步需访问 258 个专家。
无缓存时每步需读取 258 × 13 MB ≈ 3.3 GB，较全量加载（140 GB）改善约 42 倍，但仍不够理想。
本文采用 LFU（Least Frequently Used）淘汰策略，容量 90 GB（约 6900 个专家，占总池 63%）。选择 LFU 而非 LRU 的理由如下：MoE 专家访问分布呈长尾幂律分布——少数热门专家被频繁访问，大量冷门专家间歇出现。LRU 基于时间局部性假设（最近访问的将很快被再次访问），在长尾分布下会将频繁出现但暂时不在工作集的专家淘汰；LFU 基于频率累计，更适应此类分布。
| 方案 | 每步 I/O 量 | 加速比（以全量加载为基准） |
|---|---|---|
| 全量加载 256 专家/层 | 140 GB | 1× |
| 仅加载 6 激活专家（无缓存） | 3.3 GB | 42× |
| LFU 缓存（冷启动 ~57% 命中率） | 1.4 GB | 100× |
| LFU 缓存（热缓存 96.1% 命中率） | 0.13 GB | 1077× |
| CUDA Graph + LFU 热缓存 | 0.13 GB | 1077× |
2.3 内存预算
128GB 统一内存的预算分配如下：
| 用途 | 预算 | 实际 |
|---|---|---|
| 非专家权重（永久常驻） | ~10 GB | 9.77 GB |
| 专家 LFU 缓存 | ~90 GB | 按需填充 |
| KV Cache（43 层 MLA 低秩压缩） | ~0.5 GB | 0.12 GB / 步 |
| 激活值与临时缓冲 | ~5 GB | ~2 GB |
| PyTorch/CUDA 框架开销 | ~5 GB | ~5 GB |
| 余量 | ~18 GB (14%) | — |
| 合计（预算上限） | ~128 GB | — |
注：热缓存稳态 CUDA 内存峰值 ~100 GB；缓存饱和后总内存约 ~110 GB，在 128 GB 统一内存范围内有约 18 GB 余量。
三、系统技术架构
3.1 架构总览图
3.1.1 初始化流程
3.1.2 组件关系
3.2 核心组件分层架构
| 层级 | 组件名称 | 核心职责 | 关键技术 | 内存占用 |
|---|---|---|---|---|
| 应用接口层 | StreamingInferenceEngine | 对外统一接口，生命周期管理 | Facade Pattern | - |
| 推理核心层 | StreamingMoE / CapturableMoE | 三阶段异步推理 / 6-slot 图捕获 | Monkey Patch, CUDA Stream, CUDA Graph | 11.4 GB |
| 缓存优化层 | ExpertCache | LFU专家缓存，淘汰策略 | OrderedDict, Freq List | 90 GB（上限） |
| 权重管理层 | WeightStore, ExpertLoader | 权重分类型加载，格式转换 | Zero-Copy, FP4 Decode | 9.77 GB + 动态 |
| 存储抽象层 | TensorIndex | 文件抽象，元数据索引 | Per-file mmap, Offset Index | ~100 MB |
3.3 运行时数据流时序图
3.4 核心接口定义
3.4.1 TensorIndex 接口
class TensorIndex:
def __init__(self, model_path: str):
"""初始化，解析45个safetensors文件header"""
pass
def get_expert_tensors(self, layer: int, expert: int) -> List[TensorMeta]:
"""返回指定专家的6个张量元数据
Returns: [(name, offset, dtype, shape)]
"""
pass
def get_non_expert_tensor(self, name: str) -> TensorMeta:
"""获取非专家张量元数据"""
pass
3.4.2 ExpertCache 接口
class ExpertCache:
def __init__(self, capacity_bytes: int = 90 * 1024**3):
self.capacity = capacity_bytes
self.freq_list = defaultdict(OrderedDict) # LFU核心数据结构
def get(self, layer: int, expert: int) -> Optional[Dict[str, Tensor]]:
"""O(1) 查询，命中则频率递增"""
pass
def put(self, layer: int, expert: int, weights: Dict[str, Tensor]) -> None:
"""O(1) 插入，超限则从最低频bucket淘汰"""
pass
def evict(self, target_bytes: int) -> int:
"""释放指定字节数，返回释放的专家数量"""
pass
def hit_rate(self) -> float:
"""返回当前命中率"""
pass
3.4.3 StreamingMoE 接口
class StreamingMoE(nn.Module):
def __init__(self, args):
super().__init__()
self.gate = nn.Linear(...) # 仅创建门控网络
self.shared_expert = SharedExpert(...) # 仅创建共享专家
# routed experts不创建Module，按需加载
def forward(self, hidden_states: Tensor) -> Tensor:
"""三阶段异步流水线
1. Load: 计算活跃专家，异步加载缺失项
2. Compute hits: 预量化输入，计算命中专家
3. Sync + compute misses: 等待传输完成，执行FP4 GEMM
"""
pass
def _apply_expert(self, expert_weights: Dict[str, Tensor],
hidden: Tensor) -> Tensor:
"""直接调用fp4_gemm kernel，绕过linear层"""
pass
3.5 技术栈分层
四、系统组件详解
4.1 TensorIndex：张量索引
转换后的模型分布于 45 个 safetensors 文件（1 个 global.safetensors、43 个 model_L{00-42}.safetensors、1 个 mtp.safetensors）。
TensorIndex 在初始化阶段解析所有文件的 header（不加载张量数据），并对每个文件独立执行 mmap。每个 per-layer 文件仅建立 3.4 GB 的虚拟地址映射，而非对整个 158GB 模型建立单一映射。同时构建从张量名称到（文件路径、数据偏移量、dtype、形状）的完整索引，并提供 get_expert_tensors(layer, expert) 接口以快速查询指定专家的 6 个组成张量。
4.2 WeightStore：非专家权重加载器
WeightStore 负责加载全部非专家权重（1559 个张量，9.77 GB）。实现利用 TensorIndex 的 per-file mmap 进行零拷贝读取——通过 torch.frombuffer 直接从 mmap 内存区域创建张量对象。
选择 per-file mmap 而非 seek+read 的考量：单个 3.4 GB 文件的 mmap 虚拟地址开销为 3.4 GB × 45 文件 ≈ 153 GB VA，在 64 位地址空间中完全可接受；mmap 的按需调页特性保证仅实际访问的页面占用物理内存。
4.3 ExpertLoader：专家权重按需加载器
ExpertLoader.load_expert(layer, expert) 读取指定专家的 6 个 FP4 张量（w1.weight, w1.scale, w2.weight, w2.scale, w3.weight, w3.scale），每个专家约 13 MB。
性能特征：从 NVMe 批量加载 6 个专家（78 MB），page cache 命中时约 0.77 ms/专家，冷启动约 4 ms。该结果表明 I/O 延迟并非主要瓶颈——CUDA kernel 启动延迟在 M=1 decode 场景下可能更为显著。
FP4 存储格式的工程约束：逻辑维度为 [dim, inter_dim]，物理存储为 [dim, inter_dim // 2]（每字节存储两个 FP4 值）。消费侧需通过 .view(torch.float4_e2m1fn_x2) 执行 reinterpret cast。
4.4 ExpertCache：基于 LFU 的专家缓存
ExpertCache 采用 LFU 淘汰策略，数据结构为 freq_list[frequency] → OrderedDict：
get(layer, expert)：访问频率递增，将 key 移至更高频 bucket。时间复杂度 O(1)put(layer, expert, weights)：新条目频率初始化为 1；缓存超容量时从最低频 bucket 淘汰。时间复杂度 O(1)evict(target_bytes)：从_min_freqbucket 的最久未访问端弹出条目，直至释放目标字节数。
LFU 相对于 LRU 的优势已由实测验证：在长序列 MoE 推理中，部分热门专家在每一步均被激活，LRU 的时间局部性假设（最近访问的条目将很快被再次访问）在步间间隔足够长时失效；LFU 的频率累积机制则能稳定保留高频专家。热缓存稳态下命中率达 96.1%，3000+ 次淘汰场景下 LFU 仍维持约 65% 命中率。
淘汰操作显式调用 Python del 删除六个权重张量以触发 GPU 内存回收。缓存粒度为字节级——total_bytes() 精确统计每个条目大小。
4.5 StreamingMoE：运行时专家流式注入
标准的 MoE.__init__ 会创建全部 256 个 Expert 模块（mp=1 时约 129 GB）。StreamingMoE 仅创建门控网络（约 2 MB）和共享专家（约 24 MB），完全跳过 routed expert 的 ModuleList 创建。
StreamingMoE.forward() 实现为三阶段异步流水线：
- Load 阶段：基于 GPU 侧
torch.bincount+torch.sort计算活跃专家分组，避免unique().tolist()引入的 CPU-GPU 同步开销。缓存命中条目归入 ready 队列；缓存未命中条目从 NVMe 加载后，在独立 CUDA stream 上启动异步 H2D 传输。 - Compute hits 阶段：所有 cache-hit 专家共享预量化输入。
forward()中对输入x预先执行一次act_quant将其转为 FP8，w1/w3 共享此量化结果，每专家节省 2 次act_quantkernel launch。 - Sync + compute misses 阶段：
self._load_event.wait()等待异步传输完成，将新加载专家加入 LFU 缓存，随后执行计算。_apply_expert()直接调用fp4_gemmkernel，绕过linear()分发路径。必须对中间结果应用 routing weights（hidden * weights[idx, top]），否则各专家输出均匀混合导致推理失效。
Pre-quantize 策略在 M=1 decode 场景下消除约 516 次不必要的 kernel launch。
实现方式为 monkey-patching，在 Transformer 初始化前执行：
_model.MoE = StreamingMoE
transformer = Transformer(args)
模型创建耗时 0.9 秒，GPU 内存 11.4 GB。
4.6 CapturableMoE + CUDA Graph：加速解码
StreamingMoE 的三阶段流水线涉及 Python 循环、dict 查找和动态 GPU 内存分配，这些操作无法被 torch.cuda.CUDAGraph 捕获。为将推理纳入 CUDA Graph 加速框架，设计了两层架构：
CapturableMoE：固定 6-slot 权重缓冲区替代 StreamingMoE 的按需加载。每个 slot 预分配 w1/w2/w3 的 FP4 权重（uint8 存储）和 scale 缓冲区。forward() 中 M=1 decode 路径遍历 6 个 slot 执行固定次数的 GEMM，无分支、无分配——满足 CUDA Graph 捕获的静态性要求。M>1 prefill 路径保留 per-expert token grouping 以保证正确性。
CUDAGraphRuntime：Attention 路经继续在 eager 模式下执行（sparse_attn TileLang kernel + compressor 状态管理）。每步逐层执行：
1. _prepare_layer (eager):
hc_pre_attn → prepare_decode → attention → hc_post → gate (FFN-path hn) → load experts
2. Graph replay (MoE-only):
hc_pre_ffn → ffn_norm → MoE GEMM (6 slots) → hc_post
关键设计决策：
- Gate 始终实时计算（不走 graph）：消除 one-step-lag——不再使用上一步 graph 捕获的专家选择 indices，而是每步基于当前 token 的 FFN-path hidden state 精确计算 gate。
- Capture 隔离：capture 时使用安全位置（pos+500）warmup kernel，直接加载 dummy experts 到 slot，避免污染真实序列的 KV cache 和 compressor 状态。
- 压缩步回退：compress ratio 层（20 层 indexer + 20 层 dense）在
(pos+1) % ratio == 0的步上无法走 graph（动态控制流），从触发层开始 fallback 到 eager full forward。 - 多轮对话：Turn 1 执行 warmup + graph capture，Turn 2+ 跳过 warmup 直接进入 graph decode 模式。
解码性能：Attention Eager + MoE-only Graph 消除约 62% Python 开销中的 kernel launch 部分。端到端验证 diff=0.0（CUDA Graph vs eager 完全一致）。
五、开发过程
项目经历四个阶段。
5.1 方案评估
五种候选方案的评估结果如下：
- 方案 A：全量加载（158 GB > 128 GB）——不可行
- 方案 B：动态专家卸载——所有权重最终须驻留于内存，物理上限无法突破
- 方案 C：逐层流式加载（每层 3.4 GB，每步 I/O 146 GB）——理论可行，但 I/O 开销极大
- 方案 D：多机张量并行——适用于交互场景，但需多台 DGX Spark
- 方案 E：HF Transformers 原生 offload——待评估
关键点：GB10 GPU 实际最大可用约为 121.6 GB，使单卡推理策略成为可能。
5.2 权重文件转换的代码改造
将 HuggingFace 格式的 46 个 shard 转换为内部 per-layer 格式的实现如下：
转换流程（cleanup_and_convert.sh + convert.py）：
- Index 阶段：扫描 46 个 HF shard 的 safetensors header，建立 69,143 个参数的完整索引，将参数分类为 non-expert 和 expert。
- Non-expert 处理：逐个加载参数，处理完成后立即释放：
- attention 投影（wq, wk, wv, wo）以 BF16 或 FP8 格式直接转发
- wo_a 融合：FP8 存储的 wo_a（weight + scale）在转换阶段执行
weight * scale反量化为 BF16 - 每个参数保存为独立临时 safetensors 文件，按层归类至子目录
- Expert 处理：每层 256 个 routed expert，各含 6 个 FP4 张量。解码时从 int8 包装中提取高低 nibble，经 FP4 查找表（16 个离散值）映射为 FP16，与 FP4 block scale 执行反向解量化后重量化为 F8_E4M3。
- Merge 阶段：将同层所有临时文件合并为单一 per-layer safetensors 文件（如
model_L00.safetensors）。
已知陷阱：
- dtype 字符串兼容性：safetensors Rust 解析器要求格式为
F8_E4M3，而非 Python 的float8_e4m3fn。normalize_dtype()必须输出 Rust 规范格式。 - 内存峰值：FP4 → FP32 反量化过程中同时持有 int8 输入、FP16 查找结果、FP32 解量化中间结果和 F8_E4M3 输出，峰值约 200 GB。128 GB 物理内存需配合至少 128 GB 磁盘 swap。
转换完成后输出 45 个文件（150 GB），经 test_per_layer_sharding.py 的 342 项验证（包含所有张量可读性、形状正确性、wo_a 融合正确性）。
5.3 工程问题记录
早期问题（1-7）
问题 1：safetensors 全文件 mmap 导致 OOM
现象：模型加载时进程被 OOM killer 终止。
原因：safe_open 默认对全部文件执行 mmap，158 GB 文件需要 158 GB 连续虚拟地址空间，128 GB 物理内存配合有限 swap 无法满足。
解决方案：对每个 per-layer 文件（3.4 GB）独立 mmap，共 45 个独立映射。
问题 2：FP4 张量不支持 copy_ 操作
现象：load_state_dict 报错，Parameter.data.copy_() 崩溃。
原因：float4_e2m1fn_x2 类型的张量元数据为只读，不支持原地写入。
解决方案：直接构造新的 nn.Parameter 对象替换旧参数，避免使用 load_state_dict 或 .data.copy_()。
该问题否定了 PyTorch 标准权重加载路径的可用性，StreamingMoE._apply_expert() 中的权重替换必须采用参数替换而非参数拷贝。
问题 3：wo_a 融合逻辑死代码
背景：wo_a 为注意力输出投影，以 FP8 格式存储（权重 + scale）。转换时融合为 BF16 可消除运行时反量化。
融合逻辑设计为用 pending_wo_a 字典暂存先到达的组件（weight 或 scale），待两者均到达后执行融合。但 HuggingFace shard 的 key 按字母序排列：wo_a.scale 在 wo_a.weight 之前被处理。scale 到达时检查 pending_wo_a 未找到 weight，落入通用回退路径调用了 save_file，将 scale 写入独立临时文件。当 wo_a.weight 到达时，scale 已被写走，融合逻辑始终无法执行。
解决方案：在 scale 和 weight 两个处理分支中，当等待伙伴张量时均添加 continue 语句，阻止提前回退至 save_file。该修复解决了推理输出乱码问题。
问题 4：TensorIndex 偏移错位
现象：加载的专家权重数据为乱码。
原因：safetensors header 中的 data_offsets 为每个文件内部的偏移量。索引建立阶段对全部张量做了全局排序，导致将文件 A 的偏移量用于文件 B。
解决方案：直接使用 safetensors header 中的原始 data_offsets，不做全局排序。
问题 5：sparse_attn kernel 共享内存超限
现象：TileLang 编译的 CUDA kernel 在启动时报告共享内存不足。
原因：kernel 需要 141 KB 动态共享内存，而 GB10 的 opt-in 最大共享内存为 101,376 bytes（约 99 KB）。
kernel 包含四个共享内存分配：q_shared（64 KB）、o_shared（64 KB）、kv_shared（取决于 block 大小）、acc_s_cast（取决于 block 大小）。原始配置（block=64, num_stages=2, T.Pipelined）下，kv_shared 64 KB + acc_s_cast 8 KB + pipeline 双缓冲，合计约 141 KB。
优化措施：block 64→16, num_stages 2→1, threads 256→64。但 q_shared（64 KB）与 o_shared（64 KB）之和 128 KB 仍超出限制。观察到两者的生命周期不重叠——q_shared 在 GEMM 读取后释放，o_shared 在循环结束时才被写入。TileLang 编译器对其进行别名优化，复用同一 64 KB 空间。
最终有效共享内存：max(q, o) 64 KB + kv_shared 16 KB + acc_s_cast 2 KB = 约 82 KB，在 99 KB 限制内。
后续进一步优化为 head-tiled 版本：H_per_tile=32, KV_block=32, threads=128，共享内存 ~99KB（精确匹配 GB10 opt-in 上限）。
问题 6：TileLang 版本兼容性
现象：AttributeError: NestedLoopChecker instance has no attribute '_inst'。
原因：TileLang 0.1.8 中 NestedLoopChecker 存在属性未初始化缺陷。
解决方案：升级至 0.1.9。
问题 7：torch default_dtype 导致 fp4_gemm 失败
现象：fp4_gemm kernel 执行时类型不匹配崩溃。
原因：torch.set_default_dtype() 默认为 float32，而 fp4_gemm/fp8_gemm kernel 输出期望 BF16。
解决方案：在 Transformer.__init__ 中设置 torch.set_default_dtype(torch.bfloat16)，并确保所有 linear() 调用路径之前该设置已生效。
CUDA Graph 开发问题（8-14）
问题 8：CapturableMoE prefill M>1 仅用 indices[0] 加载专家
现象：多 token prefill 时输出为乱码。
原因：CapturableMoE.forward() 中 M>1 分支仅读取 indices[0]（第一个 token 的专家选择），其余 token 的专家未被激活，输出为均匀混合。
解决方案：为 M>1 实现 per-expert token grouping（与 StreamingMoE 相同逻辑），通过 torch.bincount + torch.where 分组。
问题 9：CUDA Graph capture 前未 warmup
现象：torch.cuda.graph(g) 进入时触发 cuBLAS 或 TileLang JIT 编译，CUDA 报错 “operation not permitted during capture”。
原因：cuBLAS heuristic selection 和 TileLang JIT 均为运行时懒初始化，在 capture 域内触发 CUDA 内存分配。
解决方案：warmup 至少一次完整 forward 后（所有 kernel JIT 编译完成），再进行 CUDA Graph capture。
问题 10：one-step-lag 导致专家选择偏差
现象：CUDA Graph 模式生成质量低于 StreamingMoE eager 模式。
原因：原始设计中 gate 运行在图内，输出 indices 被 graph 捕获为固定值。_prepare_layer 加载专家时读取的是上一步图重放的 indices，而非当前 token 的真实选择。
解决方案：Gate 移出 graph，每步在 _prepare_layer 中实时计算。CUDA Graph 仅捕获 MoE GEMM + hc_post（Attention Eager + MoE-only Graph 架构）。整个 one-step-lag 问题经 6 个累积 bug 的修复才完全解决。
问题 11：_prepare_layer 用 ffn_norm 导致 KV cache 不一致
现象：graph 捕获期间 kv_cache 内容与 eager 模式有差异。
原因：_prepare_layer 中计算 gate 输入的 hn 时使用了 ffn_norm，而 Block.forward 中 Attention 的 Q 计算使用 attn_norm。不同 norm 导致 kv_cache 写入值不同。
解决方案：_prepare_layer 改用 attn_norm 处理 attention 路径的 hn。
问题 12：decode_step 输出追踪错误
现象：graph 重放后 h 保持输入值，后续层处理的是嵌入输出而非隐藏状态。
原因：capture 时使用 _ = layer(s_h, ...) 丢弃了输出。修正为 s_out.copy_(layer(s_h, ...)) 将输出写入预分配 buffer。
解决方案：graph 重放后 h = s_out 读取正确的输出。
问题 13：压缩步回退从 layer 0 重算
现象：压缩步 fallback 时重新处理了所有已完成的层。
原因：_eager_forward 从 range(0, n_layers) 处理全部层，但此时 h 已包含前 li 层的 graph 输出。
解决方案：改为 range(li, n_layers) 从触发压缩的当前层开始。
问题 14：warmup / capture 污染真实序列状态
现象：capture 后推理输出异常。
原因：capture 时 _full_warmup 和 _prepare_layer 在真实序列位置写入 KV cache 和 compressor 状态。
解决方案：warmup 使用安全位置（pos+500，不超过 max_seq_len）；capture 使用 _load_dummy_experts 直接加载 slot buffer，绕过 _prepare_layer。
5.4 集成流程
StreamingInferenceEngine 的初始化顺序如下：
TensorIndex（解析 45 个文件 header + per-file mmap）
→ WeightStore（从 mmap 零拷贝加载 9.77 GB 非专家权重）
→ ExpertLoader（准备按需专家读取）
→ ExpertCache（初始化 LFU 缓存，默认容量 90 GB）
→ StreamingTransformer（monkey-patch MoE + load_state_dict）
→ warmup（prefill + 50-100 step 填充缓存 + JIT 编译）
→ CUDA Graph capture（Turn 1 only）
→ generate（prefill + token-by-token graph decode）
每个阶段记录 GPU 内存、缓存条目数与命中率。
六、实验结果
实验环境：DGX Spark（GB10, 128 GB 统一内存），模型为 DeepSeek-V4-Flash（284B 参数，158GB 权重），FP4 量化。
以下是单次执行效果：
以下是多轮会话效果
执行CUDAGraph的捕获，后续输出内容还算正常。
6.1 解码性能（热缓存稳态）
| 模式 | tok/s | ms/step | 备注 |
|---|---|---|---|
| StreamingMoE（eager baseline） | 2.31 | 433 | 三阶段流水线 |
| CUDA Graph（MoE-only） | 2.16 | 463 | Attention Eager + MoE Graph |
| 推测解码（MTP, K=4） | 1.23 | — | 0.77× baseline，受限于 attention kernel M>1 效率 |
| 首 token（冷启动） | — | 3.6s | 空 cache，加载约 300 专家 |
瓶颈分解（433ms/step, eager baseline）：
433ms decode step
├── Python/工程开销: ~270ms (62%) ← 最大瓶颈
│ ├── StreamingMoE forward 43 层 Python 循环
│ ├── expert cache 管理 (bincount/sort/where/dict)
│ └── CUDA pipeline drain
├── Attention 43 层: ~108ms (25%)
├── MoE GEMM 43 层: ~43ms (10%)
│ └── fp4_gemm × 6 expert × 43 层
├── Head 投影: ~5ms (1%)
└── 采样/其他: ~7ms (2%)
6.2 缓存性能
| 阶段 | 缓存命中率 | 缓存条目 | 每步 I/O | 备注 |
|---|---|---|---|---|
| 冷启动（首 100 token） | ~57% | 0→2230 | ~1.4 GB | 快速填充 |
| 热缓存（100+ token） | 96.1% | 6900+ | ~0.13 GB | 接近稳态 |
| 长序列（4096 token） | ~90%+ | 满（90 GB） | ~0.3 GB | 偶发冷门专家 miss |
| Evictions 后 | ~65% | 满 | ~1.2 GB | LFU 保留高频专家 |
6.3 内存占用
| 状态 | CUDA 内存 | 实际总内存 | 128GB 占比 |
|---|---|---|---|
| 初始化（仅非专家权重） | ~18 GB | ~28 GB | 22% |
| 热缓存稳态（6900+ 条目） | ~100 GB | ~110 GB | 86% |
| 缓存饱和 + OS 开销 | ~105 GB | ~118 GB | 92% |
| 剩余余量 | — | ~10 GB | 8% |
6.4 性能演进
| 阶段 | 缓存策略 | 命中率 | 解码速度 | CUDA 峰值 | 备注 |
|---|---|---|---|---|---|
| v1 实验（wo_a 损坏） | LRU | 39.4% | 1.02 tok/s | 25.7 GB | 输出异常，仅解码 4 步 |
| v1 修复（wo_a 融合） | LFU | 56.8% | 0.62 tok/s | 110+ GB | 短序列 19 token，含 workaround 回退 |
| v2 优化（head-tiling） | LFU | 72% | 1.60 tok/s | ~100 GB | sparse_attn block 64→16 优化 |
| v2 最终（head-tiled v2） | LFU | 96.1% | 2.31 tok/s | ~100 GB | H_per_tile=32, KV_block=32 |
| v3 CUDA Graph | LFU | 96.1% | 2.16 tok/s | ~100 GB | Attention Eager + MoE-only Graph |
| 预取优化（预计，还未开展） | LFU+EMA | 98%+ | 2.9-3.5 tok/s | ~100 GB | Python 开销优化 |
七、经验总结
7.1 硬件约束驱动架构创新
128 GB 的内存上限是流式推理引擎产生的主要推动力。资源约束在此场景下成为创新的催化剂。per-file mmap、LFU 缓存、三阶段异步流水线均为直接响应物理限制的设计选择。
7.2 MoE 流式推理的可行性
实验结果表明：采用 FP4 量化、LFU 缓存、per-file mmap 和 per-layer 文件布局，可在单卡上运行 284B 参数模型。热缓存稳态 CUDA 内存约 100 GB，总内存约 110 GB，在 128 GB 统一内存的硬件上可行，余量约 10 GB。
7.3 统一内存架构下的工程考量
CPU 与 GPU 共享内存池带来以下影响：
- 优势：无需显式 CPU-GPU 数据传输，
to(device)操作近乎冗余 - 注意：
torch.cuda.memory_allocated()仅反映 CUDA 分配器管理的内存，CPU-side 张量和 mmap 映射不在此统计范围内。实际总内存消耗 = CUDA tracked + CPU-side tensors + 框架开销
7.4 量化精度与工程复杂度
FP4 量化在存储效率之外引入了显著的工程复杂度：
element_size()返回 1（与 int8 一致），张量级内存节省为 0.data.copy_()不支持，load_state_dict不可用- 逻辑形状（
[d0, d1]）与物理形状（[d0, d1//2]）不一致 - 类型转换依赖
view的 reinterpret cast，而非真实数据转换
排查成本随量化激进程度非线性增长。
7.5 集成测试的必要性
wo_a 融合缺陷为典型的集成问题：各组件独立测试均通过——转换器正确读取了 scale，正确读取了 weight——但组合后功能失效。仅有端到端推理测试才能暴露该问题。建议每个新架构特性增加端到端验证步骤。
7.6 Python 开销是主力瓶颈
eager baseline 的 433ms/step 中 62%（~270ms）为 Python/CUDA 边界开销，包括 43 层 Python 循环、expert cache 管理（bincount/sort dict 操作）和 CUDA pipeline drain。CUDA Graph 部分缓解了 kernel launch 开销，但 StreamingMoE forward 的 Python 层开销仍是下一步优化的重点。
7.7 CUDA Graph 开发的教训
CUDA Graph 带来了显著的正确性挑战：
- 可捕获性：所有操作必须在 capture 域内是静态的——无分配、无分支、固定地址。CapturableMoE 的 6-slot 固定缓冲区设计是关键。
- 状态管理：KV cache、compressor state、position buffer 等可变状态必须在 eager 模式下更新，不可被 graph 捕获。
register_buffer+fill_()的 fixed-address 模式是可靠方案。 - 调试方法：graph vs eager 的 diff 比较是最有效的验证手段（
test_cudagraph.py中 multi-step diff=0.0 验证）。 - 回退路径：压缩步等动态控制流场景须设计 eager fallback 路径，且必须从正确的中间状态继续执行。
7.8 I/O 优化方向
流式推理模式下 GPU 算力非主要瓶颈，NVMe 读取延迟和 cache miss 处理是关键优化方向：
- 热门专家预取：基于 EMA 统计与 token 间连续性预测后续专家，提前加载至缓存
- 批量读取：合并同层专家 I/O 为连续读取请求
- I/O 与计算重叠：在当前层计算阶段预取后续层所需专家
本文涉及的完整代码见 AtomGit 仓库。
欢迎加入DeepSeek 技术社区。在这里，你可以找到志同道合的朋友，共同探索AI技术的奥秘。
更多推荐
5
0- 0
已为社区贡献3条内容
所有评论(0)