| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | import torch.distributed as dist |
| 7 | import torch.multiprocessing as mp |
| 8 | |
| 9 | from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config |
| 10 | from vllm.distributed import cleanup_dist_env_and_memory |
| 11 | from vllm.distributed.parallel_state import ( |
| 12 | init_distributed_environment, |
| 13 | initialize_model_parallel, |
| 14 | ) |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.utils.network_utils import get_open_port |
| 17 | from vllm.utils.system_utils import update_environment_variables |
| 18 | |
| 19 | _SHAPES = ( |
| 20 | (129, 512), |
| 21 | (129, 768), |
| 22 | (257, 768), |
| 23 | (1023, 4224), |
| 24 | (1024, 4224), |
| 25 | (8191, 4224), |
| 26 | ) |
| 27 | _N = 512 |
| 28 | |
| 29 | |
| 30 | def _reference( |
| 31 | x: torch.Tensor, |
| 32 | weight: torch.Tensor, |
| 33 | weight_scale: torch.Tensor | None, |
| 34 | world_size: int, |
| 35 | group: dist.ProcessGroup, |
| 36 | all_reduce: bool, |
| 37 | ) -> torch.Tensor: |
| 38 | M = x.shape[0] |
| 39 | padded_M = (M + world_size - 1) // world_size * world_size |
| 40 | partial = torch.empty( |
| 41 | (padded_M, weight.shape[0]), |
| 42 | dtype=x.dtype, |
| 43 | device=x.device, |
| 44 | ) |
| 45 | if weight_scale is None: |
| 46 | torch.mm(x, weight.T, out=partial[:M]) |
| 47 | else: |
| 48 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 49 | mxfp8_e4m3_quantize, |
| 50 | ) |
| 51 | |
| 52 | x_q, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True) |
| 53 | partial[:M] = torch._scaled_mm( |
| 54 | x_q, |
| 55 | weight.T, |
| 56 | x_sf.view(torch.float8_e8m0fnu), |
| 57 | weight_scale.view(torch.float8_e8m0fnu), |
| 58 | out_dtype=x.dtype, |
| 59 | ) |
| 60 | if padded_M > M: |
| 61 | partial[M:].zero_() |
| 62 | |
| 63 | if all_reduce: |
| 64 | dist.all_reduce(partial, group=group) |
| 65 | return partial[:M] |
| 66 | |
| 67 | output = torch.empty( |
| 68 | (padded_M // world_size, weight.shape[0]), |
| 69 | dtype=x.dtype, |
| 70 | device=x.device, |
| 71 | ) |
| 72 | dist.reduce_scatter_single(output, partial, group=group) |
| 73 | return output |
| 74 | |
| 75 | |
| 76 | def _assert_valid_rows_close( |
| 77 | actual: torch.Tensor, |
| 78 | expected: torch.Tensor, |
| 79 | M: int, |
| 80 | rank: int, |
| 81 | all_reduce: bool, |
| 82 | ) -> None: |
| 83 | if all_reduce: |
| 84 | torch.testing.assert_close(actual, expected, rtol=5e-2, atol=4.0) |
| 85 | return |
| 86 | rank_start = rank * actual.shape[0] |
| 87 | valid_rows = min(max(M - rank_start, 0), actual.shape[0]) |
| 88 | torch.testing.assert_close( |
| 89 | actual[:valid_rows], |
| 90 | expected[:valid_rows], |
| 91 | rtol=5e-2, |
| 92 | atol=4.0, |
| 93 | ) |
| 94 | # Sequence-parallel callers may pass an unpadded M and read the padding |
| 95 | # rows back (they flow through per-token kernels before being dropped), |
| 96 | # so RS must leave them zero like ``sp_reduce_scatter`` does. |
| 97 | assert not actual[valid_rows:].any() |
| 98 | |
| 99 | |
| 100 | def _run_mode( |
| 101 | *, |
| 102 | all_reduce: bool, |
| 103 | device: torch.device, |
| 104 | group: dist.ProcessGroup, |
| 105 | rank: int, |
| 106 | world_size: int, |
| 107 | weights: dict[int, torch.Tensor], |
| 108 | backend: str, |
| 109 | ) -> None: |
| 110 | # cute_dsl is unavailable off CUDA, so import it only inside the GPU worker. |
| 111 | from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr |
| 112 | |
| 113 | gemm_rs_ar = GemmRsAr( |
| 114 | max_M=max(M for M, _ in _SHAPES), |
| 115 | N=_N, |
| 116 | all_reduce=all_reduce, |
| 117 | ) |
| 118 | from vllm.config.quantization import QuantizationConfigArgs |
| 119 | from vllm.model_executor.layers.linear import RowParallelLinear |
| 120 | from vllm.model_executor.layers.quantization.online.base import ( |
| 121 | OnlineQuantizationConfig, |
| 122 | ) |
| 123 | |
| 124 | quant_config = ( |
| 125 | None |
| 126 | if backend == "bf16" |
| 127 | else OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8")) |
| 128 | ) |
| 129 | projections = {} |
| 130 | operands = {} |
| 131 | for K, w in weights.items(): |
| 132 | with torch.device(device): |
| 133 | linear = RowParallelLinear( |
| 134 | K * world_size, |
| 135 | _N, |
| 136 | bias=False, |
| 137 | params_dtype=torch.bfloat16, |
| 138 | quant_config=quant_config, |
| 139 | ) |
| 140 | # Select the fused path before online quantization replaces meta weights. |
| 141 | assert gemm_rs_ar.can_run(linear) |
| 142 | linear.weight = torch.nn.Parameter(w.clone(), requires_grad=False) |
| 143 | linear.quant_method.process_weights_after_loading(linear) |
| 144 | projections[K] = linear |
| 145 | weight = linear.weight |
| 146 | if backend == "flashinfer_cutedsl": |
| 147 | weight = weight.t() |
| 148 | operands[K] = (weight, getattr(linear, "weight_scale", None)) |
| 149 | |
| 150 | input_generator = torch.Generator(device=device) |
| 151 | |
| 152 | # Alternating shapes exercise producer-flag reuse across different grids |
| 153 | # and CTA-group choices. |
| 154 | for M, K in (*_SHAPES, *_SHAPES[::-1]): |
| 155 | input_generator.manual_seed(2000 + M + K) |
| 156 | x = torch.randn( |
| 157 | M, |
| 158 | K, |
| 159 | dtype=torch.bfloat16, |
| 160 | device=device, |
| 161 | generator=input_generator, |
| 162 | ) |
| 163 | expected = _reference(x, *operands[K], world_size, group, all_reduce) |
| 164 | actual = gemm_rs_ar.apply(x, projections[K]) |
| 165 | torch.accelerator.synchronize(device) |
| 166 | _assert_valid_rows_close(actual, expected, M, rank, all_reduce) |
| 167 | |
| 168 | if all_reduce: |
| 169 | # AR outputs must remain valid after the symmetric workspace is reused. |
| 170 | lifetime_M, lifetime_K = 257, 768 |
| 171 | input_generator.manual_seed(2500) |
| 172 | lifetime_x = torch.randn( |
| 173 | lifetime_M, |
| 174 | lifetime_K, |
| 175 | dtype=torch.bfloat16, |
| 176 | device=device, |
| 177 | generator=input_generator, |
| 178 | ) |
| 179 | first_output = gemm_rs_ar.apply(lifetime_x, projections[lifetime_K]) |
| 180 | first_snapshot = first_output.clone() |
| 181 | second_output = gemm_rs_ar.apply(-lifetime_x, projections[lifetime_K]) |
| 182 | torch.accelerator.synchronize(device) |
| 183 | assert first_output.data_ptr() != second_output.data_ptr() |
| 184 | torch.testing.assert_close(first_output, first_snapshot, rtol=0, atol=0) |
| 185 | |
| 186 | graph_M, graph_K = 1025, 4224 |
| 187 | input_generator.manual_seed(3000) |
| 188 | graph_x = torch.randn( |
| 189 | graph_M, |
| 190 | graph_K, |
| 191 | dtype=torch.bfloat16, |
| 192 | device=device, |
| 193 | generator=input_generator, |
| 194 | ) |
| 195 | graph_expected = _reference( |
| 196 | graph_x, |
| 197 | *operands[graph_K], |
| 198 | world_size, |
| 199 | group, |
| 200 | all_reduce, |
| 201 | ) |
| 202 | graph_output = torch.empty_like(graph_expected) |
| 203 | |
| 204 | capture_stream = torch.cuda.Stream() |
| 205 | capture_stream.wait_stream(torch.cuda.current_stream()) |
| 206 | with torch.cuda.stream(capture_stream): |
| 207 | for _ in range(3): |
| 208 | graph_output.copy_(gemm_rs_ar.apply(graph_x, projections[graph_K])) |
| 209 | capture_stream.synchronize() |
| 210 | dist.barrier(group=group) |
| 211 | |
| 212 | graph = torch.cuda.CUDAGraph() |
| 213 | with torch.cuda.graph(graph, stream=capture_stream): |
| 214 | graph_output.copy_(gemm_rs_ar.apply(graph_x, projections[graph_K])) |
| 215 | torch.cuda.current_stream().wait_stream(capture_stream) |
| 216 | dist.barrier(group=group) |
| 217 | |
| 218 | for _ in range(3): |
| 219 | graph.replay() |
| 220 | torch.accelerator.synchronize(device) |
| 221 | _assert_valid_rows_close( |
| 222 | graph_output, |
| 223 | graph_expected, |
| 224 | graph_M, |
| 225 | rank, |
| 226 | all_reduce, |
| 227 | ) |
| 228 | |
| 229 | dist.barrier(group=group) |
| 230 | del gemm_rs_ar, capture_stream, graph, graph_output |
| 231 | |
| 232 | |
| 233 | def _worker(local_rank: int, world_size: int, master_port: int, backend: str) -> None: |
| 234 | # This module pulls in cute_dsl, which is unavailable off CUDA, so importing |
| 235 | # it at module level would fail collection. Import it where it is used, as |
| 236 | # `kda.py` does. |
| 237 | |
| 238 | device = torch.device("cuda", local_rank) |
| 239 | torch.accelerator.set_device_index(device) |
| 240 | update_environment_variables( |
| 241 | { |
| 242 | "RANK": str(local_rank), |
| 243 | "LOCAL_RANK": str(local_rank), |
| 244 | "WORLD_SIZE": str(world_size), |
| 245 | "MASTER_ADDR": "localhost", |
| 246 | "MASTER_PORT": str(master_port), |
| 247 | } |
| 248 | ) |
| 249 | |
| 250 | init_distributed_environment() |
| 251 | config = VllmConfig() |
| 252 | config.model_config = ModelConfig(dtype="bfloat16") |
| 253 | config.kernel_config.linear_backend = "auto" if backend == "bf16" else backend |
| 254 | with set_current_vllm_config(config): |
| 255 | initialize_model_parallel(tensor_model_parallel_size=world_size) |
| 256 | |
| 257 | group = dist.group.WORLD |
| 258 | rank = dist.get_rank(group) |
| 259 | |
| 260 | weight_generator = torch.Generator(device=device) |
| 261 | weights = {} |
| 262 | for K in {K for _, K in _SHAPES}: |
| 263 | weight_generator.manual_seed(1000 + rank * 10 + K) |
| 264 | weights[K] = torch.randn( |
| 265 | _N, |
| 266 | K, |
| 267 | dtype=torch.bfloat16, |
| 268 | device=device, |
| 269 | generator=weight_generator, |
| 270 | ) |
| 271 | |
| 272 | # Production binds one mode per worker; exercise both mode-bound instances |
| 273 | # sequentially in this test without implying that both are initialized. |
| 274 | with set_current_vllm_config(config): |
| 275 | for all_reduce in (False, True): |
| 276 | _run_mode( |
| 277 | all_reduce=all_reduce, |
| 278 | device=device, |
| 279 | group=group, |
| 280 | rank=rank, |
| 281 | world_size=world_size, |
| 282 | weights=weights, |
| 283 | backend=backend, |
| 284 | ) |
| 285 | cleanup_dist_env_and_memory() |
| 286 | |
| 287 | |
| 288 | @pytest.mark.parametrize( |
| 289 | "backend", ["bf16", "flashinfer_cutlass", "flashinfer_cutedsl"] |
| 290 | ) |
| 291 | @pytest.mark.distributed(num_gpus=2) |
| 292 | @pytest.mark.skipif( |
| 293 | not current_platform.is_device_capability_family(100), |
| 294 | reason="GEMM-RS/AR requires SM100", |
| 295 | ) |
| 296 | def test_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, backend: str) -> None: |
| 297 | world_size = 2 |
| 298 | if torch.accelerator.device_count() < world_size: |
| 299 | pytest.skip("GEMM-RS/AR requires two GPUs") |
| 300 | |
| 301 | monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1") |
| 302 | monkeypatch.setenv("NCCL_NVLS_ENABLE", "1") |
| 303 | try: |
| 304 | mp.spawn( |
| 305 | _worker, |
| 306 | args=(world_size, get_open_port(), backend), |
| 307 | nprocs=world_size, |
| 308 | ) |
| 309 | finally: |
| 310 | cleanup_dist_env_and_memory() |