目录
在 AMD Instinct GPU上实现 SSD
AMD中国
2026年06月12日 18:40
收录于文集
共115篇

原文作者：Yixing Xu, Subhajit Dutta Chowdhury, Xuanwu Yin, Dong Li, Vikram Appia, Zhenyu Gu, Emad Barsoum

图片

Speculative Speculative Decoding（SSD）[1] 是近期提出的一种推测解码（Speculative Decoding, SD）算法，能在传统 SD 基础上进一步加速大语言模型的推理。标准 SD 的做法是：一个轻量级草稿模型先推测出多个后续 token，再由大目标模型并行验证。SD 虽然已经降低了纯自回归解码的开销，但仍存在一个顺序依赖：目标模型完成验证之后，下一轮草稿模型推测才能开始。

SSD 消除了这个依赖。在目标模型验证当前推测结果的同时，草稿模型在独立的硬件上提前计算多组可能的下一轮推测分支。如果实际验证结果匹配了其中某个预计算分支，下一轮推测立即可用，从而将草稿推测的延迟隐式地重叠到目标模型的验证时间里。

本文有三个目标：

  • 介绍 SSD 的核心思路以及它对低延迟 LLM 服务的意义

  • 总结在 AMD Instinct GPU 上通过 ROCm 启用 SSD 所需的工程工作

  • 本文报告选取在Instinct MI300X 上复现 SSD 的端到端结果，并说明这项工作对未来 ROCm 推理工作的价值

我们的实现已开源，可在此处获取[2]

图片

为什么 SSD 很重要

自回归解码一次只生成一个 token，这使得现代加速器的算力大量闲置。SD 通过引入小模型草稿推测、大模型并行验证的方式改善了这个问题。但即便如此，SD 仍需在两个阶段间交替：

1. 草稿模型生成多个候选序列

2. 目标模型并行验证

这个剩余同步点成了新的瓶颈。SSD 正是针对这个瓶颈而设计的。它将草稿生成与验证分离到不同设备上，让草稿模型为多种可能的验证结果提前推测多个分支。本质上是把原本闲置的计算资源用来提前准备未来可能需要的路径。

这对交互式 LLM 系统至关重要，这类场景往往对延迟非常敏感。如果草稿延迟能被隐藏，用户在不改变模型输出的前提下就能看到更快的 token 流。SSD 在强基线的推测解码方案上取得了更好的表现，将吞吐-延迟的 Pareto 前沿向前推进。我们在 Instinct MI300X GPU 上得出了相同的结论，这也说明 ROCm 不仅能支撑主流的推理路径，也能支持需要更深层系统集成的异步服务算法。

图 1. SSD 工作流程概览

如图 1 所示，Round T speculation 对应当前推测轮次，Actual verification outcome 对应验证结果（包含接受前缀长度和采样的 bonus token）。推测缓存以可能的验证结果为键名，预存了下一轮推测的草稿。

实验环境配置

我们在所选硬件上复现所使用的软件栈如下：

  • GPU：AMD Instinct MI300X

  • ROCm：7.2

  • Docker 镜像：基于ROCm/flashinfer v0.5.3+amd.2 构建（.devcontainer/rocm/Dockerfile）

  • PyTorch：2.9.1+rocm7.2

  • FlashInfer：amd-flashinfer 0.5.3+amd.2（包含 PR #214 [3]的所有内核修复）

  • flash-attn：上游 Dao-AILab/flash-attention，使用Composable Kernel（CK）后端的 gfx942 版本

实现遵循论文中描述的 SSD 部署模式：大目标模型跨多 GPU 分片，小草稿模型在独立 GPU 上异步运行。在 70B 基准配置中，SSD 总共使用 5 块 MI300X GPU：4 块用于目标模型，1 块用于草稿模型。

AMD GPU适配要点

这项工作最大的价值在于同时测试了多个高级推理组件：异步多设备调度、树式推测解码、分页 KV 缓存管理、自定义注意力掩码，以及图执行友好的运行时间。

1. 发现并上游修复 ROCm FlashInfer 的正确性问题

