@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f1e8fd27ae

[ROCm] Add AITER FP8 ViT encoder attention (#49937)

Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com>

yinfengLiu committed Jul 30, 2026 at 14:22 UTC f1e8fd27aea5563172a9b2559c42188419a100b2

8 files changed +563 -45

benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py new

+124

@@ -0,0 +1,124 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +"""Benchmark complete AITER BF16 and FP8 ViT attention calls on ROCm.

5 +

6 +The FP8 timing includes Q/K/V quantization and attention. Dynamic scales are

7 +calibrated once per input before timing so the measured path matches serving

8 +with a static scale file.

9 +

10 +Example:

11 + python benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py \

12 + --seq-lens 2304 4096 8192 16384 --head-dim 72

13 +"""

14 +

15 +from types import SimpleNamespace

16 +

17 +import torch

18 +

19 +from vllm.config import VllmConfig, set_current_vllm_config

20 +from vllm.config.multimodal import MultiModalConfig

21 +from vllm.model_executor.layers.attention.mm_encoder_attention import (

22 + MMEncoderAttention,

23 +)

24 +from vllm.platforms import current_platform

25 +from vllm.triton_utils import triton

26 +from vllm.utils.argparse_utils import FlexibleArgumentParser

27 +from vllm.v1.attention.backends.registry import AttentionBackendEnum

28 +

29 +

30 +def make_attention(num_heads: int, head_dim: int, fp8: bool) -> MMEncoderAttention:

31 + mm_config = MultiModalConfig(

32 + mm_encoder_attn_backend=AttentionBackendEnum.ROCM_AITER_FA,

33 + mm_encoder_attn_dtype="fp8" if fp8 else None,

34 + )

35 + vllm_config = VllmConfig()

36 + vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)

37 + with set_current_vllm_config(vllm_config):

38 + return MMEncoderAttention(num_heads, head_dim).to("cuda")

39 +

40 +

41 +def bench(

42 + seq_lens: list[int],

43 + num_heads: int,

44 + head_dim: int,

45 + warmup_ms: int,

46 + repeat_ms: int,

47 +) -> None:

48 + if not current_platform.is_rocm():

49 + raise RuntimeError("This benchmark requires ROCm and AITER.")

50 +

51 + old_dtype = torch.get_default_dtype()

52 + torch.set_default_dtype(torch.bfloat16)

53 + try:

54 + fp8_attention = make_attention(num_heads, head_dim, fp8=True)

55 + bf16_attention = make_attention(num_heads, head_dim, fp8=False)

56 + finally:

57 + torch.set_default_dtype(old_dtype)

58 +

59 + print(

60 + f"{'seq_len':>8} {'BF16 ms':>12} {'FP8 ms':>12} "

61 + f"{'speedup':>10} {'FP8/BF16':>12}"

62 + )

63 + print("-" * 60)

64 + for seq_len in seq_lens:

65 + torch.manual_seed(0)

66 + qkv = torch.randn(

67 + 1,

68 + seq_len,

69 + 3,

70 + num_heads,

71 + head_dim,

72 + device="cuda",

73 + dtype=torch.bfloat16,

74 + )

75 + q, k, v = qkv.unbind(dim=2)

76 + cu_seqlens = torch.tensor([0, seq_len], device="cuda", dtype=torch.int32)

77 + max_seqlen = torch.tensor(seq_len, device="cuda", dtype=torch.int32)

78 +

79 + # Calibrate per-tensor scales once, then benchmark the static-scale path.

80 + fp8_attention._fp8_dynamic_scale = True

81 + fp8_attention._forward_aiter_fp8(q, k, v, cu_seqlens, max_seqlen)

82 + fp8_attention._fp8_dynamic_scale = False

83 +

84 + bf16_ms = triton.testing.do_bench(

85 + lambda q=q, k=k, v=v, cu=cu_seqlens, ms=max_seqlen: (

86 + bf16_attention._forward_fa(q, k, v, cu, ms)

87 + ),

88 + warmup=warmup_ms,

89 + rep=repeat_ms,

90 + )

91 + fp8_ms = triton.testing.do_bench(

92 + lambda q=q, k=k, v=v, cu=cu_seqlens, ms=max_seqlen: (

93 + fp8_attention._forward_aiter_fp8(q, k, v, cu, ms)

94 + ),

95 + warmup=warmup_ms,

96 + rep=repeat_ms,

97 + )

