From aac306ea7d15eefa17aa088d10fc586b8e439bd8 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 16 Jun 2026 22:33:42 +0000 Subject: [PATCH 1/2] feat(hip,aiter): default silu_and_mul and rope auto backend to native Experimentation found the in-tree HIP native kernels to be the better default for silu_and_mul and the cos/sin-cache rope, so backend="auto" now resolves to native for both ops instead of routing to AITER on gfx942/gfx950. The C++ AITER path stays reachable via an explicit backend="aiter". Move the auto-selection tests to the native test modules so the native-default contract is exercised on any HIP build, independent of whether the aiter package is installed. Also reconcile the inplace rope docstring, which incorrectly claimed AITER consumes cos/sin in the query dtype (the C++ shim passes them as float32). Co-Authored-By: Claude Opus 4.7 --- README.md | 20 +++++++---- flashinfer/activation.py | 16 ++++----- flashinfer/rope.py | 34 ++++++------------- tests/rocm_tests/test_activation_aiter_hip.py | 10 ------ tests/rocm_tests/test_activation_hip.py | 12 +++++++ tests/rocm_tests/test_rope_aiter_hip.py | 15 -------- tests/rocm_tests/test_rope_hip.py | 17 ++++++++++ 7 files changed, 59 insertions(+), 65 deletions(-) diff --git a/README.md b/README.md index 46d3049695..1201942203 100644 --- a/README.md +++ b/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. diff --git a/flashinfer/activation.py b/flashinfer/activation.py index de55ba0dea..898449e3eb 100644 --- a/flashinfer/activation.py +++ b/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( `_ 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. diff --git a/flashinfer/rope.py b/flashinfer/rope.py index bfc4f08b99..3a5b3f1e58 100644 --- a/flashinfer/rope.py +++ b/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 @@ -1314,16 +1306,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 ---- diff --git a/tests/rocm_tests/test_activation_aiter_hip.py b/tests/rocm_tests/test_activation_aiter_hip.py index 5e57478b5d..fa13656f80 100644 --- a/tests/rocm_tests/test_activation_aiter_hip.py +++ b/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.""" diff --git a/tests/rocm_tests/test_activation_hip.py b/tests/rocm_tests/test_activation_hip.py index bdafabe84c..926c087350 100644 --- a/tests/rocm_tests/test_activation_hip.py +++ b/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" diff --git a/tests/rocm_tests/test_rope_aiter_hip.py b/tests/rocm_tests/test_rope_aiter_hip.py index da46b424e0..7465e5d96b 100644 --- a/tests/rocm_tests/test_rope_aiter_hip.py +++ b/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) diff --git a/tests/rocm_tests/test_rope_hip.py b/tests/rocm_tests/test_rope_hip.py index 1a61ae68c9..dcdc1cd26a 100644 --- a/tests/rocm_tests/test_rope_hip.py +++ b/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) From e07eef560ddd5d6319b321939206b5d9ab7f8b3b Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 17 Jun 2026 02:17:53 +0000 Subject: [PATCH 2/2] fix(rope): validate backend on all platforms, not just under IS_HIP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move the backend-string check and explicit-aiter validation to the top of apply_rope_with_cos_sin_cache and its inplace variant, matching silu_and_mul. Previously these checks lived inside `if IS_HIP:`, so on non-HIP platforms backend="aiter" or an unknown backend string silently fell through to the native kernel instead of raising — inconsistent with the documented contract. Addresses Copilot review feedback on PR #259. Co-Authored-By: Claude Opus 4.7 --- flashinfer/rope.py | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/flashinfer/rope.py b/flashinfer/rope.py index 3a5b3f1e58..1f3b7c3c3a 100644 --- a/flashinfer/rope.py +++ b/flashinfer/rope.py @@ -1233,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) @@ -1252,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), @@ -1320,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) @@ -1336,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(