在启动过程中，我们发现了 FlashInfer HIP 路径中的两个独立正确性 bug 并已修复：

  • bf16/fp16 MMA 类型混淆问题：rowsum 内核中 bf16 输入被错误地以错误的 MFMA 指令解释，导致草稿 token 接受率大幅下降。

  • 自定义掩码路径的行映射 bug：在AMD CDNA3 上，树解码使用自定义掩码时，自定义掩码代码路径的行映射错误会破坏分组查询注意力（GQA）的注意力输出。

这些修复至关重要，因为 SSD 高度依赖树解码和自定义掩码注意力来并行准备多个推测分支。如果不修复，算法虽然能跑通，但会产生极差的接受行为或不正确的注意力输出，导致 SSD 端到端性能不可用。

2. 适配 ROCm 的注意力栈

我们对注意力路径做了以下更新，使其兼容 ROCm 后端：

  • 将默认架构目标从 NVIDIA 9.0 改为 AMD gfx942

  • 为flash_attn_varlen_func 和flash_attn_with_kvcache 添加兼容层

  • 适配原始代码路径与 ROCm 可用注意力实现之间的 API 差异

这使得代码库更具可移植性，减少了对 NVIDIA 特定打包约定的依赖。

3. 引入双树解码后端

我们引入了一个后端开关SSD_TREE_DECODE_BACKEND，提供两条执行路径：

  • flashinfer：默认的高性能路径，支持 HIP 图执行

  • sdpa：面向正确性的回退路径，以 eager 模式运行

这种双后端设计在启动过程中非常实用。SDPA 路径提供了可靠的正确性参考，方便调试；而 FlashInfer 路径在 ROCm 内核问题修复后则用于更高性能的执行。

4. 修复 JIT 和图执行相关的运行时差异

ROCm 移植还需要几项底层运行时调整：

  • 修复 FlashInfer JIT 中与 packbits 相关的编译问题

  • 适配 NVIDIA 与 AMD FlashInfer 构建之间 plan() 参数的差异

  • 通过运行时上下文传递自定义 mask 和 wrapper 状态

  • 对支持的树解码路径条件性启用图捕获

这些细节往往决定了研究原型能否成为可复现的 ROCm 概念验证。

5. 自动化 ROCm 环境搭建

我们添加了setup_rocm.sh 脚本，用于在 ROCm 容器内自动化环境搭建流程，包括包检查、header 修复、可编辑安装、可选的 flash-attn 构建以及冒烟测试。这降低了复现成本，使 MI300X 路径更易于其他工程师验证和扩展。

快速开始

环境要求：AMD Instinct GPU、ROCm 7.2、Docker。

基本复现流程如下：

1. 克隆 SSD 仓库

2. 从源码构建 ROCm FlashInfer Docker 镜像

3. 创建并启动容器

4. 进入容器，激活 micromamba 环境，安装 FlashInfer

5. 运行 setup_rocm.sh 安装 SSD，并用 CK 后端构建 flash-attn

6. 下载目标模型和草稿模型，预处理基准数据集

7. 在 MI300X 上运行 SSD 基准测试 

步骤 1：克隆 SSD

代码块
PlainText
自动换行
复制代码
cd /home/<your-username>
git clone -b rocm-mi300x https://github.com/AMD-AGI/ssd.git
cd ssd 
复制成功

步骤 2：构建 Docker 镜像

Docker 镜像基于 FlashInfer v0.5.3+amd.2 构建，其中已包含 CDNA3 bf16 rowsum 和自定义掩码注意力的内核修复（PR #214 [3]）。将其克隆到/home/<your-username>/tmp/ 下，以便容器内可访问（容器会 bind-mount /home）：

代码块
PlainText
自动换行
复制代码
git clone --branch v0.5.3+amd.2 --depth 1 \
  https://github.com/ROCm/flashinfer.git /home/<your-username>/tmp/flashinfer-build
cd /home/<your-username>/tmp/flashinfer-build

docker build \
  --build-arg ROCM_VERSION=7.2 \
  --build-arg PY_VERSION=3.12 \
  --build-arg TORCH_VERSION=2.9.1 \
  -t flashinfer-0.5.3.amd2_rocm7.2 \
  -f .devcontainer/rocm/Dockerfile . 
