| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import json |
| 4 | import logging |
| 5 | |
| 6 | import pytest |
| 7 | import regex as re |
| 8 | |
| 9 | from tests.utils import compare_two_settings |
| 10 | from vllm._aiter_ops import is_aiter_found_and_supported |
| 11 | from vllm.config import CompilationMode |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.transformers_utils.repo_utils import hf_api |
| 14 | |
| 15 | |
| 16 | @pytest.mark.skipif( |
| 17 | not current_platform.is_rocm() or not is_aiter_found_and_supported(), |
| 18 | reason="Requires ROCm AITER attention and quantization", |
| 19 | ) |
| 20 | def test_rocm_aiter_static_attention_output_correctness(caplog_mp_spawn): |
| 21 | """Trained FP8 outputs agree when actual AITER output quantization is fused.""" |
| 22 | model = hf_api().snapshot_download( |
| 23 | "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8", |
| 24 | revision="870a1177ed082c0ba5dbefd62573ab8f91803f6a", |
| 25 | allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt", "*.jinja"], |
| 26 | ) |
| 27 | common_args = [ |
| 28 | "--dtype", |
| 29 | "bfloat16", |
| 30 | "--load-format", |
| 31 | "auto", |
| 32 | "--max-model-len", |
| 33 | "512", |
| 34 | "--max-num-batched-tokens", |
| 35 | "512", |
| 36 | "--max-num-seqs", |
| 37 | "8", |
| 38 | "--kv-cache-memory-bytes", |
| 39 | str(256 * 1024**2), |
| 40 | "--gpu-memory-utilization", |
| 41 | "0.01", |
| 42 | "--attention-backend", |
| 43 | "ROCM_AITER_UNIFIED_ATTN", |
| 44 | ] |
| 45 | |
| 46 | def settings(enabled): |
| 47 | return [ |
| 48 | *common_args, |
| 49 | "--compilation-config", |
| 50 | json.dumps( |
| 51 | { |
| 52 | "mode": CompilationMode.VLLM_COMPILE, |
| 53 | "cudagraph_mode": "NONE", |
| 54 | "splitting_ops": [], |
| 55 | "inductor_compile_config": {"force_disable_caches": True}, |
| 56 | "custom_ops": ["+quant_fp8", "-rms_norm"], |
| 57 | "pass_config": { |
| 58 | "fuse_attn_quant": enabled, |
| 59 | "fuse_norm_quant": False, |
| 60 | "fuse_act_quant": False, |
| 61 | "fuse_allreduce_rms": False, |
| 62 | }, |
| 63 | } |
| 64 | ), |
| 65 | ] |
| 66 | |
| 67 | env = { |
| 68 | "VLLM_ROCM_USE_AITER": "1", |
| 69 | "VLLM_ROCM_USE_AITER_LINEAR": "1", |
| 70 | "VLLM_DISABLE_COMPILE_CACHE": "1", |
| 71 | } |
| 72 | with caplog_mp_spawn(logging.DEBUG) as logs: |
| 73 | compare_two_settings( |
| 74 | model, |
| 75 | settings(False), |
| 76 | settings(True), |
| 77 | env1=env, |
| 78 | env2=env, |
| 79 | method="generate", |
| 80 | force_v1_runner=True, |
| 81 | ) |
| 82 | matches = [ |
| 83 | int(count) for count in re.findall(r"'attn_quant_fusion': (\d+)", logs.text) |
| 84 | ] |
| 85 | print(f"Attention output quantization rewrite counts: {matches}") |
| 86 | assert matches and all(count == 32 for count in matches), ( |
| 87 | f"Expected output quantization fused in all 32 Llama layers, found {matches}" |
| 88 | ) |