Chenxi Ma

·7 min read分钟阅读

Runtime design运行时设计

A DeepInfer-Style Runtime DesignDeepInfer 风格的运行时设计

The engineering boundaries, state invariants and verification rules mini-DeepInfer commits to, written before the implementation.mini-DeepInfer 要坚持的工程边界、状态不变量和验证方法，先于代码实现写成。

DeepInferRuntimeSchedulingDesign
Contents目录 09 sections节
  1. 1. Define the runtime’s promises first
  2. 2. Mapping the official design onto this project
  3. 3. The data contract of a step
  4. 4. V4.1 state layout
  5. 5. Scheduling policy
  6. 6. Sampling and replayability
  7. 7. Implementation order, driven by evidence
  8. 8. Verification gate
  9. 9. What is explicitly not promised
  1. 1. 先定义运行时的承诺
  2. 2. 官方设计与本项目的映射
  3. 3. 一个 step 的数据契约
  4. 4. V4.1 状态布局
  5. 5. 调度策略
  6. 6. 采样与可重放性
  7. 7. 按证据推进的实现顺序
  8. 8. 验证门槛
  9. 9. 当前明确不承诺的内容

This document was written before the code. It records the engineering boundaries, state invariants and verification methods mini-DeepInfer holds to. The goal is not to copy DeepInfer’s production code, but to keep its most instructive technical taste: separate scheduling, memory, device execution and correctness evidence, so every performance optimization maps back to a state transition that can be inspected.

DeepInfer’s public description lists optimized kernels, parallelism, DSpark, asynchronous scheduling, batch invariance, deterministic sampling and structured output as core capabilities, and its scheduler source treats SchedulerOutputs, prefill/decode, the block manager and the asynchronous step as separate boundaries:

1. Define the runtime’s promises first

mini-DeepInfer is a readable, PyTorch-first implementation that supports a single GPU and intra-node EP × DP. It must satisfy four observable promises:

  1. Every step has an explicit input and output. The scheduler output describes the token counts, sequence order, block tables and rows to sample for this iteration; device execution never changes CPU queues behind its back.
  2. Resource lifetimes can be proven. KV blocks and V4.1 state slots each have exactly one owner; preemption, completion and cancellation release both.
  3. Recomputation is equivalent. Sampling randomness depends only on the sequence seed and the token position, never on the batch row or the history of a global RNG.
  4. Performance paths come with regression evidence. Use the smallest test that exposes the state difference first, then change the implementation; a performance number must record hardware, checkpoint, batch, context, dtype and warm-up.

These four matter more than porting every production feature. EP, CP, custom kernels and multi-node communication can be swapped in later, but they must not break these boundaries.

2. Mapping the official design onto this project

Official DeepInfer conceptmini-DeepInfer implementationCurrent boundary
C++ Scheduler / SchedulerOutputsScheduler / SchedulerOutput in engine/scheduler.pySingle process, no DP-group voting
BlockManager free list and reference countsengine/block_manager.pySingle GPU, no cross-rank sharing
Iteration-level prefill/decodeScheduler._schedule_prefill() / _schedule_decode()The two q lengths are never mixed
Zero-overhead schedulingStepHandle and remap in worker/model_runner.pyThe forward pass has no host synchronization; results are collected in StepHandle.wait()
Grouped expert GEMMGroupedExperts in models/deepseek_v41/model.pyFP4 is dequantized to bf16 on demand, then three torch._grouped_mm calls; not a native FP4 GEMM
Decode CUDA GraphDecodeGraphRunner in worker/cuda_graph.pyBucketed by batch/context; padding uses reserved slot/block 0; prefill, logits and sampling stay eager
Deterministic samplingmix_seed() + Gumbel-max in layers/sampler.pyCovers token sampling only
V4.1 attention stateBlock + state slot in models/deepseek_v41/model.pyEach rank manages its own paged cache
Structured output / DSparkReserved in SamplingParams and the execution boundaryGrammar compilation and draft verification are not implemented