复制成功

步骤 3：创建并启动容器

代码块
PlainText
自动换行
复制代码
docker run -dit \
  --name ssd \
  --privileged --network=host \
  --device=/dev/kfd --device=/dev/dri \
  --ipc=host --shm-size 64G \
  --group-add video \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  -v /home:/home \
  flashinfer-0.5.3.amd2_rocm7.2 \
  /bin/bash 
复制成功

步骤 4：进入容器并配置环境

代码块
PlainText
自动换行
复制代码
docker exec -u 0 -it ssd bash

# Activate the pre-built micromamba environment
export MAMBA_EXE=/bin/micromamba
export MAMBA_ROOT_PREFIX=/opt/conda
eval "$($MAMBA_EXE shell hook --shell bash)"
micromamba activate flashinfer-py3.12-torch2.9.1-rocm7.2

# Install FlashInfer from the source used to build the Docker image
pip install --no-build-isolation -ve /home/<your-username>/tmp/flashinfer-build
复制成功

验证安装：

代码块
PlainText
自动换行
复制代码
python -c "import torch; print(torch.__version__)"          # 2.9.1+...
python -c "import flashinfer; print(flashinfer.__version__)" # 0.5.3+amd.2
复制成功

步骤 5：安装 SSD 并构建 flash-att

代码块
PlainText
自动换行
复制代码
cd /home/<your-username>/ssd
bash setup_rocm.sh
复制成功

该脚本执行pip install -e .（可编辑安装 SSD），然后从上游 Dao-AILab/flash-attention 以 CK 后端为 gfx942 构建 flash-attn。flash-attn 构建是 HIP 图模式所必需的，大约需要 10-30 分钟。 

步骤 6：下载模型和数据集

代码块
PlainText
自动换行
复制代码
export SSD_HF_CACHE=/home/<your-username>/hf_cache
export SSD_DATASET_DIR=$SSD_HF_CACHE/processed_datasets
export HSA_NO_SCRATCH_RECLAIM=1

pip install huggingface_hub datasets
huggingface-cli login        # required for the gated meta-llama/* repos

huggingface-cli download meta-llama/Llama-3.1-8B-Instruct  --cache-dir $SSD_HF_CACHE
huggingface-cli download meta-llama/Llama-3.2-1B-Instruct  --cache-dir $SSD_HF_CACHE
# For the 70B benchmark configuration (~140 GB of weights):
huggingface-cli download meta-llama/Llama-3.1-70B-Instruct --cache-dir $SSD_HF_CACHE

HF_DATASETS_CACHE=$SSD_HF_CACHE python scripts/get_data_from_hf.py
复制成功

scripts/get_data_from_hf.py 会下载 HumanEval、Alpaca、C4、GSM8K 和 UltraFeedback，并在 $SSD_DATASET_DIR 下写入预处理后的 JSONL 文件。 

步骤 7：运行快速 SSD 验证

代码块
PlainText
自动换行
复制代码
python -O bench/bench.py \
  --llama --size 8 --gpus 2 \
  --spec --async \
  --k 7 --f 3 --b 1 \
  --temp 0 \
  --numseqs 16 --output_len 128 \
  --random 
复制成功

运行 70B 基准配置

代码块
PlainText
自动换行
复制代码
python -O bench/bench.py \
  --llama --size 70 --gpus 5 \
  --spec --async \
  --k 7 --f 3 --b 1 \
  --temp 0 \
  --numseqs 128 --output_len 512 \
  --all
复制成功

如有需要，可使用SSD_TREE_DECODE_BACKEND=sdpa 作为回退后端进行调试或正确性验证。

性能结果

我们使用Llama-3.2-1B-Instruct 作为草稿模型加速Llama-3.1-70B-Instruct 目标模型，在四个基准测试上报告三种算法的平均推理速度：

AR：自回归解码

SD：标准推测解码

SSD：Speculative Speculative Decoding

表 1. AR、SD 和 SSD 在 alpaca、c4、ultrafeedback、humaneval 四个基准上的平均推理速度，TP 分别为 1、2、4。

