| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test DeepEP + DeepGEMM integration |
| 4 | DeepGEMM are gemm kernels specialized for the |
| 5 | fp8 block-quantized case. |
| 6 | """ |
| 7 | |
| 8 | import dataclasses |
| 9 | from contextlib import contextmanager |
| 10 | |
| 11 | import pytest |
| 12 | import torch.distributed |
| 13 | from torch.distributed import ProcessGroup |
| 14 | from typing_extensions import ParamSpec |
| 15 | |
| 16 | import vllm.envs as envs |
| 17 | from vllm.config import VllmConfig, set_current_vllm_config |
| 18 | from vllm.forward_context import set_forward_context |
| 19 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 20 | from vllm.model_executor.layers.fused_moe.config import ( |
| 21 | FusedMoEQuantConfig, |
| 22 | fp8_w8a8_moe_quant_config, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts |
| 25 | from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel |
| 26 | from vllm.utils.deep_gemm import ( |
| 27 | get_mk_alignment_for_contiguous_layout, |
| 28 | is_deep_gemm_e8m0_used, |
| 29 | is_deep_gemm_supported, |
| 30 | ) |
| 31 | from vllm.utils.import_utils import has_deep_ep, has_deep_gemm |
| 32 | from vllm.utils.math_utils import next_power_of_2 |
| 33 | from vllm.utils.torch_utils import set_random_seed |
| 34 | from vllm.v1.worker.workspace import init_workspace_manager |
| 35 | |
| 36 | from ...utils import multi_gpu_test |
| 37 | from .parallel_utils import ProcessGroupInfo, parallel_launch |
| 38 | from .utils import make_dummy_moe_config, make_test_weights |
| 39 | |
| 40 | if has_deep_ep(): |
| 41 | from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import ( |
| 42 | DeepEPHTPrepareAndFinalize, |
| 43 | ) |
| 44 | from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( |
| 45 | DeepEPLLPrepareAndFinalize, |
| 46 | ) |
| 47 | |
| 48 | from .parallel_utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a |
| 49 | |
| 50 | if has_deep_gemm(): |
| 51 | from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import ( |
| 52 | BatchedDeepGemmExperts, |
| 53 | ) |
| 54 | from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( |
| 55 | DeepGemmExperts, |
| 56 | ) |
| 57 | |
| 58 | requires_deep_ep = pytest.mark.skipif( |
| 59 | not has_deep_ep(), |
| 60 | reason="Requires deep_ep kernels", |
| 61 | ) |
| 62 | |
| 63 | requires_deep_gemm = pytest.mark.skipif( |
| 64 | not is_deep_gemm_supported(), |
| 65 | reason="Requires deep_gemm kernels", |
| 66 | ) |
| 67 | |
| 68 | P = ParamSpec("P") |
| 69 | |
| 70 | |
| 71 | @contextmanager |
| 72 | def with_dp_metadata(M: int, world_size: int): |
| 73 | num_tokens_across_dp = torch.tensor([M] * world_size, device="cpu", dtype=torch.int) |
| 74 | |
| 75 | vllm_config = VllmConfig() |
| 76 | vllm_config.parallel_config.data_parallel_size = world_size |
| 77 | vllm_config.parallel_config.enable_expert_parallel = True |
| 78 | |
| 79 | with set_forward_context( |
| 80 | None, |
| 81 | vllm_config, |
| 82 | num_tokens=M, |
| 83 | num_tokens_across_dp=num_tokens_across_dp, |
| 84 | ): |
| 85 | yield |
| 86 | |
| 87 | |
| 88 | def make_block_quant_fp8_weights( |
| 89 | e: int, |
| 90 | n: int, |
| 91 | k: int, |
| 92 | block_size: list[int], |
| 93 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 94 | """Return weights w1q, w2q, w1_scale, w2_scale.""" |
| 95 | (_, w1q, w1_scale, _), (_, w2q, w2_scale, _) = make_test_weights( |
| 96 | e, n, k, torch.bfloat16, torch.float8_e4m3fn, block_shape=block_size |
| 97 | ) |
| 98 | return w1q, w2q, w1_scale, w2_scale |
| 99 | |
| 100 | |
| 101 | @dataclasses.dataclass |
| 102 | class TestConfig: |
| 103 | topk: int |
| 104 | m: int |
| 105 | k: int |
| 106 | n: int |
| 107 | num_experts: int |
| 108 | per_act_token_quant: bool |
| 109 | block_size: list[int] |
| 110 | # configs for testing low-latency kernels |
| 111 | low_latency: bool |
| 112 | use_fp8_dispatch: bool | None = False |
| 113 | |
| 114 | |
| 115 | @dataclasses.dataclass |
| 116 | class TestTensors: |
| 117 | rank_tokens: torch.Tensor # all ranks make this many tokens |
| 118 | rank_token_scales: torch.Tensor | None |
| 119 | topk: torch.Tensor |
| 120 | topk_weights: torch.Tensor |
| 121 | config: TestConfig |
| 122 | |
| 123 | @staticmethod |
| 124 | def make(config: TestConfig, rank) -> "TestTensors": |
| 125 | dtype = torch.bfloat16 |
| 126 | topk, m, k = (config.topk, config.m, config.k) |
| 127 | |
| 128 | fp8_info = torch.finfo(torch.float8_e4m3fn) |
| 129 | fp8_max, fp8_min = fp8_info.max, fp8_info.min |
| 130 | device = torch.accelerator.current_device_index() |
| 131 | rank_tokens = torch.randn((m, k), device=device, dtype=dtype) / 10.0 |
| 132 | rank_tokens = rank_tokens.clamp(min=fp8_min, max=fp8_max) |
| 133 | rank_token_scales = None |
| 134 | |
| 135 | topk_ids = torch.randint( |
| 136 | low=0, |
| 137 | high=config.num_experts, |
| 138 | size=(m, topk), |
| 139 | device=device, |
| 140 | ).to(dtype=torch.int64) |
| 141 | |
| 142 | topk_weights = torch.randn( |
| 143 | topk_ids.shape, |
| 144 | dtype=torch.float32, |
| 145 | device=device, |
| 146 | ) |
| 147 | |
| 148 | return TestTensors( |
| 149 | rank_tokens=rank_tokens, |
| 150 | rank_token_scales=rank_token_scales, |
| 151 | topk=topk_ids, |
| 152 | topk_weights=topk_weights, |
| 153 | config=config, |
| 154 | ) |
| 155 | |
| 156 | |
| 157 | def make_ll_modular_kernel( |
| 158 | pg: ProcessGroup, |
| 159 | pgi: ProcessGroupInfo, |
| 160 | max_tokens_per_rank: int, |
| 161 | dp_size: int, |
| 162 | hidden_size: int, |
| 163 | q_dtype: torch.dtype | None, |
| 164 | test_config: TestConfig, |
| 165 | quant_config: FusedMoEQuantConfig, |
| 166 | ) -> FusedMoEKernel: |
| 167 | assert test_config.low_latency |
| 168 | assert test_config.use_fp8_dispatch is not None |
| 169 | |
| 170 | a2a: DeepEPLLPrepareAndFinalize = make_deepep_a2a( |
| 171 | pg=pg, |
| 172 | pgi=pgi, |
| 173 | dp_size=dp_size, |
| 174 | deepep_ht_args=None, |
| 175 | deepep_ll_args=DeepEPLLArgs( |
| 176 | max_tokens_per_rank=max_tokens_per_rank, |
| 177 | hidden_size=hidden_size, |
| 178 | num_experts=test_config.num_experts, |
| 179 | use_fp8_dispatch=test_config.use_fp8_dispatch, |
| 180 | ), |
| 181 | q_dtype=q_dtype, |
| 182 | block_shape=test_config.block_size, |
| 183 | ) |
| 184 | |
| 185 | fused_experts = BatchedDeepGemmExperts( |
| 186 | max_num_tokens=max_tokens_per_rank, |
| 187 | num_dispatchers=pgi.world_size // dp_size, |
| 188 | quant_config=quant_config, |
| 189 | moe_config=make_dummy_moe_config(), |
| 190 | ) |
| 191 | return FusedMoEKernel( |
| 192 | prepare_finalize=a2a, |
| 193 | fused_experts=fused_experts, |
| 194 | ) |
| 195 | |
| 196 | |
| 197 | def make_ht_modular_kernel( |
| 198 | pg: ProcessGroup, |
| 199 | pgi: ProcessGroupInfo, |
| 200 | dp_size: int, |
| 201 | num_local_experts: int, |
| 202 | q_dtype: torch.dtype | None, |
| 203 | test_config: TestConfig, |
| 204 | quant_config: FusedMoEQuantConfig, |
| 205 | ) -> FusedMoEKernel: |
| 206 | assert not test_config.low_latency |
| 207 | assert test_config.use_fp8_dispatch is None |
| 208 | |
| 209 | a2a: DeepEPHTPrepareAndFinalize = make_deepep_a2a( |
| 210 | pg=pg, |
| 211 | pgi=pgi, |
| 212 | dp_size=dp_size, |
| 213 | deepep_ht_args=DeepEPHTArgs(num_local_experts=num_local_experts), |
| 214 | deepep_ll_args=None, |
| 215 | q_dtype=q_dtype, |
| 216 | block_shape=test_config.block_size, |
| 217 | ) |
| 218 | |
| 219 | fused_experts = DeepGemmExperts( |
| 220 | moe_config=make_dummy_moe_config(), |
| 221 | quant_config=quant_config, |
| 222 | ) |
| 223 | return FusedMoEKernel( |
| 224 | prepare_finalize=a2a, |
| 225 | fused_experts=fused_experts, |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | def make_modular_kernel( |
| 230 | pg: ProcessGroup, |
| 231 | pgi: ProcessGroupInfo, |
| 232 | dp_size: int, |
| 233 | num_local_experts: int, |
| 234 | test_tensors: TestTensors, |
| 235 | quant_config: FusedMoEQuantConfig, |
| 236 | ) -> FusedMoEKernel: |
| 237 | q_dtype = torch.float8_e4m3fn |
| 238 | test_config = test_tensors.config |
| 239 | |
| 240 | mk: FusedMoEKernel |
| 241 | # Make modular kernel |
| 242 | if test_config.low_latency: |
| 243 | max_tokens_per_rank = max(64, next_power_of_2(test_tensors.rank_tokens.size(0))) |
| 244 | hidden_size = test_tensors.rank_tokens.size(-1) |
| 245 | |
| 246 | mk = make_ll_modular_kernel( |
| 247 | pg=pg, |
| 248 | pgi=pgi, |
| 249 | max_tokens_per_rank=max_tokens_per_rank, |
| 250 | dp_size=dp_size, |
| 251 | hidden_size=hidden_size, |
| 252 | q_dtype=q_dtype, |
| 253 | test_config=test_config, |
| 254 | quant_config=quant_config, |
| 255 | ) |
| 256 | else: |
| 257 | mk = make_ht_modular_kernel( |
| 258 | pg, |
| 259 | pgi, |
| 260 | dp_size, |
| 261 | num_local_experts, |
| 262 | q_dtype, |
| 263 | test_config, |
| 264 | quant_config=quant_config, |
| 265 | ) |
| 266 | |
| 267 | return mk |
| 268 | |
| 269 | |
| 270 | def deepep_deepgemm_moe_impl( |
| 271 | pg: ProcessGroup, |
| 272 | pgi: ProcessGroupInfo, |
| 273 | dp_size: int, |
| 274 | test_tensors: TestTensors, |
| 275 | w1: torch.Tensor, |
| 276 | w2: torch.Tensor, |
| 277 | w1_scale: torch.Tensor | None, |
| 278 | w2_scale: torch.Tensor | None, |
| 279 | ) -> torch.Tensor: |
| 280 | test_config = test_tensors.config |
| 281 | num_experts = test_config.num_experts |
| 282 | num_local_experts = w1.size(0) |
| 283 | |
| 284 | def build_expert_map(): |
| 285 | num_local_experts = w1.size(0) |
| 286 | expert_map = torch.full((num_experts,), fill_value=-1, dtype=torch.int32) |
| 287 | s = pgi.rank * num_local_experts |
| 288 | e = s + num_local_experts |
| 289 | expert_map[s:e] = torch.tensor(list(range(num_local_experts))) |
| 290 | device = torch.accelerator.current_device_index() |
| 291 | return expert_map.to(device=device, dtype=torch.int32) |
| 292 | |
| 293 | quant_config = fp8_w8a8_moe_quant_config( |
| 294 | w1_scale=w1_scale, |
| 295 | w2_scale=w2_scale, |
| 296 | # Low-Latency kernels can't dispatch scales. |
| 297 | a1_scale=(None if test_config.low_latency else test_tensors.rank_token_scales), |
| 298 | block_shape=test_config.block_size, |
| 299 | ) |
| 300 | |
| 301 | # Make modular kernel |
| 302 | mk: FusedMoEKernel = make_modular_kernel( |
| 303 | pg=pg, |
| 304 | pgi=pgi, |
| 305 | dp_size=dp_size, |
| 306 | num_local_experts=num_local_experts, |
| 307 | test_tensors=test_tensors, |
| 308 | quant_config=quant_config, |
| 309 | ) |
| 310 | |
| 311 | with with_dp_metadata( |
| 312 | M=test_tensors.rank_tokens.size(0), world_size=pgi.world_size |
| 313 | ): |
| 314 | out = mk.apply( |
| 315 | hidden_states=test_tensors.rank_tokens, |
| 316 | w1=w1, |
| 317 | w2=w2, |
| 318 | topk_weights=test_tensors.topk_weights, |
| 319 | topk_ids=test_tensors.topk, |
| 320 | activation=MoEActivation.SILU, |
| 321 | global_num_experts=num_experts, |
| 322 | expert_map=build_expert_map(), |
| 323 | apply_router_weight_on_input=False, |
| 324 | ) |
| 325 | return out |
| 326 | |
| 327 | |
| 328 | def triton_impl( |
| 329 | a: torch.Tensor, |
| 330 | topk_ids: torch.Tensor, |
| 331 | topk_weights: torch.Tensor, |
| 332 | w1: torch.Tensor, |
| 333 | w2: torch.Tensor, |
| 334 | w1_scale: torch.Tensor, |
| 335 | w2_scale: torch.Tensor, |
| 336 | a1_scale: torch.Tensor, |
| 337 | block_shape: list[int], |
| 338 | ): |
| 339 | quant_config = fp8_w8a8_moe_quant_config( |
| 340 | w1_scale=w1_scale, |
| 341 | w2_scale=w2_scale, |
| 342 | a1_scale=a1_scale, |
| 343 | block_shape=block_shape, |
| 344 | ) |
| 345 | |
| 346 | return fused_experts( |
| 347 | hidden_states=a, |
| 348 | w1=w1, |
| 349 | w2=w2, |
| 350 | topk_weights=topk_weights, |
| 351 | topk_ids=topk_ids, |
| 352 | quant_config=quant_config, |
| 353 | ) |
| 354 | |
| 355 | |
| 356 | def _test_deepep_deepgemm_moe( |
| 357 | pgi: ProcessGroupInfo, |
| 358 | dp_size: int, |
| 359 | config: TestConfig, |
| 360 | w1: torch.Tensor, |
| 361 | w2: torch.Tensor, |
| 362 | w1_scale: torch.Tensor, |
| 363 | w2_scale: torch.Tensor, |
| 364 | ): |
| 365 | device = torch.device(f"cuda:{pgi.local_rank}") |
| 366 | init_workspace_manager(device) |
| 367 | |
| 368 | set_random_seed(pgi.rank) |
| 369 | |
| 370 | device = torch.accelerator.current_device_index() |
| 371 | w1 = w1.to(device=device) |
| 372 | w2 = w2.to(device=device) |
| 373 | w1_scale = w1_scale.to(device=device) |
| 374 | w2_scale = w2_scale.to(device=device) |
| 375 | |
| 376 | if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP: |
| 377 | pg = torch.distributed.split_group( |
| 378 | split_ranks=[list(range(pgi.world_size))], |
| 379 | group_desc="deepep_deepgemm_test", |
| 380 | ) |
| 381 | else: |
| 382 | pg = torch.distributed.new_group(list(range(pgi.world_size))) |
| 383 | test_tensors = TestTensors.make(config, pgi.rank) |
| 384 | block_shape = [w1.size(1) // w1_scale.size(1), w1.size(2) // w1_scale.size(2)] |
| 385 | |
| 386 | with set_current_vllm_config(VllmConfig()): |
| 387 | # Reference |
| 388 | triton_moe = triton_impl( |
| 389 | a=test_tensors.rank_tokens, |
| 390 | topk_ids=test_tensors.topk, |
| 391 | topk_weights=test_tensors.topk_weights, |
| 392 | w1=w1, |
| 393 | w2=w2, |
| 394 | w1_scale=w1_scale, |
| 395 | w2_scale=w2_scale, |
| 396 | a1_scale=test_tensors.rank_token_scales, |
| 397 | block_shape=block_shape, |
| 398 | ) |
| 399 | |
| 400 | # Slice experts for this rank. |
| 401 | num_local_experts = config.num_experts // pgi.world_size |
| 402 | e_start = num_local_experts * pgi.rank |
| 403 | e_end = e_start + num_local_experts |
| 404 | w1_ep = w1[e_start:e_end] |
| 405 | w2_ep = w2[e_start:e_end] |
| 406 | w1_scale_ep = w1_scale[e_start:e_end] |
| 407 | w2_scale_ep = w2_scale[e_start:e_end] |
| 408 | |
| 409 | deepep_moe = deepep_deepgemm_moe_impl( |
| 410 | pg, |
| 411 | pgi, |
| 412 | dp_size, |
| 413 | test_tensors, |
| 414 | w1_ep, |
| 415 | w2_ep, |
| 416 | w1_scale_ep, |
| 417 | w2_scale_ep, |
| 418 | ) |
| 419 | |
| 420 | torch.testing.assert_close( |
| 421 | triton_moe, |
| 422 | deepep_moe, |
| 423 | atol=6e-2, |
| 424 | rtol=6e-2, |
| 425 | ) |
| 426 | |
| 427 | |
| 428 | MNKs = [ |
| 429 | (8, 128, 128), |
| 430 | (8, 128, 512), |
| 431 | (3, 1024, 2048), |
| 432 | (32, 128, 1024), |
| 433 | (45, 512, 2048), |
| 434 | (64, 1024, 1024), |
| 435 | (129, 128, 256), |
| 436 | (129, 1024, 2048), |
| 437 | (222, 1024, 2048), |
| 438 | ] |
| 439 | |
| 440 | TOPKS = [2, 6] |
| 441 | NUM_EXPERTS = [32] |
| 442 | |
| 443 | |
| 444 | @pytest.mark.parametrize("mnk", MNKs) |
| 445 | @pytest.mark.parametrize("num_experts", NUM_EXPERTS) |
| 446 | @pytest.mark.parametrize("topk", TOPKS) |
| 447 | @pytest.mark.parametrize("world_dp_size", [(2, 1)]) |
| 448 | @multi_gpu_test(num_gpus=2) |
| 449 | @requires_deep_ep |
| 450 | @requires_deep_gemm |
| 451 | def test_ht_deepep_deepgemm_moe( |
| 452 | mnk: tuple[int, int, int], |
| 453 | num_experts: int, |
| 454 | topk: int, |
| 455 | world_dp_size: tuple[int, int], |
| 456 | disable_deepgemm_ue8m0, |
| 457 | workspace_init, |
| 458 | ): |
| 459 | """Tests for High-Throughput DeepEP + DeepGemm integration.""" |
| 460 | m, n, k = mnk |
| 461 | set_random_seed(7) |
| 462 | |
| 463 | if topk > num_experts: |
| 464 | pytest.skip(f"Skipping test: topk={topk} > E={num_experts}") |
| 465 | |
| 466 | block_m = get_mk_alignment_for_contiguous_layout()[0] |
| 467 | block_size = [block_m, block_m] |
| 468 | |
| 469 | world_size, dp_size = world_dp_size |
| 470 | config = TestConfig( |
| 471 | topk=topk, |
| 472 | m=m, |
| 473 | k=k, |
| 474 | n=n, |
| 475 | num_experts=num_experts, |
| 476 | per_act_token_quant=False, |
| 477 | block_size=block_size, |
| 478 | low_latency=False, |
| 479 | use_fp8_dispatch=None, |
| 480 | ) |
| 481 | |
| 482 | w1, w2, w1_scale, w2_scale = make_block_quant_fp8_weights( |
| 483 | num_experts, n, k, block_size |
| 484 | ) |
| 485 | |
| 486 | parallel_launch( |
| 487 | world_size, |
| 488 | _test_deepep_deepgemm_moe, |
| 489 | dp_size, |
| 490 | config, |
| 491 | w1, |
| 492 | w2, |
| 493 | w1_scale, |
| 494 | w2_scale, |
| 495 | ) |
| 496 | |
| 497 | |
| 498 | MNKs = [ |
| 499 | (1, 128, 2560), |
| 500 | (2, 128, 2560), |
| 501 | (3, 1024, 2560), |
| 502 | (32, 128, 2560), |
| 503 | (45, 512, 2560), |
| 504 | (64, 1024, 2560), |
| 505 | (222, 1024, 2560), |
| 506 | ] |
| 507 | # Fix tests for USE_FP8_DISPATCH=True |
| 508 | USE_FP8_DISPATCH = [False] |
| 509 | |
| 510 | |
| 511 | @pytest.mark.parametrize("mnk", MNKs) |
| 512 | @pytest.mark.parametrize("num_experts", NUM_EXPERTS) |
| 513 | @pytest.mark.parametrize("topk", TOPKS) |
| 514 | @pytest.mark.parametrize("use_fp8_dispatch", USE_FP8_DISPATCH) |
| 515 | @pytest.mark.parametrize("block_size", [[128, 128]]) |
| 516 | @pytest.mark.parametrize("world_dp_size", [(2, 1)]) |
| 517 | @multi_gpu_test(num_gpus=2) |
| 518 | @requires_deep_ep |
| 519 | @requires_deep_gemm |
| 520 | def test_ll_deepep_deepgemm_moe( |
| 521 | mnk: tuple[int, int, int], |
| 522 | num_experts: int, |
| 523 | topk: int, |
| 524 | use_fp8_dispatch: bool, |
| 525 | block_size: list[int], |
| 526 | world_dp_size: tuple[int, int], |
| 527 | disable_deepgemm_ue8m0, |
| 528 | workspace_init, |
| 529 | ): |
| 530 | """Tests for Low-Latency DeepEP + DeepGemm integration.""" |
| 531 | assert not is_deep_gemm_e8m0_used() |
| 532 | |
| 533 | m, n, k = mnk |
| 534 | set_random_seed(7) |
| 535 | |
| 536 | if topk > num_experts: |
| 537 | pytest.skip(f"Skipping test: topk={topk} > E={num_experts}") |
| 538 | |
| 539 | world_size, dp_size = world_dp_size |
| 540 | config = TestConfig( |
| 541 | topk=topk, |
| 542 | m=m, |
| 543 | k=k, |
| 544 | n=n, |
| 545 | num_experts=num_experts, |
| 546 | per_act_token_quant=False, |
| 547 | block_size=block_size, |
| 548 | low_latency=True, |
| 549 | use_fp8_dispatch=use_fp8_dispatch, |
| 550 | ) |
| 551 | |
| 552 | w1, w2, w1_scale, w2_scale = make_block_quant_fp8_weights( |
| 553 | num_experts, n, k, block_size |
| 554 | ) |
| 555 | |
| 556 | parallel_launch( |
| 557 | world_size, |
| 558 | _test_deepep_deepgemm_moe, |
| 559 | dp_size, |
| 560 | config, |
| 561 | w1, |
| 562 | w2, |
| 563 | w1_scale, |
| 564 | w2_scale, |
| 565 | ) |