DSpark 的实现和测评

发表于 分类于 Disqus：

DSpark 的实现和测评

DSpark = DFlash 的并行 backbone forward（1 次）+ N 步轻量 Markov 序列修正，全部在 CUDA Graph 内。 本文结合 vLLM 源码分析 DSpark 的实现细节，并在 Qwen3-8B 上实测 deepseek-ai 官方 draft 和社区 Dogacel draft 的效果差异。

1. 背景：投机解码与并行起草

投机解码（Speculative Decoding, SD）用一个小 draft 模型并行猜测 N 个 token，再由 target 模型一次 verify，通过 rejection sampling 保证输出分布不变。SD 的收益来自把 decode 阶段的 memory-bound 转为 compute-bound–bs=1 时 GPU 利用率极低，draft 的轻量 GEMM + target 的 batched verify 填充了 GPU 空闲。

vLLM v1 的 SD 框架支持多种 method：eagle、eagle3、dflash、dspark、medusa、ngram、mtp 等。DSpark 继承自 DFlash，核心改进是序列马尔可夫采样。

继承链：

1
2
3
4
BaseSpeculator (ABC)
  └─ DraftModelSpeculator
      └─ DFlashSpeculator
          └─ DSparkSpeculator   ← 本文主角

模型类继承链：

1
2
3
Qwen3ForCausalLM
  └─ DFlashQwen3ForCausalLM
      └─ Qwen3DSparkForCausalLM

2. DSpark vs DFlash：两个核心差异

DSpark 的 docstring 写得非常清楚，和 DFlash 的差异只有两点。

2.1 Anchor-as-first-prediction（锚位即首预测）

DFlash：每个 request 发 1 + N 个 query token（1 个 anchor/bonus + N 个 mask token）。anchor 是上一步验证通过的 token，只有 N 个 mask 位置做预测：

1
2
3
4
5
DFlash query layout (1+N=9, N=8):
  [anchor] [mask] [mask] [mask] [mask] [mask] [mask] [mask] [mask]
     ↑       ↑      ↑      ↑      ↑      ↑      ↑      ↑      ↑
   bonus    pred   pred   pred   pred   pred   pred   pred   pred
   (不采样)

DSpark：anchor 本身也是预测位置，每个 request 只发 N 个 query token：

1
2
3
4
5
DSpark query layout (N=8):
  [anchor] [noise] [noise] [noise] [noise] [noise] [noise] [noise]
     ↑       ↑       ↑       ↑       ↑       ↑       ↑       ↑
   pred     pred    pred    pred    pred    pred    pred    pred
   (采样)

代码（DSparkSpeculator.__init__）：

1
2
3
4
5
6
7
self.sample_from_anchor = getattr(
    self.draft_model_config.hf_config, "sample_from_anchor", True
)
if self.sample_from_anchor:
    self.num_query_per_req = self.num_speculative_steps      # N
else:
    self.num_query_per_req = 1 + self.num_speculative_steps  # 1+N (兼容旧格式)

在 Triton kernel _prepare_dflash_inputs_kernel 中，通过 SAMPLE_FROM_ANCHOR 编译常量控制采样行为：

1
2
3
4
# DSpark: 所有 N 个位置都采样，sample_pos = query_pos + 1（标准 next-token）
sample_off = 0 if SAMPLE_FROM_ANCHOR else 1
is_sample = is_query & (query_off >= sample_off)
sample_pos = query_pos + 1 if SAMPLE_FROM_ANCHOR else query_pos

2.2 Sequential Markov Sampling（序列马尔可夫采样）

这是 DSpark 的核心创新。

DFlash：N 个 mask 位置的 hidden states 一次性并行采样，各位置之间无依赖。

DSpark：先并行 forward 得到所有 N 个位置的 hidden states，然后从左到右逐个采样，每步用前一个采样出的 token 注入一个 Markov bias：

1
2
3
4
5
6
7
8
9
10
11
12
并行 backbone forward → [h₀, h₁, h₂, ..., h₇]
                           │     │     │         │
                           ▼     ▼     ▼         ▼
base_logits[0]  base_logits[1]  ...    base_logits[7]
     +                +                   +
markov_bias(     markov_bias(        markov_bias(
  anchor)          sample₀)            sample₆)
     │                │                   │
     ▼                ▼                   ▼
  sample₀         sample₁     ...      sample₇
     │
     └──────────────→ 传给下一步作为 prev