以TP=4 为例，在此 MI300X 配置下，SSD 达到 225.86 tokens/s，而 SD 为 138.06 tokens/s，自回归解码仅为52.32 tokens/s。具体来说：

  • 相比自回归解码加速4.32 倍

  • 相比标准推测解码加速 1.64 倍

这些结果的意义体现在两个方面。第一，它表明 SSD 在 AMD 硬件上保留了其核心优势：将草稿和验证工作重叠的做法在 MI300X 上是有效的。第二，SD 与 SSD 之间的差距证实了额外的适配工作是有价值的。单纯移植 SD 已经能降低延迟，而 SSD 则证明了 ROCm 可以支持更新一类异步解码算法，带来更大的性能提升。

这项工作为何重要

这项工作超越了一张简单的基准测试表格。

  • 它将最新的研究解码算法快速引入 AMD 硬件，有助于展示 AMD 软件栈对新兴 LLM 服务技术的就绪程度。

  • 它产出了可复用的移植模式，包括后端抽象、兼容层、图执行感知的回退机制、以及自动化环境搭建流程。

  • 它提供了一个具体的概念验证，表明AMD 主流GPU不仅适用于标准Transformer 

    推理，也能承载更先进的多设备推理调度策略。

简而言之，这既是一次算法复现，也是一次系统工程落地。性能结果是可见的产出，但更重要的影响是降低了未来 ROCm 在推测解码、树解码以及其他延迟优化的 LLM 服务创新上的门槛。

回顾之前的 ROCm SD 系列博客

AMD 团队在多个方向上推进推测解码：本文聚焦于在主流Instinct GPU 上启用最新的 SSD 方法（通过独立硬件分别运行草稿生成器和验证器来加速 SD）；此前的 ROCm SD 博客则侧重于设计新算法 [4-6]，以及介绍原始 SD 方法并将其应用到 AMD GPU [7-9]。

总结

本文介绍了 SSD，一种新的解码算法。它通过将验证与未来推测任务分离到不同设备上，消除了推测解码中残留的同步瓶颈。我们还详细描述了在 Instinct MI300X GPU 上通过 ROCm 启用 SSD 所需的关键工程工作，包括注意力栈适配、双树解码后端、运行时集成修改以及自动化环境搭建支持。

复现结果表明，SSD 在 AMD 硬件上能够带来显著的加速效果。在 LLaMA 系列模型、TP=4 的配置下，SSD 相比自回归解码实现了 4.32 倍加速，相比标准推测解码实现了1.64 倍加速。更广泛地看，本项目证明了 ROCm 不仅能支撑成熟的推理栈，也能承载新一代异步化、系统侧更重的 LLM 解码方法。

参考链接

1. SSD 论文：https://arxiv.org/pdf/2603.03251

2. SSD 开源实现：https://github.com/AMD-AGI/ssd/tree/rocm-upstream

3. ROCm FlashInfer PR #214：https://github.com/ROCm/flashinfer/pull/214

4. FLy：语义级推测解码新范式：https://rocm.blogs.amd.com/artificial-intelligence/fly/README.html

5. Gumiho：草稿序列中早期 token 更重要：https://rocm.blogs.amd.com/software-tools-optimization/gumiho/README.html

6. 多模态 AI 推理加速：https://rocm.blogs.amd.com/software-tools-optimization/multimodal-spec-dec/README.html

7. 推测解码深度解析：https://rocm.blogs.amd.com/software-tools-optimization/speculative-decoding---deep-dive/README.html

8. 在 MI300X 上实现 LLM 推理加速（最高 3 倍）：https://rocm.blogs.amd.com/artificial-intelligence/spec_decode_mi300x/README.html

9. 使用推测采样在 AMD GPU 上加速文本生成：https://rocm.blogs.amd.com/artificial-intelligence/speculative-decoding/README.html

免责声明：第三方内容由相关第三方许可给您，而非 AMD 许可。所有链接的第三方内容均"按原样"提供，不附带任何形式的保证。使用此类第三方内容由您自行决定，AMD 在任何情况下均不对因使用第三方内容而产生的任何直接、间接、特殊或其他后果性损害承担责任。您承担所有风险，并对因使用第三方内容可能产生的任何损害负全部责任。

cv50475848