ssd：AI Agent 工具实践指南

工作中遇到相关需求时，ssd值得先读说明，因为它主要用于支持推测推测解码的轻量级推理引擎（SSD）。从日常自动化的使用方式看，输入边界、依赖和失败处理如果不清楚就很难稳定复用是采用前必须回答的问题。短测时我会用一项范围明确的真实任务完成最小试跑，并保留配置时间、输出质量、异常信息和维护痕迹的结果，方便团队复盘。它更像给愿意先做小范围验证并复查原始文档的团队准备的可审查方案，是否长期使用应由试跑数据决定。

Speculative 投机 Decoding

Paper

“在所有小说中，每当一个人遇到不同的选择时，他都会选择其中一个并排除其他选择；在几乎深不可测的彭翠的作品中，他同时选择了所有选择。”

——豪尔赫·路易斯·博尔赫斯，《小径分岔的花园》（1941）

SSD 是一种新的 LLM 推理算法。它非常准确，而且速度非常快。

SSD 是一种新型推测解码（SD）。在正常的 SD 中，一个小型且快速的模型会猜测较大且较慢的模型可能生成的接下来的几个令牌，然后大型模型在一次前向传递中验证它们：起草和验证在同一硬件上相继发生。

在 SSD 中，它们在不同的硬件上并行发生。小模型提前预测可能的验证结果，并立即推测所有结果。如果猜测正确，则可以立即返回推测，从而完全消除起草开销。

此自定义推理引擎支持：

• SSD 算法的参考实现
• 优化的 SD 和自回归基线
• Qwen3 + Llama3 型号系列
• 张量并行性
• PagedAttention、CUDAgraphs、torch编译、前缀缓存

设置

要求：Python 3.11+，CUDA >= 12.8。此代码是在 H100 上编写和测试的。

如果未安装uv：

curl -LsSf https://astral.sh/uv/install.sh | sh
# if `uv` is not found in this shell:
export PATH="$HOME/.local/bin:$PATH"

然后：

git clone https://github.com/tanishqkumar/ssd && cd ssd
uv sync                    # core SSD deps
# uv sync --extra scripts  # add deps used by scripts/
source .venv/bin/activate
python -c "from ssd import LLM; print('ok')"

通过环境变量设置路径。 SSD_HF_CACHE 应指向 HuggingFace hub 目录 - 这是包含 models--org--name/ 子目录的目录（e.g./data/huggingface/hub，而不是 /data/huggingface/）。 SSD_DATASET_DIR 应指向包含数据集子目录的目录（humaneval/、alpaca/ 等）。

export SSD_HF_CACHE=/path/to/huggingface/hub
export SSD_DATASET_DIR=/path/to/processed_datasets
export SSD_CUDA_ARCH=9.0   # 9.0=H100, 8.0=A100, 8.9=L40/4090

下载模型+数据集

如果您已经通过 huggingface-cli 或类似方式下载了模型，则可以直接跳到数据集 - 只需确保 SSD_HF_CACHE 指向正确的位置即可。下载脚本需要额外的 scripts：uv sync --extra scripts。

# models (uses SSD_HF_CACHE)
python scripts/download_from_hf.py llama

# datasets (writes to $HF_DATASETS_CACHE/processed_datasets)
export HF_DATASETS_CACHE=/path/to  # parent of SSD_DATASET_DIR
python scripts/get_data_from_hf.py --num-samples 10000

用途

下面的所有命令都从 bench/ 目录内部运行。大型型号（Llama-3 70B、Qwen-3 32B）在生成开始之前需要几分钟时间来生成 load/warmup/compile。始终使用 python -O 来禁用调试开销。

基准测试

使用 --all 对四个数据集进行全面评估。由于不同的数据分布的可预测程度不同，因此 SD/SSD 的速度很大程度上取决于数据集。对多种类型数据集的多种提示进行平均 给出一个整体情况。 --numseqs 是针对每个数据集的，因此 --numseqs 128 --all 总共运行 128 × 4 = 512 个提示。

cd bench

# AR — Llama 70B, 4 GPUs
python -O bench.py --llama --size 70 --gpus 4 --b 1 --temp 0 --numseqs 128 --output_len 512 --all

# Sync spec decode — 70B target + 1B draft, 4 GPUs, k=6
python -O bench.py --llama --size 70 --gpus 4 --spec --k 6 --b 1 --temp 0 --numseqs 128 --output_len 512 --all

# Async spec decode (SSD) — 70B target (4 GPUs) + 1B draft (1 GPU), k=7, f=3
python -O bench.py --llama --size 70 --gpus 5 --spec --async --k 7 --f 3 --b 1 --temp 0 --numseqs 128 --output_len 512 --all

对于 Qwen 型号，请使用 --qwen --size 32。有关完整参数，请参阅 bench/bench.py。对于 SGLang/vLLM 基线，请参阅 bench/README.md。

聊天

仅与 Llama-3.1 70B 进行交互式流媒体聊天。支持AR、同步SD和异步SD (SSD)。通过 --metrics 在每次响应后打印令牌计数、速度和 TTFT。

cd bench

# AR — 4 GPUs
python -O chat.py --ssd --gpus 4

# Sync spec decode — 4 GPUs, k=6
python -O chat.py --ssd --spec --k 6 --gpus 4

# Async spec decode (SSD) — 5 GPUs, k=7, f=3
python -O chat.py --ssd --spec --async --k 7 --f 3 --gpus 5 --metrics

还支持 SGLang 和 vLLM 聊天后端（自动启动其服务器）进行比较：

python -O chat.py --sglang        # spec decode
python -O chat.py --sglang --ar   # autoregressive
python -O chat.py --vllm          # spec decode

路线图

近期将支持的功能：

• 在最多 4 个设备上并行草稿数据（增加推测缓存大小）以避免计算限制
• OpenAI-compatible 对 HTTP 的推理
• 新型号和MoE支持：GPT-OSS和Kimi-K2.5。

欢迎投稿！

历史