代码在 _sample_sequential：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def _sample_sequential(self, num_reqs, head_hidden):
    n_spec = self.num_speculative_steps
    # 1. 一次性算出所有 N 个位置的 base logits
    base_logits = self.model.compute_draft_logits(sample_hidden)  # [B, N, V]

    # 2. anchor token 作为初始 prev
    prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]]

    # 3. 逐位置采样
    for i in range(n_spec):
        markov_embed = self.model.markov_embed(prev)    # [B, r]
        bias = self.model.markov_bias(markov_embed)     # [B, V]
        logits_i = base_logits[:, i] + bias             # 加上 Markov 偏置
        draft_sampled_i = gumbel_sample(logits_i, ...)   # 采样
        self.draft_tokens[:num_reqs, i] = draft_sampled_i
        prev = draft_sampled_i                           # 传给下一步

一句话总结：并行 forward 拿到所有位置的 base prediction，再用 N 步轻量 Markov 修正注入序列依赖–把「N 个独立预测」变成「N 个有依赖的预测」。

3. Markov Head 结构

DSparkMarkovHead 是一个 low-rank 转移偏置头：

1
2
3
4
5
6
7
8
9
prev_token_id
    │
    │ markov_w1: Embedding(V, r)     ← V 是 vocab_size，r 是 markov_rank
    ▼
markov_embed [B, r]
    │
    │ markov_w2: ParallelLMHead(r, V)  ← r → V 的线性投影
    ▼
markov_bias [B, V]  ← 加到 base_logits 上

代码（qwen3_dspark.py）：

1
2
3
4
5
class DSparkMarkovHead(nn.Module):
    def __init__(self, vocab_size, draft_vocab_size, markov_rank, ...):
        self.markov_w1 = nn.Embedding(vocab_size, markov_rank)       # V×r
        self.markov_w2 = ParallelLMHead(
            draft_vocab_size, markov_rank, bias=False, disable_tp=True)  # r×V

两个权重都是 replicated（disable_tp=True），因为 Markov head 每步都跑，分片会引入 all-reduce 和 full-vocab gather。

参数量 = 2×V×r2 \times V \times r。当 V=151936V=151936（Qwen3 词表）、r=64r=64 时约 19.4M 参数，相比 8B backbone 可以忽略。

4. 完整的 Draft 一步流程

DSparkSpeculator._generate_draft 只有两行：

1
2
3
def _generate_draft(self, num_reqs, num_tokens_padded, ...):
    head_hidden = self._run_model(...)                    # 1. 并行 backbone forward
    self._sample_sequential(num_reqs, head_hidden)        # 2. 序列 Markov 采样

Step 1：并行 Backbone Forward（继承自 DFlash）

  • 输入：N 个 query token（anchor + mask/noise），position 已对齐
  • 上下文 KV 已在 precompute_and_store_context_kv 中预填充
  • 非因果 attention：N 个 query 位置可以互相 attend
  • 整个 forward 被 CUDA Graph 捕获

Step 2：Sequential Markov Sampling（DSpark 独有）

  • 取出 N 个位置的 hidden states
  • 一次性算出 base logits（compute_draft_logits）
  • 逐位置：base_logits[i] + markov_bias(prev) -> gumbel_sample
  • 这个循环也被 CUDA Graph 捕获（所有 buffer 预分配固定地址）

Context KV 预计算（DFlash 的关键优化）

避免逐层跑 target 的 forward 来填充 draft KV cache，而是用 target 的中间层 hidden states 一次性投影：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
target aux hidden states [num_ctx, H_target]
    │
    │ fc 层投影到 draft hidden size
    ▼
context_states [num_ctx, H_draft]
    │
    │ ① Fused GEMM（所有层的 KV projection 合成一个矩阵乘法）
    ▼
all_kv_flat [num_ctx, L×2×kv_size]
    │
    │ ② Grouped RMSNorm（所有层的 K-norm 一次算完）
    ▼
all_k_normed [L, num_ctx, nkv, hd]
    │
    │ ③ Fused RoPE（所有层一次应用）
    ▼
all_k_final [L, num_ctx, nkv, hd]  → per-layer 写入 KV cache

代码核心（DFlashQwen3Model.precompute_and_store_context_kv）：

1
2
3
4
5
# 融合所有层的 KV 权重做一次大 GEMM
all_kv_flat = F.linear(normed_context_states, self._fused_kv_weight, self._fused_kv_bias)
# 分离 K/V，per-layer 写入 cache
all_kv = all_kv_flat.view(num_ctx, L, 2, nkv, hd).permute(2, 1, 0, 3, 4).contiguous()
all_k, all_v = all_kv[0], all_kv[1]