98 + speedup = bf16_ms / fp8_ms

99 + ratio = fp8_ms / bf16_ms

100 + print(

101 + f"{seq_len:>8} {bf16_ms:>12.3f} {fp8_ms:>12.3f} "

102 + f"{speedup:>9.2f}x {ratio:>12.3f}"

103 + )

104 +

105 +

106 +if __name__ == "__main__":

107 + parser = FlexibleArgumentParser(

108 + description="Benchmark AITER BF16 vs FP8 ViT attention."

109 + )

110 + parser.add_argument(

111 + "--seq-lens", type=int, nargs="+", default=[2304, 4096, 8192, 16384]

112 + )

113 + parser.add_argument("--num-heads", type=int, default=16)

114 + parser.add_argument("--head-dim", type=int, default=72)

115 + parser.add_argument("--warmup-ms", type=int, default=100)

116 + parser.add_argument("--repeat-ms", type=int, default=500)

117 + args = parser.parse_args()

118 + bench(

119 + args.seq_lens,

120 + args.num_heads,

121 + args.head_dim,

122 + args.warmup_ms,

123 + args.repeat_ms,

124 + )

docs/features/quantization/fp8_vit_attn.md

+27 -8

@@ -3,9 +3,9 @@

3 For visual understanding workloads with large images (e.g. QHD, 4K) and relatively

4 short text prompts/generation, the ViT encoder attention can become a significant

5 bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM

6 -supports optional FP8 quantization for the ViT encoder attention via the

7 -FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the

8 -cuDNN attention call.

6 +supports optional FP8 quantization for ViT encoder attention via FlashInfer

7 +cuDNN on NVIDIA GPUs and AITER on AMD GPUs. Q/K/V are quantized on-the-fly to

8 +FP8 before the attention call.

10 !!! note

11 - Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,

@@ -15,20 +15,29 @@ cuDNN attention call.

15 requests. Smaller images may see no speedup due to quantization overhead

16 (3 quantization kernel launches + un-padding).

17 - FP8 tensor-core speedup is more pronounced on GB300 than GB200.

18 + - On ROCm, packed variable-length image and video batches are supported by

19 + AITER's native varlen FP8 attention.

21 ## Requirements

21 -- FlashInfer cuDNN backend with cuDNN >= 9.17.1.

23 +- NVIDIA: FlashInfer cuDNN backend with cuDNN >= 9.17.1.

24 +- AMD: AITER with `flash_attn_varlen_fp8_pertensor_func` support on gfx942

25 + (MI300 series) or gfx950 (MI350 series).

27 ## Usage

25 -Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together

26 -with `--mm-encoder-attn-backend FLASHINFER`:

29 +Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` and selecting

30 +the backend for the current platform:

32 ```bash

33 vllm serve $MODEL \

34 --mm-encoder-attn-backend FLASHINFER \

35 --mm-encoder-attn-dtype fp8

36 +

37 +# AMD ROCm

38 +vllm serve $MODEL \

39 + --mm-encoder-attn-backend ROCM_AITER_FA \

40 + --mm-encoder-attn-dtype fp8

41 ```

43 By default (no scale file), **dynamic scaling** is used: a 16-entry circular

@@ -45,7 +54,7 @@ reuse them to avoid the dynamic overhead:

54 # Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,

55 # then dumps the learned scales to JSON).

56 vllm bench mm-processor \

48 - --model $MODEL --mm-encoder-attn-backend FLASHINFER \

57 + --model $MODEL --mm-encoder-attn-backend $MM_ATTN_BACKEND \

58 --mm-encoder-attn-dtype fp8 \

59 --mm-encoder-fp8-scale-save-path /path/to/scales.json \

60 --dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \

@@ -53,7 +62,7 @@ vllm bench mm-processor \

63 # Step 2: serve with static scales (no dynamic overhead).

64 vllm serve $MODEL \

56 - --mm-encoder-attn-backend FLASHINFER \

65 + --mm-encoder-attn-backend $MM_ATTN_BACKEND \

66 --mm-encoder-attn-dtype fp8 \

67 --mm-encoder-fp8-scale-path /path/to/scales.json

68 ```

@@ -94,6 +103,16 @@ Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.

103

104 Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.

105

