Skip to content

Commit 3331ced

Browse files
gshtrasyewentao256
authored andcommitted
[ROCm][Bugfix] Only enable +rms_norm based on aiter if not explicitly disabled (#25275)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Signed-off-by: yewentao256 <zhyanwentao@126.com>
1 parent b614e0f commit 3331ced

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

vllm/platforms/rocm.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -340,7 +340,8 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
340340
else:
341341
parallel_config.worker_cls = "vllm.worker.worker.Worker"
342342
# Aiter rms norm perform best when CUDA Graph capture is enabled.
343-
if use_v1 and use_aiter_rms_norm and not is_eager_execution:
343+
if (use_v1 and use_aiter_rms_norm and not is_eager_execution
344+
and "-rms_norm" not in compilation_config.custom_ops):
344345
compilation_config.custom_ops.append("+rms_norm")
345346

346347
@classmethod

0 commit comments

Comments
 (0)