Chenxi Ma

·6 min read分钟阅读

Code tour代码导览

mini-DeepInfer Technical Notes: From a Request to DeepSeek-V4.1mini-DeepInfer 技术说明：从请求到 DeepSeek-V4.1

A code tour of how mini-DeepInfer turns one request into reproducible V4.1 inference steps: lifecycle, attention state, paged blocks, ZOS, sampling and evidence.代码导览：mini-DeepInfer 如何把一个请求变成可复现的 V4.1 推理步骤——生命周期、注意力状态、分页 block、ZOS、采样与正确性证据。

DeepInferDeepSeek-V4.1RuntimeCode tour
Contents目录 09 sections节
  1. 1. Scope and operating boundary
  2. 2. The lifecycle of a request
  3. 3. V4.1’s attention state
  4. 4. Paged blocks and state slots
  5. 5. Zero-overhead scheduling (ZOS)
  6. 6. Sampling and determinism
  7. 7. Correctness evidence
  8. 8. Where to start changing things
  9. 9. Known boundaries and future work
  1. 1. 范围与运行边界
  2. 2. 一次请求的生命周期
  3. 3. V4.1 的注意力状态
  4. 4. 分页 block 与 state slot
  5. 5. Zero-overhead scheduling（ZOS）
  6. 6. 采样与确定性
  7. 7. 正确性证据
  8. 8. 从哪里开始改
  9. 9. 已知边界与后续工作

This document explains how mini-DeepInfer turns one request into a reproducible V4.1 inference step. It follows the current repository code and focuses on implementation boundaries that can be read, verified and changed; it is not a promise about DeepInfer’s production-cluster performance. Read A DeepInfer-Style Runtime Design first for the mapping to the official design, the step data contract and the implementation order, then follow the tour below.

1. Scope and operating boundary

mini-DeepInfer is a PyTorch-first teaching runtime that supports a single GPU and intra-node EP × DP. It contains:

  • DeepSeek-V4.1’s compressed KV, sliding window, two-level sparse indexer, hyper-connections, Engram and MoE;
  • DeepInfer-style paged KV blocks, per-sequence state slots, continuous batching, chunked prefill, preemption and deterministic sampling;
  • component-level and end-to-end comparison tests against the official V4.1 reference.

It does not replicate the production C++ scheduler, CP and multi-node communication, the official custom CUDA/TileLang kernels, or the DSpark/MTP serving path. The default checkpoint is a small random-weight model: good for testing shapes and execution semantics, useless for judging generation quality.

2. The lifecycle of a request

The entry point is minideepinfer/entrypoints/llm.py. LLM.generate() creates a Sequence for each prompt and hands it to LLMEngine; every step() goes through three phases:

  1. Scheduler.schedule() decides on the CPU whether this step is prefill or decode, and reserves KV blocks and state slots for the sequences;
  2. ModelRunner.execute() sends the batch, block tables, positions and sampling parameters to the GPU asynchronously;
  3. Scheduler.process_outputs() commits tokens once the results are readable, handles EOS/stop tokens/length limits, and releases finished or preempted state. Explicit cancellation goes through LLMEngine.abort_request(), which first drains the in-flight step, then lets the scheduler invalidate the epoch and reclaim both kinds of cache.

The scheduler never mixes prefill and decode in one step. Prefill can be cut into chunks by max_num_batched_tokens; decode advances every running sequence by one token. When blocks run out, the sequence with the longer output is preempted first, goes back to the waiting queue and is recomputed later.

sequenceDiagram
    participant U as Caller
    participant E as LLMEngine
    participant S as Scheduler
    participant R as ModelRunner
    participant G as GPU cache/model
    U->>E: generate(prompt, sampling_params)
    loop until all sequences finish
        E->>S: schedule()
        S-->>E: prefill/decode batch + reservations
        E->>R: execute(batch)
        R->>G: async forward + sample
        G-->>R: sampled token buffer
        E->>S: process_outputs(previous step)
    end
    E-->>U: token ids / text

3. V4.1’s attention state

V4.1 maintains two kinds of state at once, and cannot be understood through a traditional [layer, batch, head, time] KV cache alone:

StateRuntime locationWritten byReleased by
Compressed KV and indexer keysPaged blocks, PAGE_TOKENS=128 tokens eachKV source layers, per compress_ratioBlocks released when the sequence ends or is preempted
Sliding window and compressor tailOne state slot per sequenceEvery layer updates the window; KV source layers update the tailThe slot is released together with the sequence’s blocks

V41Cache.allocate() allocates these tensors in one go from the model configuration. BlockManager only manages ownership and reference counts of blocks; state slots are managed by a second BlockManager, so scheduling must satisfy the capacity of both resources at once.