106 +**Complete AITER attention call on MI300X** (BF16 input, 16 heads,

107 +head_dim=72; FP8 includes Q/K/V quantization):

108 +

110 +| --------------- | ---------- | --------- | ------- |

111 +| 2304 | 0.467 ms | 0.337 ms | **1.38x** |

112 +| 4096 | 0.812 ms | 0.764 ms | **1.06x** |

113 +| 8192 | 2.555 ms | 2.364 ms | **1.08x** |

114 +| 16384 | 9.769 ms | 8.655 ms | **1.13x** |

115 +

116 ## Accuracy

117

118 ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated

tests/kernels/attention/test_mha_attn.py

+135 -1

@@ -20,7 +20,7 @@ from vllm.platforms import current_platform

20 from vllm.platforms.cpu import CpuPlatform

21 from vllm.platforms.cuda import CudaPlatform

22 from vllm.platforms.interface import DeviceCapability

23 -from vllm.platforms.rocm import RocmPlatform

23 +from vllm.platforms.rocm import RocmPlatform, on_mi3xx

24 from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed

25 from vllm.v1.attention.backends.registry import AttentionBackendEnum

26 from vllm.v1.attention.selector import _cached_get_attn_backend

@@ -346,3 +346,137 @@ def test_mha_attn_varlen_forward_flashinfer(

346 torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)

347 finally:

348 vllm_config.model_config = old_model_config

349 +

350 +

351 +@pytest.mark.skipif(

352 + not current_platform.is_rocm() or not on_mi3xx(),

353 + reason="AITER FP8 attention requires MI300/MI350",

354 +)

355 +@pytest.mark.parametrize("var_seq_len", [[128, 193], [256, 384, 511]])

356 +@pytest.mark.parametrize("head_size", [64, 72, 80])

357 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half])

358 +def test_mha_attn_varlen_forward_aiter_fp8(

359 + default_vllm_config,

360 + var_seq_len: list[int],

361 + head_size: int,

362 + dtype: torch.dtype,

363 +):

364 + """Compare packed AITER FP8 ViT attention with a BF16/FP16 reference."""

365 + aiter = pytest.importorskip("aiter")

366 + if not hasattr(aiter, "flash_attn_varlen_fp8_pertensor_func"):

367 + pytest.skip("Installed AITER does not provide varlen FP8 attention")

368 +

369 + num_heads = 4

370 + set_random_seed(0)

371 + torch.set_default_device("cuda")

372 + torch.set_default_dtype(dtype)

373 +

374 + vllm_config = get_current_vllm_config()

375 + old_model_config = getattr(vllm_config, "model_config", None)

376 + minimal_model_config = type(

377 + "MinimalModelConfig",

378 + (),

379 + {

380 + "multimodal_config": MultiModalConfig(

381 + mm_encoder_attn_backend=AttentionBackendEnum.ROCM_AITER_FA,

382 + mm_encoder_attn_dtype="fp8",

383 + ),

384 + },

385 + )()

386 + vllm_config.model_config = minimal_model_config

387 + try:

388 + total_len = sum(var_seq_len)

389 + # Keep Q/K/V as non-contiguous views, matching interleaved ViT QKV.

390 + qkv = torch.randn(1, total_len, 3, num_heads, head_size)

391 + q, k, v = qkv.unbind(dim=2)

392 + cu_seqlens = torch.tensor(

393 + [0] + list(itertools.accumulate(var_seq_len)), dtype=torch.int32

394 + )

395 + max_seqlen = torch.tensor(max(var_seq_len), dtype=torch.int32)

396 + scale = 1.0 / head_size**0.5

397 +

398 + attn = MMEncoderAttention(num_heads, head_size, scale=scale)

399 + assert attn.attn_backend == AttentionBackendEnum.ROCM_AITER_FA

400 + assert attn.fp8_enabled

401 + output = attn(

402 + q,

403 + k,

404 + v,

405 + cu_seqlens=cu_seqlens,

406 + max_seqlen=max_seqlen,

407 + )

408 +

409 + ref_output = torch.cat(

410 + [

411 + ref_attention(q_i, k_i, v_i, scale=scale)

412 + for q_i, k_i, v_i in zip(

413 + torch.split(q, var_seq_len, dim=1),

414 + torch.split(k, var_seq_len, dim=1),

415 + torch.split(v, var_seq_len, dim=1),

416 + )

417 + ],

418 + dim=1,

419 + )

