| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test DeepEP v2 (ElasticBuffer) dispatch-combine logic. |
| 4 | Compares against a pure-PyTorch reference MoE implementation. |
| 5 | """ |
| 6 | |
| 7 | import dataclasses |
| 8 | |
| 9 | import pytest |
| 10 | import torch.distributed |
| 11 | from torch.distributed import ProcessGroup |
| 12 | |
| 13 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 14 | from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights |
| 15 | from tests.kernels.utils import torch_experts |
| 16 | from vllm.config import VllmConfig, set_current_vllm_config |
| 17 | from vllm.forward_context import set_forward_context |
| 18 | from vllm.model_executor.layers.fused_moe import TritonExperts |
| 19 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 20 | from vllm.model_executor.layers.fused_moe.config import ( |
| 21 | FusedMoEQuantConfig, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.utils.flashinfer import has_flashinfer |
| 26 | from vllm.utils.import_utils import has_deep_ep_v2 |
| 27 | from vllm.utils.torch_utils import set_random_seed |
| 28 | from vllm.v1.worker.workspace import init_workspace_manager |
| 29 | |
| 30 | from ...utils import multi_gpu_test |
| 31 | from .parallel_utils import ProcessGroupInfo, parallel_launch |
| 32 | |
| 33 | if has_deep_ep_v2(): |
| 34 | from .parallel_utils import DeepEPV2Args, make_deepep_v2_a2a |
| 35 | |
| 36 | requires_deep_ep_v2 = pytest.mark.skipif( |
| 37 | not has_deep_ep_v2(), |
| 38 | reason="Requires DeepEP v2 (ElasticBuffer)", |
| 39 | ) |
| 40 | |
| 41 | |
| 42 | def assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None: |
| 43 | close = torch.isclose(actual, expected, atol=2e-1, rtol=2e-1) |
| 44 | close_fraction = close.float().mean().item() |
| 45 | assert close_fraction > 0.99, ( |
| 46 | f"Only {close_fraction:.1%} of FP8 outputs are within tolerance" |
| 47 | ) |
| 48 | |
| 49 | |
| 50 | @dataclasses.dataclass |
| 51 | class TestConfig: |
| 52 | dtype: torch.dtype |
| 53 | topk: int |
| 54 | m: int |
| 55 | k: int |
| 56 | n: int |
| 57 | num_experts: int |
| 58 | |
| 59 | |
| 60 | @dataclasses.dataclass |
| 61 | class TestTensors: |
| 62 | rank_tokens: torch.Tensor |
| 63 | rank_token_scales: torch.Tensor | None |
| 64 | intermediate_scales: torch.Tensor | None |
| 65 | topk: torch.Tensor |
| 66 | topk_weights: torch.Tensor |
| 67 | config: TestConfig |
| 68 | |
| 69 | @staticmethod |
| 70 | def make(config: TestConfig) -> "TestTensors": |
| 71 | assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn] |
| 72 | token_dtype = ( |
| 73 | torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype |
| 74 | ) |
| 75 | rank_tokens = ( |
| 76 | torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10 |
| 77 | ) |
| 78 | if config.dtype == torch.float8_e4m3fn: |
| 79 | rank_token_scales = torch.tensor(1 / 448, device="cuda") |
| 80 | intermediate_scales = torch.tensor(8 / 448, device="cuda") |
| 81 | else: |
| 82 | rank_token_scales = None |
| 83 | intermediate_scales = None |
| 84 | |
| 85 | topk = torch.stack( |
| 86 | [ |
| 87 | torch.randperm(config.num_experts, device="cuda")[: config.topk] |
| 88 | for _ in range(config.m) |
| 89 | ] |
| 90 | ).to(dtype=torch.int64) |
| 91 | topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda") |
| 92 | return TestTensors( |
| 93 | rank_tokens=rank_tokens, |
| 94 | rank_token_scales=rank_token_scales, |
| 95 | intermediate_scales=intermediate_scales, |
| 96 | topk=topk, |
| 97 | topk_weights=topk_weights, |
| 98 | config=config, |
| 99 | ) |
| 100 | |
| 101 | |
| 102 | def make_modular_kernel( |
| 103 | pg: ProcessGroup, |
| 104 | pgi: ProcessGroupInfo, |
| 105 | dp_size: int, |
| 106 | hidden_size: int, |
| 107 | num_experts: int, |
| 108 | num_local_experts: int, |
| 109 | topk: int, |
| 110 | q_dtype: torch.dtype | None, |
| 111 | use_fp8_dispatch: bool, |
| 112 | quant_config: FusedMoEQuantConfig, |
| 113 | use_cudagraph: bool = False, |
| 114 | ) -> FusedMoEKernel: |
| 115 | v2_args = DeepEPV2Args( |
| 116 | num_local_experts=num_local_experts, |
| 117 | num_experts=num_experts, |
| 118 | num_topk=topk, |
| 119 | hidden_size=hidden_size, |
| 120 | max_tokens_per_rank=8192, |
| 121 | use_fp8_dispatch=use_fp8_dispatch, |
| 122 | ) |
| 123 | |
| 124 | a2a = make_deepep_v2_a2a( |
| 125 | pg=pg, |
| 126 | pgi=pgi, |
| 127 | dp_size=dp_size, |
| 128 | v2_args=v2_args, |
| 129 | use_cudagraph=use_cudagraph, |
| 130 | ) |
| 131 | |
| 132 | moe_config = make_dummy_moe_config( |
| 133 | num_experts=num_local_experts, |
| 134 | experts_per_token=topk, |
| 135 | hidden_dim=hidden_size, |
| 136 | ) |
| 137 | |
| 138 | fused_experts = TritonExperts( |
| 139 | moe_config=moe_config, |
| 140 | quant_config=quant_config, |
| 141 | ) |
| 142 | |
| 143 | mk = FusedMoEKernel( |
| 144 | prepare_finalize=a2a, |
| 145 | fused_experts=fused_experts, |
| 146 | ) |
| 147 | return mk |
| 148 | |
| 149 | |
| 150 | def deepep_v2_moe_impl( |
| 151 | pg: ProcessGroup, |
| 152 | pgi: ProcessGroupInfo, |
| 153 | dp_size: int, |
| 154 | test_tensors: TestTensors, |
| 155 | w1: torch.Tensor, |
| 156 | w2: torch.Tensor, |
| 157 | w1_scale: torch.Tensor | None, |
| 158 | w2_scale: torch.Tensor | None, |
| 159 | num_experts: int, |
| 160 | topk: int, |
| 161 | use_fp8_dispatch: bool, |
| 162 | per_act_token_quant: bool, |
| 163 | ) -> torch.Tensor: |
| 164 | num_local_experts = w1.size(0) |
| 165 | |
| 166 | def build_expert_map(): |
| 167 | expert_map = torch.full((num_experts,), fill_value=-1, dtype=torch.int32) |
| 168 | s = pgi.rank * num_local_experts |
| 169 | e = s + num_local_experts |
| 170 | expert_map[s:e] = torch.tensor(list(range(num_local_experts))) |
| 171 | device = torch.accelerator.current_device_index() |
| 172 | return expert_map.to(device=device, dtype=torch.int32) |
| 173 | |
| 174 | is_quantized = w1.dtype == torch.float8_e4m3fn |
| 175 | q_dtype = torch.float8_e4m3fn if is_quantized else None |
| 176 | |
| 177 | quant_config = FusedMoEQuantConfig.make( |
| 178 | q_dtype, |
| 179 | w1_scale=w1_scale, |
| 180 | w2_scale=w2_scale, |
| 181 | per_act_token_quant=per_act_token_quant, |
| 182 | a1_scale=test_tensors.rank_token_scales, |
| 183 | a2_scale=test_tensors.intermediate_scales, |
| 184 | ) |
| 185 | |
| 186 | hidden_size = test_tensors.rank_tokens.size(1) |
| 187 | |
| 188 | mk: FusedMoEKernel = make_modular_kernel( |
| 189 | pg, |
| 190 | pgi, |
| 191 | dp_size, |
| 192 | hidden_size, |
| 193 | num_experts, |
| 194 | num_local_experts, |
| 195 | topk, |
| 196 | q_dtype, |
| 197 | use_fp8_dispatch, |
| 198 | quant_config, |
| 199 | ) |
| 200 | |
| 201 | out = mk.apply( |
| 202 | hidden_states=test_tensors.rank_tokens, |
| 203 | w1=w1, |
| 204 | w2=w2, |
| 205 | topk_weights=test_tensors.topk_weights, |
| 206 | topk_ids=test_tensors.topk, |
| 207 | activation=MoEActivation.SILU, |
| 208 | global_num_experts=num_experts, |
| 209 | expert_map=build_expert_map(), |
| 210 | apply_router_weight_on_input=False, |
| 211 | ) |
| 212 | |
| 213 | return out |
| 214 | |
| 215 | |
| 216 | def _deep_ep_v2_moe( |
| 217 | pgi: ProcessGroupInfo, |
| 218 | dp_size: int, |
| 219 | config: TestConfig, |
| 220 | w1: torch.Tensor, |
| 221 | w2: torch.Tensor, |
| 222 | w1_scale: torch.Tensor | None, |
| 223 | w2_scale: torch.Tensor | None, |
| 224 | use_fp8_dispatch: bool, |
| 225 | per_act_token_quant: bool, |
| 226 | ): |
| 227 | device = torch.device(f"cuda:{pgi.local_rank}") |
| 228 | init_workspace_manager(device) |
| 229 | |
| 230 | is_quantized = w1.dtype == torch.float8_e4m3fn |
| 231 | device_idx = torch.accelerator.current_device_index() |
| 232 | w1 = w1.to(device=device_idx) |
| 233 | w2 = w2.to(device=device_idx) |
| 234 | if is_quantized: |
| 235 | assert w1_scale is not None and w2_scale is not None |
| 236 | w1_scale = w1_scale.to(device=device_idx) |
| 237 | w2_scale = w2_scale.to(device=device_idx) |
| 238 | |
| 239 | pg = torch.distributed.new_group(list(range(pgi.world_size))) |
| 240 | # The caller's set_random_seed() only seeds the parent; spawn() gives each |
| 241 | # worker a fresh unseeded RNG, so seed here too or the inputs below differ |
| 242 | # every run. Offset by rank to keep the ranks' data distinct. |
| 243 | set_random_seed(7 + pgi.rank) |
| 244 | test_tensors = TestTensors.make(config) |
| 245 | |
| 246 | with set_current_vllm_config(VllmConfig()): |
| 247 | # Reference |
| 248 | q_dtype = torch.float8_e4m3fn if is_quantized else None |
| 249 | torch_combined = torch_experts( |
| 250 | test_tensors.rank_tokens, |
| 251 | w1, |
| 252 | w2, |
| 253 | test_tensors.topk_weights, |
| 254 | test_tensors.topk, |
| 255 | w1_scale=w1_scale, |
| 256 | w2_scale=w2_scale, |
| 257 | a1_scale=test_tensors.rank_token_scales, |
| 258 | a2_scale=test_tensors.intermediate_scales, |
| 259 | quant_dtype=q_dtype, |
| 260 | per_act_token_quant=per_act_token_quant, |
| 261 | ) |
| 262 | |
| 263 | # Splice experts for this rank |
| 264 | num_local_experts = config.num_experts // pgi.world_size |
| 265 | e_start = num_local_experts * pgi.rank |
| 266 | e_end = e_start + num_local_experts |
| 267 | w1_ep = w1[e_start:e_end] |
| 268 | w2_ep = w2[e_start:e_end] |
| 269 | |
| 270 | w1_scale_ep, w2_scale_ep = None, None |
| 271 | if is_quantized: |
| 272 | w1_scale_ep = w1_scale[e_start:e_end] # type: ignore |
| 273 | w2_scale_ep = w2_scale[e_start:e_end] # type: ignore |
| 274 | |
| 275 | deepep_combined = deepep_v2_moe_impl( |
| 276 | pg, |
| 277 | pgi, |
| 278 | dp_size, |
| 279 | test_tensors, |
| 280 | w1_ep, |
| 281 | w2_ep, |
| 282 | w1_scale_ep, |
| 283 | w2_scale_ep, |
| 284 | config.num_experts, |
| 285 | config.topk, |
| 286 | use_fp8_dispatch, |
| 287 | per_act_token_quant, |
| 288 | ) |
| 289 | |
| 290 | if is_quantized: |
| 291 | assert_fp8_close(torch_combined, deepep_combined) |
| 292 | else: |
| 293 | torch.testing.assert_close( |
| 294 | torch_combined, |
| 295 | deepep_combined, |
| 296 | atol=6e-2, |
| 297 | rtol=6e-2, |
| 298 | ) |
| 299 | |
| 300 | |
| 301 | MNKs = [ |
| 302 | (1, 256, 256), |
| 303 | (2, 256, 512), |
| 304 | (3, 1024, 2048), |
| 305 | (32, 256, 1024), |
| 306 | (45, 512, 2048), |
| 307 | (64, 1024, 1024), |
| 308 | (222, 1024, 2048), |
| 309 | ] |
| 310 | |
| 311 | DTYPES = [torch.bfloat16, torch.float8_e4m3fn] |
| 312 | |
| 313 | |
| 314 | @pytest.mark.parametrize("dtype", DTYPES) |
| 315 | @pytest.mark.parametrize("m,n,k", MNKs) |
| 316 | @pytest.mark.parametrize("num_experts", [32]) |
| 317 | @pytest.mark.parametrize("topk", [6]) |
| 318 | @pytest.mark.parametrize("world_dp_size", [(2, 1)]) |
| 319 | @multi_gpu_test(num_gpus=2) |
| 320 | @requires_deep_ep_v2 |
| 321 | def test_deep_ep_v2_moe( |
| 322 | dtype: torch.dtype, |
| 323 | m: int, |
| 324 | n: int, |
| 325 | k: int, |
| 326 | num_experts: int, |
| 327 | topk: int, |
| 328 | world_dp_size: tuple[int, int], |
| 329 | workspace_init, |
| 330 | ): |
| 331 | per_act_token_quant = False |
| 332 | use_fp8_dispatch = False |
| 333 | |
| 334 | set_random_seed(7) |
| 335 | world_size, dp_size = world_dp_size |
| 336 | config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts) |
| 337 | |
| 338 | quant_dtype = dtype if dtype == torch.float8_e4m3fn else None |
| 339 | (_, w1, w1_scale, _), (_, w2, w2_scale, _) = make_test_weights( |
| 340 | num_experts, |
| 341 | n, |
| 342 | k, |
| 343 | quant_dtype=quant_dtype, |
| 344 | per_out_ch_quant=True, |
| 345 | ) |
| 346 | |
| 347 | parallel_launch( |
| 348 | world_size, |
| 349 | _deep_ep_v2_moe, |
| 350 | dp_size, |
| 351 | config, |
| 352 | w1, |
| 353 | w2, |
| 354 | w1_scale, |
| 355 | w2_scale, |
| 356 | use_fp8_dispatch, |
| 357 | per_act_token_quant, |
| 358 | ) |
| 359 | |
| 360 | |
| 361 | EXPERTS_BACKENDS = [ |
| 362 | "flashinfer_trtllm", |
| 363 | "flashinfer_cutlass", |
| 364 | "trtllm_fp8", |
| 365 | ] |
| 366 | |
| 367 | |
| 368 | def _make_experts( |
| 369 | experts_backend: str, |
| 370 | config: TestConfig, |
| 371 | moe_config, |
| 372 | num_local_experts: int, |
| 373 | rank: int, |
| 374 | w1_bf16: torch.Tensor, |
| 375 | w2_bf16: torch.Tensor, |
| 376 | test_tensors: TestTensors, |
| 377 | ): |
| 378 | e_start = num_local_experts * rank |
| 379 | e_end = e_start + num_local_experts |
| 380 | |
| 381 | if experts_backend != "trtllm_fp8": |
| 382 | from vllm.model_executor.layers.fused_moe.config import ( |
| 383 | FUSED_MOE_UNQUANTIZED_CONFIG, |
| 384 | ) |
| 385 | from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( |
| 386 | backend_to_kernel_cls, |
| 387 | convert_to_unquantized_kernel_format, |
| 388 | map_unquantized_backend, |
| 389 | ) |
| 390 | |
| 391 | torch_combined = torch_experts( |
| 392 | test_tensors.rank_tokens, |
| 393 | w1_bf16, |
| 394 | w2_bf16, |
| 395 | test_tensors.topk_weights, |
| 396 | test_tensors.topk, |
| 397 | ) |
| 398 | backend = map_unquantized_backend(experts_backend) |
| 399 | w1_ep, w2_ep = convert_to_unquantized_kernel_format( |
| 400 | backend, |
| 401 | moe_config, |
| 402 | w1_bf16[e_start:e_end], |
| 403 | w2_bf16[e_start:e_end], |
| 404 | ) |
| 405 | experts_cls = next( |
| 406 | cls |
| 407 | for cls in backend_to_kernel_cls(backend) |
| 408 | if issubclass(cls, mk.FusedMoEExpertsModular) |
| 409 | ) |
| 410 | fused_experts = experts_cls( |
| 411 | moe_config=moe_config, |
| 412 | quant_config=FUSED_MOE_UNQUANTIZED_CONFIG, |
| 413 | ) |
| 414 | return fused_experts, w1_ep, w2_ep, torch_combined, 1e-1, 2e-1 |
| 415 | |
| 416 | from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves |
| 417 | from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( |
| 418 | TrtLlmFp8ExpertsModular, |
| 419 | ) |
| 420 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 421 | Fp8MoeBackend, |
| 422 | convert_to_fp8_moe_kernel_format, |
| 423 | ) |
| 424 | |
| 425 | w1_ref = w1_bf16.to(torch.float8_e4m3fn).to(torch.bfloat16) |
| 426 | w2_ref = w2_bf16.to(torch.float8_e4m3fn).to(torch.bfloat16) |
| 427 | |
| 428 | block_shape = [128, 128] |
| 429 | qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape) |
| 430 | assert qw.w13_weight_scale is not None |
| 431 | assert qw.w2_weight_scale is not None |
| 432 | |
| 433 | reference_topk_weights = test_tensors.topk_weights.to(torch.bfloat16).to( |
| 434 | torch.float32 |
| 435 | ) |
| 436 | torch_combined = torch_experts( |
| 437 | test_tensors.rank_tokens, |
| 438 | qw.w13_weight, |
| 439 | qw.w2_weight, |
| 440 | reference_topk_weights, |
| 441 | test_tensors.topk, |
| 442 | w1_scale=qw.w13_weight_scale, |
| 443 | w2_scale=qw.w2_weight_scale, |
| 444 | quant_dtype=torch.float8_e4m3fn, |
| 445 | block_shape=block_shape, |
| 446 | ) |
| 447 | |
| 448 | class _MockLayer: |
| 449 | weight_block_size = block_shape |
| 450 | |
| 451 | class moe_config: |
| 452 | is_act_and_mul = True |
| 453 | intermediate_size_per_partition = config.n |
| 454 | |
| 455 | class activation: |
| 456 | is_gated = True |
| 457 | |
| 458 | w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = convert_to_fp8_moe_kernel_format( |
| 459 | fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 460 | layer=_MockLayer(), |
| 461 | w13=qw.w13_weight[e_start:e_end], |
| 462 | w2=qw.w2_weight[e_start:e_end], |
| 463 | w13_scale=qw.w13_weight_scale[e_start:e_end], |
| 464 | w2_scale=qw.w2_weight_scale[e_start:e_end], |
| 465 | w13_input_scale=None, |
| 466 | w2_input_scale=None, |
| 467 | ) |
| 468 | |
| 469 | fused_experts = TrtLlmFp8ExpertsModular( |
| 470 | moe_config=moe_config, |
| 471 | quant_config=FusedMoEQuantConfig.make( |
| 472 | torch.float8_e4m3fn, |
| 473 | block_shape=block_shape, |
| 474 | w1_scale=w1_scale_ep, |
| 475 | w2_scale=w2_scale_ep, |
| 476 | ), |
| 477 | ) |
| 478 | return fused_experts, w1_ep, w2_ep, torch_combined, 6e-2, 6e-2 |
| 479 | |
| 480 | |
| 481 | def _deep_ep_v2_moe_backends( |
| 482 | pgi: ProcessGroupInfo, |
| 483 | dp_size: int, |
| 484 | config: TestConfig, |
| 485 | use_cudagraph: bool, |
| 486 | experts_backend: str, |
| 487 | ): |
| 488 | import tempfile |
| 489 | |
| 490 | from vllm.config import KernelConfig |
| 491 | from vllm.distributed import ( |
| 492 | init_distributed_environment, |
| 493 | initialize_model_parallel, |
| 494 | ) |
| 495 | |
| 496 | device = torch.device(f"cuda:{pgi.local_rank}") |
| 497 | init_workspace_manager(device) |
| 498 | |
| 499 | pg = torch.distributed.new_group(list(range(pgi.world_size))) |
| 500 | set_random_seed(7 + pgi.rank) |
| 501 | test_tensors = TestTensors.make(config) |
| 502 | num_local_experts = config.num_experts // pgi.world_size |
| 503 | hidden_size = config.k |
| 504 | |
| 505 | # All ranks must use the same global weights before taking their EP slice. |
| 506 | w1_bf16 = ( |
| 507 | torch.randn( |
| 508 | (config.num_experts, 2 * config.n, config.k), |
| 509 | device="cuda", |
| 510 | dtype=torch.bfloat16, |
| 511 | ) |
| 512 | / 15 |
| 513 | ) |
| 514 | w2_bf16 = ( |
| 515 | torch.randn( |
| 516 | (config.num_experts, config.k, config.n), |
| 517 | device="cuda", |
| 518 | dtype=torch.bfloat16, |
| 519 | ) |
| 520 | / 15 |
| 521 | ) |
| 522 | torch.distributed.broadcast(w1_bf16, src=0, group=pg) |
| 523 | torch.distributed.broadcast(w2_bf16, src=0, group=pg) |
| 524 | |
| 525 | vllm_cfg = VllmConfig() |
| 526 | vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm") |
| 527 | |
| 528 | with set_current_vllm_config(vllm_cfg): |
| 529 | temp_file = tempfile.mktemp() |
| 530 | init_distributed_environment( |
| 531 | world_size=pgi.world_size, |
| 532 | rank=pgi.rank, |
| 533 | distributed_init_method=f"file://{temp_file}", |
| 534 | local_rank=pgi.local_rank, |
| 535 | backend="nccl", |
| 536 | ) |
| 537 | initialize_model_parallel(tensor_model_parallel_size=1) |
| 538 | |
| 539 | moe_config = make_dummy_moe_config( |
| 540 | num_experts=config.num_experts, |
| 541 | num_local_experts=num_local_experts, |
| 542 | experts_per_token=config.topk, |
| 543 | hidden_dim=hidden_size, |
| 544 | intermediate_size=config.n, |
| 545 | ) |
| 546 | moe_config = dataclasses.replace( |
| 547 | moe_config, |
| 548 | moe_parallel_config=dataclasses.replace( |
| 549 | moe_config.moe_parallel_config, |
| 550 | ep_size=pgi.world_size, |
| 551 | ep_rank=pgi.rank, |
| 552 | use_ep=True, |
| 553 | all2all_backend="deepep_v2", |
| 554 | ), |
| 555 | ) |
| 556 | |
| 557 | ( |
| 558 | fused_experts, |
| 559 | w1_ep, |
| 560 | w2_ep, |
| 561 | torch_combined, |
| 562 | atol, |
| 563 | rtol, |
| 564 | ) = _make_experts( |
| 565 | experts_backend, |
| 566 | config, |
| 567 | moe_config, |
| 568 | num_local_experts, |
| 569 | pgi.rank, |
| 570 | w1_bf16, |
| 571 | w2_bf16, |
| 572 | test_tensors, |
| 573 | ) |
| 574 | |
| 575 | v2_args = DeepEPV2Args( |
| 576 | num_local_experts=num_local_experts, |
| 577 | num_experts=config.num_experts, |
| 578 | num_topk=config.topk, |
| 579 | hidden_size=hidden_size, |
| 580 | max_tokens_per_rank=8192, |
| 581 | use_fp8_dispatch=False, |
| 582 | ) |
| 583 | a2a = make_deepep_v2_a2a( |
| 584 | pg=pg, |
| 585 | pgi=pgi, |
| 586 | dp_size=dp_size, |
| 587 | v2_args=v2_args, |
| 588 | use_cudagraph=use_cudagraph, |
| 589 | ) |
| 590 | mk_kernel = FusedMoEKernel( |
| 591 | prepare_finalize=a2a, |
| 592 | fused_experts=fused_experts, |
| 593 | ) |
| 594 | |
| 595 | with set_forward_context(None, vllm_cfg): |
| 596 | for _ in range(3): |
| 597 | out = mk_kernel.apply( |
| 598 | hidden_states=test_tensors.rank_tokens, |
| 599 | w1=w1_ep, |
| 600 | w2=w2_ep, |
| 601 | topk_weights=test_tensors.topk_weights, |
| 602 | topk_ids=test_tensors.topk, |
| 603 | activation=MoEActivation.SILU, |
| 604 | global_num_experts=config.num_experts, |
| 605 | expert_map=None, |
| 606 | apply_router_weight_on_input=False, |
| 607 | ) |
| 608 | |
| 609 | torch.testing.assert_close(torch_combined, out, atol=atol, rtol=rtol) |
| 610 | |
| 611 | |
| 612 | @pytest.mark.parametrize("m,n,k", [(32, 256, 1024)]) |
| 613 | @pytest.mark.parametrize("num_experts", [32]) |
| 614 | @pytest.mark.parametrize("topk", [6]) |
| 615 | @pytest.mark.parametrize("world_dp_size", [(2, 1)]) |
| 616 | @pytest.mark.parametrize("experts_backend", EXPERTS_BACKENDS) |
| 617 | @pytest.mark.parametrize("use_cudagraph", [True, False]) |
| 618 | @multi_gpu_test(num_gpus=2) |
| 619 | @requires_deep_ep_v2 |
| 620 | @pytest.mark.skipif( |
| 621 | not has_flashinfer() or not current_platform.has_device_capability(100), |
| 622 | reason="Requires FlashInfer TRTLLM fused MoE (SM100)", |
| 623 | ) |
| 624 | def test_deep_ep_v2_moe_backends( |
| 625 | m: int, |
| 626 | n: int, |
| 627 | k: int, |
| 628 | num_experts: int, |
| 629 | topk: int, |
| 630 | world_dp_size: tuple[int, int], |
| 631 | experts_backend: str, |
| 632 | use_cudagraph: bool, |
| 633 | workspace_init, |
| 634 | ): |
| 635 | set_random_seed(7) |
| 636 | world_size, dp_size = world_dp_size |
| 637 | config = TestConfig( |
| 638 | dtype=torch.float8_e4m3fn |
| 639 | if experts_backend == "trtllm_fp8" |
| 640 | else torch.bfloat16, |
| 641 | topk=topk, |
| 642 | m=m, |
| 643 | k=k, |
| 644 | n=n, |
| 645 | num_experts=num_experts, |
| 646 | ) |
| 647 | |
| 648 | parallel_launch( |
| 649 | world_size, |
| 650 | _deep_ep_v2_moe_backends, |
| 651 | dp_size, |
| 652 | config, |
| 653 | use_cudagraph, |
| 654 | experts_backend, |
| 655 | ) |