引言
前沿大语言模型（LLM）规模持续扩大，对 GPU 的计算力和内存带宽需求日益高涨。GPU 厂商和模型开发者正转向低精度浮点格式，其中 FP4（4 位浮点）量化尤为引人注目。例如，FP4 量化的Llama 3.3 70B模型体积缩小 3.5 倍，在MMLU等基准测试中质量损失微乎其微。
然而，当前硬件支持存在明显短板。虽然 NVIDIA GB200 和 AMD MI350 等新一代 GPU 原生支持 FP4 矩阵乘法，但广泛部署的 AMD Instinct MI250 和 MI300 系列 GPU 尚无此能力，导致用户无法在现有 AMD 硬件上高效运行 FP4 模型。
为填补这一空白，我们开发了 Petit——专为 AMD GPU 设计的 FP16/BF16 × FP4 混合精度 GPU 内核集合。Petit 无需硬件升级，即可在 MI200 和 MI300 系列上服务 FP4 模型，带来显著性能提升：
- 使用SGLang时，Llama 3.3 70B 端到端推理性能提升 1.74 倍；
- 等效矩阵乘法执行速度比 AMD 最先进的 GEMM 库 hipBLASLt 快至 3.7 倍。
Petit 以 BSD 许可开源，已集成至 SGLang 0.4.10。用户可通过以下命令在 AMD MI250/MI300X 上启动 Llama 3.3 70B FP4 模型服务：
python -m sglang.launch_server --model-path nvidia/Llama-3.3-70B-Instruct-FP4 --host 0.0.0.0 --port 30000本文详述优化历程和技术细节，Petit 充分利用 AMD 开源软件生态，并引入离线重排和硬件特定低级优化等创新。
与硬件架构协同设计高效 GPU 内核
现代 GPU 通过堆叠紧凑的计算单元（CU）实现海量计算吞吐，但需应用与底层架构深度协同设计，方能发挥峰值性能。如图 1 所示，Petit 的开发遵循多项关键协同设计原则。
通过预处理实现高效去量化
Petit 高效利用 AMD GPU 的专用 MatrixCore 硬件加速矩阵乘法。一个 wavefront（64 个线程组）可高效集体乘以两个 BF16/FP16 16×16 矩阵。但 MI300X GPU 无原生 FP4 权重 MatrixCore 支持，因此需将 FP4 权重去量化至 BF16/FP16，同时保持内存加载和 MatrixCore 准备的高效性。
这带来核心挑战：内存加载和 MatrixCore 准备需不同数据布局。内存高效需连续加载 1024 字节块，而 MatrixCore 期望 16×16 瓦片分布于 wavefront。传统 GPU 端重排开销巨大。
NVIDIA GPU 的Marlin实现通过磁盘预排元素避免此问题。我们将 8 个连续 FP4 值打包至 32 位整数，需 31 条指令去量化。Petit 进一步定制位打包格式：前 4 个 FP4 以 BF8 布局重排，其余存于剩余位。借助 AMD 独有的 v_bfrev_b32 和 v_cvt_pk_f32_bf8 指令（支持子双字寻址 SDWA），仅用 15 条指令去量化 8 个 FP4 值，乘法性能提升 30%。
掌握内存层次结构
MI300X 等 GPU 算术密度极高（>500），CU 需每字节执行数百操作以达峰值 FLOPS，因此最大化有效内存带宽至关重要。Petit 采用瓦片化和双缓冲等成熟技术，并针对 AMD 特性优化：
- 避免 LDS Bank 冲突：AMD GPU LDS 分 32 银行，每周期允许 32 个唯一银行并发访问。冲突会导致序列化瓶颈，尤其 wavefront 有 64 线程。Petit 基于银行设计实现置换数据布局，实现无冲突 LDS 使用。
- Chiplet 和互连：每个 MI300 chiplet（XCD）有 4MB 本地 L2 缓存，全 GPU 共享 256MB L3 缓存。互连带宽高但延迟大。Petit 实现拓扑感知工作负载分区，减少互连流量，优先朴素网格分区而非全局条带分区（当剖析显示互连开销更大时）。
生成高质量机器码
GPU 使用简单乱序执行单元最大化 CU 密度，但分支和管线停顿代价高昂。AMD GPU 提供条件移动和边界内存指令，完全消除分支。例如，Petit 利用指定范围的缓冲加载/存储指令，GPU 自动丢弃越界访问；LDS 超 64KB 访问亦自动处理，无性能罚则。此外，Petit 提供编译器提示，重叠 MFMA（矩阵融合乘加）指令与内存访问，有效隐藏内存延迟。
标准编译器可能未充分利用高级 GPU ISA（如有意越界访问为未定义行为）。这些优化需手动构建和验证。
性能结果
端到端推理性能
我们通过 FP4 与 BF16 模型端到端推理性能对比评估 Petit 实际效果。测试使用 Llama 3.3 70B 两种变体与 SGLang v0.4.10，测量批次大小 10 和 64 的输入/输出 token 吞吐量。环境为 AMD 开发者云 VM（1× MI300X GPU、240 GB RAM、5 TB SSD），运行 ROCm 6.4.2 于 Ubuntu 24.04.1。
图 2 显示离线生成基准结果（使用真实 ShareGPT 痕迹，反映生产性能）。总体上，Petit 服务 Llama 3.3 70B FP4 模型比 SGLang 原生 BF16 模型快 1.74 倍（批次 10）和 1.60 倍（批次 64）。在小批次内存带宽受限的生产场景，Petit 高效利用 3.5 倍小体积的 FP4 模型，实现更高吞吐。可用以下命令复现：
python -m sglang.bench_offline_throughput --model-path nvidia/Llama-3.3-70B-Instruct-FP4 --num-prompts 10
python -m sglang.bench_offline_throughput --model-path nvidia/Llama-3.3-70B-Instruct-FP4 --num-prompts 64详细性能分析
我们进一步对比 Petit 与 hipBLASLt（AMD 最先进的低级汇编 GEMM 库）性能。注意两库目标略异：
- Petit：BF16 矩阵 × NVFP4 矩阵（16 元素共享 1 FP8 尺度）。
- hipBLASLt：两个 BF16 矩阵。
虽非完全相同，但结果提供量化洞见。我们考察 Llama 3 70B 服务时实际权重矩阵尺寸，测量 m=16（解码工作负载）和 m=256（预填充工作负载）性能，平均 100 次运行（50 次预热后）。两库均调优至最佳配置。
图 3a 和 3b 显示 GEMM 性能。对于 m=16（解码主导），Petit 比 hipBLASLt 快至 3.7 倍，平均提升 2.56 倍。对于 m=256（预填充），Petit 快至 1.09 倍。