420 + diff = (output.float() - ref_output.float()).abs()

421 + cosine = torch.nn.functional.cosine_similarity(

422 + output.float().flatten(),

423 + ref_output.float().flatten(),

424 + dim=0,

425 + )

426 +

427 + assert output.shape == ref_output.shape

428 + assert output.dtype == dtype

429 + assert cosine.item() >= 0.99

430 + assert diff.mean().item() <= 0.03

431 + assert diff.max().item() <= 0.30

432 + finally:

433 + vllm_config.model_config = old_model_config

434 +

435 +

436 +@pytest.mark.skipif(

437 + not current_platform.is_rocm(), reason="AITER FP8 attention requires ROCm"

438 +)

439 +def test_mha_attn_aiter_fp8_rejects_unsupported_arch(default_vllm_config):

440 + vllm_config = get_current_vllm_config()

441 + old_model_config = getattr(vllm_config, "model_config", None)

442 + vllm_config.model_config = type(

443 + "MinimalModelConfig",

444 + (),

445 + {

446 + "multimodal_config": MultiModalConfig(

447 + mm_encoder_attn_backend=AttentionBackendEnum.ROCM_AITER_FA,

448 + mm_encoder_attn_dtype="fp8",

449 + ),

450 + },

451 + )()

452 + try:

453 + with (

454 + patch("vllm.platforms.rocm.on_mi3xx", return_value=False),

455 + pytest.raises(ValueError, match="gfx942 or gfx950"),

456 + ):

457 + MMEncoderAttention(4, 64)

458 + finally:

459 + vllm_config.model_config = old_model_config

460 +

461 +

462 +@pytest.mark.skipif(

463 + not current_platform.is_rocm(), reason="AITER FP8 attention requires ROCm"

464 +)

465 +def test_mha_attn_fp8_rejects_wrong_backend(default_vllm_config):

466 + vllm_config = get_current_vllm_config()

467 + old_model_config = getattr(vllm_config, "model_config", None)

468 + vllm_config.model_config = type(

469 + "MinimalModelConfig",

470 + (),

471 + {

472 + "multimodal_config": MultiModalConfig(

473 + mm_encoder_attn_backend=AttentionBackendEnum.FLASH_ATTN,

474 + mm_encoder_attn_dtype="fp8",

475 + ),

476 + },

477 + )()

478 + try:

479 + with pytest.raises(ValueError, match="requires either"):

480 + MMEncoderAttention(4, 64)

481 + finally:

482 + vllm_config.model_config = old_model_config

tests/kernels/core/test_vit_fp8_scaling.py

+44 -3

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.attention.mm_encoder_attention import (

14 _FP8_AMAX_HISTORY_LEN,

15 _FP8_MAX,

16 )

17 +from vllm.platforms import current_platform

18 from vllm.utils.flashinfer import (

19 is_flashinfer_cudnn_fp8_prefill_attn_supported,

20 )

@@ -29,8 +30,17 @@ NUM_HEADS = 16

30 HEAD_DIM = 72

33 +def _is_mi3xx() -> bool:

34 + if not current_platform.is_rocm():

35 + return False

36 +

37 + from vllm.platforms.rocm import on_mi3xx

38 +

39 + return on_mi3xx()

40 +

41 +

42 @contextlib.contextmanager

33 -def _build_attention(mm_config):

43 +def _build_attention(mm_config, attn_backend=None):

