[ROCm][AITER] Fix AITER import regression for explicit backend selection

tests/kernels/attention/test_aiter_flash_attn.py

@@ -5,10 +5,17 @@
 import pytest
 import torch
 
-import vllm.v1.attention.backends.rocm_aiter_fa  # noqa: F401
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
-from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
+
+# Import AITER backend if on ROCm and aiter is available
+if current_platform.is_rocm():
+    from vllm._aiter_ops import is_aiter_found_and_supported
+
+    if is_aiter_found_and_supported():
+        import aiter
+
+        from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
 
 NUM_HEADS = [(4, 4), (8, 2)]
 HEAD_SIZES = [128, 256]
@@ -102,8 +109,11 @@ def test_varlen_with_paged_kv(
     num_blocks: int,
     q_dtype: torch.dtype | None,
 ) -> None:
-    if not is_flash_attn_varlen_func_available():
-        pytest.skip("flash_attn_varlen_func required to run this test.")
+    from vllm._aiter_ops import is_aiter_found_and_supported
+
+    if not is_aiter_found_and_supported():
+        pytest.skip("aiter package required for this test.")
+
     torch.set_default_device("cuda")
     set_random_seed(0)
     num_seqs = len(seq_lens)
@@ -129,6 +139,8 @@ def test_varlen_with_paged_kv(
     cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(
         dim=0, dtype=torch.int32
     )
+    # Save kv_lens as list before converting to tensor
+    kv_lens_list = kv_lens
     kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
 
     max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
@@ -141,41 +153,91 @@ def test_varlen_with_paged_kv(
     maybe_quantized_query = query
     maybe_quantized_key_cache = key_cache
     maybe_quantized_value_cache = value_cache
-    k_descale = None
-    v_descale = None
+    k_scale_tensor = None
+    v_scale_tensor = None
+    dequant = False
+
     if q_dtype is not None:
         # QKV are drawn from N(0, 1): no need for a fp8 scaling factor
         maybe_quantized_query = query.to(q_dtype)
         maybe_quantized_key_cache = key_cache.to(q_dtype)
         maybe_quantized_value_cache = value_cache.to(q_dtype)
-
+        dequant = True
         scale_shape = (num_seqs, num_kv_heads)
-        k_descale = torch.ones(scale_shape, dtype=torch.float32)
-        v_descale = torch.ones(scale_shape, dtype=torch.float32)
 
-    torch.ops.vllm.flash_attn_varlen_func(
-        maybe_quantized_query,
-        maybe_quantized_key_cache,
-        maybe_quantized_value_cache,
-        out=output,
+        # For per-seq-per-head scales (matching AITER backend expectation)
+        k_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)
+        v_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)
+
+    # Prepare metadata for cp_mha_gather_cache
+    # token_to_batch: maps each token to its batch index
+    token_to_batch = torch.zeros(sum(kv_lens_list), dtype=torch.int32)
+    seq_starts = torch.zeros(num_seqs, dtype=torch.int32)
+
+    token_idx = 0
+    for batch_idx, kv_len in enumerate(kv_lens_list):
+        token_to_batch[token_idx : token_idx + kv_len] = batch_idx
+        seq_starts[batch_idx] = 0  # Assuming all sequences start at 0 in their blocks
+        token_idx += kv_len
+
+    # Allocate buffers for gathered KV
+    total_kv_tokens = sum(kv_lens_list)
+    gathered_key = torch.empty(
+        total_kv_tokens, num_kv_heads, head_size, dtype=maybe_quantized_key_cache.dtype
+    )
+    gathered_value = torch.empty(
+        total_kv_tokens,
+        num_kv_heads,
+        head_size,
+        dtype=maybe_quantized_value_cache.dtype,
+    )
+
+    # Gather paged KV cache into contiguous tensors using triton kernel
+    cp_mha_gather_cache(
+        key_cache=maybe_quantized_key_cache,
+        value_cache=maybe_quantized_value_cache,
+        key=gathered_key,
+        value=gathered_value,
+        block_tables=block_tables,
+        k_scales=k_scale_tensor
+        if k_scale_tensor is not None
+        else torch.ones(1, dtype=torch.float32),
+        v_scales=v_scale_tensor
+        if v_scale_tensor is not None
+        else torch.ones(1, dtype=torch.float32),
+        cu_seqlens_kv=cu_seq_lens,
+        token_to_batch=token_to_batch,
+        seq_starts=seq_starts,
+        dequant=dequant,
+        kv_cache_layout="NHD",
+        total_tokens=total_kv_tokens,
+    )
+
+    # Call aiter flash attention with gathered KV
+    aiter.flash_attn_varlen_func(
+        q=maybe_quantized_query,
+        k=gathered_key,
+        v=gathered_value,
         cu_seqlens_q=cu_query_lens,
+        cu_seqlens_k=cu_seq_lens,
         max_seqlen_q=max_query_len,
         max_seqlen_k=max_kv_len,
+        min_seqlen_q=1,
+        dropout_p=0.0,
         softmax_scale=scale,
-        alibi_slopes=None,
+        causal=True,
         window_size=window_size,
-        block_table=block_tables,
-        cu_seqlens_k=cu_seq_lens,
-        k_scale=k_descale,
-        v_scale=v_descale,
+        alibi_slopes=None,
+        return_lse=False,
+        out=output,
     )
 
     ref_output = ref_paged_attn(
         query=query,
         key_cache=key_cache,
         value_cache=value_cache,
         query_lens=query_lens,
-        kv_lens=kv_lens,
+        kv_lens=kv_lens_list,
         block_tables=block_tables,
         scale=scale,
         sliding_window=sliding_window,
@@ -189,3 +251,8 @@ def test_varlen_with_paged_kv(
         torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
         f"{torch.max(torch.abs(output - ref_output))}",
     )
+
+    # Log diff stats for tracking changes
+    print(f"Max abs diff: {torch.max(torch.abs(output - ref_output))}")
+    print(f"Mean diff: {torch.mean(torch.abs(output - ref_output))}")
+    print(f"Min diff: {torch.std(torch.abs(output - ref_output))}")

vllm/_aiter_ops.py

@@ -31,12 +31,22 @@ def is_aiter_found() -> bool:
 
 
 def is_aiter_found_and_supported() -> bool:
-    """Check if AITER is available AND enabled via environment variable.
+    """Check if AITER library is available and platform supports it.
 
-    Checks: platform (ROCm), device arch (gfx9), library existence,
-    and VLLM_ROCM_USE_AITER env variable.
+    Checks: platform (ROCm), device arch (gfx9), and library existence.
+    Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled().
+
+    This function determines if aiter CAN be used, not if it SHOULD be used.
+
+    Separation of concerns:
+    - This function: Can aiter work on this system? (platform + library availability)
+    - rocm_aiter_ops.is_enabled(): Should aiter be used by default? (adds env var check)
+    - Backend selection: Can explicitly request aiter regardless of env var
+
+    This allows explicit backend selection via attention_config to work even when
+    VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery.
     """
-    if current_platform.is_rocm() and IS_AITER_FOUND and envs.VLLM_ROCM_USE_AITER:
+    if current_platform.is_rocm() and IS_AITER_FOUND:
         from vllm.platforms.rocm import on_gfx9
 
         return on_gfx9()
@@ -62,14 +72,23 @@ def wrapper(*args, **kwargs):
 # because it returns wrong result on gfx942.
 # This is a workaround to get the correct FP8 dtype.
 # This might because that the get_gfx() is wrapped as a custom op.
-if is_aiter_found_and_supported():
+#
+# Import strategy to avoid unwanted JIT warnings:
+# - Only import aiter dtypes at module load if VLLM_ROCM_USE_AITER=1 (env var set)
+# - This prevents JIT warnings during backend auto-discovery when aiter is not preferred
+# - Explicit backend selection (via attention_config) still works because:
+#   1. Backend modules (rocm_aiter_fa.py) import aiter directly when loaded
+#   2. Individual op implementations import aiter locally when called
+#   3. This module's ops are only called when an aiter backend is actually in use
+if is_aiter_found_and_supported() and envs.VLLM_ROCM_USE_AITER:
     from aiter import dtypes
 
     AITER_FP8_DTYPE = dtypes.fp8
 else:
-    # Placeholder when AITER is disabled - prevents NameError during module load.
-    # Note: When AITER is disabled, ops are not registered, so fake implementations
-    # referencing this variable won't actually be called at runtime.
+    # Placeholder when AITER is not the default - prevents NameError during module load.
+    # Note: This fallback is used for fake implementations and type checking.
+    # If an AITER backend is explicitly selected (even with env var=0),
+    # the backend module will import aiter directly (rocm_aiter_fa.py line 35).
     AITER_FP8_DTYPE = _FP8_DTYPE
 
 
@@ -1654,5 +1673,87 @@ def shuffle_weights(
 
         return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)
 
+    @staticmethod
+    def flash_attn_varlen_func(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        min_seqlen_q: int | None = None,
+        dropout_p: float = 0.0,
+        softmax_scale: float | None = None,
+        causal: bool = False,
+        window_size: tuple[int, int] | None = None,
+        alibi_slopes: torch.Tensor | None = None,
+        return_lse: bool = False,
+        out: torch.Tensor | None = None,
+    ):
+        """
+        Flash attention with variable length sequences.
+
+        This function is NOT wrapped with @is_aiter_supported decorator
+        to allow explicit backend selection via attention_config to work
+        even when VLLM_ROCM_USE_AITER=0.
+
+        Note: This performs lazy import of aiter.flash_attn_varlen_func
+        """
+        from aiter import flash_attn_varlen_func
+
+        return flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            min_seqlen_q=min_seqlen_q,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            alibi_slopes=alibi_slopes,
+            return_lse=return_lse,
+            out=out,
+        )
+
+    @staticmethod
+    def pa_fwd_asm(
+        Q: torch.Tensor,
+        K: torch.Tensor,
+        V: torch.Tensor,
+        block_tables: torch.Tensor,
+        context_lens: torch.Tensor,
+        block_tables_stride0: int,
+        K_QScale: torch.Tensor,
+        V_QScale: torch.Tensor,
+        out_: torch.Tensor,
+    ):
+        """
+        Paged attention forward pass using assembly kernel.
+
+        This function is NOT wrapped with @is_aiter_supported decorator
+        to allow explicit backend selection via attention_config to work
+        even when VLLM_ROCM_USE_AITER=0.
+
+        Note: This performs lazy import of aiter.pa_fwd_asm
+        """
+        from aiter import pa_fwd_asm
+
+        return pa_fwd_asm(
+            Q=Q,
+            K=K,
+            V=V,
+            block_tables=block_tables,
+            context_lens=context_lens,
+            block_tables_stride0=block_tables_stride0,
+            K_QScale=K_QScale,
+            V_QScale=V_QScale,
+            out_=out_,
+        )
+
 
 rocm_aiter_ops.register_ops_once()

vllm/v1/attention/backends/fa_utils.py

@@ -34,6 +34,15 @@ def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any:  # type: ignore[no
                 "to be installed. Please install flash-attn first."
             )
 
+    # ROCm doesn't use scheduler metadata (FA3 feature), provide stub
+    def get_scheduler_metadata(*args: Any, **kwargs: Any) -> None:  # type: ignore[misc]
+        return None
+
+    # ROCm uses the C++ custom op for reshape_and_cache
+    from vllm import _custom_ops as ops
+
+    reshape_and_cache_flash = ops.reshape_and_cache_flash
+
 
 def get_flash_attn_version(requires_alibi: bool = False) -> int | None:
     # import here to avoid circular dependencies
@@ -128,4 +137,19 @@ def flash_attn_supports_mla():
 
 
 def is_flash_attn_varlen_func_available() -> bool:
-    return current_platform.is_cuda() or current_platform.is_xpu()
+    if current_platform.is_cuda() or current_platform.is_xpu():
+        return True
+
+    # On ROCm, flash_attn_varlen_func is provided by the aiter package
+    # not vllm_flash_attn extension. This function checks CAPABILITY
+    # (can it work?) not PREFERENCE (should it be used by default?).
+    # Tests and other code use this to determine if the functionality
+    # exists on the system, regardless of VLLM_ROCM_USE_AITER.
+    # We use is_aiter_found_and_supported() instead of importing aiter directly
+    # to avoid triggering JIT compilation warnings during the availability check.
+    if current_platform.is_rocm():
+        from vllm._aiter_ops import is_aiter_found_and_supported
+
+        return is_aiter_found_and_supported()
+
+    return False

vllm/v1/attention/backends/rocm_aiter_fa.py

@@ -34,9 +34,6 @@
 if current_platform.is_rocm():
     from vllm.triton_utils import tl, triton
 
-    if rocm_aiter_ops.is_enabled():
-        import aiter
-
     def block_size(x, head_dim):
         return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))
 
@@ -798,7 +795,7 @@ def extend_for_sliding_window(
             total_tokens=swa_total_tokens,
         )
 
-        aiter.flash_attn_varlen_func(
+        rocm_aiter_ops.flash_attn_varlen_func(
             q=query,
             k=key_fetched,
             v=value_fetched,
@@ -848,7 +845,7 @@ def extend_forward(
                 v_scale,
             )
             return
-        out, lse = aiter.flash_attn_varlen_func(
+        out, lse = rocm_aiter_ops.flash_attn_varlen_func(
             q=query,
             k=key,
             v=value,
@@ -895,7 +892,7 @@ def extend_forward(
                 total_tokens=total_token_per_batch[chunk_idx],
             )
 
-            suf_out, suf_lse = aiter.flash_attn_varlen_func(
+            suf_out, suf_lse = rocm_aiter_ops.flash_attn_varlen_func(
                 q=query,
                 k=key_fetched,
                 v=value_fetched,
@@ -1053,7 +1050,7 @@ def forward(
                 prefill_key = key[num_decode_tokens + num_extend_tokens :]
                 prefill_value = value[num_decode_tokens + num_extend_tokens :]
 
-                aiter.flash_attn_varlen_func(
+                rocm_aiter_ops.flash_attn_varlen_func(
                     q=prefill_query,
                     k=prefill_key,
                     v=prefill_value,
@@ -1159,7 +1156,7 @@ def forward(
                     )
                     new_key_cache = key_cache.view_as(k_cache_template)
                     new_value_cache = value_cache.view_as(v_cache_template)
-                    aiter.pa_fwd_asm(
+                    rocm_aiter_ops.pa_fwd_asm(
                         Q=query[:num_decode_tokens],
                         K=new_key_cache,
                         V=new_value_cache,