The mapping table is a maintenance constraint. A new feature must first say which boundary it belongs to: scheduling policy does not belong in the model forward, and the model must not touch the waiting/running queues directly.

3. The data contract of a step

One engine iteration runs in this order:

schedule(N)
  -> SchedulerOutput(prompt/decode, seqs, seqlens_q, sample_seqs, reservations)
execute(N)
  -> ModelRunner builds the batch on the GPU stream, runs the model, samples
process(N-1)
  -> wait for the previous StepHandle, commit tokens, handle stop/length/abort, release state

With ZOS enabled, process(N-1) can overlap execute(N). “Overlap” here describes a dependency relationship; it does not mean every Python operator is already synchronization-free. Anything that must read a GPU result has to be concentrated in StepHandle.wait().

3.1 The scheduler output must be a complete snapshot

SchedulerOutput must at least answer:

  • Is this round prefill or decode?
  • How many query tokens does each sequence in the batch take?
  • Which sequences are sampled, and which epoch does a sampled result belong to?
  • Which sequences are ignored, finished, preempted or cancelled?

ModelRunner only consumes this snapshot and never re-derives scheduling policy. That lets scheduling be tested on CPU alone, and lets a fake executor replay the entire lifecycle.

3.2 CPU view and GPU view

Under ZOS a sequence has two lengths:

CPU-known tokens:        len(token_ids)
GPU-reserved tokens:     get_len() = len(token_ids) + num_inflight_tokens
Written/reserved KV:     num_cached_tokens

Decode is legal only when num_cached_tokens == get_len() - 1. The scheduler adds a reservation when it schedules and appends the real token only when the result returns; a stale result whose epoch is no longer valid must be discarded and can never be written back to the sequence.

4. V4.1 state layout

V4.1’s cache is not one traditional KV table; it is two kinds of state with different lifetimes:

  1. Paged compressed state: compressed KV and indexer keys are written per compressed entry into 128-token blocks. A block table addresses them in batches.
  2. Per-sequence state: the sliding-window ring and the compressor tail each take one state slot per sequence. It must move, be preempted and be released together with the sequence.

So a sequence’s resources must be released in pairs:

block_manager.free(seq.block_table)
state_manager.free([seq.state_slot])

Releasing only the blocks makes the scheduler see free KV while the next request leaks a state slot.

5. Scheduling policy

Prefill

  • The waiting queue stays FCFS.
  • If the prompt at the head cannot fit, stop trying instead of skipping it, which would cause unpredictable starvation.
  • max_num_batched_tokens limits only the q tokens of this round; chunked prefill lets one prompt span several steps.
  • Only the last chunk of a prompt produces a sampling reservation.
  • Before it first enters running, the blocks for the whole prompt and one state slot are allocated together.

Decode

  • Each running sequence advances one token at a time.
  • Candidate order is decided by output length first.
  • When blocks run out, the sequence with the longest output is released and put back at the head of the waiting queue.
  • Preemption increments the epoch, so stale GPU results cannot contaminate the recomputation.
  • After EOS, max_tokens or max_model_len, resources are released at the point where the result is committed.

Cancellation

Cancellation is an explicit terminal state. It cannot masquerade as EOS and it is not just deleting an object from a queue. The safe order is:

  1. Let any StepHandle still on the GPU finish.
  2. Let the scheduler invalidate the request’s epoch.
  3. Remove it from waiting/running.
  4. Release both the blocks and the state slot.
  5. Return the abort finish reason.

This way a cancellation cannot reuse a block while the GPU is still reading the old one.

6. Sampling and replayability

The official documentation treats deterministic sampling and batch invariance as separate capabilities. mini-DeepInfer implements a teaching-grade equivalent with these constraints:

sample_seed = mix_seed(sequence_seed, output_position)
token = GumbelMax(logits / temperature, sample_seed)

Sampling must not use the batch row as the random seed, and top-k/top-p selection must not depend on other rows. Verification must cover:

  • A single-row batch equals a multi-row batch.
  • Forward order equals reversed order.
  • Token sequences are identical after preemption and recomputation.
  • Cancelling one request does not change other requests’ results.