Inside one attention layer, the hidden state passes through a low-rank query, a single-KV-head quantization round trip, the compressor and the indexer:

flowchart LR
    X["hidden x"] --> Q["low-rank query"]
    X --> K["KV projection + FP8 round trip"]
    K --> W["window ring"]
    X --> C["compressor"]
    C --> CKV["compressed KV / FP4"]
    C --> IK["indexer keys / FP4"]
    Q --> I["two-level top-k indexer"]
    IK --> I
    Q --> A["sparse attention + sink"]
    I --> A
    W --> A
    CKV --> A
    A --> O["grouped low-rank output"]

Only layers configured as KV source write the compressed cache, and only index-source layers compute top-k; the other layers reuse the nearest source layer’s result through kv_source_of() and index_source_of(). This keeps V4.1’s cross-layer reuse semantics while letting the cache layout be accessed by a continuous batch.

4. Paged blocks and state slots

Each Sequence keeps its own block_table, state_slot, cached token count and in-flight token count. Block 0 is a reserved padding block; real blocks are allocated from 1. On release, a block goes to the tail of the free list and can be reused only after its reference count reaches zero.

The scheduler’s invariants:

  • num_cached_tokens is the position already written to KV/state;
  • num_inflight_tokens is the results already submitted to the GPU but not yet processed by the CPU;
  • decode advances only when num_cached_tokens == sequence_length - 1;
  • preemption releases blocks and slot and increments epoch, invalidating old in-flight results;
  • the same seq_seed and position give the same random sample, so recomputation can reproduce a discarded token.

5. Zero-overhead scheduling (ZOS)

With enable_zos, the CPU does not wait for the previous round’s sampled results before scheduling the next one. ModelRunner keeps the previous round’s GPU token buffer; if a token needed by the next decode has not returned to the CPU, it indexes that buffer directly on the GPU through remap. The CPU still makes capacity decisions from reservation counts, while the GPU reads tokens by dependency on the same stream.

step N:     schedule(N) -> enqueue forward/sample -------------------->
step N+1:   schedule(N+1) -> remap previous sampled buffer -> enqueue ->
CPU:        process(N-1) -------------------------------> commit token

The forward pass no longer contains host synchronization; everything that must read a GPU result is concentrated in StepHandle.wait(), so the CPU work of step N can overlap the GPU work of step N-1.

6. Sampling and determinism

layers/sampler.py implements temperature, top-k and top-p sampling on a Gumbel-max path. The random seed is mixed from the sequence seed and the token position, not from the batch row. So the same input with the same seq_seed should produce the same token sequence under any batch composition and preemption history.

This determinism is part of preemption correctness: preemption throws away the cache and the prefix must be recomputed; if sampling depended on global RNG state, recomputation would change later tokens and the scheduler could not verify that “release and replay” is equivalent.

7. Correctness evidence

Tests come in three layers:

LayerCommand/filePurpose
CPU resourcestests/test_block_manager.py, tests/test_scheduler.pyVerify free list, reference counts, chunked prefill, preemption and boundaries
Samplingtests/test_sampler.pyVerify seed, temperature, top-k/top-p and batch independence
V4.1 comparisontests/test_v41.pyCompare quantization, indexer, sparse attention and end-to-end logits against the official reference

Run the full check:

cd mini-DeepInfer
pytest tests/test_block_manager.py tests/test_scheduler.py tests/test_sampler.py
pytest

GPU and official-reference tests need a working CUDA environment; without a GPU run the CPU subset and mark the checks that did not run as unverified, not as passed.

8. Where to start changing things

GoalEntry point
Change scheduling policyminideepinfer/engine/scheduler.py
Change block/slot lifecycleminideepinfer/engine/block_manager.py, sequence.py
Change ZOS batch constructionminideepinfer/worker/model_runner.py
Change V4.1 attention / MoEminideepinfer/models/deepseek_v41/model.py
Change low-precision stateminideepinfer/models/deepseek_v41/quant.py, weights.py
Change reproducible experimentstests/, tools/v41_reference.py

The recommended order is: write a minimal test that exposes the state or scheduling difference, change the implementation, then regress with reference logits and the preemption / batch-composition variants. A performance number must state GPU, checkpoint, batch, context length, dtype and whether warm-up is included.

9. Known boundaries and future work

