A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 11:00 UTC
[ROCm] fix num_stages for default moe config to avoid triton OutOfResource error (#17744)
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>
Hongxia Yang committed
May 6, 2025 at 20:39 UTC
ed3a1d2106a42b1522013b304e5dc9ca172385ec
1 file changed
+3
-1
vllm/model_executor/layers/fused_moe/fused_moe.py
+3
-1
@@ -747,13 +747,15 @@ def get_default_config(
747
if dtype == "fp8_w8a8" and block_shape is not None:
748
# Block-wise quant: BLOCK_SIZE_N must be divisible by block_shape[0]
749
# BLOCK_SIZE_K must be divisible by block_shape[1]
750
+ # num_stages=3 can cause triton.runtime.errors.OutOfResources
751
+ # on ROCm, set it to 2 instead.
752
config = {
753
"BLOCK_SIZE_M": 64,
754
"BLOCK_SIZE_N": block_shape[0],
755
"BLOCK_SIZE_K": block_shape[1],
756
"GROUP_SIZE_M": 32,
757
"num_warps": 4,
756
- "num_stages": 3,
758
+ "num_stages": 3 if not current_platform.is_rocm() else 2,
759
}
760
elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None:
761
# moe wna16 kernels