5. Probabilistic Rejection Sampling 与 Reduced Vocab

DSpark 支持 draft_sample_method="probabilistic"（Gumbel-based rejection sampling）。Draft 采样时把 logits 通过 Gumbel max trick 得到 draft_logits，Target verify 时用相同 Gumbel seed 验证，保证输出分布不变。

支持 reduced draft vocab：draft 在小词表上算 logits，然后 scatter 到 target vocab 位置：

1
2
3
4
if self._d2t_scatter_index is not None:
    buf = self._draft_scatter_buf[:num_reqs]  # [-inf, -inf, ...]
    buf.index_copy_(1, self._d2t_scatter_index, logits_i)  # 只填 draft vocab 列
    logits_i = buf  # 变成 target vocab 大小

6. CUDA Graph 覆盖

DFlash/DSpark 的 CUDA Graph 是 FULL mode，覆盖整个 draft step：

1
2
3
# DFlashSpeculator.init_cudagraph_manager
if wants_full and supports_full:
    cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY

为了让 Markov 循环能被 CG 捕获，所有 buffer 都是预分配的固定地址：

Buffer 用途 CG 兼容性
draft_tokens 输出 token ✅ 固定地址
draft_logits probabilistic 模式的 processed logits ✅ 固定地址
_draft_scatter_buf reduced vocab scatter buffer ✅ 固定地址
_anchor_idx 每个 request 的 anchor 位置索引 ✅ 固定地址
input_buffers.input_ids anchor token 读取 ✅ 固定地址

7. 模型加载与权重共享

load_dspark_model（dspark/utils.py）做了几件事：

  1. 创建 draft config，设置非因果注意力
  2. 加载 draft 模型
  3. Embed tokens 共享：如果 draft 没有自己的 embedding，用 target 的
  4. LM head 共享：同理
1
2
3
if _should_share(draft_model, "has_own_embed_tokens", draft_embed, target_embed):
    del draft_inner.embed_tokens
    draft_inner.embed_tokens = target_embed

权重加载（Qwen3DSparkForCausalLM.load_weights）：

  • 跳过 t2d（训练用映射，推理不需要）
  • d2t -> draft_id_to_target_id（推理用的 draft→target 映射）
  • 跳过 mask_embedding（DSpark 通过 vocab row 做 mask，不用单独参数）和 confidence_head（未接入推理）
  • 调用 _build_fused_kv_buffers() 构建 fused KV 权重

8. 实验环境

项目 配置
Target Model Qwen/Qwen3-8B
推理引擎 vLLM v0.26.0
conda 环境 dspark-vllm
nsys 版本 2026.1.3（vLLM traces）/ 2025.3.0（DeepSpec trace）
采集参数 -t cuda,nvtx,osrt,cudnn,cublas --python-backtrace=cuda --cudabacktrace=all
Benchmark SPEED-Bench（qualitative split, coding category）

Draft Model 配置

配置名 Draft Model 架构 来源
Baseline 无（纯 Qwen3-8B） Qwen3 -
DSpark（deepseek-ai） deepseek-ai/dspark_qwen3_8b_block7 Qwen3DSparkSt deepseek-ai 官方
DSpark（Dogacel） Dogacel/Qwen3-8B-DSpark EAGLE3 社区训练

Dogacel 的 vLLM 启动参数：speculative 开启、acceptance: 0.85、num_spec_tokens: 7、max_model_len: 2048。

9. 性能对比

9.1 端到端性能（3 prompts, 各 64 tokens）

配置 耗时 vs Baseline 加速比
Baseline 0.75s - 1.00x
DSpark（deepseek-ai） 0.50s -33% 1.50x
DSpark（Dogacel） 0.77s +3% 0.97x

9.2 投机解码指标（DeepSpec evaluator trace）

指标 DSpark（deepseek-ai）
verify_steps 12
mean_accept_len 7.1
推测 每次提议 ~7 tokens，几乎全部被接受

Dogacel 的 trace 中未发现 dspark_propose / target_verify 的 NVTX range，推测 acceptance rate 极低。

mean_accept_len=7.1 意味着 N=8 时几乎全部接受–backbone 的并行预测质量极高，Markov head 的序列修正有效。verify_steps=12 表示 12 步验证共接受约 85 个 token（12×7.112 \times 7.1）。

10. Trace 分析

10.1 Trace 文件清单

