Skip to content

[CI][BugFix][AMD] Add check for model_config being None and update conftest.py to load AITER of available to fix Kernels MoE Test % N #33952

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
rasmith wants to merge 15 commits into from
Closed

[CI][BugFix][AMD] Add check for model_config being None and update conftest.py to load AITER of available to fix Kernels MoE Test % N #33952

Show file tree
Hide file tree
Changes from 14 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
9737831
Fixing the skinny gemm dispatch logic. Weights can be padded for it t…
gshtras Jan 29, 2026
5c1f660
Added weight padding into fp8 skinny gemm tests
gshtras Jan 29, 2026
515e0c5
add default_vllm_config so tests pass
rasmith Feb 3, 2026
a965c65
Merge branch 'main' of github.com:rasmith/vllm
rasmith Feb 4, 2026
e8bc3c2
Merge branch 'main' of github.com:rasmith/vllm
rasmith Feb 4, 2026
4cecd78
Merge branch 'main' of github.com:rasmith/vllm
rasmith Feb 4, 2026
cd2f3b6
enable aiter in tests
rasmith Feb 5, 2026
1eb74e2
Merge branch 'main' of github.com:rasmith/vllm
rasmith Feb 5, 2026
7b09163
merge main
rasmith Feb 5, 2026
32caeb5
typo
rasmith Feb 6, 2026
afe4b5c
undo mori
rasmith Feb 6, 2026
1cfd671
update comment
rasmith Feb 6, 2026
57d5660
pre-commit
rasmith Feb 6, 2026
9d0c5f8
fix typo
rasmith Feb 6, 2026
09d804c
only load ops and ensure env var not st when tests run
rasmith Feb 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions tests/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,19 @@
from torch._inductor.utils import fresh_cache


def use_aiter_if_available():

@robertgshaw2-redhat robertgshaw2-redhat Feb 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hm, not sure if this is a good idea. I think many tests explicity set or dont set this env

@rasmith rasmith Feb 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I updated it so that only the function pointers are loaded, which was happening before, but the environment variable is set to 0.

@rasmith rasmith Feb 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@robertgshaw2-redhat I could also do this in _aiter_ops.py, which also works. Basically, the functions get loaded but the env var will be unset:

@@ -36,7 +36,7 @@ def is_aiter_found_and_supported() -> bool:
     Checks: platform (ROCm), device arch (gfx9), library existence,
     and VLLM_ROCM_USE_AITER env variable.
     """
-    if current_platform.is_rocm() and IS_AITER_FOUND and envs.VLLM_ROCM_USE_AITER:
+    if current_platform.is_rocm() and IS_AITER_FOUND:
         from vllm.platforms.rocm import on_gfx9

from vllm.platforms import current_platform
from importlib.util import find_spec

# NOTE: it's not possible to use vllm._aiter_ops.is_aiter_found
# because the aiter ops won't load and all tests that want
# to use aiter will fail because no aiter ops will be loaded.
if current_platform.is_rocm() and find_spec("aiter") is not None:
os.environ["VLLM_ROCM_USE_AITER"] = "1"


use_aiter_if_available()

if TYPE_CHECKING:
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
from transformers.generation.utils import GenerateOutput
Expand Down
4 changes: 0 additions & 4 deletions tests/kernels/moe/test_rocm_aiter_topk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
# and the platform is not ROCm.

import importlib.util
import os

import pytest
import torch
Expand All @@ -20,9 +19,6 @@
if not current_platform.is_rocm():
pytest.skip("This test can only run on ROCm.", allow_module_level=True)

# This environment variable must be set so ops will be registered.
os.environ["VLLM_ROCM_USE_AITER"] = "1"

# this import statement is needed to ensure the ops are registered
import vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe # noqa: F401

Expand Down
1 change: 1 addition & 0 deletions vllm/config/vllm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
envs.VLLM_ROCM_USE_AITER
and envs.VLLM_ROCM_USE_AITER_RMSNORM
and envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
and cfg.model_config is not None
and cfg.model_config.get_hidden_size() == 2880
)

Expand Down