44 """Yield an MMEncoderAttention with the given multimodal config.

46 The VllmConfig context stays active while the test runs so that

@@ -44,7 +54,13 @@ def _build_attention(mm_config):

54 )

55 from vllm.v1.attention.backends.registry import AttentionBackendEnum

47 - if not is_flashinfer_cudnn_fp8_prefill_attn_supported():

57 + if attn_backend is None:

58 + attn_backend = AttentionBackendEnum.FLASHINFER

59 +

60 + if (

61 + attn_backend == AttentionBackendEnum.FLASHINFER

62 + and not is_flashinfer_cudnn_fp8_prefill_attn_supported()

63 + ):

64 yield None

65 return

@@ -56,7 +72,7 @@ def _build_attention(mm_config):

72 patch(

73 "vllm.model_executor.layers.attention.mm_encoder_attention"

74 ".get_vit_attn_backend",

59 - return_value=AttentionBackendEnum.FLASHINFER,

75 + return_value=attn_backend,

76 ),

77 ):

78 attn = MMEncoderAttention(

@@ -169,6 +185,31 @@ def test_static_scales_loaded(_make_static_attention) -> None:

185 assert not hasattr(attn, "_fp8_q_amax")

186

187

188 +@pytest.mark.skipif(

189 + not _is_mi3xx(),

190 + reason="AITER FP8 attention requires MI300/MI350",

191 +)

192 +def test_aiter_static_scales_loaded(tmp_path) -> None:

193 + """Verify AITER reuses the existing static FP8 scale loading path."""

194 + from vllm.config.multimodal import MultiModalConfig

195 + from vllm.v1.attention.backends.registry import AttentionBackendEnum

196 +

197 + scale_file = tmp_path / "aiter_scales.json"

198 + scale_file.write_text(json.dumps({LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0}}))

199 + mm_config = MultiModalConfig(

200 + mm_encoder_attn_dtype="fp8",

201 + mm_encoder_fp8_scale_path=str(scale_file),

202 + )

203 +

204 + with _build_attention(mm_config, AttentionBackendEnum.ROCM_AITER_FA) as attn:

205 + assert attn is not None

206 + assert attn.fp8_enabled

207 + assert not attn._fp8_dynamic_scale

208 + assert attn._fp8_q_scale.item() == 224.0

209 + assert attn._fp8_k_scale.item() == 198.0

210 + assert attn._fp8_v_scale.item() == 210.0

211 +

212 +

213 def test_static_scales_missing_layer(tmp_path) -> None:

214 """Verify error when requested layer is not in the scale file."""

215 from vllm.config import VllmConfig, set_current_vllm_config

vllm/_aiter_ops.py

+95

@@ -1404,6 +1404,65 @@ def _triton_rotary_embedding_fake(

1404 return

1405

1406

1407 +def _rocm_aiter_fp8_attn_impl(

1408 + q: torch.Tensor,

1409 + k: torch.Tensor,

1410 + v: torch.Tensor,

1411 + q_descale: torch.Tensor,

1412 + k_descale: torch.Tensor,

1413 + v_descale: torch.Tensor,

1414 + batch_size: int,

1415 + output_dtype: torch.dtype,

1416 + scale: float | None = None,

1417 + cu_seqlens: torch.Tensor | None = None,

1418 + max_seqlen: torch.Tensor | None = None,

1419 +) -> torch.Tensor:

1420 + """Run AITER FP8 attention for fixed or packed inputs."""

1421 + from aiter import flash_attn_varlen_fp8_pertensor_func

1422 +

1423 + q_len = q.size(1)

1424 + if cu_seqlens is None:

1425 + cu_seqlens = torch.arange(

1426 + 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device

1427 + )

1428 + max_seqlen_value = q_len if max_seqlen is None else max_seqlen.item()

1429 +

1430 + q, k, v = (x.flatten(0, 1) for x in (q, k, v))

1431 + output = flash_attn_varlen_fp8_pertensor_func(

1432 + q,

1433 + k,

1434 + v,

1435 + q_descale=q_descale,

1436 + k_descale=k_descale,

1437 + v_descale=v_descale,

1438 + cu_seqlens_q=cu_seqlens,

1439 + cu_seqlens_k=cu_seqlens,

1440 + max_seqlen_q=max_seqlen_value,

1441 + max_seqlen_k=max_seqlen_value,

1442 + causal=False,

1443 + softmax_scale=scale,

1444 + )

1445 + return output.to(output_dtype).reshape(batch_size, q_len, *output.shape[1:])

1446 +

1447 +

1448 +def _rocm_aiter_fp8_attn_fake(

1449 + q: torch.Tensor,

1450 + k: torch.Tensor,

1451 + v: torch.Tensor,

1452 + q_descale: torch.Tensor,

1453 + k_descale: torch.Tensor,

1454 + v_descale: torch.Tensor,

1455 + batch_size: int,

1456 + output_dtype: torch.dtype,

1457 + scale: float | None = None,

1458 + cu_seqlens: torch.Tensor | None = None,

1459 + max_seqlen: torch.Tensor | None = None,

1460 +) -> torch.Tensor:

1461 + return torch.empty(

1462 + (*q.shape[:-1], v.shape[-1]), device=q.device, dtype=output_dtype

1463 + )

1464 +

1465 +

1466 # Global flag to ensure ops are registered only once

1467 _OPS_REGISTERED = False

1468

@@ -1959,6 +2018,14 @@ class rocm_aiter_ops:

2018 dispatch_key=current_platform.dispatch_key,

2019 )

2020

2021 + direct_register_custom_op(

2022 + op_name="aiter_fp8_attn_wrapper",

2023 + op_func=_rocm_aiter_fp8_attn_impl,

2024 + mutates_args=[],

2025 + fake_impl=_rocm_aiter_fp8_attn_fake,

2026 + dispatch_key=current_platform.dispatch_key,

2027 + )

2028 +

2029 direct_register_custom_op(

2030 op_name="rocm_aiter_gemm_a8wfp4",

2031 op_func=_rocm_aiter_gemm_a8wfp4_impl,

@@ -2830,6 +2897,34 @@ class rocm_aiter_ops:

2897 sink_ptr=sink_ptr,

2898 )

2899

2900 + @staticmethod

2901 + def fp8_attn_wrapper(

2902 + q: torch.Tensor,

2903 + k: torch.Tensor,

2904 + v: torch.Tensor,

2905 + q_descale: torch.Tensor,

2906 + k_descale: torch.Tensor,

2907 + v_descale: torch.Tensor,

2908 + batch_size: int,

2909 + output_dtype: torch.dtype,

2910 + scale: float | None = None,

2911 + cu_seqlens: torch.Tensor | None = None,

2912 + max_seqlen: torch.Tensor | None = None,

2913 + ) -> torch.Tensor:

2914 + return torch.ops.vllm.aiter_fp8_attn_wrapper(

2915 + q,

2916 + k,

2917 + v,

2918 + q_descale,

2919 + k_descale,

2920 + v_descale,

2921 + batch_size,

2922 + output_dtype,

2923 + scale,

2924 + cu_seqlens,

2925 + max_seqlen,

2926 + )

2927 +

2928 @staticmethod

2929 def pa_fwd_asm(

2930 Q: torch.Tensor,

vllm/model_executor/layers/attention/mm_encoder_attention.py

+108 -32

@@ -24,6 +24,7 @@ from vllm.model_executor.models.vision import (

24 get_multimodal_config,

25 get_vit_attn_backend,

26 )

27 +from vllm.platforms import current_platform

28 from vllm.utils.flashinfer import (

29 is_flashinfer_cudnn_fp8_prefill_attn_supported,

30 )

@@ -32,6 +33,7 @@ from vllm.utils.torch_utils import async_tensor_h2d

33 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version

34 from vllm.v1.attention.backends.registry import AttentionBackendEnum

35 from vllm.v1.attention.ops.vit_attn_wrappers import (

36 + vit_aiter_fp8_attn_wrapper,

37 vit_flash_attn_wrapper,

38 vit_flashinfer_wrapper,

39 vit_torch_sdpa_wrapper,

@@ -380,8 +382,8 @@ class MMEncoderAttention(CustomOp):

382 No-op if FP8 is not requested. Raises ``ValueError`` if FP8 is

383 requested but the platform does not support it.

384 """

