@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 350ca72c04

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 17:45 UTC

[ROCm][AITER] Fix AITER import regression for explicit backend selection (#33749)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Feb 6, 2026 at 09:08 UTC 350ca72c0423f2b094723dfea479a8f0eb7a46d5

5 files changed +262 -66

tests/kernels/attention/test_aiter_flash_attn.py

+87 -20

@@ -5,10 +5,17 @@

5 import pytest

6 import torch

8 -import vllm.v1.attention.backends.rocm_aiter_fa # noqa: F401

8 from vllm.platforms import current_platform

9 from vllm.utils.torch_utils import set_random_seed

11 -from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available

10 +

11 +# Import AITER backend if on ROCm and aiter is available

12 +if current_platform.is_rocm():

13 + from vllm._aiter_ops import is_aiter_found_and_supported

14 +

15 + if is_aiter_found_and_supported():

16 + import aiter

17 +

18 + from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache

20 NUM_HEADS = [(4, 4), (8, 2)]

21 HEAD_SIZES = [128, 256]

@@ -102,8 +109,11 @@ def test_varlen_with_paged_kv(

109 num_blocks: int,

110 q_dtype: torch.dtype | None,

111 ) -> None:

105 - if not is_flash_attn_varlen_func_available():

106 - pytest.skip("flash_attn_varlen_func required to run this test.")

112 + from vllm._aiter_ops import is_aiter_found_and_supported

113 +

114 + if not is_aiter_found_and_supported():

115 + pytest.skip("aiter package required for this test.")

116 +

117 torch.set_default_device("cuda")

118 set_random_seed(0)

119 num_seqs = len(seq_lens)

@@ -129,6 +139,8 @@ def test_varlen_with_paged_kv(

139 cu_seq_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(

140 dim=0, dtype=torch.int32

141 )

142 + # Save kv_lens as list before converting to tensor

143 + kv_lens_list = kv_lens

144 kv_lens = torch.tensor(kv_lens, dtype=torch.int32)

145

146 max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size

@@ -141,33 +153,83 @@ def test_varlen_with_paged_kv(

153 maybe_quantized_query = query

154 maybe_quantized_key_cache = key_cache

155 maybe_quantized_value_cache = value_cache

144 - k_descale = None

145 - v_descale = None

156 + k_scale_tensor = None

157 + v_scale_tensor = None

158 + dequant = False

159 +

160 if q_dtype is not None:

161 # QKV are drawn from N(0, 1): no need for a fp8 scaling factor

162 maybe_quantized_query = query.to(q_dtype)

163 maybe_quantized_key_cache = key_cache.to(q_dtype)

164 maybe_quantized_value_cache = value_cache.to(q_dtype)

151 -

165 + dequant = True

166 scale_shape = (num_seqs, num_kv_heads)

153 - k_descale = torch.ones(scale_shape, dtype=torch.float32)

154 - v_descale = torch.ones(scale_shape, dtype=torch.float32)

167

156 - torch.ops.vllm.flash_attn_varlen_func(

157 - maybe_quantized_query,

158 - maybe_quantized_key_cache,

159 - maybe_quantized_value_cache,

160 - out=output,

168 + # For per-seq-per-head scales (matching AITER backend expectation)

169 + k_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)

170 + v_scale_tensor = torch.ones(scale_shape, dtype=torch.float32)

171 +

172 + # Prepare metadata for cp_mha_gather_cache

173 + # token_to_batch: maps each token to its batch index

174 + token_to_batch = torch.zeros(sum(kv_lens_list), dtype=torch.int32)

175 + seq_starts = torch.zeros(num_seqs, dtype=torch.int32)

176 +

177 + token_idx = 0

178 + for batch_idx, kv_len in enumerate(kv_lens_list):

179 + token_to_batch[token_idx : token_idx + kv_len] = batch_idx

180 + seq_starts[batch_idx] = 0 # Assuming all sequences start at 0 in their blocks

181 + token_idx += kv_len

182 +

183 + # Allocate buffers for gathered KV

184 + total_kv_tokens = sum(kv_lens_list)

185 + gathered_key = torch.empty(

186 + total_kv_tokens, num_kv_heads, head_size, dtype=maybe_quantized_key_cache.dtype

187 + )

188 + gathered_value = torch.empty(

189 + total_kv_tokens,

190 + num_kv_heads,

191 + head_size,

192 + dtype=maybe_quantized_value_cache.dtype,

193 + )

194 +

195 + # Gather paged KV cache into contiguous tensors using triton kernel

196 + cp_mha_gather_cache(

197 + key_cache=maybe_quantized_key_cache,

198 + value_cache=maybe_quantized_value_cache,

199 + key=gathered_key,

200 + value=gathered_value,

201 + block_tables=block_tables,

202 + k_scales=k_scale_tensor

203 + if k_scale_tensor is not None

204 + else torch.ones(1, dtype=torch.float32),

205 + v_scales=v_scale_tensor

206 + if v_scale_tensor is not None

207 + else torch.ones(1, dtype=torch.float32),

208 + cu_seqlens_kv=cu_seq_lens,

209 + token_to_batch=token_to_batch,

210 + seq_starts=seq_starts,

211 + dequant=dequant,

212 + kv_cache_layout="NHD",

213 + total_tokens=total_kv_tokens,

214 + )

215 +

216 + # Call aiter flash attention with gathered KV

217 + aiter.flash_attn_varlen_func(

218 + q=maybe_quantized_query,

219 + k=gathered_key,

220 + v=gathered_value,

221 cu_seqlens_q=cu_query_lens,

222 + cu_seqlens_k=cu_seq_lens,

223 max_seqlen_q=max_query_len,

224 max_seqlen_k=max_kv_len,

225 + min_seqlen_q=1,

226 + dropout_p=0.0,

227 softmax_scale=scale,

165 - alibi_slopes=None,

228 + causal=True,

229 window_size=window_size,

167 - block_table=block_tables,

168 - cu_seqlens_k=cu_seq_lens,

169 - k_scale=k_descale,

170 - v_scale=v_descale,

230 + alibi_slopes=None,

231 + return_lse=False,

232 + out=output,

233 )

234

235 ref_output = ref_paged_attn(

@@ -175,7 +237,7 @@ def test_varlen_with_paged_kv(

237 key_cache=key_cache,

238 value_cache=value_cache,

239 query_lens=query_lens,

178 - kv_lens=kv_lens,

240 + kv_lens=kv_lens_list,

241 block_tables=block_tables,

242 scale=scale,

243 sliding_window=sliding_window,

@@ -189,3 +251,8 @@ def test_varlen_with_paged_kv(

251 torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),

252 f"{torch.max(torch.abs(output - ref_output))}",

253 )

254 +

255 + # Log diff stats for tracking changes

256 + print(f"Max abs diff: {torch.max(torch.abs(output - ref_output))}")

257 + print(f"Mean diff: {torch.mean(torch.abs(output - ref_output))}")

258 + print(f"Min diff: {torch.std(torch.abs(output - ref_output))}")

vllm/_aiter_ops.py

+114 -34

@@ -14,7 +14,10 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (

14 rocm_aiter_sparse_attn_indexer_fake,

15 )

17 -_FP8_DTYPE = current_platform.fp8_dtype()

17 +# fp8_dtype is not cached.

18 +# on ROCm the fp8_dtype always calls is_fp8_fnuz

19 +# which is a host op, so we cache it once here.

20 +FP8_DTYPE = current_platform.fp8_dtype()

23 def is_aiter_found() -> bool:

@@ -31,12 +34,22 @@ IS_AITER_FOUND = is_aiter_found()

36 def is_aiter_found_and_supported() -> bool:

34 - """Check if AITER is available AND enabled via environment variable.

37 + """Check if AITER library is available and platform supports it.

36 - Checks: platform (ROCm), device arch (gfx9), library existence,

37 - and VLLM_ROCM_USE_AITER env variable.

39 + Checks: platform (ROCm), device arch (gfx9), and library existence.

40 + Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled().

41 +

42 + This function determines if aiter CAN be used, not if it SHOULD be used.

43 +

44 + Separation of concerns:

45 + - This function: Can aiter work on this system? (platform + library availability)

46 + - rocm_aiter_ops.is_enabled(): Should aiter be used by default? (adds env var check)

47 + - Backend selection: Can explicitly request aiter regardless of env var

48 +

49 + This allows explicit backend selection via attention_config to work even when

50 + VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery.

51 """

39 - if current_platform.is_rocm() and IS_AITER_FOUND and envs.VLLM_ROCM_USE_AITER:

52 + if current_platform.is_rocm() and IS_AITER_FOUND:

53 from vllm.platforms.rocm import on_gfx9

55 return on_gfx9()

@@ -58,21 +71,6 @@ def if_aiter_supported(func: Callable) -> Callable:

71 return wrapper

61 -# Can't use dtypes.fp8 directly inside an op

62 -# because it returns wrong result on gfx942.

63 -# This is a workaround to get the correct FP8 dtype.

64 -# This might because that the get_gfx() is wrapped as a custom op.

65 -if is_aiter_found_and_supported():

66 - from aiter import dtypes

67 -

68 - AITER_FP8_DTYPE = dtypes.fp8

69 -else:

70 - # Placeholder when AITER is disabled - prevents NameError during module load.

71 - # Note: When AITER is disabled, ops are not registered, so fake implementations

72 - # referencing this variable won't actually be called at runtime.

73 - AITER_FP8_DTYPE = _FP8_DTYPE

74 -

75 -

74 def _rocm_aiter_fused_moe_impl(

75 hidden_states: torch.Tensor,

76 w1: torch.Tensor,

@@ -539,7 +537,7 @@ def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl(

537 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

538 import aiter as rocm_aiter

539

542 - assert quant_dtype in [torch.int8, _FP8_DTYPE]

540 + assert quant_dtype in [torch.int8, FP8_DTYPE]

541

542 y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)

543 out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)

@@ -581,7 +579,7 @@ def _rocm_aiter_rmsnorm_fused_dynamic_quant_impl(

579 ) -> tuple[torch.Tensor, torch.Tensor]:

580 import aiter as rocm_aiter

581

584 - assert quant_dtype in [torch.int8, _FP8_DTYPE]

582 + assert quant_dtype in [torch.int8, FP8_DTYPE]

583

584 y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)

585 out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)

@@ -630,10 +628,10 @@ def _rocm_aiter_per_token_quant_impl(

628 ) -> tuple[torch.Tensor, torch.Tensor]:

629 from aiter.ops.quant import dynamic_per_token_scaled_quant

630

633 - assert quant_dtype in [torch.int8, _FP8_DTYPE]

631 + assert quant_dtype in [torch.int8, FP8_DTYPE]

632

633 out_shape = x.shape

636 - out = torch.empty(x.shape, dtype=_FP8_DTYPE, device=x.device)

634 + out = torch.empty(x.shape, dtype=FP8_DTYPE, device=x.device)

635 if scale is None:

636 scale = torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device)

637 dynamic_per_token_scaled_quant(

@@ -653,7 +651,7 @@ def _rocm_aiter_per_token_quant_fake(

651 ) -> tuple[torch.Tensor, torch.Tensor]:

652 out_shape = x.shape

653 return (

656 - torch.empty(x.shape, dtype=_FP8_DTYPE, device=x.device),

654 + torch.empty(x.shape, dtype=FP8_DTYPE, device=x.device),

655 torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device),

656 )

657

@@ -675,7 +673,7 @@ def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl(

673 None,

674 None,

675 group_size=group_size,

678 - dtype_quant=AITER_FP8_DTYPE,

676 + dtype_quant=FP8_DTYPE,

677 res1=residual,

678 )

679 return (

@@ -695,7 +693,7 @@ def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake(

693 M, N = x.shape

694 scale_shape = (M, (N + group_size - 1) // group_size)

695 return (

698 - torch.empty_like(x, dtype=AITER_FP8_DTYPE, device=x.device),

696 + torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),

697 torch.empty_like(residual, device=residual.device),

698 torch.empty(scale_shape, dtype=torch.float32, device=x.device),

699 )

@@ -717,7 +715,7 @@ def _rocm_aiter_rmsnorm_fp8_group_quant_impl(

715 None,

716 None,

717 group_size=group_size,

720 - dtype_quant=AITER_FP8_DTYPE,

718 + dtype_quant=FP8_DTYPE,

719 res1=None,

720 )

721 return (x_quant, x_quant_scales)

@@ -732,7 +730,7 @@ def _rocm_aiter_rmsnorm_fp8_group_quant_fake(

730 M, N = x.shape

731 scale_shape = (M, (N + group_size - 1) // group_size)

732 return (

735 - torch.empty_like(x, dtype=AITER_FP8_DTYPE, device=x.device),

733 + torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),

734 torch.empty(scale_shape, dtype=torch.float32, device=x.device),

735 )

736

@@ -745,7 +743,7 @@ def _rocm_aiter_group_fp8_quant_impl(

743 from aiter import QuantType, get_hip_quant

744

745 aiter_per1x128_quant = get_hip_quant(QuantType.per_1x128)

748 - return aiter_per1x128_quant(x.contiguous(), quant_dtype=AITER_FP8_DTYPE)

746 + return aiter_per1x128_quant(x.contiguous(), quant_dtype=FP8_DTYPE)

747

748

749 def _rocm_aiter_group_fp8_quant_fake(

@@ -753,7 +751,7 @@ def _rocm_aiter_group_fp8_quant_fake(

751 group_size: int,

752 ) -> tuple[torch.Tensor, torch.Tensor]:

753 M, N = x.shape

756 - x_fp8 = torch.empty((M, N), dtype=AITER_FP8_DTYPE, device=x.device)

754 + x_fp8 = torch.empty((M, N), dtype=FP8_DTYPE, device=x.device)

755 out_bs = torch.empty(

756 (

757 M,

@@ -775,7 +773,7 @@ def _rocm_aiter_act_mul_and_fp8_group_quant_impl(

773 x,

774 activation="silu",

775 group_size=group_size,

778 - dtype_quant=AITER_FP8_DTYPE,

776 + dtype_quant=FP8_DTYPE,

777 )

778

779

@@ -786,7 +784,7 @@ def _rocm_aiter_act_mul_and_fp8_group_quant_fake(

784 M, N = x.shape

785 assert N % 2 == 0

786 N_half = N // 2

789 - x_fp8 = torch.empty((M, N_half), dtype=AITER_FP8_DTYPE, device=x.device)

787 + x_fp8 = torch.empty((M, N_half), dtype=FP8_DTYPE, device=x.device)

788 out_bs = torch.empty(

789 (

790 M,

@@ -986,7 +984,7 @@ class rocm_aiter_ops:

984 @classmethod

985 @if_aiter_supported

986 def is_shuffle_kv_cache_enabled(cls) -> bool:

989 - return cls._AITER_ENABLED and cls._SHUFFLE_KV_CACHE_ENABLED

987 + return cls._SHUFFLE_KV_CACHE_ENABLED

988

989 @classmethod

990 @if_aiter_supported

@@ -1654,5 +1652,87 @@ class rocm_aiter_ops:

1652

1653 return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)

1654

1655 + @staticmethod

1656 + def flash_attn_varlen_func(

1657 + q: torch.Tensor,

1658 + k: torch.Tensor,

1659 + v: torch.Tensor,

1660 + cu_seqlens_q: torch.Tensor,

1661 + cu_seqlens_k: torch.Tensor,

1662 + max_seqlen_q: int,

1663 + max_seqlen_k: int,

1664 + min_seqlen_q: int | None = None,

1665 + dropout_p: float = 0.0,

1666 + softmax_scale: float | None = None,

1667 + causal: bool = False,

1668 + window_size: tuple[int, int] | None = None,

1669 + alibi_slopes: torch.Tensor | None = None,

1670 + return_lse: bool = False,

1671 + out: torch.Tensor | None = None,

1672 + ):

1673 + """

1674 + Flash attention with variable length sequences.

1675 +

1676 + This function is NOT wrapped with @is_aiter_supported decorator

1677 + to allow explicit backend selection via attention_config to work

1678 + even when VLLM_ROCM_USE_AITER=0.

1679 +

1680 + Note: This performs lazy import of aiter.flash_attn_varlen_func

1681 + """

1682 + from aiter import flash_attn_varlen_func

1683 +

1684 + return flash_attn_varlen_func(

1685 + q=q,

1686 + k=k,

1687 + v=v,

1688 + cu_seqlens_q=cu_seqlens_q,

1689 + cu_seqlens_k=cu_seqlens_k,

1690 + max_seqlen_q=max_seqlen_q,

1691 + max_seqlen_k=max_seqlen_k,

1692 + min_seqlen_q=min_seqlen_q,

1693 + dropout_p=dropout_p,

1694 + softmax_scale=softmax_scale,

1695 + causal=causal,

1696 + window_size=window_size,

1697 + alibi_slopes=alibi_slopes,

1698 + return_lse=return_lse,

1699 + out=out,

1700 + )

1701 +

1702 + @staticmethod

1703 + def pa_fwd_asm(

1704 + Q: torch.Tensor,

1705 + K: torch.Tensor,

1706 + V: torch.Tensor,

1707 + block_tables: torch.Tensor,

1708 + context_lens: torch.Tensor,

1709 + block_tables_stride0: int,

1710 + K_QScale: torch.Tensor,

1711 + V_QScale: torch.Tensor,

1712 + out_: torch.Tensor,

1713 + ):

1714 + """

1715 + Paged attention forward pass using assembly kernel.

1716 +

1717 + This function is NOT wrapped with @is_aiter_supported decorator

1718 + to allow explicit backend selection via attention_config to work

1719 + even when VLLM_ROCM_USE_AITER=0.

1720 +

1721 + Note: This performs lazy import of aiter.pa_fwd_asm

1722 + """

1723 + from aiter import pa_fwd_asm

1724 +

1725 + return pa_fwd_asm(

1726 + Q=Q,

1727 + K=K,

1728 + V=V,

1729 + block_tables=block_tables,

1730 + context_lens=context_lens,

1731 + block_tables_stride0=block_tables_stride0,

1732 + K_QScale=K_QScale,

1733 + V_QScale=V_QScale,

1734 + out_=out_,

1735 + )

1736 +

1737

1738 rocm_aiter_ops.register_ops_once()

vllm/v1/attention/backends/fa_utils.py

+45 -1

@@ -8,6 +8,12 @@ from vllm.platforms import current_platform

9 logger = init_logger(__name__)

11 +# Track whether upstream flash-attn is available on ROCm.

12 +# Set during module initialization and never modified afterwards.

13 +# This module-level flag avoids repeated import attempts and ensures

14 +# consistent behavior (similar to IS_AITER_FOUND in _aiter_ops.py).

15 +_ROCM_FLASH_ATTN_AVAILABLE = False

16 +

17 if current_platform.is_cuda():

18 from vllm._custom_ops import reshape_and_cache_flash

19 from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]

@@ -26,6 +32,9 @@ elif current_platform.is_xpu():

32 elif current_platform.is_rocm():

33 try:

34 from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]

35 +

36 + # Mark that upstream flash-attn is available on ROCm

37 + _ROCM_FLASH_ATTN_AVAILABLE = True

38 except ImportError:

40 def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-redef,misc]

@@ -34,6 +43,15 @@ elif current_platform.is_rocm():

43 "to be installed. Please install flash-attn first."

44 )

46 + # ROCm doesn't use scheduler metadata (FA3 feature), provide stub

47 + def get_scheduler_metadata(*args: Any, **kwargs: Any) -> None: # type: ignore[misc]

48 + return None

49 +

50 + # ROCm uses the C++ custom op for reshape_and_cache

51 + from vllm import _custom_ops as ops

52 +

53 + reshape_and_cache_flash = ops.reshape_and_cache_flash

54 +

56 def get_flash_attn_version(requires_alibi: bool = False) -> int | None:

57 # import here to avoid circular dependencies

@@ -128,4 +146,30 @@ def flash_attn_supports_mla():

146

147

148 def is_flash_attn_varlen_func_available() -> bool:

131 - return current_platform.is_cuda() or current_platform.is_xpu()

149 + """Check if flash_attn_varlen_func is available.

150 +

151 + This function determines whether the flash_attn_varlen_func imported at module

152 + level is a working implementation or a stub.

153 +

154 + Platform-specific sources:

155 + - CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func

156 + - XPU: ipex_ops.flash_attn_varlen_func

157 + - ROCm: upstream flash_attn.flash_attn_varlen_func (if available)

158 +

159 + Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py)

160 + which uses rocm_aiter_ops.flash_attn_varlen_func. The condition to use AITER is

161 + handled separately via _aiter_ops.is_aiter_found_and_supported().

162 +

163 + Returns:

164 + bool: True if a working flash_attn_varlen_func implementation is available.

165 + """

166 + if current_platform.is_cuda() or current_platform.is_xpu():

167 + # CUDA and XPU always have flash_attn_varlen_func available

168 + return True

169 +

170 + if current_platform.is_rocm():

171 + # Use the flag set during module import to check if

172 + # upstream flash-attn was successfully imported

173 + return _ROCM_FLASH_ATTN_AVAILABLE

174 +

175 + return False

vllm/v1/attention/backends/rocm_aiter_fa.py

+9 -8

@@ -34,9 +34,6 @@ _CP_TOKENS_PER_ITER_ROCM = 32 * 1024

34 if current_platform.is_rocm():

35 from vllm.triton_utils import tl, triton

37 - if rocm_aiter_ops.is_enabled():

38 - import aiter

39 -

37 def block_size(x, head_dim):

38 return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))

@@ -798,7 +795,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

795 total_tokens=swa_total_tokens,

796 )

797

801 - aiter.flash_attn_varlen_func(

798 + rocm_aiter_ops.flash_attn_varlen_func(

799 q=query,

800 k=key_fetched,

801 v=value_fetched,

@@ -848,7 +845,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

845 v_scale,

846 )

847 return

851 - out, lse = aiter.flash_attn_varlen_func(

848 + out, lse = rocm_aiter_ops.flash_attn_varlen_func(

849 q=query,

850 k=key,

851 v=value,

@@ -895,7 +892,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

892 total_tokens=total_token_per_batch[chunk_idx],

893 )

894

898 - suf_out, suf_lse = aiter.flash_attn_varlen_func(

895 + suf_out, suf_lse = rocm_aiter_ops.flash_attn_varlen_func(

896 q=query,

897 k=key_fetched,

898 v=value_fetched,

@@ -1053,7 +1050,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1050 prefill_key = key[num_decode_tokens + num_extend_tokens :]

1051 prefill_value = value[num_decode_tokens + num_extend_tokens :]

1052

1056 - aiter.flash_attn_varlen_func(

1053 + rocm_aiter_ops.flash_attn_varlen_func(

1054 q=prefill_query,

1055 k=prefill_key,

1056 v=prefill_value,

@@ -1159,7 +1156,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1156 )

1157 new_key_cache = key_cache.view_as(k_cache_template)

1158 new_value_cache = value_cache.view_as(v_cache_template)

1162 - aiter.pa_fwd_asm(

1159 + rocm_aiter_ops.pa_fwd_asm(

1160 Q=query[:num_decode_tokens],

1161 K=new_key_cache,

1162 V=new_value_cache,

@@ -1188,6 +1185,10 @@ class AiterFlashAttentionImpl(AttentionImpl):

1185 device=output.device,

1186 )

1187

1188 + # import so that aiter register the op to the namespace of

1189 + # torch.ops.aiter

1190 + import aiter # noqa: F401

1191 +

1192 torch.ops.aiter.paged_attention_v1(

1193 output[:num_decode_tokens],

1194 workspace_buffer,

vllm/v1/spec_decode/eagle.py

+7 -3

@@ -222,9 +222,13 @@ class SpecDecodeBaseProposer:

222 RocmAttentionMetadata,

223 ]

224 # ROCM_AITER_FA is an optional backend

225 - from vllm._aiter_ops import rocm_aiter_ops

226 -

227 - if rocm_aiter_ops.is_enabled() and find_spec(

225 + # We check is_enabled() here to avoid importing the backend module during

226 + # auto-discovery when VLLM_ROCM_USE_AITER=0, which would trigger aiter

227 + # import and JIT compilation warnings. Explicit backend selection via

228 + # attention_config still works because the backend module is loaded

229 + # directly when selected, not through this auto-discovery path.

230 + # Check if backend module exists to allow explicit selection

231 + if find_spec(

232 AttentionBackendEnum.ROCM_AITER_FA.get_path(include_classname=False)

233 ):

234 from vllm.v1.attention.backends.rocm_aiter_fa import (