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 7, 00:15 UTC
[ROCm][CI] Fix logprob divergence for TitanML/tiny-mixtral under AITER rms_norm (#36101)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 9, 2026 at 12:07 UTC
1e0f917b349338ac09377dd277ded5e1e62df77e
2 files changed
+39
tests/models/language/generation/test_common.py
+4
@@ -126,6 +126,10 @@ def test_models(
126
127
if use_rocm_aiter and (model in AITER_MODEL_LIST):
128
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
129
+ if model == "TitanML/tiny-mixtral":
130
+ # Untrained model: near-uniform logits make argmax sensitive to
131
+ # AITER's bfloat16 rounding error in plain rms_norm.
132
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0")
133
elif use_rocm_aiter and model not in AITER_MODEL_LIST:
134
# Skip model that are not using AITER tests.
135
# When more AITER kernels are added, this list will not be
tests/utils.py
+35
@@ -1602,6 +1602,41 @@ def override_cutlass_fp8_supported(value: bool):
1602
yield
1603
1604
1605
+def disable_aiter_plain_rmsnorm(monkeypatch) -> None:
1606
+ """Patch dispatch_rocm_rmsnorm_func so the plain (non-fused) rms_norm path
1607
+ always uses the native float32 kernel for the duration of a test.
1608
+
1609
+ The fused path (rms_norm2d_with_add, selected when with_fused_add=True) is
1610
+ left on AITER -- only the plain path is redirected to native.
1611
+
1612
+ AITER's plain rms_norm accumulates variance in bfloat16 (~1 ULP/call),
1613
+ which drifts the KV cache over many decode steps. This drift is irrelevant
1614
+ for a trained model (rank-1/rank-2 gap ~1-3 nats >> 1 ULP), but breaks
1615
+ logprob comparison tests with randomly-initialised models like
1616
+ TitanML/tiny-mixtral whose rank-1/rank-2 gap is only O(1/sqrt(V)) ~0.006
1617
+ nats -- smaller than the accumulated per-step error.
1618
+ """
1619
+ import torch
1620
+
1621
+ import vllm.model_executor.layers.layernorm as _ln_mod
1622
+ from vllm.model_executor.layers.layernorm import rms_norm as _native
1623
+
1624
+ _orig = _ln_mod.dispatch_rocm_rmsnorm_func
1625
+
1626
+ def _native_plain(
1627
+ with_fused_add: bool, dtype: torch.dtype, use_aiter: bool = False
1628
+ ):
1629
+ if (
1630
+ use_aiter
1631
+ and not with_fused_add
1632
+ and dtype in (torch.float16, torch.bfloat16)
1633
+ ):
1634
+ return _native
1635
+ return _orig(with_fused_add, dtype, use_aiter)
1636
+
1637
+ monkeypatch.setattr(_ln_mod, "dispatch_rocm_rmsnorm_func", _native_plain)
1638
+
1639
+
1640
def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)):
1641
"""
1642
Generate prompts which a bunch of assignments,