383 - # Populate defaults so ``_forward_flashinfer`` can

384 - # check ``self.fp8_enabled`` and others without AttributeError.

385 + # Populate defaults so backend forward methods can check

386 + # ``self.fp8_enabled`` and related state without AttributeError.

387 self.fp8_enabled = False

388 self._fp8_dynamic_scale = False

389 self.fp8_quant: QuantFP8 | None = None

@@ -393,13 +395,38 @@ class MMEncoderAttention(CustomOp):

395 if mm_cfg is None or mm_cfg.mm_encoder_attn_dtype != "fp8":

396 return

397

396 - # FP8 path

397 - if not is_flashinfer_cudnn_fp8_prefill_attn_supported():

398 + if self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA:

399 + if not current_platform.is_rocm():

400 + raise ValueError("AITER FP8 ViT attention requires ROCm.")

401 +

402 + from vllm.platforms.rocm import on_mi3xx

403 +

404 + if not on_mi3xx():

405 + raise ValueError(

406 + "AITER FP8 ViT attention requires an MI300-series or "

407 + "MI350-series GPU (gfx942 or gfx950)."

408 + )

409 + try:

410 + from aiter import flash_attn_varlen_fp8_pertensor_func # noqa: F401

411 + except ImportError as exc:

412 + raise ValueError(

413 + "mm_encoder_attn_dtype='fp8' with ROCM_AITER_FA requires "

414 + "an AITER build that provides "

415 + "flash_attn_varlen_fp8_pertensor_func."

416 + ) from exc