7. Implementation order, driven by evidence

Later code advances in this order, each step as an independently reviewable commit:

  1. A complete request lifecycle: request id, cancellation, finish reason, resource snapshots and CPU regression tests.
  2. Observable scheduling cost: record prefill/decode tokens, preemption count and block/slot watermarks, keeping metrics separate from scheduler output.
  3. Sampling boundary extension: a stop-token set and an optional logits mask, staying batch-invariant.
  4. Device execution optimization (implemented): remove host-visible branches from forward and use grouped GEMM and decode CUDA Graph; next, evaluate weight dequantization, logits and sampling overhead.
  5. Advanced paths: DSpark draft/verify, SWA bounded replay, structured output and end-to-end verification on the real checkpoint.

Steps 4 and 5 may not be declared “done” without a benchmark and a reference comparison. If an optimization changes the cache layout or the scheduler contract, update this document’s mapping table and invariants first.

The tiny-checkpoint verification of step 4 is in PRs #13, #14, #15 and #16: regression tests for empty experts, EP sharding, compression-group boundaries and graph padding, and a comparison against the official reference logits (LOGIT_ATOL=0.05). Tiny offline output throughput on 1 × B200 rose from 478 to 2,643 tok/s (median of 3 runs); that does not mean the real-checkpoint multi-GPU prefill numbers were updated. The workload, TTFT, TPOT, goodput and capture overhead are as recorded in each PR’s experiment notes.

8. Verification gate

Every change must at least pass:

pytest tests/test_block_manager.py tests/test_scheduler.py
pytest
git diff --check

When CUDA and the official reference are available, also run tests/test_sampler.py, tests/test_v41.py and the end-to-end preemption matrix. Without the matching hardware, report skipped/unverified; never write it up as passed.

9. What is explicitly not promised

  • A single-GPU PyTorch implementation is not official multi-GPU throughput.
  • A random-weight checkpoint verifies execution semantics, not model capability.
  • Without a grammar compiler, “structured output” cannot be written as supported.
  • Without a draft model and acceptance traces, DSpark cannot be written as supported.
  • Every benchmark must state hardware, model, batch, length, dtype and warm-up.

This document describes the skeleton of the implementation. The code can get faster step by step, but these boundaries must stay readable, testable and replayable.

这份文档先于代码实现，记录 mini-DeepInfer 要坚持的工程边界、状态不变量和验证方法。目标不是复制 DeepInfer 的生产代码，而是保留它最值得学习的技术品味：把调度、内存、设备执行和正确性证据拆开，让每一个性能优化都能落回一个可检查的状态转移。

官方 DeepInfer 的公开说明将优化 kernel、并行、DSpark、异步调度、batch-invariant、确定性采样和结构化输出列为核心能力；官方调度器源码则把 SchedulerOutputs、prefill/decode、block manager 和异步 step 作为独立边界：

1. 先定义运行时的承诺

mini-DeepInfer 当前是以 PyTorch 为主的可读实现，支持单卡和机内 EP × DP。它需要满足四个可观察的承诺：

  1. 每一个 step 都有明确输入和输出。 调度输出描述本轮 token 数、序列顺序、block table 和需要采样的行；设备执行不偷偷改变 CPU 队列。
  2. 资源生命周期可以被证明。 KV block 和 V4.1 state slot 都有唯一的所有者；抢占、完成、取消都会释放两类资源。
  3. 重算是等价的。 采样随机性只由序列 seed 和 token position 决定，不能依赖 batch 行号或全局 RNG 的历史。
  4. 性能路径有可回归的正确性证据。 先用最小测试暴露状态差异，再改实现；性能数字必须同时记录硬件、checkpoint、batch、上下文、dtype 和 warm-up。

这四条比“把所有生产特性搬过来”更重要。EP、CP、定制 kernel 和多机通信都可以以后替换进来，但不能破坏这些边界。

2. 官方设计与本项目的映射

