TokenSpeed在AMD GPU上跑起来：ROCm部署LLM推理的完整路径-育师

TokenSpeed在AMD GPU上跑起来：ROCm部署LLM推理的完整路径

【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeed

TokenSpeed 是一款光速级 LLM 推理引擎，不仅支持 NVIDIA GPU，也完整支持 AMD GPU 的 ROCm 部署。本文带你从零开始，在 AMD Instinct 系列（MI350X/MI355X，gfx950）及 RDNA4 显卡上完成 TokenSpeed 的安装、验证与推理服务启动，掌握 ROCm 环境部署 LLM 推理的完整路径。

一、先了解 TokenSpeed 的 AMD 支持能力

TokenSpeed 采用"内核插件化"架构：核心引擎与硬件无关，真正的加速由tokenspeed-kernel内核注册表按 GPU 架构自动选择实现。针对 AMD，项目内置了独立的 tokenspeed-kernel-amd 内核包，用 Gluon 编写的高性能内核覆盖：

算子族	gfx950（CDNA4）	gfx1250 / gfx1201
Attention	MHA、MLA、DSA、KDA	MHA、MLA、KDA（gfx1201 走通用 Triton 路径）
MoE	BF16、MXFP4 权重	—
GEMM / Sampling	专用调优内核	—

在 MI355X 上实测，Gluon Attention 内核比 Triton 基线快 1.4–2.3 倍，小 batch 下 MoE 比 Triton 快 1.7–2.1 倍。性能细节见 tokenspeed-kernel-amd/README.md。

二、环境准备：ROCm 7.2 + PyTorch

AMD 路径的官方验证环境是ROCm 7.2，要求 Linux x86_64 主机、Python 3.10–3.13。关键点：必须先安装 ROCm 版 PyTorch，再安装 TokenSpeed，否则 pip 会选到 CUDA 依赖。

python -m pip install "torch==2.14.0" "torchvision==0.29.0" \ --index-url https://download.pytorch.org/whl/rocm7.2

安装完成后，一行断言即可确认 PyTorch 是 HIP 后端：

python -c 'import torch; assert torch.version.hip'

三、安装 TokenSpeed：Nightly 一行命令

在上面的环境基础上，直接使用 ROCm 专用 nightly 索引：

python -m pip install --upgrade tokenspeed \ --extra-index-url https://lightseek.org/whl/nightly/rocm7.2 python -m pip check

这里有三个新手容易踩的坑：

不要同时配置 CUDA 与 ROCm 两个索引——内核 wheel 版本号相同但后端依赖不同，混用会装坏；
想固定版本可指定tokenspeed==0.1.0.postYYYYMMDD（该日期需同时存在于两个索引）；
升级旧版本时用--no-cache-dir --force-reinstall。

完整说明见 docs/guides/getting-started.md。

四、验证与首次启动

安装完成后用两条命令做冒烟验证：

tokenspeed env tokenspeed serve --help

然后启动第一个推理服务。对于 RDNA4（gfx1201，如 Radeon AI PRO R9700）这类消费级/工作站显卡，初始支持 Dense BF16 路径，推荐从 Triton 注意力 + eager 执行开始：

tokenspeed serve Qwen/Qwen3-0.6B \ --host 127.0.0.1 \ --dtype bfloat16 \ --world-size 1 \ --attention-backend triton \ --enforce-eager \ --disable-prefill-graph

服务起来后，标准 OpenAI 兼容 API 即可在http://127.0.0.1:8000上使用。

五、大模型实战：8 卡 gfx950 部署旗舰模型

如果你有 8 张 MI350X/MI355X，AMD 的标准部署路径已经过完整调优。以 Kimi K3 为例（docs/recipes/models.md 中有全部模型的启动配方）：

tokenspeed serve moonshotai/Kimi-K3 \ --served-model-name kimi-k3 \ --trust-remote-code \ --max-model-len 8192 \ --kv-cache-dtype fp8 \ --tensor-parallel-size 8 \ --enable-expert-parallel \ --attention-backend mla \ --moe-backend auto \ --max-num-seqs 32 \ --disable-kvstore \ --host 0.0.0.0 \ --port 8000

这条命令的精髓在于两个auto/自动选择：

--attention-backend mla自动选用 gfx950 上的 MLA 内核；如需强制 Gluon 内核（含投机解码 drafter），换成--attention-backend gluon --drafter-attention-backend gluon，注册表找不到对应内核时会快速失败而非静默降级；
--moe-backend auto在 TP8/EP8 下自动选用专用 Gluon SiTU 内核，单 token 解码还会自动融合路由 MXFP4 专家与共享专家。

追求极限吞吐时，还有一键开启 Gluon Petit MegaMoE 的 DP8/EP8 组合，配方同样在 docs/recipes/models.md 中。

六、进阶：Docker 容器与源码安装

容器方式：项目提供了现成的 AMD 镜像构建文件 docker/Dockerfile.amd，基于rocm7.2.4-pytorch-2.13.0-gfx950基础镜像，按顺序安装 tokenspeed-kernel-amd、tokenspeed-kernel（TOKENSPEED_KERNEL_BACKEND=rocm）、tokenspeed-scheduler 和主运行时，适合团队统一环境。

源码方式：在仓库根目录依次安装即可，内核包的 Python 元数据会自动带上 ROCm 后端依赖：

pip install --force-reinstall --no-deps ./tokenspeed-kernel-amd --no-build-isolation TOKENSPEED_KERNEL_BACKEND=rocm pip install tokenspeed-kernel/python/ --no-build-isolation pip install tokenspeed-scheduler/ pip install ./python --no-build-isolation

⚠️ 顺序很重要：tokenspeed-kernel-amd必须先于tokenspeed-kernel就位，后者才能解析出 ROCm 需求。

七、常见架构速查

GPU	架构代号	推荐路径
MI350X / MI355X	gfx950（CDNA4）	Gluon 专用内核（MLA/DSA/KDA + MoE + GEMM），旗舰性能
Radeon AI PRO R9700	gfx1201（RDNA4）	通用 Triton 内核 + ROCm 库实现，Dense BF16
gfx1250	—	Attention/MoE 专用内核已就位

TokenSpeed 会自动检测架构：gfx1201与 CDNA4/CDNA5 分开识别，CDNA 专属的矩阵与异步拷贝能力在 RDNA4 上保持关闭，避免误用。