| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Bounded ROCm model coverage for actual sequence-parallel AsyncTP rewrites.""" |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from tests.utils import multi_gpu_test |
| 8 | from vllm.config import PassConfig |
| 9 | from vllm.platforms import current_platform |
| 10 | |
| 11 | from .common import INDUCTOR_GRAPH_PARTITION, custom_ops_combos |
| 12 | from .models import TRITON_ATTN, llama3_8b |
| 13 | |
| 14 | |
| 15 | @multi_gpu_test(num_gpus=2) |
| 16 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") |
| 17 | @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) |
| 18 | @pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm"))) |
| 19 | @pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) |
| 20 | def test_tp2_rocm_async_tp_bf16_fusions( |
| 21 | inductor_graph_partition, custom_ops, attn_backend, run_e2e_fusion_test |
| 22 | ): |
| 23 | n_layers = 4 |
| 24 | model_kwargs = dict( |
| 25 | hf_overrides=dict( |
| 26 | num_hidden_layers=n_layers, |
| 27 | hidden_size=512, |
| 28 | intermediate_size=1024, |
| 29 | num_attention_heads=8, |
| 30 | num_key_value_heads=2, |
| 31 | head_dim=64, |
| 32 | ), |
| 33 | load_format="dummy", |
| 34 | max_model_len=1024, |
| 35 | kv_cache_memory_bytes=64 * 1024**2, |
| 36 | gpu_memory_utilization=0.01, |
| 37 | max_num_batched_tokens=1024, |
| 38 | max_num_seqs=32, |
| 39 | ) |
| 40 | run_e2e_fusion_test( |
| 41 | llama3_8b.model_name, |
| 42 | llama3_8b.matches(n_layers), |
| 43 | model_kwargs, |
| 44 | attn_backend, |
| 45 | dict( |
| 46 | use_inductor_graph_partition=inductor_graph_partition, |
| 47 | custom_ops=custom_ops.split(","), |
| 48 | pass_config=PassConfig( |
| 49 | enable_qk_norm_rope_fusion=True, |
| 50 | enable_sp=True, |
| 51 | fuse_gemm_comms=True, |
| 52 | fuse_allreduce_rms=False, |
| 53 | sp_min_token_num=512, |
| 54 | ), |
| 55 | ), |
| 56 | ["norm_rope_fusion", "sequence_parallel", "async_tp"], |
| 57 | tp_size=2, |
| 58 | ) |