The current implementation puts “readable and matching” first:

  • The forward pass is mostly PyTorch / Triton; host synchronization inside forward has been removed and results are collected in StepHandle.wait(), so ZOS can truly overlap;
  • MoE uses grouped GEMM (FP4 weights dequantized to bf16 on demand, not a native FP4 GEMM), and decode uses a CUDA Graph bucketed by batch / context; prefill, logits and sampling are still eager;
  • Real V4.1-Flash prefill on 4 × B200 is still only 11.9% of the official engine; the remaining gap is in performance report C1;
  • DSpark speculative decoding, a prefill cost model and end-to-end batch invariance are still on the roadmap;
  • A random-weight checkpoint verifies execution semantics only; it cannot prove DeepSeek-V4.1’s task quality or production throughput.

For the roadmap, license and third-party notices see the README and THIRD_PARTY_NOTICES.

这份文档说明 mini-DeepInfer 如何把一个请求变成可复现的 V4.1 推理步骤。它对应当前仓库代码，重点是可以读懂、可以验证和可以改动的实现边界；它不是对 DeepInfer 生产集群性能的承诺。先阅读 DeepInfer 风格的运行时设计，了解官方设计与本项目的映射、step 数据契约和后续实现顺序，再阅读下面的代码导览。

1. 范围与运行边界

mini-DeepInfer 是以 PyTorch 为主的教学型运行时，支持单卡和机内 EP × DP，包含：

  • DeepSeek-V4.1 的压缩 KV、滑动窗口、两级稀疏 indexer、hyper-connections、Engram 和 MoE；
  • DeepInfer 风格的分页 KV block、每序列 state slot、连续批处理、chunked prefill、抢占和确定性采样；
  • 与官方 V4.1 reference 的组件和端到端对照测试。

它没有复刻生产系统的 C++ 调度器、CP 与多机通信、官方的定制 CUDA/TileLang kernel 或 DSpark/MTP 服务路径。默认 checkpoint 是随机权重的小型模型，适合测试形状和执行语义，不能用来评价生成质量。

2. 一次请求的生命周期

入口在 minideepinfer/entrypoints/llm.py。LLM.generate() 为每个 prompt 创建 Sequence，交给 LLMEngine；每次 step() 都经历三个阶段：

  1. Scheduler.schedule() 在 CPU 上决定本步是 prefill 还是 decode，并为序列预留 KV block 与 state slot；
  2. ModelRunner.execute() 将 batch、block table、位置和采样参数异步送入 GPU；
  3. Scheduler.process_outputs() 在结果可读后提交 token，处理 EOS/stop token/长度上限，并释放已完成或被抢占的状态。 显式取消通过 LLMEngine.abort_request() 先排空在途 step，再让 scheduler 使 epoch 失效并回收两类缓存。

调度器不把 prefill 和 decode 混在同一步。prefill 可以按 max_num_batched_tokens 切块；decode 每条运行中的序列前进一个 token。当 block 不足时，较长输出优先被抢占，释放的序列回到 waiting 队列并在之后重算。

sequenceDiagram
    participant U as Caller
    participant E as LLMEngine
    participant S as Scheduler
    participant R as ModelRunner
    participant G as GPU cache/model
    U->>E: generate(prompt, sampling_params)
    loop until all sequences finish
        E->>S: schedule()
        S-->>E: prefill/decode batch + reservations
        E->>R: execute(batch)
        R->>G: async forward + sample
        G-->>R: sampled token buffer
        E->>S: process_outputs(previous step)
    end
    E-->>U: token ids / text

3. V4.1 的注意力状态

V4.1 同时维护两类状态，不能只用传统的 [layer, batch, head, time] KV cache 来理解：

状态运行时位置写入方式释放方式
压缩 KV 与 indexer key分页 block，每块 PAGE_TOKENS=128 个 tokenKV source layer 按 compress_ratio 写入序列结束或抢占时释放 block
滑动窗口与 compressor tail每个序列一个 state slot每层更新窗口；KV source layer 更新 tail与序列 block 一起释放 slot

V41Cache.allocate() 按模型配置一次性分配这些张量。BlockManager 只管理 block 的所有权和引用计数；state slot 由第二个 BlockManager 管理，因此调度时必须同时满足两种资源的容量。

在一个注意力层中，hidden state 经过低秩 query、单 KV 头的量化往返、压缩器和 indexer：

flowchart LR
    X["hidden x"] --> Q["low-rank query"]
    X --> K["KV projection + FP8 round trip"]
    K --> W["window ring"]
    X --> C["compressor"]
    C --> CKV["compressed KV / FP4"]
    C --> IK["indexer keys / FP4"]
    Q --> I["two-level top-k indexer"]
    IK --> I
    Q --> A["sparse attention + sink"]
    I --> A
    W --> A
    CKV --> A
    A --> O["grouped low-rank output"]

只有配置为 KV source 的层写压缩缓存，只有 index source 的层计算 top-k；其余层通过 kv_source_of() 和 index_source_of() 复用最近的源层结果。这样既保留 V4.1 的层间复用语义，又让缓存布局能被连续 batch 访问。

