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 5, 23:30 UTC
[ROCm][CI] Fix AITER state leak in shared_fused_moe_routed_transform test (#38137)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 26, 2026 at 11:26 UTC
bdc1719eb9b27ef4a30c1f753b406bf05c8edd3a
3 files changed
+21
-3
tests/kernels/moe/test_routing_simulator.py
+9
-2
@@ -125,8 +125,15 @@ def test_routing_strategy_integration(monkeypatch, device):
125
env_name = "VLLM_MOE_ROUTING_SIMULATION_STRATEGY"
126
monkeypatch.setenv(env_name, strategy)
127
128
- # Force reload of environment variable
129
- envs.environment_variables[env_name] = lambda s=strategy: s
128
+ # Temporarily override the envs lookup so the router factory
129
+ # reads the monkeypatched value instead of the module-load-time
130
+ # default. Use monkeypatch.setitem so the original lambda is
131
+ # restored automatically at teardown.
132
+ monkeypatch.setitem(
133
+ envs.environment_variables,
134
+ env_name,
135
+ lambda s=strategy: s,
136
+ )
137
138
# Test the select_experts method
139
topk_weights, topk_ids = fused_moe.router.select_experts(
tests/kernels/moe/test_shared_fused_moe_routed_transform.py
+1
-1
@@ -137,7 +137,7 @@ def test_routed_input_transform_inside_vs_outside(
137
Method A (inside): SharedFusedMoE with routed_input_transform
138
Method B (outside): Manually transform, then SharedFusedMoE without transform
139
"""
140
- if current_platform.is_rocm() and use_rocm_aiter:
140
+ if current_platform.is_rocm():
141
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
142
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1" if use_rocm_aiter else "0")
143
from vllm._aiter_ops import rocm_aiter_ops
vllm/distributed/parallel_state.py
+11
@@ -1905,6 +1905,17 @@ def destroy_distributed_environment():
1905
def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
1906
# Reset environment variable cache
1907
envs.disable_envs_cache()
1908
+
1909
+ # Reset rocm_aiter_ops class variables to match current os.environ.
1910
+ # These are class-level attributes that persist across tests and are
1911
+ # NOT restored by monkeypatch (which only restores os.environ).
1912
+ from vllm.platforms import current_platform
1913
+
1914
+ if current_platform.is_rocm():
1915
+ from vllm._aiter_ops import rocm_aiter_ops
1916
+
1917
+ rocm_aiter_ops.refresh_env_variables()
1918
+
1919
# Ensure all objects are not frozen before cleanup
1920
gc.unfreeze()
1921