feat(hip,aiter): use CK rmsnorm2d for the plain rmsnorm AITER backend

[demandal25]

Summary

Switch the standalone rmsnorm AITER backend from AITER's vLLM-style rms_norm C++ symbol to its CK rmsnorm2d symbol (the rmsnorm2d_fwd entry point), aligning it with the fused_add_rmsnorm path which already uses the CK rmsnorm2d_with_add kernel. The Python API and JIT wiring are unchanged.

What changed

Kernel

• flashinfer/csrc_rocm/norm_aiter.cu — rmsnorm_aiter now forward-declares and calls CK rmsnorm2d(out, input, weight2d, eps, 0) instead of rms_norm. Reshapes weight to [1, n] as CK requires (same as the fused-add path).

Dtype gating (regression fix)

• flashinfer/norm.py — CK rmsnorm2d only accepts fp16/bf16, whereas the old rms_norm dispatch also accepted fp32. backend="auto" is now gated to 2D fp16/bf16 inputs so fp32 routes to native instead of hitting a CK-specific error, and backend="aiter" rejects fp32 with a clear Python-level ValueError. (fp32 rmsnorm is unsupported by the native ROCm kernel too, so this only changes which error surfaces — it does not remove working functionality.)

Tests / docs

• tests/rocm_tests/test_rmsnorm_aiter_hip.py — added fp32 coverage (auto→native selection, explicit-aiter rejection); updated kernel name in the header.
• README.md — feature matrix, C++-integration note, and backend-exceptions bullet now reference CK rmsnorm2d; auto criteria note fp16/bf16.

Notes on correctness verification

• Non-contiguous tensors are handled correctly. CK rmsnorm2d passes input.stride(0) / out.stride(0) to the kernel, so sliced/non-contiguous inputs and outputs produce correct results (verified: max abs err 0.0, untouched regions preserved). This is actually more stride-aware than the old packed-layout rms_norm kernel.
• Precision is unchanged from the old path (fp16 ~3.9e-3, bf16 ~6.25e-2 vs an fp32 reference), so the existing AITER tolerances stay.

Test plan

• pytest tests/rocm_tests/test_norm_hip.py tests/rocm_tests/test_rmsnorm_aiter_hip.py -n auto --reruns 2 — 767 passed
• Verified the rebuilt JIT shim links the CK rmsnorm2d symbol (nm)
• pre-commit run --files on all changed files — clean

[Copilot]

Pull request overview

This PR updates the ROCm AITER RMSNorm backend to call AITER’s CK rmsnorm2d (via the rmsnorm2d_fwd entry point) instead of the older rms_norm symbol, aligning the standalone RMSNorm path with the existing fused-add RMSNorm path that already uses CK kernels. It also adjusts Python-side backend selection/validation and updates tests/docs accordingly.

Changes:

• Switch AITER plain rmsnorm C++ entry point to reshape weight to [1, n] and call CK rmsnorm2d.
• Tighten backend="auto" selection to 2D fp16/bf16 (and add a clearer error for backend="aiter" with unsupported dtypes).
• Add ROCm tests for fp32 routing/rejection and update README references to the CK kernel name.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
flashinfer/csrc_rocm/norm_aiter.cu	Routes plain AITER RMSNorm through CK rmsnorm2d with required weight reshaping.
flashinfer/norm.py	Updates auto-backend gating and adds Python-level dtype validation for backend="aiter".
tests/rocm_tests/test_rmsnorm_aiter_hip.py	Adds coverage for fp32 auto-routing and explicit AITER fp32 rejection; updates header text.
README.md	Updates documentation references from rms_norm to CK rmsnorm2d.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.