4. 分页 block 与 state slot

每条 Sequence 保存自己的 block_table、state_slot、已缓存 token 数和在途 token 数。block 0 是保留的 padding block；真实 block 从 1 开始分配。释放时 block 进入 free list 尾部，引用计数归零后才可再次使用。

调度器的几个不变量是：

  • num_cached_tokens 表示已经写入 KV/state 的位置；
  • num_inflight_tokens 表示已经提交到 GPU 但 CPU 尚未处理的结果；
  • decode 只在 num_cached_tokens == sequence_length - 1 时推进；
  • 抢占会释放 block 和 slot，并增加 epoch，使旧的在途结果失效；
  • 相同 seq_seed 和位置会产生相同的随机采样，因此重算可以复现被丢弃的 token。

5. Zero-overhead scheduling（ZOS）

开启 enable_zos 时，CPU 不必等待上一轮采样结果才开始安排下一轮。ModelRunner 会保留上一轮 GPU token buffer；如果下一轮 decode 所需 token 尚未回到 CPU，就通过 remap 在 GPU 上直接索引该 buffer。CPU 仍然依据预留计数做容量判断，GPU 则在同一 stream 上按依赖读取 token。

step N:     schedule(N) -> enqueue forward/sample -------------------->
step N+1:   schedule(N+1) -> remap previous sampled buffer -> enqueue ->
CPU:        process(N-1) -------------------------------> commit token

forward 内的 host 同步已经去除，需要读取 GPU 结果的地方都集中在 StepHandle.wait()，因此第 N 步的 CPU 工作可以与第 N-1 步的 GPU 工作重叠。

6. 采样与确定性

layers/sampler.py 使用 Gumbel-max 路径实现 temperature、top-k 和 top-p 采样。随机种子由序列 seed 与 token position 混合得到，而不是由 batch 行号决定。因此，同一输入、同一 seq_seed 在不同 batch 组成和抢占历史下仍应得到同一 token 序列。

这项确定性是抢占正确性的组成部分：抢占本身丢掉了缓存，需要重算前缀；如果采样依赖全局 RNG 状态，重算会改变后续 token，调度器就无法验证“释放后重放”等价。

7. 正确性证据

测试分三层：

层级命令/文件目的
CPU 资源tests/test_block_manager.py, tests/test_scheduler.py验证 free list、引用计数、chunked prefill、抢占和边界
采样tests/test_sampler.py验证 seed、温度、top-k/top-p 和 batch 独立性
V4.1 对照tests/test_v41.py对比官方 reference 的量化、indexer、稀疏 attention 和端到端 logits

运行完整检查：

cd mini-DeepInfer
pytest tests/test_block_manager.py tests/test_scheduler.py tests/test_sampler.py
pytest

GPU 和官方 reference 测试需要可用的 CUDA 环境；没有 GPU 时应运行 CPU 子集，并把未执行的检查标为未验证，而不是把它们当成通过。

8. 从哪里开始改

目标入口
改调度策略minideepinfer/engine/scheduler.py
改 block/slot 生命周期minideepinfer/engine/block_manager.py、sequence.py
改 ZOS batch 构造minideepinfer/worker/model_runner.py
改 V4.1 attention / MoEminideepinfer/models/deepseek_v41/model.py
改低精度状态minideepinfer/models/deepseek_v41/quant.py、weights.py
改可复现实验tests/、tools/v41_reference.py

推荐的修改顺序是先写一个能暴露状态或调度差异的最小测试，再改实现，最后用 reference logits 和抢占/批组成变体回归。性能数字必须注明 GPU、checkpoint、batch、上下文长度、dtype 和是否包含 warm-up。

9. 已知边界与后续工作

当前实现优先保证“读得懂且对得上”：

  • forward 主要使用 PyTorch / Triton；forward 内的 host 同步已经去除，结果统一在 StepHandle.wait() 回收，ZOS 因此可以真正重叠；
  • MoE 已使用 grouped GEMM（FP4 权重按需反量化为 bf16，不是原生 FP4 GEMM），decode 使用按 batch / 上下文分桶的 CUDA Graph；prefill、logits 与采样仍是 eager；
  • 真实 V4.1-Flash 的 prefill 在 4 × B200 上仍只有官方的 11.9%，剩余差距见性能报告 C1；
  • DSpark 投机解码、prefill 成本模型和端到端 batch-invariant 仍在路线图中；
  • 随机权重 checkpoint 只能验证执行语义，不能证明 DeepSeek-V4.1 的任务质量或生产吞吐。

相关路线图、许可和第三方声明见 README 与 THIRD_PARTY_NOTICES。