Source-Level Compilation and Tuning of vLLM on ROCm: A Practical Guide for Domestic AI Accelerators

Source-level compilation and tuning of vLLM on ROCm centers on selecting a matching ROCm version, enabling the Flash Attention ROCm backend, and tuning parameters for memory and communication bottlenecks. Based on Mingxin Technology's measured experience on the AMD MI308X platform (ROCm 7.2), detailed in reports R1–R4, this article systematically outlines key compilation and tuning points, providing a transferable reference framework for deployment on domestic accelerators such as Ascend.

Why Source-Level Compilation? Current State and Challenges of the ROCm Ecosystem

vLLM's official support is most mature for NVIDIA CUDA. However, on AMD ROCm platforms, precompiled PyPI packages often fail to fully leverage ROCm-specific optimizations (e.g., HIP backend, ROCm's Flash Attention implementation). Source-level compilation allows developers to:

• Enable ROCm-specific acceleration: For example, -DUSE_ROCM=ON and -DROCM_ARCH=gfx942 (for MI300 series) automatically enable ROCm's GEMM optimizations and the hipBLASLt library.
• Control Flash Attention version: vLLM's flash-attn library on ROCm requires manual specification of the version (e.g., flash-attn==2.6.1+rocm); otherwise, the default CUDA version leads to compilation failure or performance regression.
• Avoid ABI compatibility issues: Different ROCm versions (e.g., 7.1 vs. 7.2) have ABI-incompatible HIP runtime libraries. Source compilation ensures consistency with the system's ROCm version.

Key steps (using ROCm 7.2 + vLLM 0.20.1 as an example):

1. Install dependencies: pip install ninja setuptools wheel, and ensure ROCm's hipcc and rocm-smi are in the PATH.
2. Clone the vLLM repository: git clone --branch v0.20.1-rocm https://github.com/vllm-project/vllm.git (use the ROCm branch).
3. Compile: python setup.py build_ext --inplace, and set the environment variable VLLM_ROCM_USE_FLASH_ATTN_V2=1 to enable the ROCm-optimized Flash Attention.

Tuning Parameters: Breaking Through Memory and Communication Bottlenecks

On ROCm platforms, vLLM tuning must focus on the following parameters, which directly impact KV Cache utilization efficiency and inference latency. In Mingxin FX100's measured R2 report, parameter tuning achieved a 26–32% reduction in TTFT and a 29–40% increase in throughput (480B model, TP8 configuration).

1. --max-model-len and --gpu-memory-utilization

• Function: Controls the model context length and GPU memory allocation ratio. Due to differences in memory bandwidth between ROCm and CUDA (MI308X: 3.5 TB/s vs. H100: 3.35 TB/s), it is advisable to lower gpu-memory-utilization (recommended 0.85–0.90) to avoid OOM.
• Measured recommendation: On MI308X ×8, --max-model-len 4096 combined with --gpu-memory-utilization 0.88 achieves a balance between KV Cache usage and inference latency (R2 baseline).

2. --block-size and --enable-prefix-caching

• Function: block-size controls the granularity of the KV Cache. On ROCm, 16 (default) is recommended to match ROCm's page size. --enable-prefix-caching allows reuse of prefix KV Caches, significantly reducing first-token latency in long-context scenarios.
• Measured data: With prefix caching enabled, the 480B model's TTFT at concurrency 8 dropped from 10.17s to 7.53s (a 26% reduction), and throughput increased from 4.1 tok/s to 74.9 tok/s (measured, report R2).

3. --num-scheduler-steps and --max-num-batched-tokens

• Function: Controls the scheduler step count and batch token count. ROCm's scheduler is more sensitive to num-scheduler-steps than CUDA; setting it to 4–8 is recommended to balance CPU overhead and GPU utilization.
• Tuning tip: In measured R3, --num-scheduler-steps 4 combined with --max-num-batched-tokens 8192 increased throughput by 35% (TP4×2 full-system metric).

4. Communication Optimization: --distributed-executor-backend and NCCL Replacement

• Function: ROCm uses an NCCL variant compiled with hipcc (rccl) by default. For cross-node deployment, ensure the RCCL_IB_HCA environment variable points to the correct InfiniBand device and set RCCL_MSCCL_ENABLE=1 to enable MSCCL optimizations.
• Measured comparison: In a single-node 8-GPU scenario, --distributed-executor-backend mp (multiprocess) showed 12% lower latency than the ray mode (measured, report R1).

Differentiated Tuning for Ascend (Huawei) Platforms

Domestic accelerator platforms (e.g., Ascend 910B) face similar challenges in vLLM adaptation, but the following differences must be noted:

• Flash Attention support: Ascend currently relies on torch_npu's npu_fusion_attention, not ROCm's flash-attn. In measured R9, Ascend 910B enabled this via the --use-npu-fusion-attention parameter, but the Ascend driver version must be ≥ 23.0.rc1.
• Memory management: Ascend's HBM bandwidth (910B: 1.2 TB/s) is lower than MI308X. Therefore, gpu-memory-utilization should be reduced to 0.80–0.85, and --enable-prefix-caching should be prioritized to reduce memory fragmentation.
• Compilation differences: Ascend requires pip install torch_npu instead of torch, and the --use-npu flag must be added in setup.py. In Mingxin's measured R9, vLLM compilation time on Ascend 910B increased by approximately 30% (due to custom operators in torch_npu).

Conclusion

Source-level compilation and tuning of vLLM on ROCm is a critical step in unlocking the potential of domestic AI accelerators. By selecting the correct ROCm version, enabling the Flash Attention ROCm backend, and finely adjusting memory and communication parameters, teams can achieve a 26–32% reduction in TTFT and a 29–40% increase in throughput on 480B-scale models (measured, reports R2/R3). Mingxin Technology has deep adaptation experience on both ROCm and Ascend platforms. We welcome teams interested in source-level tuning or KV Cache acceleration to participate in joint testing, driving the deployment efficiency of domestic accelerators.

Key Points Q&A

Q: How can Flash Attention compilation failures be avoided when compiling vLLM on ROCm?
A: Use vLLM's ROCm branch (e.g., v0.20.1-rocm) and set the environment variable VLLM_ROCM_USE_FLASH_ATTN_V2=1. Also ensure the flash-attn version is 2.6.1+rocm, not the CUDA version.

Q: What are the core tuning parameters for vLLM on ROCm?
A: Key parameters include --gpu-memory-utilization (recommended 0.85–0.90), --block-size 16, --enable-prefix-caching (for long-context scenarios), --num-scheduler-steps 4-8, and the communication backend --distributed-executor-backend mp.

Q: What are the main differences in vLLM tuning between ROCm and Ascend platforms?
A: ROCm relies on the flash-attn ROCm backend, while Ascend uses npu_fusion_attention. Due to Ascend's lower HBM bandwidth (1.2 TB/s vs. MI308X's 3.5 TB/s), gpu-memory-utilization should be reduced to 0.80–0.85, and prefix caching should be prioritized.