SGLang 适配 AMD 平台的开发工作全面梳理

SGLang 的核心分层是纯 Python 的调度/前端层硬件无关→ 算子与后端层sgl-kernel、attention backend、MoE、量化、通信→ 构建/运行环境层。从 NVIDIA 迁到 AMD绝大部分工作集中在后两层。下面按模块逐项梳理。一、构建系统与运行环境NVIDIA 路径官方 pip wheel 开箱即用CUDA toolkit flashinfer 等预编译依赖直接装。AMD 路径需要额外开发sgl-kernel 需用专门的 ROCm 构建脚本源码编译cd sgl-kernel python setup_rocm.py installPython 包也要切换pyproject_other.toml并执行pip install -e python[all_hip]维护独立的 ROCm Dockerfile 与镜像rocm/sgl-dev、lmsysorg/sglang:*-rocm*-mi30x并按 GPU 架构分镜像gfx942MI300X/MI325X、gfx950MI350/MI355X、gfx1151Strix Halowave32 架构需专门镜像且关闭 AITER容器与宿主适配暴露/dev/kfd、/dev/dri加入 video/render 组设置HIP_FORCE_DEV_KERNARG、RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES等环境变量条件依赖管理pyproject.toml需要按平台分支CUDA 版依赖 flashinfer/flash-attnROCm 版依赖 aiter/triton二、底层算子库sgl-kernel / sglang-kernel移植这是工作量最大的部分。sgl-kernel 覆盖 MLA decode、GEMMFP8/INT8/AWQ/GPTQ、MoE topk、RMSNorm、RoPE、采样、投机解码等大量自定义 CUDA kernel。CUDA→HIP 转换hipify 工具可自动处理约 95% 的语法但必须手工处理的关键点包括warp size 32→64AMD wavefront 为 64 线程所有依赖 warp 内 shuffle/ballot 语义的 kernel 都要改删除/重写 PTX 内联汇编如 PDL programmatic dependent launchNCCL→RCCL API 替换双后端条件编译C 头文件用#ifdef USE_ROCM选择 CUDA/ROCm 两套实现Python 侧按torch.version.hip分支导入不同算子NVIDIA 专有库的替代flashinfer 的 silu_and_mul / rmsnorm / RoPE / sampling 等算子在 AMD 上要么用 PyTorch 回退实现要么接入 AITER/CK 等价算子架构检测从is_sm90_supported换成is_gfx942/is_gfx950NVTX 标记换成 roctxCUTLASS/cuBLAS 类依赖替换换成 Composable KernelCK、hipBLASLt并需要按 (M, N, K) shape 做 CSV 驱动的 GEMM tuningAITER 提供 shape 采集和调优工作流三、Attention 后端开发NVIDIA默认有 flashinfer、FlashAttention 3、trtllm_mha/gen、FlashMLA 等多套后端这些全部不可用FA3 是 Hopper 专属、TRT-LLM 无 ROCm 移植计划。AMD 需要开发 Triton attention backendHIP 版 Triton和AITER backend新增aiter_backend.py遵循BaseAttnBackend接口封装 AITER 的 attention API约 300–400 行处理硬件约束差异如 CK batch-prefill 路径 head_dim 上限 256head_dim 512 的模型如 Gemma-4prefill 要路由到 Triton unified attention某些路径禁止 fp8_e4m3 KV cachebatched decode 输出错误须强制 bf16 KVMLA 专项NVIDIA 用 FlashMLA/cutlass_mla_decodeAMD 用 AITER 的 MLA decode kernelgfx950 上走 gluon sparse kernel其他架构走 tilelang partialcombine还要做 fused decode MLASGLANG_ROCM_FUSED_DECODE_MLA等开关长上下文 decode 的分区策略如SGLANG_AITER_DECODE_PARTITION_BUCKETS按序列长度桶生成多个 decode graph四、MoE 与专家并行EPAll-to-All 通信后端替换NVIDIA 用 DeepEP基于 NVSHMEM/IBGDA、pplx-kernels、NIXLAMD 需要接入MORI-EP--moe-a2a-backend mori即 AMD 原生的 dispatch/combine 实现MoE 计算后端替换NVIDIA 的 DeepGEMM/CUTLASS/TRT-LLM runner 不可用AMD 走 Triton runner 或 AITER fused MoECK 内核还要处理 MoE padding、gate/up 权重 tile 布局interleave vs separated等差异计算通信重叠适配TBOtwo-batch overlap在 AMD 上要基于 MORI 的异步 API 重做双流流水线并利用 SDMA copy engine 做零算力开销的数据搬运ROCm 特有的多流限制HIP 默认硬件队列只有 4 个共享专家与路由专家重叠需设置GPU_MAX_HW_QUEUES5并开SGLANG_ROCM_USE_MULTI_STREAM还需修多节点稳定性问题MORI inter_node 模式曾出现 CUDA graph capture 挂起和高并发下 dispatch/combine RDMA 死锁五、量化路径重建量化是平台差异最明显的领域之一NVIDIA 专有路径全部失效Marlinawq_marlin、gptq_marlin、modelopt_fp8/fp4、gguf 在 AMD 上不支持AMD 替代路径FP8 走 AITER 或 TritonDeepSeek-V3/R1 预量化模型开箱即用AWQ 用 Triton 反量化 kernel无 Marlin 加速MXFP4 需要 CDNA3/CDNA4 SGLANG_USE_AITER1NVFP4 模型需通过 Petit kernelpetit_nvfp4另有 AMD 专属的在线量化quark_int4fp8_moe新量化格式的适配常留尾巴例如 MiniMax-M3 的 MXFP8 需要 PR 把mxfp8加进 ROCm 支持列表并绕过 bf16 router GEMM 缺口block-scale FP8 GEMM 还需在 CK 与 Triton 实现间做精度/速度权衡gfx942 默认 CK但 GSM8K 等评测有小幅精度差六、CUDA Graph 与投机解码HIP graph 基本可用但坑更多EAGLE 投机解码在 MI300X 上开启 CUDA graph 时因tree speculative sampling kernel 在 HIP 上不可用而回退到 greedy verification且该回退路径存在不确定性死锁需要移植/重写投机解码相关 kernelverify、draft_extendAITER 侧有 unified verify / draft_extend 的独立开关SGLANG_AITER_UNIFIED_VERIFY等内存策略差异AITER 注意力在长上下文8K下要把mem-fraction-static乘 0.85 预留非静态 workspace否则 OOM——这类默认值都要为 AMD 单独调校七、通信与 PD 分离NCCL→RCCLAPI 基本兼容但环境变量、拓扑感知、allreduce 性能调优要重做SGLang 的 custom all-reduce CUDA kernel 需要 HIP 移植或改用 AITER 的 all-gather 实现SGLANG_USE_AITER_AGPD 分离的 KV cache 传输NVIDIA 走 Mooncake/NIXLAMD 用 MORI-IO 做 p2p KV 传输八、工具链、调试与 CI性能分析Nsight Systems/Compute → rocprof / roctx / rpd_tracertorch profiler 的 kineto 在 ROCm 上的行为差异要适配CI 基础设施需要 MI300/MI350 的 GPU runner为 ROCm 单独跑测试集很多 kernel 路径依赖特定 gfx 架构无法进常规 CI只能人工覆盖精度与性能验证单 kernel 正确性验证策略warp 语义变化、原子操作、共享内存 bank conflict 模式在 wave64 下都不同移植后需逐一跑数值对齐测试九、平台检测与代码组织约定整个 Python 层要建立统一的设备抽象所有硬件分支收拢到if torch.version.hip:条件或独立文件中保证 NVIDIA 路径零改动工具函数层is_cuda()/is_hip()、显存查询、设备属性要封装好因为torch.cuda.*在 ROCm 下虽是别名可用但显存上报、SM/CU 数量、wave size 等语义不同。工作量参考以精简版 mini-sglang 的评估为例完整 AMD 适配约1000–1200 行改动、涉及约 20 个文件可拆成 4 个 PR设备抽象与构建系统~200 行、HIP kernel 移植 RCCL~400 行、AITER attention backend~300–400 行、Docker/文档/CI~150 行。完整版 SGLang 则是长期持续投入量化矩阵、每个新模型的 per-arch tuning、MORI 多节点稳定性、投机解码 kernel 补齐都需要随版本迭代。总结Python 调度层几乎零成本真正的开发工作在「kernel 移植hipify 手工修 warp/汇编→ 用 AITER/Triton/CK 重建 attention、MoE、量化后端 → 用 MORI/RCCL 替换 NVSHMEM/NCCL 通信 → 处理 CUDA graph 与投机解码的 HIP 特有缺陷 → 按 gfx 架构做调优与 CI」这条链路上。回报是显著的——AITER 集成后 DeepSeek-R1 在 MI300X 上端到端吞吐从 6484 提升到 13704 tok/s约 2.1 倍。参考1工作内容梳理一一、构建系统与安装对应「一、构建系统与运行环境」SGLang 官方文档 · AMD GPU 平台安装指南https://docs.sglang.ai/_sources/platforms/amd_gpu.md 中文版镜像https://sglang.org/zh/platforms/amd_gpu —— 支撑源码安装流程setup_rocm.py install、切换pyproject_other.tomlpip install -e python[all_hip]、rocm.Dockerfile、Docker 运行参数/dev/kfd、/dev/dri、--group-add video等、GRUB/NUMA 系统级调优。AMD ROCm 官方文档 · SGLang inference and serving on ROCmhttps://rocm.docs.amd.com/projects/ai-ecosystem/en/latest/inference/sglang.html—— 支撑AMD 官方维护的 SGLang on ROCm 支持矩阵ROCm 版本 × Instinct/Radeon × gfx 架构。sgl-project/sglang PR #33939[AMD] Add gfx1151 (Strix Halo) Docker imagehttps://github.com/sgl-project/sglang/pull/33939 镜像解析https://prhub.com.cn/sgl-project/sglang/pull/33939 —— 支撑gfx1151 专用镜像、setup_rocm.py仅接受 gfx942/gfx950/gfx1250 需打补丁、RDNA 上SGLANG_USE_AITER0且默认--attention-backend triton等细节。二、sgl-kernel 内核移植对应「二、底层算子库移植」sgl-project/mini-sglang Issue #90Interest in adding AMD GPU (ROCm) supporthttps://github.com/sgl-project/mini-sglang/issues/90—— 这是「工作量参考」的直接出处~1000–1200 行 / ~20 个文件 / 4 个 PR 的拆分hipify-perl 覆盖约 95% 语法warp 32→64、PDL 内联汇编、NCCL→RCCL 需手工处理flashinfer 算子替换与 NVTX→roctx新增attention/aiter.py后端~300–400 行等。sglang-kernelsgl-kernel内核库完全指南GitCode 博客https://blog.gitcode.com/73c0a33a7f2453365058fcff9bc3fca1.html—— 支撑sgl-kernel 按平台拆分pyproject_rocm.toml/setup_rocm.py、按torch.version.hip动态加载 ROCm 算子的机制。SGLang 源码阅读 · sgl-kernel “five wheels, one tree”https://www.jinnpan.com/sources/sglang.html—— 支撑一份 C/Triton 源码、五个平台 wheelCUDA/ROCm/CPU/MUSA/Metal的工程组织方式ROCm 用 CMake hipcc Composable Kernel 构建。三、Attention 后端与 AITER对应「三、Attention 后端开发」AMD ROCm 博客 · Accelerate DeepSeek-R1 Inference: Integrate AITER into SGLanghttps://rocm.blogs.amd.com/artificial-intelligence/aiter-intergration-s/README.html—— 支撑AITER 替换 MoE / attention / GEMM / all-reduce 的集成点、AITER 的 JIT ASM/CK/HIP 多级 kernel 层次以及 prefill 延迟降 52%、decode 降 47%、吞吐翻倍的实测数据。AMD ROCm 博客 · AITER: AI Tensor Engine For ROCmhttps://rocm.blogs.amd.com/software-tools-optimization/aiter-ai-tensor-engine/README.html—— 支撑结尾引用的性能数字来源SGLang DeepSeek V3/R1 在 MI300X 上 6484.76 → 13704.36 tok/s2 倍以上提升。AMD ROCm 博客 · Supercharge DeepSeek-R1 Inference on AMD Instinct MI300Xhttps://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html—— 支撑AITER kernel 级加速幅度GEMM 2x、MoE 3x、MLA decode 17x、MHA prefill 14x、MI300X vs H200 的对比、SGLang 服务端超参调优chunked_prefill_size 等。四、MoE / EP 与通信对应「四、MoE 与专家并行」「七、通信与 PD 分离」SGLang 文档 · Expert ParallelismHardware-Specific Configurationhttps://sgl-project-sglang-93.mintlify.app/distributed/expert-parallelism—— 支撑AMD 上--moe-a2a-backend mori的配置方式及「MORI 目前只支持 normal mode」。GitHub · ROCm/moriModular RDMA Interfacehttps://github.com/rocm/mori—— 支撑MORI-EP / MORI-IO 已集成进 SGLang 和 vLLM 用于 MoE EP 与 PD 分离、MORI-SHMEM、扩展到 64 GPU 等项目动态。SGLang 文档 · Environment VariablesMORI Configurationhttps://sgl-project.github.io/references/environment_variables.html—— 支撑SGLANG_MORI_DISPATCH_DTYPE、MORI_DISABLE_AUTO_XGMI等 MORI 调优环境变量。AMD 招聘 JD · AI Software Engineer (Shanghai)https://bebee.com/cn/jobs/ai-software-engineer-high-performance-amd-shanghai-shanghai–t7xk-751205891—— 支撑AMD 官方对「把 MORI-EP/IO 落进 SGLang/vLLM」工作内容的描述dispatch/combine、KVCache transfer、EPLB、8–64 GPU EP 维护可作为「该方向是持续工程投入」的旁证。五、PD 分离 KV 传输对应「七、通信与 PD 分离」Infera 文档 · PD disaggregationMooncake vs MoRIIO 对比https://rocm.docs.amd.com/projects/infera/en/main/features/pd_disaggregation.html—— 支撑SGLang 侧--disaggregation-transfer-backend mooncake/mori的差异write/push vs read/pull、bootstrap server vs ZMQ。SGLang 文档 · DeepSeek-V4 Cookbook · PD Disaggregation on MI355X (MORI-IO)https://lmsysorg.mintlify.app/cookbook/autoregressive/DeepSeek/DeepSeek-V4—— 支撑GPU_MAX_HW_QUEUES5、MORI_SHMEM_HEAP_SIZE、MegaMoEv2、FP8 KV pool 拆分等 AMD 侧 PD 调优细节。sgl-project/sglang PR #26288[PD][AMD] incremental KV transfer with decode radix cachehttps://github.com/sgl-project/sglang/pull/26288 镜像解析https://prhub.com.cn/sgl-project/sglang/pull/26288 —— 支撑mori 后端的增量 KV 传输与 decode radix cache 适配工作mori/conn.py、参数白名单等。Prefill/decode disaggregation reaches production: state of playlecompute.frhttps://lecompute.fr/en/runtimes/disaggregation-prefill-decode-production/—— 支撑MORI-IO 作为 AMD 的 NIXL 对位、SGLang 同时支持 NIXL / Mooncake / MORI-IO 三种 KV 传输。六、CUDA Graph 与投机解码对应「六、CUDA Graph 与投机解码」sgl-project/sglang Issue #29347[Bug][AMD][ROCm/MI300X] EAGLE CUDA graph 不确定性死锁https://github.com/sgl-project/sglang/issues/29347—— 支撑tree speculative sampling kernel 在 HIP 上不可用 → 回退 greedy verification → 约 1/3 概率挂死--disable-cuda-graph或去掉 EAGLE 作为 workaround但均牺牲吞吐。七、补充阅读SGLang 遇上 ROCm 首测体验CSDNhttps://blog.csdn.net/2600_96323197/article/details/162310474—— ROCm 7.x 下源码编译实践PYTORCH_ROCM_ARCHgfx942、Triton ROCm 分支锁定等踩坑记录。AMD ROCm 文档 · SGLang inference performance testingMAD 基准https://rocm.docs.amd.com/en/docs-6.4.2/how-to/rocm-for-ai/inference/benchmark-docker/sglang.html—— 官方基准复现流程可用于移植后的性能回归验证。参考2: 工作内容梳理2作为一家新的 GPGPU 公司要让 SGLang 在你的硬件上高效运行核心工作不是修改 SGLang 源码而是通过其 Out-of-TreeOOT插件体系以旁加载side-installed的方式补齐硬件特有的实现。SGLang 近年来的架构演进正是为了系统性地降低新硬件厂商的接入成本。适配工作可以归纳为四个核心层面平台插件、算子库、通信库以及可选的通用插件钩子。一、硬件平台插件Hardware Platform Plugin声明“你的硬件是什么”这是最基础的一层。你需要发布一个 Python 包在其中实现一个SRTPlatform的子类并通过setuptools入口点sglang.srt.platforms注册它。SGLang 启动时会通过activate()回调自动探测你的硬件是否可用并加载对应的平台类。平台类需要覆盖 SGLang 运行时依赖的设备侧核心抽象主要包括需要实现的接口/能力作用说明DeviceMixin设备操作设备身份识别如is_gpgpu()的等价逻辑、set_device()、get_device_name()、显存查询等基础操作KV Cache 池通过build_kv_pool(request)或工厂方法返回适配你显存布局的 KV Cache 池类需覆盖 MHA、MLA、稀疏注意力等不同架构的池实现Attention 后端提供 Prefill 和 Decode 阶段的 Attention 执行后端使 SGLang 的注意力调度能调用你的硬件实现Graph 捕获若你的硬件支持图执行类似 CUDA Graph需要提供图捕获与回放的后端这对降低解码阶段 CPU 开销至关重要编译后端若你的软件栈提供torch.compile的等效后端需要在此声明否则 SGLang 会回退到 eager 模式这一层的设计原则是fail-fast未实现的方法默认抛出NotImplementedError迫使你在适配初期就暴露缺失项而不是静默产生错误结果。二、算子库提供“在你的硬件上怎么算”这是性能差异的最大来源。SGLang 的核心热点算子Attention、GEMM、归一化、激活、MoE 路由等在 NVIDIA 生态中依赖flashinfer、sgl_kernel等 CUDA 专用库在你的 GPGPU 上需要用你的软件栈重新实现这些算子。有两种并存的算子集成路径建议同时使用自定义算子库静态注册将高性能 Kernel 编译成.so或 Python 扩展通过 PyTorch 的TORCH_LIBRARY_FRAGMENT机制注册为torch.ops.你的命名空间.op_name。当 SGLang 检测到你的平台激活时自动加载该算子库并在热路径中替换 CUDA 实现。昇腾 NPU 的sgl_kernel_npu就是这一模式的参考。OOT 算子分发动态优先级替换SGLang 提供了MultiPlatformOp的 OOT 注册机制。你的插件可以通过register_oot_impl()为指定的算子如SiluAndMul、RMSNorm注册VENDOR类型的实现并指定优先级。运行时dispatch_forward()会优先查找 OOT 注册表若你的算子可用则调用你的实现否则回退到内置或其他后端。可以通过SGLANG_OOT_PREFERvendor全局设定你的算子优先或通过SGLANG_OOT_OP_PREFER按算子精确控制。需要优先覆盖的算子类别参考摩尔线程 MATE 库的实践Attention对接 FlashAttention / FlashMLA 等接口语义、GEMM / DeepGEMM、归一化RMSNorm、LayerNorm、激活函数、MoE 相关路由与专家计算、采样Sampling算子。三、分布式通信让多卡跑起来SGLang 的分布式推理依赖集合通信。你的 GPGPU 通常不会直接运行 NCCL因此需要在 OOT 框架下注册通信器。具体做法是通过GroupCoordinator.register_oot_communicator(factory_fn)注册一个通信器工厂并提供all_reduce、reduce_scatter_tensor、all_gather_into_tensor等集合操作的实现。你的通信库可以基于自研的集合通信库类似 HCCL、FlagCX 的定位或封装torch.distributed中你已支持的后端。通信器工厂的优先级高于 SGLang 内置的硬件通信逻辑确保你的路径被选中。四、通用插件钩子可选但强烈建议用于长尾兼容SGLang 还提供了sglang.srt.plugins入口点的通用插件机制支持before / after / around / replace类型的钩子可以拦截任意函数或方法。当某些模型或路径中存在平台接口无法优雅覆盖的硬件相关分支时可以用通用插件做临时性兼容避免 fork SGLang。但应优先通过平台接口和算子注册解决问题通用钩子更适合作为模型支持初期的过渡手段。参考路径以摩尔线程 MUSA 为镜摩尔线程的 MUSA 后端从“代码共建”走到 SGLang 主线“原生支持”提供了一条可参照的路线图在算子层开源了 MATE 算子库对接 FlashAttention、FlashMLA、DeepGEMM 等主流接口并通过多平台分发机制将 MUSA 实现路由到 SGLang 执行路径中同时完成了从环境构建、算子覆盖到分布式推理的全链路打通。你的适配工作本质上是在重复这条已被验证的路径只是硬件和软件栈不同。建议的推进顺序第一步最小可运行闭环——实现SRTPlatform的设备操作与 KV 池提供最基本的 Attention 和 GEMM 算子确保单卡上能跑通一个主流模型如 Qwen 或 Llama 系列的generate请求。第二步性能关键路径——覆盖 Prefill / Decode Attention、MoE、采样等热点算子接入图捕获验证吞吐与延迟是否进入可用区间。第三步分布式与长尾模型——注册集合通信器跑通 TP / PP / EP 配置针对目标客户的重点模型如 DeepSeek、Kimi 等补齐模型特有的算子与缓存路径。第四步上游化可选但有利于客户信心——像摩尔线程那样将适配代码合入 SGLang 主线获得“官方原生支持”的身份会显著降低客户在选型时的采用顾虑。