From b430f3b3f240f04415214626a166667ed018f717 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 16 Jun 2026 17:42:39 +0000 Subject: [PATCH 1/2] feat(hip,aiter): use CK rmsnorm2d for the plain rmsnorm AITER backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch the standalone rmsnorm AITER shim from AITER's vLLM-style rms_norm symbol to its CK rmsnorm2d symbol (the rmsnorm2d_fwd entry point), aligning it with the fused_add_rmsnorm path. The shim reshapes weight to [1, n] as CK requires; the Python API and JIT wiring are unchanged. CK rmsnorm2d only accepts fp16/bf16 (the old rms_norm dispatch also accepted fp32), so gate backend="auto" to 2D fp16/bf16 — fp32 now routes to native instead of hitting a CK-specific error — and reject fp32 with a clear Python error under backend="aiter". Note fp32 rmsnorm is unsupported by the native ROCm kernel too, so this only changes which error surfaces. Co-Authored-By: Claude Opus 4.7 --- README.md | 6 ++--- flashinfer/csrc_rocm/norm_aiter.cu | 17 ++++++++----- flashinfer/norm.py | 29 +++++++++++++++------- tests/rocm_tests/test_rmsnorm_aiter_hip.py | 23 ++++++++++++++--- 4 files changed, 54 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 80a09a07b8..b47fa3ea6b 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,7 @@ kernel for non-attention ops). **AITER** = ROCm AITER backend. | **POD attention** | ✅ `fa2` | — | HIP | MHA / GQA / MQA; single + batch variants (`PODWithPagedKVCacheWrapper`, `BatchPODWithPagedKVCacheWrapper`); JIT-only (excluded from AOT, same as upstream CUDA) | | **RoPE (positional encoding)** | ✅ `native` | ✅ | **AITER** for the cos/sin-cache path on gfx942/gfx950; else **HIP `native`** | LLaMA-style + LLaMA 3.1 scaling; fused RoPE + fp8 quant + paged-KV append (E4M3FNUZ, E5M2FNUZ). AITER backend covers `apply_rope_with_cos_sin_cache` and its inplace variant via AITER's C++ `rope_cached_positions_2c_fwd_impl` (linked at the C++ level, no runtime `import aiter`); cos/sin passed as float32 | | **Paged KV-cache append** | ✅ `native` | ✅ | **AITER** when `fp16/bf16` + `NHD` + gfx942/gfx950 + AITER importable; else **HIP `native`** | `append_paged_kv_cache`; fp8 KV-cache supported on the HIP path | -| **RMSNorm** | ✅ `native` | ✅ | **AITER** for 2-D inputs on gfx942/gfx950; else **HIP `native`** (3-D inputs or AITER unavailable) | `rmsnorm`; AITER's C++ `rms_norm` (linked at the C++ level, no runtime `import aiter`); fp16/bf16, 2-D only, slightly lower precision at `hidden_size >= 1024` | +| **RMSNorm** | ✅ `native` | ✅ | **AITER** for 2-D inputs on gfx942/gfx950; else **HIP `native`** (3-D inputs or AITER unavailable) | `rmsnorm`; AITER's C++ CK `rmsnorm2d` (the `rmsnorm2d_fwd` entry point, linked at the C++ level, no runtime `import aiter`); fp16/bf16, 2-D only, slightly lower precision at `hidden_size >= 1024` | | **Fused add RMSNorm** | ✅ `native` | ✅ | **AITER** on gfx942/gfx950; else **HIP `native`** | `fused_add_rmsnorm`; AITER's C++ CK `rmsnorm2d_with_add` (linked at the C++ level, no runtime `import aiter`); 2-D only, slightly lower precision at `hidden_size >= 1024` | | **LayerNorm / Gemma RMSNorm** | ✅ | — | HIP | | | **Sampling** | ✅ | — | HIP | Top-K / Top-P / Min-P / OnlineSoftmax / SamplingFromLogits | @@ -328,7 +328,7 @@ prefill/decode), `backend="native"` for non-attention ops The `rmsnorm`, `fused_add_rmsnorm`, `silu_and_mul`, and `rope` (cos/sin-cache) AITER backends are integrated at the **C++ level**: FlashInfer's JIT compiles a small HIP shim that calls AITER's C++ kernels -(`rms_norm`, `rmsnorm2d_with_add`, `aiter::silu_and_mul`, +(`rmsnorm2d`, `rmsnorm2d_with_add`, `aiter::silu_and_mul`, `rope_cached_positions_2c_fwd_impl`) directly and links a symbol-visible AITER `.so` — there is no runtime `import aiter` on these paths. The first JIT build of each op builds the corresponding AITER module once with @@ -342,7 +342,7 @@ shape-based gating to the other ops. Backend-specific exceptions to "auto picks AITER when supported": -* `rmsnorm`: `backend="auto"` picks the AITER C++ path (`rms_norm`) only +* `rmsnorm`: `backend="auto"` picks the AITER C++ path (CK `rmsnorm2d`) only for 2-D inputs; 3-D inputs fall back to the HIP `native` kernel. * `batch_decode`: `use_cuda_graph=True` or `use_tensor_cores=True` force `auto` back to `fa2` (AITER decode does not support either), diff --git a/flashinfer/csrc_rocm/norm_aiter.cu b/flashinfer/csrc_rocm/norm_aiter.cu index 84a6181301..fde0e4f750 100644 --- a/flashinfer/csrc_rocm/norm_aiter.cu +++ b/flashinfer/csrc_rocm/norm_aiter.cu @@ -1,10 +1,10 @@ // SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 // -// PyTorch entry point for AITER's CK fused-add RMSNorm (rmsnorm2d_with_add). -// FlashInfer links the symbol-visible AITER module (see -// flashinfer/jit/aiter_source.py) and calls the kernel directly — no Python -// `import aiter` at runtime. +// PyTorch entry points for AITER's CK RMSNorm kernels: plain rmsnorm +// (rmsnorm2d) and fused-add rmsnorm (rmsnorm2d_with_add). FlashInfer links the +// symbol-visible AITER module (see flashinfer/jit/aiter_source.py) and calls the +// kernels directly — no Python `import aiter` at runtime. // // FlashInfer's fused_add_rmsnorm is in-place: // residual = input + residual; input = rmsnorm(residual) * weight @@ -22,7 +22,10 @@ void rmsnorm2d_with_add(at::Tensor& out, at::Tensor& input, at::Tensor& residual_in, at::Tensor& residual_out, at::Tensor& weight, double epsilon, int use_model_sensitive_rmsnorm); -void rms_norm(at::Tensor& out, at::Tensor& input, at::Tensor& weight, double epsilon); +// CK 2D forward (the symbol the AITER `rmsnorm2d_fwd` pybind name binds to); +// the in-place overload writes `out` directly. +void rmsnorm2d(at::Tensor& out, at::Tensor& input, at::Tensor& weight, double epsilon, + int use_model_sensitive_rmsnorm); void fused_add_rmsnorm_aiter(at::Tensor input, at::Tensor residual, at::Tensor weight, double eps) { const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); @@ -36,5 +39,7 @@ void fused_add_rmsnorm_aiter(at::Tensor input, at::Tensor residual, at::Tensor w void rmsnorm_aiter(at::Tensor out, at::Tensor input, at::Tensor weight, double eps) { const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); - rms_norm(out, input, weight, eps); + // CK expects weight as [1, n]; FlashInfer passes [n] (see fused-add note above). + at::Tensor weight2d = weight.reshape({1, -1}); + rmsnorm2d(out, input, weight2d, eps, /*use_model_sensitive_rmsnorm=*/0); } diff --git a/flashinfer/norm.py b/flashinfer/norm.py index 433fd1b68f..7ea1c446cb 100644 --- a/flashinfer/norm.py +++ b/flashinfer/norm.py @@ -38,14 +38,21 @@ def get_norm_aiter_module(): return gen_norm_aiter_module().build_and_load() def _auto_select_norm_backend(input: torch.Tensor) -> str: - # auto routes plain rmsnorm to the C++ AITER kernel on supported devices - # (2D inputs only) and falls back to native everywhere else (3D inputs, or - # when AITER is not installed, so auto never raises). Note: AITER's rms_norm - # uses lower-precision reductions that exceed the flashinfer test tolerance - # at hidden_size >= 1024 (fp16 atol ~4e-3, bf16 ~7e-2). + # auto routes plain rmsnorm to the C++ AITER kernel only when the input is + # 2D fp16/bf16 (the CK rmsnorm2d kernel rejects other ranks/dtypes) and + # AITER is available; everything else falls back to native. (fp32 is + # unsupported by both the AITER and native ROCm kernels, so it still raises + # downstream — routing it to native just keeps the error consistent with the + # default kernel rather than exposing a CK-specific message.) Note: AITER's + # CK rmsnorm2d uses lower-precision reductions that exceed the flashinfer + # test tolerance at hidden_size >= 1024 (fp16 atol ~4e-3, bf16 ~7e-2). from .aiter_utils import is_aiter_available - if input.ndim == 2 and is_aiter_available(input.device): + if ( + input.ndim == 2 + and input.dtype in (torch.float16, torch.bfloat16) + and is_aiter_available(input.device) + ): return "aiter" return "native" @@ -86,10 +93,10 @@ def rmsnorm( `_ backend: str Kernel backend to use. ``"auto"`` (default) selects the best available backend: - the AITER C++ kernel for 2D inputs on supported ROCm devices, else native. + the AITER C++ kernel for 2D fp16/bf16 inputs on supported ROCm devices, else native. ``"native"`` uses the FlashInfer JIT kernel on all platforms. - ``"aiter"`` uses AMD AITER's ``rms_norm`` C++ kernel — ROCm (gfx942/gfx950) - only, 2D inputs only. Precision is slightly lower than ``"native"`` at + ``"aiter"`` uses AMD AITER's CK ``rmsnorm2d`` C++ kernel — ROCm (gfx942/gfx950) + only, 2D fp16/bf16 inputs only. Precision is slightly lower than ``"native"`` at ``hidden_size >= 1024`` (fp16 atol ~4e-3, bf16 ~7e-2). Returns @@ -108,6 +115,10 @@ def rmsnorm( f"AITER rmsnorm only supports 2D inputs; got {input.ndim}D. " "Use backend='native' for 3D inputs." ) + if input.dtype not in (torch.float16, torch.bfloat16): + raise ValueError( + f"AITER rmsnorm only supports float16/bfloat16 inputs; got {input.dtype}." + ) if out is None: out = torch.empty_like(input) get_norm_aiter_module().rmsnorm_aiter(out, input, weight, eps) diff --git a/tests/rocm_tests/test_rmsnorm_aiter_hip.py b/tests/rocm_tests/test_rmsnorm_aiter_hip.py index 0dc6902146..8cbf24dc0f 100644 --- a/tests/rocm_tests/test_rmsnorm_aiter_hip.py +++ b/tests/rocm_tests/test_rmsnorm_aiter_hip.py @@ -1,9 +1,9 @@ # SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 # -# Tests for the AITER rms_norm backend exposed via flashinfer.rmsnorm(backend="aiter"). +# Tests for the AITER CK rmsnorm2d backend exposed via flashinfer.rmsnorm(backend="aiter"). # -# Note on tolerances: AITER rms_norm uses lower-precision reductions than the +# Note on tolerances: AITER's CK rmsnorm2d uses lower-precision reductions than the # flashinfer native JIT kernel. For production inference these differences are # negligible, but they exceed the native kernel's test tolerance (fp16 atol=1e-3, # bf16 atol=1.6e-2). The tolerances below reflect AITER's actual precision. @@ -43,14 +43,31 @@ def test_rmsnorm_aiter_vs_ref(dtype, hidden_size, batch_size): @requires_aiter def test_rmsnorm_auto_backend_selects_aiter_for_2d(): - """auto backend on gfx942/950 routes 2D inputs to AITER and 3D inputs to native.""" + """auto backend on gfx942/950 routes 2D fp16/bf16 to AITER; 3D or fp32 to native.""" from flashinfer.norm import _auto_select_norm_backend device = torch.device("cuda:0") x2d = torch.randn(8, 128, dtype=torch.float16, device=device) x3d = torch.randn(8, 4, 128, dtype=torch.float16, device=device) + x2d_fp32 = torch.randn(8, 128, dtype=torch.float32, device=device) assert _auto_select_norm_backend(x2d) == "aiter" assert _auto_select_norm_backend(x3d) == "native" + # CK rmsnorm2d rejects fp32, so auto must not select it (routes to native). + assert _auto_select_norm_backend(x2d_fp32) == "native" + + +@requires_aiter +def test_rmsnorm_aiter_rejects_fp32(): + """backend='aiter' with fp32 raises a clear FlashInfer error, not the deep CK one. + + fp32 is unsupported by both ROCm rmsnorm kernels; the Python-level guard + rejects it before reaching AITER so the message is actionable. + """ + device = torch.device("cuda:0") + x = torch.randn(8, 128, dtype=torch.float32, device=device) + w = torch.randn(128, dtype=torch.float32, device=device) + with pytest.raises(ValueError, match="float16/bfloat16"): + flashinfer.rmsnorm(x, w, backend="aiter") @requires_aiter From c8b5ea30ce61748e65c026ab4e7676b1d7071e1b Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 16 Jun 2026 19:34:32 +0000 Subject: [PATCH 2/2] fix(hip,aiter): guard rmsnorm weight dtype on the AITER path Address Copilot review on #258: - CK rmsnorm2d derives a single dtype from input and reads weight with it, so a mismatched weight dtype silently produced NaN/garbage. Reject weight.dtype != input.dtype under backend="aiter", and have backend="auto" fall back to native on a mismatch (native handles fp16 input + fp32 weight fine). - Update the README RMSNorm matrix row and backend-exception bullet to reflect the fp16/bf16 + matching-weight-dtype auto gating. Co-Authored-By: Claude Opus 4.7 --- README.md | 5 ++-- flashinfer/norm.py | 29 +++++++++++++++------- tests/rocm_tests/test_rmsnorm_aiter_hip.py | 24 +++++++++++++++--- 3 files changed, 43 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index b47fa3ea6b..46d3049695 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,7 @@ kernel for non-attention ops). **AITER** = ROCm AITER backend. | **POD attention** | ✅ `fa2` | — | HIP | MHA / GQA / MQA; single + batch variants (`PODWithPagedKVCacheWrapper`, `BatchPODWithPagedKVCacheWrapper`); JIT-only (excluded from AOT, same as upstream CUDA) | | **RoPE (positional encoding)** | ✅ `native` | ✅ | **AITER** for the cos/sin-cache path on gfx942/gfx950; else **HIP `native`** | LLaMA-style + LLaMA 3.1 scaling; fused RoPE + fp8 quant + paged-KV append (E4M3FNUZ, E5M2FNUZ). AITER backend covers `apply_rope_with_cos_sin_cache` and its inplace variant via AITER's C++ `rope_cached_positions_2c_fwd_impl` (linked at the C++ level, no runtime `import aiter`); cos/sin passed as float32 | | **Paged KV-cache append** | ✅ `native` | ✅ | **AITER** when `fp16/bf16` + `NHD` + gfx942/gfx950 + AITER importable; else **HIP `native`** | `append_paged_kv_cache`; fp8 KV-cache supported on the HIP path | -| **RMSNorm** | ✅ `native` | ✅ | **AITER** for 2-D inputs on gfx942/gfx950; else **HIP `native`** (3-D inputs or AITER unavailable) | `rmsnorm`; AITER's C++ CK `rmsnorm2d` (the `rmsnorm2d_fwd` entry point, linked at the C++ level, no runtime `import aiter`); fp16/bf16, 2-D only, slightly lower precision at `hidden_size >= 1024` | +| **RMSNorm** | ✅ `native` | ✅ | **AITER** for 2-D fp16/bf16 inputs on gfx942/gfx950; else **HIP `native`** (3-D inputs, fp32, or AITER unavailable) | `rmsnorm`; AITER's C++ CK `rmsnorm2d` (the `rmsnorm2d_fwd` entry point, linked at the C++ level, no runtime `import aiter`); fp16/bf16, 2-D only, weight dtype must match input, slightly lower precision at `hidden_size >= 1024` | | **Fused add RMSNorm** | ✅ `native` | ✅ | **AITER** on gfx942/gfx950; else **HIP `native`** | `fused_add_rmsnorm`; AITER's C++ CK `rmsnorm2d_with_add` (linked at the C++ level, no runtime `import aiter`); 2-D only, slightly lower precision at `hidden_size >= 1024` | | **LayerNorm / Gemma RMSNorm** | ✅ | — | HIP | | | **Sampling** | ✅ | — | HIP | Top-K / Top-P / Min-P / OnlineSoftmax / SamplingFromLogits | @@ -343,7 +343,8 @@ shape-based gating to the other ops. Backend-specific exceptions to "auto picks AITER when supported": * `rmsnorm`: `backend="auto"` picks the AITER C++ path (CK `rmsnorm2d`) only - for 2-D inputs; 3-D inputs fall back to the HIP `native` kernel. + for 2-D fp16/bf16 inputs whose weight dtype matches; 3-D inputs, fp32, or a + mismatched weight dtype fall back to the HIP `native` kernel. * `batch_decode`: `use_cuda_graph=True` or `use_tensor_cores=True` force `auto` back to `fa2` (AITER decode does not support either), and `pos_encoding_mode != "NONE"` raises under `backend="aiter"`. diff --git a/flashinfer/norm.py b/flashinfer/norm.py index 7ea1c446cb..c19ac92592 100644 --- a/flashinfer/norm.py +++ b/flashinfer/norm.py @@ -37,20 +37,22 @@ def get_norm_aiter_module(): return gen_norm_aiter_module().build_and_load() - def _auto_select_norm_backend(input: torch.Tensor) -> str: + def _auto_select_norm_backend(input: torch.Tensor, weight: torch.Tensor) -> str: # auto routes plain rmsnorm to the C++ AITER kernel only when the input is - # 2D fp16/bf16 (the CK rmsnorm2d kernel rejects other ranks/dtypes) and - # AITER is available; everything else falls back to native. (fp32 is - # unsupported by both the AITER and native ROCm kernels, so it still raises - # downstream — routing it to native just keeps the error consistent with the - # default kernel rather than exposing a CK-specific message.) Note: AITER's - # CK rmsnorm2d uses lower-precision reductions that exceed the flashinfer - # test tolerance at hidden_size >= 1024 (fp16 atol ~4e-3, bf16 ~7e-2). + # 2D fp16/bf16 with a matching weight dtype (the CK rmsnorm2d kernel rejects + # other ranks/dtypes and reads weight with the input dtype) and AITER is + # available; everything else falls back to native. (fp32 is unsupported by + # both the AITER and native ROCm kernels, so it still raises downstream — + # routing it to native just keeps the error consistent with the default + # kernel rather than exposing a CK-specific message.) Note: AITER's CK + # rmsnorm2d uses lower-precision reductions that exceed the flashinfer test + # tolerance at hidden_size >= 1024 (fp16 atol ~4e-3, bf16 ~7e-2). from .aiter_utils import is_aiter_available if ( input.ndim == 2 and input.dtype in (torch.float16, torch.bfloat16) + and weight.dtype == input.dtype and is_aiter_available(input.device) ): return "aiter" @@ -105,7 +107,9 @@ def rmsnorm( Normalized tensor, 2D shape (batch_size, hidden_size) or 3D shape (batch_size, num_heads, hidden_size). """ if IS_HIP: - _backend = backend if backend != "auto" else _auto_select_norm_backend(input) + _backend = ( + backend if backend != "auto" else _auto_select_norm_backend(input, weight) + ) if _backend == "aiter": from .aiter_utils import require_aiter @@ -119,6 +123,13 @@ def rmsnorm( raise ValueError( f"AITER rmsnorm only supports float16/bfloat16 inputs; got {input.dtype}." ) + if weight.dtype != input.dtype: + # CK rmsnorm2d derives a single dtype from input and reads weight + # bytes with it; a mismatched weight dtype silently yields NaN/garbage. + raise ValueError( + f"AITER rmsnorm requires weight.dtype == input.dtype; got " + f"weight {weight.dtype} vs input {input.dtype}." + ) if out is None: out = torch.empty_like(input) get_norm_aiter_module().rmsnorm_aiter(out, input, weight, eps) diff --git a/tests/rocm_tests/test_rmsnorm_aiter_hip.py b/tests/rocm_tests/test_rmsnorm_aiter_hip.py index 8cbf24dc0f..17ea3a5219 100644 --- a/tests/rocm_tests/test_rmsnorm_aiter_hip.py +++ b/tests/rocm_tests/test_rmsnorm_aiter_hip.py @@ -43,17 +43,22 @@ def test_rmsnorm_aiter_vs_ref(dtype, hidden_size, batch_size): @requires_aiter def test_rmsnorm_auto_backend_selects_aiter_for_2d(): - """auto backend on gfx942/950 routes 2D fp16/bf16 to AITER; 3D or fp32 to native.""" + """auto routes 2D fp16/bf16 (matching weight) to AITER; 3D, fp32, or a + mismatched weight dtype to native.""" from flashinfer.norm import _auto_select_norm_backend device = torch.device("cuda:0") x2d = torch.randn(8, 128, dtype=torch.float16, device=device) + w = torch.randn(128, dtype=torch.float16, device=device) x3d = torch.randn(8, 4, 128, dtype=torch.float16, device=device) x2d_fp32 = torch.randn(8, 128, dtype=torch.float32, device=device) - assert _auto_select_norm_backend(x2d) == "aiter" - assert _auto_select_norm_backend(x3d) == "native" + w_fp32 = torch.randn(128, dtype=torch.float32, device=device) + assert _auto_select_norm_backend(x2d, w) == "aiter" + assert _auto_select_norm_backend(x3d, w) == "native" # CK rmsnorm2d rejects fp32, so auto must not select it (routes to native). - assert _auto_select_norm_backend(x2d_fp32) == "native" + assert _auto_select_norm_backend(x2d_fp32, w_fp32) == "native" + # CK reads weight with the input dtype, so a mismatch must not select AITER. + assert _auto_select_norm_backend(x2d, w_fp32) == "native" @requires_aiter @@ -70,6 +75,17 @@ def test_rmsnorm_aiter_rejects_fp32(): flashinfer.rmsnorm(x, w, backend="aiter") +@requires_aiter +def test_rmsnorm_aiter_rejects_weight_dtype_mismatch(): + """backend='aiter' with weight.dtype != input.dtype raises rather than silently + producing NaN/garbage (CK rmsnorm2d reads weight bytes with the input dtype).""" + device = torch.device("cuda:0") + x = torch.randn(8, 128, dtype=torch.float16, device=device) + w = torch.randn(128, dtype=torch.float32, device=device) + with pytest.raises(ValueError, match="weight.dtype == input.dtype"): + flashinfer.rmsnorm(x, w, backend="aiter") + + @requires_aiter def test_rmsnorm_aiter_with_out_tensor(): """backend='aiter' respects the out= argument."""