417 + elif self.attn_backend == AttentionBackendEnum.FLASHINFER:

418 + if not is_flashinfer_cudnn_fp8_prefill_attn_supported():

419 + raise ValueError(

420 + "mm_encoder_attn_dtype='fp8' requires the FlashInfer "

421 + "cuDNN backend with cuDNN >= 9.17.1 on Blackwell (SM 100) "

422 + "or newer. cuDNN's FP8 SDPA path with bf16/fp16 output is "

423 + "not available on Hopper (H100/H200) or earlier."

424 + )

425 + else:

426 raise ValueError(

399 - "mm_encoder_attn_dtype='fp8' requires the FlashInfer "

400 - "cuDNN backend with cuDNN >= 9.17.1 on Blackwell (SM 100) "

401 - "or newer. cuDNN's FP8 SDPA path with bf16/fp16 output is "

402 - "not available on Hopper (H100/H200) or earlier."

427 + "mm_encoder_attn_dtype='fp8' requires either the "

428 + "ROCM_AITER_FA backend on ROCm or the FlashInfer cuDNN "

429 + f"backend on CUDA, got {self.attn_backend}."

430 )

431

432 self.fp8_enabled = True

@@ -642,6 +669,37 @@ class MMEncoderAttention(CustomOp):

669 buffer_wrapped,

670 )

671