官方 DeepInfer 概念mini-DeepInfer 的实现当前边界
C++ Scheduler / SchedulerOutputsengine/scheduler.py 的 Scheduler / SchedulerOutput单进程、无 DP 组投票
BlockManager free list 与引用计数engine/block_manager.py单 GPU，无跨 rank 共享
iteration-level prefill/decodeScheduler._schedule_prefill() / _schedule_decode()不混合两种 q length
zero-overhead schedulingworker/model_runner.py 的 StepHandle 与 remapforward 已去除 host 同步；结果在 StepHandle.wait() 回收
grouped expert GEMMmodels/deepseek_v41/model.py 的 GroupedExpertsFP4 按需反量化为 bf16，三次 torch._grouped_mm，非原生 FP4 GEMM
decode CUDA Graphworker/cuda_graph.py 的 DecodeGraphRunner按 batch/context 分桶；padding 使用保留 slot/block 0；prefill、logits 与采样仍 eager
deterministic samplinglayers/sampler.py 的 mix_seed() + Gumbel-max只覆盖 token 采样
V4.1 attention statemodels/deepseek_v41/model.py 的 block + state slot每个 rank 独立管理分页缓存
structured output / DSpark预留在 SamplingParams 和执行边界尚未实现 grammar 编译或 draft verify

映射表是维护约束。新增功能必须先说明属于哪个边界，不能把调度策略塞进模型 forward，也不能让模型直接操作 waiting/running 队列。

3. 一个 step 的数据契约

一次引擎迭代按下面的顺序进行：

schedule(N)
  -> SchedulerOutput(prompt/decode, seqs, seqlens_q, sample_seqs, reservations)
execute(N)
  -> ModelRunner 在 GPU stream 上构造 batch、运行 model、采样
process(N-1)
  -> 等待上一个 StepHandle，提交 token，处理 stop/length/abort，释放状态

开启 ZOS 时，process(N-1) 可以和 execute(N) 重叠。这里的“重叠”只描述依赖关系，不代表 Python 的每个算子都已经无同步；任何需要读取 GPU 结果的地方都必须集中在 StepHandle.wait()。

3.1 调度输出必须是完整快照

SchedulerOutput 至少要能回答：

  • 本轮是 prefill 还是 decode？
  • batch 中每个序列占多少 query token？
  • 哪些序列要采样，采样结果对应哪一个 epoch？
  • 哪些序列被忽略、完成、抢占或取消？

ModelRunner 只消费这份快照，不重新推断调度策略。这样可以在 CPU 上单独测试调度，也可以用假执行器重放整个生命周期。

3.2 CPU 视图与 GPU 视图

ZOS 下同一条序列有两个长度：

CPU 已知 token:       len(token_ids)
GPU 已预留 token:     get_len() = len(token_ids) + num_inflight_tokens
已写入/预留 KV:       num_cached_tokens

decode 的合法条件是 num_cached_tokens == get_len() - 1。调度时先增加 reservation，结果返回时才追加真实 token；旧结果携带的 epoch 如果已经失效，就只能丢弃，不能重新写入序列。

4. V4.1 状态布局

V4.1 的缓存不是一张传统 KV 表，而是两种不同生命周期的状态：

  1. 分页压缩状态： 压缩 KV 和 indexer key 按压缩后的 entry 写进 128-token block。它适合用 block table 做批量寻址。
  2. 序列状态： sliding-window ring 和 compressor tail 每条序列各占一个 state slot。它必须跟着序列一起迁移、抢占和释放。

因此，一个序列的资源释放必须成对发生：

block_manager.free(seq.block_table)
state_manager.free([seq.state_slot])

不能只释放 block；否则 scheduler 看起来有空闲 KV，但下一条请求会在 state slot 上泄漏。

5. 调度策略

Prefill

  • waiting 队列保持 FCFS；
  • 队头 prompt 不可容纳时停止继续尝试，避免越过队头造成不可预测的饥饿；
  • max_num_batched_tokens 只限制本轮 q token，chunked prefill 允许一个 prompt 跨多个 step；
  • prompt 最后一块才产生采样 reservation；
  • 首次进入 running 前，同时分配完整 prompt 所需的 block 和一个 state slot。

