xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-24 14:56:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

[ROCm] fix num_stages for default moe config to avoid triton OutOfResource error (#17744)

Browse Source 
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>
This commit is contained in:
Hongxia Yang 2025-05-06 20:39:48 -04:00 committed by GitHub
parent 022afbeb4e
commit ed3a1d2106
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/model_executor/layers/fused_moe/fused_moe.py
View File

@ -747,13 +747,15 @@ def get_default_config(
if dtype == "fp8_w8a8" and block_shape is not None:
# Block-wise quant: BLOCK_SIZE_N must be divisible by block_shape[0]
# BLOCK_SIZE_K must be divisible by block_shape[1]
# num_stages=3 can cause triton.runtime.errors.OutOfResources
# on ROCm, set it to 2 instead.
config = {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": block_shape[0],
"BLOCK_SIZE_K": block_shape[1],
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3,
"num_stages": 3 if not current_platform.is_rocm() else 2,
}
elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None:
# moe wna16 kernels