biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix(rocm): Use refresh_env_variables() for rocm_aiter_ops in test_moe (#31711)

Browse Source 
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-01-14 00:41:54 +05:30
committed by GitHub
parent f28125d87b
commit 69f8a0ea37
1 changed files with 6 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
tests/kernels/moe/test_moe.py
View File

@@ -6,8 +6,6 @@ Run `pytest tests/kernels/test_moe.py`.
"""
import functools
import importlib
import sys
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
@@ -592,15 +590,13 @@ def test_mixtral_moe(
"""Make sure our Mixtral MoE implementation agrees with the one from
huggingface."""
# clear the cache before every test
# Force reload aiter_ops to pick up the new environment variables.
if "rocm_aiter_ops" in sys.modules:
importlib.reload(rocm_aiter_ops)
# Explicitly set AITER env var based on test parameter to ensure
# consistent behavior regardless of external environment
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
rocm_aiter_ops.refresh_env_variables()
if use_rocm_aiter:
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
if dtype == torch.float32:
pytest.skip("AITER ROCm test skip for float32")
if use_rocm_aiter and dtype == torch.float32:
pytest.skip("AITER ROCm test skip for float32")
monkeypatch.setenv("RANK", "0")
monkeypatch.setenv("LOCAL_RANK", "0")