文件 大小 来源 CUDA Kernel 数据
baseline_trace.nsys-rep 1.5 MB vLLM profile_baseline.py ❌ 无
dogacel_trace.nsys-rep 1.8 MB vLLM profile_dogacel.py ❌ 无
trace.nsys-rep（DeepSpec） 3.5 MB DeepSpec evaluator ✅ 有

10.2 CUDA Kernel 缺失原因

vLLM 的 EngineCore 在子进程中运行，nsys 默认只 trace 主进程。三个 vLLM trace 均无 GPU kernel 数据。

解决方案：重新采集时添加 --trace-fork 参数：

1
2
3
4
5
6
nsys profile -t cuda,nvtx,osrt,cudnn,cublas \
  --python-backtrace=cuda --cudabacktrace=all \
  --trace-fork \
  --force-overwrite=true \
  -o baseline_trace_v2 \
  bash -c '...'

10.3 DeepSpec Evaluator Kernel 分布

Kernel 耗时占比 Instances 说明
CUTLASS GEMM (16×16) 66.3% 7,177 主要 matmul（Q/K/V/O + MLP）
elementwise_kernel 3.9% 8,522 RoPE、残差等
reduce_kernel (mean) 2.8% 4,298 RMSNorm
CUTLASS GEMM (32×32) 2.7% 382 大块矩阵乘法
Flash Attention 1.6% 864 Attention 计算
Softmax forward 1.0% 168 Softmax

关键观察：GEMM 占 66.3%，但 bs=1 decode 时本质是 memory-bound（M=1 瘦矩阵乘）。kernel launch 开销显著（~20000 次 launch）。vLLM 的 CUDA Graph 会消除大部分 launch 开销，fused kernel 会压缩 elementwise/reduce 占比。预期 vLLM 路径下 GEMM 占比升至 80%+。

10.4 NVTX Range 对比

NVTX Range Baseline Dogacel deepseek-ai dspark
dspark_propose ❌ ❌ ✅
target_verify ❌ ❌ ✅
decode_sample ❌ ❌ ✅
warmup ❌ ❌ ✅
VLLM::EngineCore ✅ ✅ ❌

Dogacel 缺少 dspark_propose/target_verify 说明其 draft forward 未走标准 dspark 代码路径。

11. Dogacel 无效原因：架构不匹配

维度 deepseek-ai（有效） Dogacel（无效）
Draft 架构 Qwen3DSparkSt EAGLE3
与 vLLM dspark 实现兼容 ✅ 完全对齐 ❌ 不匹配
NVTX range 存在 ✅ propose + verify ❌ 无
Mean accept len 7.1 推测极低
端到端加速 1.50x 0.97x（负优化）

根因：Dogacel 用 EAGLE3 架构训练 draft，中间层 hidden state 接口与 dspark 实现不兼容。即使 draft 能加载运行，acceptance rate 极低，draft 开销 > SD 收益。即使模型本身学得不差，接口不对也白搭。

12. SPEED-Bench 数据集

12.1 整体结构

SPEED-Bench（SPEculative Evaluation Dataset）是 NVIDIA 出的投机解码评测基准。

Split 样本数 用途
qualitative 880（11 类×80） 测 SD 质量（acceptance rate）
throughput_1k/2k/8k/16k/32k 1536×5 测系统吞吐（高并发）

12.2 Qualitative Split（质量评测）

从 18 个公开数据源聚合，分成 11 个 category：Coding、Math、Humanities、STEM、Writing、Summarization、Roleplay、RAG、Multilingual、Reasoning、QA。每类 80 个样本，用 OpenAI text-embedding-3-small 做嵌入，greedy 选择 + swap 优化最大化语义多样性（平均 pairwise cosine similarity 从 SpecBench 的 0.22 降到 0.14）。

12.3 Throughput Split（吞吐评测）

固定输入长度桶（1K/2K/8K/16K/32K），每桶 1536 条（512×3），分 3 个难度类别：low_entropy（coding 类）、high_entropy（creative writing 类）、mixed_entropy。用 tiktoken 精确 pad/truncate，不用 random token（会扭曲 MoE routing 和 acceptance behavior）。

12.4 为什么选 coding 类做 benchmark

  1. Coding 是低熵任务–token 可预测性高，SD 的 acceptance rate 天然高，是 best-case 场景
  2. 语义多样性好–80 条 prompt 覆盖 Python(27)、C++(9)、Java(10)、Go(13)、JS(11)、Rust(3) 等，来自 LiveCodeBench、Code Contests、HumanEvalPack
  3. 固定输出长度（--speed-bench-output-len 2048）–隔离 prefill 影响，纯测 decode
  4. 两种并发对比：--max-concurrency 32（batched，模拟生产环境）vs --max-concurrency 1（单流，测纯 decode 延迟）
  5. --disable-shuffle 保证可复现，--temperature 1.0 高温采样更反映真实使用场景

