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 10, 05:15 UTC
fix(rocm): Use refresh_env_variables() for rocm_aiter_ops in test_moe (#31711)
Signed-off-by: rabi <ramishra@redhat.com>
Rabi Mishra committed
Jan 14, 2026 at 00:41 UTC
69f8a0ea37a018b2e54897c44f8832de51a2bf5c
1 file changed
+6
-10
tests/kernels/moe/test_moe.py
+6
-10
@@ -6,8 +6,6 @@ Run `pytest tests/kernels/test_moe.py`.
6
"""
7
8
import functools
9
-import importlib
10
-import sys
9
from collections.abc import Callable
10
from dataclasses import dataclass
11
from typing import Any
@@ -592,15 +590,13 @@ def test_mixtral_moe(
590
"""Make sure our Mixtral MoE implementation agrees with the one from
591
huggingface."""
592
595
- # clear the cache before every test
596
- # Force reload aiter_ops to pick up the new environment variables.
597
- if "rocm_aiter_ops" in sys.modules:
598
- importlib.reload(rocm_aiter_ops)
593
+ # Explicitly set AITER env var based on test parameter to ensure
594
+ # consistent behavior regardless of external environment
595
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
596
+ rocm_aiter_ops.refresh_env_variables()
597
600
- if use_rocm_aiter:
601
- monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
602
- if dtype == torch.float32:
603
- pytest.skip("AITER ROCm test skip for float32")
598
+ if use_rocm_aiter and dtype == torch.float32:
599
+ pytest.skip("AITER ROCm test skip for float32")
600
601
monkeypatch.setenv("RANK", "0")
602
monkeypatch.setenv("LOCAL_RANK", "0")