| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm projection/collective contract corresponding to the SM100 GEMM-RS/AR. |
| 4 | |
| 5 | Exercise the production RowParallelLinear and sequence-parallel reduction |
| 6 | paths: native BF16/MXFP8 GEMM, AITER all-reduce, and RCCL reduce-scatter. This |
| 7 | checks the same observable contract, not NVIDIA's fused GEMM/communication |
| 8 | implementation or its performance. |
| 9 | """ |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | import torch.distributed as dist |
| 14 | import torch.multiprocessing as mp |
| 15 | |
| 16 | from tests.kernels.test_gemm_rs_ar import _N, _SHAPES |
| 17 | from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config |
| 18 | from vllm.config.quantization import QuantizationConfigArgs |
| 19 | from vllm.distributed import cleanup_dist_env_and_memory, get_tp_group |
| 20 | from vllm.distributed.parallel_state import ( |
| 21 | graph_capture, |
| 22 | init_distributed_environment, |
| 23 | initialize_model_parallel, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.linear import RowParallelLinear |
| 26 | from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig |
| 27 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 28 | _mxfp8_e4m3_quantize_torch, |
| 29 | dequant_mxfp8_to_bf16, |
| 30 | ) |
| 31 | from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter |
| 32 | from vllm.platforms import current_platform |
| 33 | from vllm.utils.network_utils import get_open_port |
| 34 | from vllm.utils.system_utils import update_environment_variables |
| 35 | |
| 36 | # Retain the expanded parent's shape transitions, plus fully padded RS ranks. |
| 37 | _ROCM_SHAPES = ((1, 512), (7, 768), *_SHAPES) |
| 38 | |
| 39 | |
| 40 | def _reference( |
| 41 | x: torch.Tensor, |
| 42 | layer: RowParallelLinear, |
| 43 | all_reduce: bool, |
| 44 | ) -> torch.Tensor: |
| 45 | weight = layer.weight |
| 46 | if hasattr(layer, "weight_scale"): |
| 47 | # The production ROCm path uses the Triton quantizer and dot_scaled; |
| 48 | # use the independent torch quantizer and FP32 matmul as the oracle. |
| 49 | x_q, x_scale = _mxfp8_e4m3_quantize_torch(x) |
| 50 | x = dequant_mxfp8_to_bf16(x_q, x_scale) |
| 51 | weight = dequant_mxfp8_to_bf16(weight, layer.weight_scale) |
| 52 | partial = (x.float() @ weight.float().T).to(torch.bfloat16).float() |
| 53 | group = get_tp_group() |
| 54 | # Always use FP32 RCCL AR for the oracle, including the RS candidate. |
| 55 | dist.all_reduce(partial, group=group.device_group) |
| 56 | expected = partial.to(torch.bfloat16) |
| 57 | if all_reduce: |
| 58 | return expected |
| 59 | padding = -x.shape[0] % group.world_size |
| 60 | expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)) |
| 61 | return expected.chunk(group.world_size)[group.rank_in_group] |
| 62 | |
| 63 | |
| 64 | def _apply(x: torch.Tensor, layer: RowParallelLinear) -> torch.Tensor: |
| 65 | output, _ = layer(x) |
| 66 | return output if layer.reduce_results else sp_reduce_scatter(output) |
| 67 | |
| 68 | |
| 69 | def _assert_output( |
| 70 | actual: torch.Tensor, expected: torch.Tensor, M: int, all_reduce: bool |
| 71 | ) -> None: |
| 72 | assert actual.dtype == torch.bfloat16 |
| 73 | assert actual.shape == expected.shape |
| 74 | assert torch.isfinite(actual).all() |
| 75 | # Match the parent's tolerance at 0.1 weight scale, then bound total error. |
| 76 | # BF16 partial/reduction rounding can dominate near a cancelled output. |
| 77 | torch.testing.assert_close(actual, expected, rtol=5e-2, atol=0.4) |
| 78 | error = (actual.float() - expected.float()).norm() |
| 79 | assert error <= 5e-3 * expected.float().norm() + 1e-6 |
| 80 | if not all_reduce: |
| 81 | rank_start = get_tp_group().rank_in_group * actual.shape[0] |
| 82 | valid_rows = min(max(M - rank_start, 0), actual.shape[0]) |
| 83 | assert not actual[valid_rows:].any(), "RS padding must remain zero" |
| 84 | |
| 85 | |
| 86 | def _run_mode( |
| 87 | device: torch.device, backend: str, all_reduce: bool, world_size: int |
| 88 | ) -> None: |
| 89 | group = get_tp_group() |
| 90 | generator = torch.Generator(device=device) |
| 91 | quant_config = ( |
| 92 | OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8")) |
| 93 | if backend == "mxfp8" |
| 94 | else None |
| 95 | ) |
| 96 | layers = {} |
| 97 | for K in sorted({K for _, K in _ROCM_SHAPES}): |
| 98 | with torch.device(device): |
| 99 | layer = RowParallelLinear( |
| 100 | K * world_size, |
| 101 | _N, |
| 102 | bias=False, |
| 103 | params_dtype=torch.bfloat16, |
| 104 | quant_config=quant_config, |
| 105 | reduce_results=all_reduce, |
| 106 | ) |
| 107 | generator.manual_seed(1000 + group.rank_in_group * 100 + K) |
| 108 | weight = ( |
| 109 | torch.randn(_N, K, dtype=torch.bfloat16, device=device, generator=generator) |
| 110 | * 0.1 |
| 111 | ) |
| 112 | layer.weight = torch.nn.Parameter(weight, requires_grad=False) |
| 113 | layer.quant_method.process_weights_after_loading(layer) |
| 114 | if backend == "mxfp8": |
| 115 | from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( |
| 116 | RocmDotScaledMxfp8LinearKernel, |
| 117 | ) |
| 118 | |
| 119 | assert isinstance(layer.quant_method.kernel, RocmDotScaledMxfp8LinearKernel) |
| 120 | assert layer.weight.dtype == torch.float8_e4m3fn |
| 121 | layers[K] = layer |
| 122 | |
| 123 | def make_input(M: int, K: int, seed: int) -> torch.Tensor: |
| 124 | generator.manual_seed(seed + group.rank_in_group * 100) |
| 125 | return torch.randn( |
| 126 | M, K, dtype=torch.bfloat16, device=device, generator=generator |
| 127 | ) |
| 128 | |
| 129 | comm = group.device_communicator |
| 130 | assert comm is not None |
| 131 | assert comm.use_aiter_allreduce |
| 132 | assert comm.aiter_ar_comm is not None and not comm.aiter_ar_comm.disabled |
| 133 | assert comm.qr_comm is None or comm.qr_comm.disabled |
| 134 | assert comm.pynccl_comm is not None and not comm.pynccl_comm.disabled |
| 135 | for M, K in (*_ROCM_SHAPES, *_ROCM_SHAPES[::-1]): |
| 136 | x = make_input(M, K, 2000 + M + K) |
| 137 | if all_reduce: |
| 138 | probe = torch.empty(M, _N, dtype=x.dtype, device=device) |
| 139 | assert comm.aiter_ar_comm.should_custom_ar(probe) |
| 140 | expected = _reference(x, layers[K], all_reduce) |
| 141 | actual = _apply(x, layers[K]) |
| 142 | _assert_output(actual, expected, M, all_reduce) |
| 143 | |
| 144 | # Retained outputs must survive the next call's collective buffer reuse. |
| 145 | x = make_input(257, 768, 2500) |
| 146 | first = _apply(x, layers[768]) |
| 147 | snapshot = first.clone() |
| 148 | second = _apply(-x, layers[768]) |
| 149 | torch.accelerator.synchronize(device) |
| 150 | assert first.data_ptr() != second.data_ptr() |
| 151 | torch.testing.assert_close(first, snapshot, rtol=0, atol=0) |
| 152 | |
| 153 | M, K = 1025, 4224 |
| 154 | graph_x = make_input(M, K, 3000) |
| 155 | for _ in range(3): |
| 156 | _apply(graph_x, layers[K]) |
| 157 | torch.accelerator.synchronize(device) |
| 158 | dist.barrier(group=group.device_group) |
| 159 | graph = torch.cuda.CUDAGraph() |
| 160 | # vLLM's context registers the AITER IPC addresses used by graph replay. |
| 161 | with ( |
| 162 | graph_capture(device=device) as context, |
| 163 | torch.cuda.graph(graph, stream=context.stream), |
| 164 | ): |
| 165 | graph_output = _apply(graph_x, layers[K]) |
| 166 | torch.cuda.current_stream().wait_stream(context.stream) |
| 167 | dist.barrier(group=group.device_group) |
| 168 | for replay in range(3): |
| 169 | graph_x.copy_(make_input(M, K, 3001 + replay)) |
| 170 | expected = _reference(graph_x, layers[K], all_reduce) |
| 171 | graph.replay() |
| 172 | torch.accelerator.synchronize(device) |
| 173 | _assert_output(graph_output, expected, M, all_reduce) |
| 174 | dist.barrier(group=group.device_group) |
| 175 | |
| 176 | |
| 177 | def _worker(rank: int, world_size: int, port: int, backend: str) -> None: |
| 178 | device = torch.device("cuda", rank) |
| 179 | torch.accelerator.set_device_index(device) |
| 180 | update_environment_variables( |
| 181 | { |
| 182 | "RANK": str(rank), |
| 183 | "LOCAL_RANK": str(rank), |
| 184 | "WORLD_SIZE": str(world_size), |
| 185 | "MASTER_ADDR": "localhost", |
| 186 | "MASTER_PORT": str(port), |
| 187 | } |
| 188 | ) |
| 189 | config = VllmConfig() |
| 190 | config.model_config = ModelConfig(dtype="bfloat16") |
| 191 | try: |
| 192 | with set_current_vllm_config(config), torch.inference_mode(): |
| 193 | init_distributed_environment() |
| 194 | initialize_model_parallel(tensor_model_parallel_size=world_size) |
| 195 | for all_reduce in (False, True): |
| 196 | _run_mode(device, backend, all_reduce, world_size) |
| 197 | finally: |
| 198 | cleanup_dist_env_and_memory() |
| 199 | |
| 200 | |
| 201 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm GEMM-RS/AR contract") |
| 202 | @pytest.mark.parametrize("world_size", [2, 4]) |
| 203 | @pytest.mark.parametrize( |
| 204 | "backend", |
| 205 | [ |
| 206 | "bf16", |
| 207 | pytest.param( |
| 208 | "mxfp8", |
| 209 | marks=pytest.mark.skipif( |
| 210 | not current_platform.supports_mx(), reason="native MXFP8 requires CDNA4" |
| 211 | ), |
| 212 | ), |
| 213 | ], |
| 214 | ) |
| 215 | @pytest.mark.distributed(num_gpus=4) |
| 216 | def test_rocm_gemm_rs_ar( |
| 217 | monkeypatch: pytest.MonkeyPatch, world_size: int, backend: str |
| 218 | ) -> None: |
| 219 | """Preserve rank reduction, tails, shape reuse, lifetime, and graph replay.""" |
| 220 | if torch.accelerator.device_count() < world_size: |
| 221 | pytest.skip(f"GEMM-RS/AR requires {world_size} GPUs") |
| 222 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 223 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1") |
| 224 | monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE") |
| 225 | try: |
| 226 | mp.spawn( |
| 227 | _worker, args=(world_size, get_open_port(), backend), nprocs=world_size |
| 228 | ) |
| 229 | finally: |
| 230 | cleanup_dist_env_and_memory() |