Decode

  • 每条 running 序列一次前进一个 token；
  • 先按输出长度决定候选顺序；
  • block 不足时释放最长输出的序列并把它放回 waiting 队列头；
  • 抢占会递增 epoch，后续旧的 GPU 结果不会污染重算；
  • 到达 EOS、max_tokens 或 max_model_len 后，在结果提交点释放资源。

取消请求

取消是一个显式的终态。它不能伪装成 EOS，也不能只从队列中删除对象。安全顺序是：

  1. 让仍在 GPU 上的 StepHandle 完成；
  2. 让 scheduler 使请求的 epoch 失效；
  3. 从 waiting/running 中移除；
  4. 同时释放 block 和 state slot；
  5. 返回 abort finish reason。

这样取消不会在 GPU 仍读取旧 block 时提前复用同一块存储。

6. 采样与可重放性

官方文档把 deterministic sampling 和 batch-invariant 作为独立能力。mini-DeepInfer 使用以下约束实现教学版等价：

sample_seed = mix_seed(sequence_seed, output_position)
token = GumbelMax(logits / temperature, sample_seed)

采样不能使用 batch row 作为随机种子，也不能让 top-k/top-p 的筛选依赖其他行。验证必须覆盖：

  • 单行与多行 batch 相同；
  • 正序与逆序 batch 相同；
  • 抢占重算后的 token 序列相同；
  • 取消一个请求不改变其他请求的结果。

7. 按证据推进的实现顺序

后续代码按下面的顺序推进，每一步都要有一个可以独立审阅的提交：

  1. 请求生命周期完整化： request id、取消、finish reason、资源快照和 CPU 回归测试。
  2. 调度成本可观察： 记录 prefill/decode token、抢占次数、block/slot 水位，并将指标与调度输出分离。
  3. 采样边界扩展： stop token 集合和可选的 logits mask，保持 batch-invariant。
  4. 设备执行优化（已实现）： 去除 forward 中的 host-visible 分支，使用 grouped GEMM 和 decode CUDA Graph；下一步评估权重反量化、logits 与采样开销。
  5. 高级路径： DSpark draft/verify、SWA bounded replay、结构化输出和真实 checkpoint 的端到端验证。

第 4、5 步都不能在没有基准和 reference 对照的情况下直接宣称“完成”。如果一个优化改变了 cache layout 或 scheduler contract，必须先更新本文件的映射表和不变量。

第 4 步的 tiny checkpoint 验证见 PR #13、 #14、#15、 #16：包含空专家、EP 分片、压缩组边界和 graph padding 回归测试， 以及官方 reference logits 对照（LOGIT_ATOL=0.05）。其中 1×B200 的 tiny 离线输出吞吐从 478 提升至 2643 tok/s（3 次取中位数）；这不代表真实 checkpoint 的多卡 prefill 数据已更新。具体 workload、 TTFT、TPOT、goodput 和 capture 开销以各 PR 的实验记录为准。

8. 验证门槛

每次修改至少通过：

pytest tests/test_block_manager.py tests/test_scheduler.py
pytest
git diff --check

有 CUDA 和官方 reference 时，再运行 tests/test_sampler.py、tests/test_v41.py 和端到端抢占矩阵。没有对应硬件时，报告为 skipped/未验证，不把它写成通过。

9. 当前明确不承诺的内容

  • 单卡 PyTorch 实现不等于官方多卡吞吐；
  • 随机权重 checkpoint 只能验证执行语义，不能证明模型能力；
  • 没有 grammar compiler 时，structured output 不能写成已支持；
  • 没有 draft model 和 acceptance trace 时，DSpark 不能写成已支持；
  • 任何 benchmark 都必须注明硬件、模型、batch、长度、dtype 和 warm-up。

这份文档描述的是实现的“骨架”。代码可以逐步变快，但这些边界必须保持可读、可测、可重放。

mini-DeepInfer · Next in this project接着读mini-DeepInfer Technical Notes: From a Request to DeepSeek-V4.1mini-DeepInfer 技术说明：从请求到 DeepSeek-V4.1 →