672 + def _quantize_qkv_fp8(

673 + self,

674 + query: torch.Tensor,

675 + key: torch.Tensor,

676 + value: torch.Tensor,

677 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

678 + assert self.fp8_quant is not None

679 +

680 + if self._fp8_dynamic_scale:

681 + self._record_amax_and_update_scales(query, key, value)

682 +

683 + query = quantize_fp8_maybe_pad_head_dim(

684 + query,

685 + self._fp8_q_scale,

686 + skip_scale=self.skip_scale_q,

687 + fp8_quant=self.fp8_quant,

688 + )

689 + key = quantize_fp8_maybe_pad_head_dim(

690 + key,

691 + self._fp8_k_scale,

692 + skip_scale=self.skip_scale_k,

693 + fp8_quant=self.fp8_quant,

694 + )

695 + value = quantize_fp8_maybe_pad_head_dim(

696 + value,

697 + self._fp8_v_scale,

698 + skip_scale=self.skip_scale_v,

699 + fp8_quant=self.fp8_quant,

700 + )

701 + return query, key, value

702 +

703 def _forward_flashinfer(

704 self,

705 query: torch.Tensor,

@@ -653,29 +711,7 @@ class MMEncoderAttention(CustomOp):

711 | None = None, # Only used for FlashInfer CuDNN backend

712 ) -> torch.Tensor:

713 if self.fp8_enabled:

656 - assert self.fp8_quant is not None

657 -

658 - if self._fp8_dynamic_scale:

659 - self._record_amax_and_update_scales(query, key, value)

660 -

661 - query = quantize_fp8_maybe_pad_head_dim(

662 - query,

663 - self._fp8_q_scale,

664 - skip_scale=self.skip_scale_q,

665 - fp8_quant=self.fp8_quant,

666 - )

667 - key = quantize_fp8_maybe_pad_head_dim(

668 - key,

669 - self._fp8_k_scale,

670 - skip_scale=self.skip_scale_k,

671 - fp8_quant=self.fp8_quant,

672 - )

673 - value = quantize_fp8_maybe_pad_head_dim(

674 - value,

675 - self._fp8_v_scale,

676 - skip_scale=self.skip_scale_v,

677 - fp8_quant=self.fp8_quant,

678 - )

714 + query, key, value = self._quantize_qkv_fp8(query, key, value)

715

716 output = vit_flashinfer_wrapper(

717 q=query,

@@ -697,6 +733,44 @@ class MMEncoderAttention(CustomOp):

733

734 return output

735

736 + def _forward_aiter_fp8(

737 + self,

738 + query: torch.Tensor,

739 + key: torch.Tensor,

740 + value: torch.Tensor,

741 + cu_seqlens: torch.Tensor | None = None,

742 + max_seqlen: torch.Tensor | None = None,

743 + ) -> torch.Tensor:

744 + assert (cu_seqlens is not None and max_seqlen is not None) or (

745 + cu_seqlens is None and max_seqlen is None

746 + ), "cu_seqlens and max_seqlen should be both set or both None."

747 +

748 + bsz, q_len = query.size()[:2]

749 + kv_len = key.size(1)

750 + is_reshaped = query.dim() != 4

751 + query, key, value = self.view_qkv_to_4d(query, key, value, bsz, q_len, kv_len)

752 +

753 + query, key, value = self._quantize_qkv_fp8(query, key, value)

754 +

755 + output = vit_aiter_fp8_attn_wrapper(

756 + q=query,

757 + k=key,

758 + v=value,

759 + q_descale=self._fp8_q_scale,

760 + k_descale=self._fp8_k_scale,

761 + v_descale=self._fp8_v_scale,

762 + batch_size=bsz,

763 + output_dtype=self.dtype,

764 + scale=self.scale,

765 + cu_seqlens=cu_seqlens,

766 + max_seqlen=max_seqlen,

767 + )

768 + if output.shape[-1] != self.head_size:

769 + output = output[..., : self.head_size].contiguous()

770 + if is_reshaped:

771 + output = output.reshape(bsz, q_len, -1)

772 + return output

773 +

774 def forward_native(

775 self,

776 query: torch.Tensor,

@@ -719,7 +793,9 @@ class MMEncoderAttention(CustomOp):

793 sequence_lengths: torch.Tensor

794 | None = None, # Only used for FlashInfer CuDNN backend

795 ) -> torch.Tensor:

722 - if self.is_flash_attn_backend:

796 + if self.fp8_enabled and self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA:

797 + return self._forward_aiter_fp8(query, key, value, cu_seqlens, max_seqlen)

798 + elif self.is_flash_attn_backend:

799 return self._forward_fa(query, key, value, cu_seqlens, max_seqlen)

800 elif self.attn_backend == AttentionBackendEnum.TRITON_ATTN:

801 return self._forward_triton(query, key, value, cu_seqlens, max_seqlen)

vllm/platforms/rocm.py

+1 -1

@@ -657,7 +657,7 @@ class RocmPlatform(Platform):

657

658 from vllm._aiter_ops import rocm_aiter_ops

659

660 - if rocm_aiter_ops.is_enabled() and on_gfx9():

660 + if rocm_aiter_ops.is_mha_enabled() and on_gfx9():

661 logger.info_once("Using AITER Flash Attention backend for ViT model.")

662 return AttentionBackendEnum.ROCM_AITER_FA

663

vllm/v1/attention/ops/vit_attn_wrappers.py

+29

@@ -18,6 +18,7 @@ import einops

18 import torch

19 import torch.nn.functional as F

21 +from vllm._aiter_ops import rocm_aiter_ops

22 from vllm.platforms import current_platform

23 from vllm.utils.torch_utils import direct_register_custom_op

@@ -114,6 +115,34 @@ def vit_flash_attn_wrapper(

115 )

116

117

118 +def vit_aiter_fp8_attn_wrapper(

119 + q: torch.Tensor,

120 + k: torch.Tensor,

121 + v: torch.Tensor,

122 + q_descale: torch.Tensor,

123 + k_descale: torch.Tensor,

124 + v_descale: torch.Tensor,

125 + batch_size: int,

126 + output_dtype: torch.dtype,

127 + scale: float | None = None,

128 + cu_seqlens: torch.Tensor | None = None,

129 + max_seqlen: torch.Tensor | None = None,

130 +) -> torch.Tensor:

131 + return rocm_aiter_ops.fp8_attn_wrapper(

132 + q,

133 + k,

134 + v,

135 + q_descale,

136 + k_descale,

137 + v_descale,

138 + batch_size,

139 + output_dtype,

140 + scale,

141 + cu_seqlens,

142 + max_seqlen,

143 + )

144 +

145 +

146 def triton_attn_wrapper(

147 q: torch.Tensor,

148 k: torch.Tensor,