feat(hip,aiter): default silu_and_mul and rope auto backend to native

README.md

@@ -61,14 +61,14 @@ kernel for non-attention ops). **AITER** = ROCm AITER backend.
 | **Cascade attention** | ✅ | — | HIP | Two-level shared-prefix attention; a fused single-kernel HIP variant is gated behind `FLASHINFER_HIP_FUSED_CASCADE=1` |
 | **MLA (Multi-Latent Attention)** | — | ✅ | **AITER** (no HIP fallback) | DeepSeek-style 192/128 head-dim split; bf16 + `page_size=1`; `backend="auto"` (default) resolves to `"aiter"` |
 | **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 |
+| **RoPE (positional encoding)** | ✅ `native` | ✅ | **HIP `native`** (AITER is opt-in via `backend="aiter"`) | LLaMA-style + LLaMA 3.1 scaling; fused RoPE + fp8 quant + paged-KV append (E4M3FNUZ, E5M2FNUZ). The 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. `backend="auto"` resolves to HIP `native` — the in-tree kernel is the better default; request AITER explicitly with `backend="aiter"` |
 | **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 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 |
 | **Logits processor** | ✅ | — | HIP | Composable processor pipeline (cap, mask, temperature, …) |
-| **Activation** | ✅ `native` | ✅ | **AITER** for `silu_and_mul` on gfx942/gfx950; else **HIP `native`** | SiLU / GELU with fused gating. AITER path (`silu_and_mul` only) via AITER's C++ `aiter::silu_and_mul` (linked at the C++ level, no runtime `import aiter`); matches native precision in fp16, lower in bf16 |
+| **Activation** | ✅ `native` | ✅ | **HIP `native`** (AITER is opt-in via `backend="aiter"`) | SiLU / GELU with fused gating. AITER path (`silu_and_mul` only) via AITER's C++ `aiter::silu_and_mul` (linked at the C++ level, no runtime `import aiter`); matches native precision in fp16, lower in bf16. `backend="auto"` resolves to HIP `native` — the in-tree kernel is the better default; request AITER explicitly with `backend="aiter"` |
 | **Quantization** | ✅ | — | HIP | `packbits`, `segment_packbits` |
 | **`torch.compile`** | ✅ (opt-in) | n/a | n/a | Set `FLASHINFER_USE_TORCH_CUSTOM_OPS=1` **before** importing `flashinfer`; requires PyTorch ≥ 2.4. Without it, `torch.compile` raises a clear error if it traces into a flashinfer op |
 
@@ -334,14 +334,20 @@ 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
 `AITER_SYMBOL_VISIBLE=1` and caches it under
 `~/.cache/flashinfer/aiter_libs/` (the CK `module_rmsnorm` build is large
-and can take many minutes the first time). For these ops,
-`backend="auto"` resolves to AITER on gfx942/gfx950 and to HIP `native`
-elsewhere, subject to op-specific constraints (e.g. `rmsnorm` auto only
-routes 2-D inputs to AITER); a later performance pass may extend this
-shape-based gating to the other ops.
+and can take many minutes the first time). Auto-routing differs per op:
+for `rmsnorm` and `fused_add_rmsnorm`, `backend="auto"` resolves to AITER
+on gfx942/gfx950 (subject to op-specific constraints — `rmsnorm` auto only
+routes 2-D fp16/bf16 inputs to AITER) and to HIP `native` elsewhere. For
+`silu_and_mul` and `rope` (cos/sin-cache), experimentation found the
+in-tree HIP `native` kernel to be the better default, so `backend="auto"`
+resolves to `native` on all platforms; the AITER C++ path for these two
+ops is opt-in via an explicit `backend="aiter"`.
 
 Backend-specific exceptions to "auto picks AITER when supported":
 
+* `silu_and_mul` and `rope` (cos/sin-cache): `backend="auto"` always resolves
+  to the HIP `native` kernel — experimentation found it the better default.
+  The AITER C++ path is reachable only via an explicit `backend="aiter"`.
 * `rmsnorm`: `backend="auto"` picks the AITER C++ path (CK `rmsnorm2d`) only
   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.

