| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm MoE coverage for the TP2 GEMM/collective fusion contract.""" |
| 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 |
| 12 | from .models import TRITON_ATTN, qwen3_a3b |
| 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[:1]) |
| 18 | @pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) |
| 19 | def test_tp2_rocm_moe_async_tp_bf16_fusions( |
| 20 | inductor_graph_partition, attn_backend, run_e2e_fusion_test |
| 21 | ): |
| 22 | """MoE dense projections retain the Qwen SP/AsyncTP rewrite counts.""" |
| 23 | n_layers = 4 |
| 24 | run_e2e_fusion_test( |
| 25 | qwen3_a3b.model_name, |
| 26 | qwen3_a3b.matches(n_layers), |
| 27 | dict( |
| 28 | hf_overrides=dict( |
| 29 | num_hidden_layers=n_layers, |
| 30 | hidden_size=512, |
| 31 | intermediate_size=1024, |
| 32 | moe_intermediate_size=128, |
| 33 | num_attention_heads=8, |
| 34 | num_key_value_heads=2, |
| 35 | head_dim=64, |
| 36 | num_experts=8, |
| 37 | num_experts_per_tok=2, |
| 38 | ), |
| 39 | load_format="dummy", |
| 40 | max_model_len=1024, |
| 41 | kv_cache_memory_bytes=64 * 1024**2, |
| 42 | gpu_memory_utilization=0.01, |
| 43 | max_num_batched_tokens=1024, |
| 44 | max_num_seqs=32, |
| 45 | ), |
| 46 | attn_backend, |
| 47 | dict( |
| 48 | use_inductor_graph_partition=inductor_graph_partition, |
| 49 | custom_ops=["-rms_norm"], |
| 50 | pass_config=PassConfig( |
| 51 | enable_qk_norm_rope_fusion=True, |
| 52 | enable_sp=True, |
| 53 | fuse_gemm_comms=True, |
| 54 | fuse_allreduce_rms=False, |
| 55 | sp_min_token_num=512, |
| 56 | ), |
| 57 | ), |
| 58 | ["norm_rope_fusion", "sequence_parallel", "async_tp"], |
| 59 | tp_size=2, |
| 60 | ) |