[ROCm][CI] Fix logprob divergence for TitanML/tiny-mixtral under AITER rms_norm (#36101)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Andreas Karatzas
2026-03-09 12:07:44 -05:00
committed by GitHub
parent c174d54f86
commit 1e0f917b34
2 changed files with 39 additions and 0 deletions

View File

@@ -126,6 +126,10 @@ def test_models(
if use_rocm_aiter and (model in AITER_MODEL_LIST):
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
if model == "TitanML/tiny-mixtral":
# Untrained model: near-uniform logits make argmax sensitive to
# AITER's bfloat16 rounding error in plain rms_norm.
monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0")
elif use_rocm_aiter and model not in AITER_MODEL_LIST:
# Skip model that are not using AITER tests.
# When more AITER kernels are added, this list will not be