flashinfer/activation.py

@@ -42,13 +42,10 @@ def get_silu_and_mul_aiter_module():
         return gen_silu_and_mul_aiter_module().build_and_load()
 
     def _auto_select_silu_and_mul_backend(input: torch.Tensor) -> str:
-        # auto routes to the C++ AITER kernel on supported gfx942/gfx950 devices
-        # and falls back to native everywhere else (incl. when AITER is not
-        # installed, so auto never raises). (Shape/precision tuning is deferred to
-        # a later performance pass.)
-        from .aiter_utils import is_aiter_available
-
-        return "aiter" if is_aiter_available(input.device) else "native"
+        # Experimentation found the in-tree native kernel to be the better default
+        # for silu_and_mul, so auto always resolves to native. The C++ AITER kernel
+        # remains reachable via an explicit backend="aiter".
+        return "native"
 
 
 @functools.cache
@@ -110,9 +107,8 @@ def silu_and_mul(
         <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
 
     backend: str
-        Kernel backend to use. ``"auto"`` (default) routes to the C++ AITER kernel
-        on ROCm (gfx942/gfx950) and to the native FlashInfer JIT kernel everywhere
-        else.
+        Kernel backend to use. ``"auto"`` (default) resolves to the native
+        FlashInfer JIT kernel on all platforms.
         ``"native"`` uses the FlashInfer JIT kernel on all platforms.
         ``"aiter"`` uses AMD AITER's ``silu_and_mul`` C++ kernel — ROCm
         (gfx942/gfx950) only; raises ``ValueError`` on any other platform.

flashinfer/rope.py

@@ -38,16 +38,10 @@ def get_rope_aiter_module():
         return gen_rope_aiter_module().build_and_load()
 
     def _auto_select_rope_backend(query: torch.Tensor, key: torch.Tensor) -> str:
-        # auto routes to the C++ AITER kernel on supported devices and falls back
-        # to native everywhere else. auto must never raise, so fall back to native
-        # for cases the AITER kernel rejects: mixed query/key dtypes (it rotates
-        # both with one cos/sin table), and when AITER is unavailable. (Shape /
-        # precision tuning is deferred to a later performance pass.)
-        from .aiter_utils import is_aiter_available
-
-        if key.dtype != query.dtype:
-            return "native"
-        return "aiter" if is_aiter_available(query.device) else "native"
+        # Experimentation found the in-tree native kernel to be the better default
+        # for the cos/sin-cache rope, so auto always resolves to native. The C++
+        # AITER kernel remains reachable via an explicit backend="aiter".
+        return "native"
 
     def _apply_rope_cos_sin_cache_aiter(
         query: torch.Tensor,
@@ -1218,10 +1212,8 @@ def apply_rope_with_cos_sin_cache(
           we rotate the even dimensions ``([..., ::2])`` and odd dimensions ``([..., 1::2])``.
 
     backend : str
-        Kernel backend to use. ``"auto"`` (default) routes to the C++ AITER kernel
-        on ROCm (gfx942/gfx950) and to the native FlashInfer JIT kernel everywhere
-        else; it also falls back to native for mixed query/key dtypes and when the
-        AITER package is unavailable, so it never raises.
+        Kernel backend to use. ``"auto"`` (default) resolves to the native
+        FlashInfer JIT kernel on all platforms.
         ``"native"`` uses the FlashInfer JIT kernel on all platforms.
         ``"aiter"`` uses AMD AITER's C++ ``rope_cached_positions_2c_fwd_impl``
         kernel — ROCm (gfx942/gfx950) only; requires the ``aiter`` package and
@@ -1241,6 +1233,18 @@ def apply_rope_with_cos_sin_cache(
     if cos_sin_cache.dtype != torch.float32:
         raise ValueError("cos_sin_cache should be float32")
 
+    if backend not in ("auto", "native", "aiter"):
+        raise ValueError(
+            f"Unknown backend {backend!r}; expected one of 'auto', 'native', 'aiter'."
+        )
+    if backend == "aiter":
+        # Validate the explicit opt-in up front so a misconfiguration (unsupported
+        # device or missing aiter package) surfaces as a clear ValueError here
+        # instead of silently falling through to native on non-HIP platforms.
+        from .aiter_utils import require_aiter
+
+        require_aiter(query.device, "rope")
+
     query_out = torch.empty_like(query)
     key_out = torch.empty_like(key)
 
@@ -1260,10 +1264,6 @@ def apply_rope_with_cos_sin_cache(
                 is_neox=is_neox,
             )
             return query_out, key_out
-        if _backend != "native":
-            raise ValueError(
-                f"Unknown backend {backend!r}; expected one of 'auto', 'native', 'aiter'."
-            )
 
     _apply_rope_pos_ids_cos_sin_cache(
         q=query.view(query.shape[0], -1, head_size),
@@ -1314,16 +1314,12 @@ def apply_rope_with_cos_sin_cache_inplace(
           we rotate the even dimensions ``([..., ::2])`` and odd dimensions ``([..., 1::2])``.
 
     backend : str
-        Kernel backend to use. ``"auto"`` (default) selects the best backend for
-        the call: on ROCm (gfx942/gfx950) it picks AITER for fp16 inputs with at
-        least ~2048 tokens (where AITER's kernel is measurably faster), and stays
-        on ``"native"`` otherwise — for bf16 (precision), small token counts
-        (launch overhead), and non-ROCm platforms.
+        Kernel backend to use. ``"auto"`` (default) resolves to the native
+        FlashInfer JIT kernel on all platforms.
         ``"native"`` uses the FlashInfer JIT kernel on all platforms.
-        ``"aiter"`` uses AMD AITER's rope_cached kernel — ROCm (gfx942/gfx950) only;
-        requires the ``aiter`` package. Precision is slightly lower than ``"native"``
-        for bfloat16 (max abs error ~5e-2 vs ~3e-2) because AITER consumes the cos/sin
-        tables in the query dtype rather than float32.
+        ``"aiter"`` uses AMD AITER's C++ ``rope_cached_positions_2c_fwd_impl``
+        kernel — ROCm (gfx942/gfx950) only; requires the ``aiter`` package and
+        query/key to share a dtype.
 
     Note
     ----
@@ -1332,6 +1328,18 @@ def apply_rope_with_cos_sin_cache_inplace(
     if cos_sin_cache.dtype != torch.float32:
         raise ValueError("cos_sin_cache should be float32")
 
+    if backend not in ("auto", "native", "aiter"):
+        raise ValueError(
+            f"Unknown backend {backend!r}; expected one of 'auto', 'native', 'aiter'."
+        )
+    if backend == "aiter":
+        # Validate the explicit opt-in up front so a misconfiguration (unsupported
+        # device or missing aiter package) surfaces as a clear ValueError here
+        # instead of silently falling through to native on non-HIP platforms.
+        from .aiter_utils import require_aiter
+
+        require_aiter(query.device, "rope")
+
     if IS_HIP:
         _backend = (
             backend if backend != "auto" else _auto_select_rope_backend(query, key)
@@ -1348,10 +1356,6 @@ def apply_rope_with_cos_sin_cache_inplace(
                 is_neox=is_neox,
             )
             return
-        if _backend != "native":
-            raise ValueError(
-                f"Unknown backend {backend!r}; expected one of 'auto', 'native', 'aiter'."
-            )
 
     # pass q_rope and k_rope as q and k to perform inplace operation
     _apply_rope_pos_ids_cos_sin_cache(

tests/rocm_tests/test_activation_aiter_hip.py

@@ -40,16 +40,6 @@ def test_silu_and_mul_aiter_vs_ref(dtype, d, num_tokens):
     torch.testing.assert_close(got.float(), ref.float(), rtol=rtol, atol=atol)
 
 
-@requires_aiter
-def test_silu_and_mul_auto_backend_selection():
-    """auto routes to the C++ AITER kernel on supported gfx942/gfx950 devices."""
-    from flashinfer.activation import _auto_select_silu_and_mul_backend
-
-    device = torch.device("cuda:0")
-    x = torch.empty(8, 256, dtype=torch.float16, device=device)
-    assert _auto_select_silu_and_mul_backend(x) == "aiter"
-
-
 @requires_aiter
 def test_silu_and_mul_aiter_with_out_tensor():
     """backend='aiter' writes the correct result into the supplied out= tensor."""

tests/rocm_tests/test_activation_hip.py

@@ -79,3 +79,15 @@ def test_gelu_and_mul(num_tokens, d, dtype):
     out = flashinfer.activation.gelu_and_mul(x)
     rtol, atol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (1e-3, 1e-3)
     torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
+
+
+def test_silu_and_mul_auto_backend_selection():
+    """auto resolves to the in-tree native kernel; AITER is opt-in via backend='aiter'.
+
+    Independent of the aiter package, so it lives in the native test module to run
+    on any HIP build (the contract must hold even when aiter is not installed).
+    """
+    from flashinfer.activation import _auto_select_silu_and_mul_backend
+
+    x = torch.empty(8, 256, dtype=torch.float16, device="cuda")
+    assert _auto_select_silu_and_mul_backend(x) == "native"

tests/rocm_tests/test_rope_aiter_hip.py

@@ -122,21 +122,6 @@ def test_rope_cos_sin_cache_aiter_inplace(is_neox_style, dtype):
     torch.testing.assert_close(key_inplace, key_out, rtol=0, atol=0)
 
 
-def test_rope_auto_backend_selection():
-    """auto routes to AITER on supported devices, but never raises: it falls back
-    to native for mixed query/key dtypes."""
-    from flashinfer.rope import _auto_select_rope_backend
-
-    device = torch.device("cuda:0")
-    q = torch.randn(2048, 128, dtype=torch.float16, device=device)
-    k = torch.randn(2048, 128, dtype=torch.float16, device=device)
-    assert _auto_select_rope_backend(q, k) == "aiter"
-
-    # Mixed q/k dtype must fall back to native, not route into AITER and raise.
-    k_bf16 = torch.randn(2048, 128, dtype=torch.bfloat16, device=device)
-    assert _auto_select_rope_backend(q, k_bf16) == "native"
-
-
 def test_rope_unknown_backend_raises():
     device = torch.device("cuda:0")
     cos_sin_cache = torch.randn(64, 64, dtype=torch.float32, device=device)

tests/rocm_tests/test_rope_hip.py

@@ -1329,6 +1329,23 @@ def _make_gqa_inputs(n):
             )
 
 
+def test_rope_auto_backend_selection():
+    """auto resolves to the in-tree native kernel; AITER is opt-in via backend='aiter'.
+
+    Independent of the aiter package, so it lives in the native test module to run
+    on any HIP build (the contract must hold even when aiter is not installed).
+    """
+    from flashinfer.rope import _auto_select_rope_backend
+
+    q = torch.randn(2048, 128, dtype=torch.float16, device="cuda")
+    k = torch.randn(2048, 128, dtype=torch.float16, device="cuda")
+    assert _auto_select_rope_backend(q, k) == "native"
+
+    # Mixed q/k dtype must also resolve to native (it would be rejected by AITER).
+    k_bf16 = torch.randn(2048, 128, dtype=torch.bfloat16, device="cuda")
+    assert _auto_select_rope_backend(q, k_bf16) == "native"
+
+
 if __name__ == "__main__":
     test_rope(2, 1, 8, 8, 1, 128, "llama", 1.0, False)
     test_rope_pos_ids(2, 1, 8, 8, 1, 128, "llama", 1.0, False, True)