13. 接受率与训练效果的关系

Acceptance rate 的天花板由 draft 训练质量决定，工程实现决定能打到多少天花板。

训练侧决定上限

  • Draft 的 hidden state 和 target 的中间层对齐越好，token 分布越接近，accept 越高
  • deepseek-ai 的 block7 专门按 dspark 接口训练，hidden state 严格对齐 Qwen3-8B 第 7 层，所以 mean_accept_len=7.1
  • Dogacel 用 EAGLE3 方式训练，hidden state 映射方式不同，接口不对

工程侧决定下限

  • vLLM dspark 的 propose → verify pipeline 是否正确对接 draft
  • KV cache 的 layout、position ID 对齐、temperature sampling 一致性
  • CUDA Graph 是否覆盖 draft forward（没覆盖的话 launch overhead 会吃掉 SD 收益）
维度 deepseek-ai（训练+工程都对） Dogacel（工程接口不对）
Draft 架构 Qwen3DSparkSt EAGLE3
Hidden state 接口 ✅ 正确对接 ❌ 不匹配
NVTX range ✅ propose + verify ❌ 无
Mean accept len 7.1 推测极低
端到端加速 1.50x 0.97x

一句话总结：训练决定 draft 能不能猜对，工程决定猜对的部分能不能高效用上。Dogacel 的情况是工程接口就不对，猜得再准也走不进去。

14. vLLM 推理引擎优化对 Kernel 分布的影响

无 vLLM 优化的 kernel 分布（DeepSpec evaluator）

Kernel 占比 说明
CUTLASS GEMM (16×16) 66.3% bs=1 时是 memory-bound
elementwise 3.9% RoPE、残差等，未融合
reduce (mean) 2.8% RMSNorm，未融合
Flash Attention 1.6% decode 时计算量小
Softmax 1.0% 未融合
总 kernel launch ~20000 次 launch 开销显著

vLLM 优化后的预期变化

  1. CUDA Graph：20000 次 kernel launch → 1 次 graph launch
  2. Fused kernel：RMSNorm + residual + RoPE 融合为 1 个 kernel
  3. FlashInfer/FlashAttention decode-optimized：attention kernel 更高效
  4. GEMM 占比升至 80%+：其他开销被压缩后，GEMM 成为绝对瓶颈

对投机解码的启示

Baseline 的 decode 在 vLLM 下 GEMM 占 80%+，本质是 memory-bound（M=1 瘦矩阵乘，GPU 利用率低）。SD 的价值在于用 draft 的轻量 GEMM + target 的 batched verify 填充 GPU 空闲。当 batch size 增大（高并发），decode 从 memory-bound 转向 compute-bound，SD 收益下降–这也是 SPEED-Bench throughput split 存在的意义。

15. 总结

  1. DSpark = DFlash 并行 backbone forward + N 步 Markov 序列修正，全部在 CUDA Graph 内，用极小的开销把并行预测的「无依赖」缺陷补上
  2. 实测 deepseek-ai 官方 draft 在 Qwen3-8B 上实现 1.50x 加速，mean_accept_len=7.1（N=8 几乎全接受）
  3. Dogacel 社区 draft 因架构不匹配（EAGLE3 vs DSpark）完全无效，0.97x 负优化
  4. 接受率天花板由训练决定，工程决定下限–hidden state 接口对齐是前提
  5. SPEED-Bench coding 类是 SD 的 best-case 场景，低熵任务下 acceptance rate 天然高

下一步

  • 重新采集 vLLM trace（加 --trace-fork），获取真实 kernel 数据
  • 检查 vLLM v0.26.0 是否支持 EAGLE3 method（Dogacel 应走 EAGLE3 而非 dspark）
  • 用 throughput split + --max-concurrency 32 测高并发下 SD 效果
  • 尝试 ngram、medusa 等 baseline 对比

参考：

  • vLLM 源码：vllm/v1/worker/gpu/spec_decode/dspark/speculator.py、vllm/model_executor/models/qwen3_dspark.py
  • vLLM PR：#50138、#50694、#50737
  • 模型：deepseek-ai/dspark_qwen3_8b_block7、Dogacel/Qwen3-8B-DSpark
  • 数据集：nvidia/SPEED-Bench，arXiv: 2604.09557