| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test deepep dispatch-combine logic.""" |
| 4 | |
| 5 | import dataclasses |
| 6 | |
| 7 | import pytest |
| 8 | import torch.distributed |
| 9 | from torch.distributed import ProcessGroup |
| 10 | |
| 11 | import vllm.envs as envs |
| 12 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 13 | from vllm import _custom_ops as ops |
| 14 | from vllm.config import VllmConfig, set_current_vllm_config |
| 15 | from vllm.model_executor.layers.activation import SiluAndMul |
| 16 | from vllm.model_executor.layers.fused_moe import TritonExperts |
| 17 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 18 | from vllm.model_executor.layers.fused_moe.config import ( |
| 19 | FusedMoEQuantConfig, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( |
| 22 | BatchedTritonExperts, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel |
| 25 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 26 | per_token_group_quant_fp8, |
| 27 | ) |
| 28 | from vllm.platforms import current_platform |
| 29 | from vllm.utils.import_utils import has_deep_ep |
| 30 | from vllm.utils.torch_utils import set_random_seed |
| 31 | from vllm.v1.worker.workspace import init_workspace_manager |
| 32 | |
| 33 | from ...utils import multi_gpu_test |
| 34 | from .parallel_utils import ProcessGroupInfo, parallel_launch |
| 35 | |
| 36 | if has_deep_ep(): |
| 37 | from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import ( |
| 38 | DeepEPHTPrepareAndFinalize, |
| 39 | ) |
| 40 | from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( |
| 41 | DeepEPLLPrepareAndFinalize, |
| 42 | ) |
| 43 | |
| 44 | from .parallel_utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a |
| 45 | |
| 46 | requires_deep_ep = pytest.mark.skipif( |
| 47 | not has_deep_ep(), |
| 48 | reason="Requires deep_ep kernels", |
| 49 | ) |
| 50 | |
| 51 | MAX_TOKENS_PER_RANK = 64 |
| 52 | |
| 53 | |
| 54 | def make_weights( |
| 55 | e, n, k, dtype |
| 56 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 57 | """Return weights w1, w2, w1_scale, w2_scale.""" |
| 58 | if dtype in [torch.float16, torch.bfloat16]: |
| 59 | w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 |
| 60 | w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 |
| 61 | return w1, w2, None, None |
| 62 | |
| 63 | # per-out-channel weight quantization |
| 64 | assert dtype == current_platform.fp8_dtype() |
| 65 | # Keep FP8 inputs finite and bounded. torch.empty made this test depend on |
| 66 | # allocator contents, while larger values exceed its fixed W8A8 tolerance. |
| 67 | w1 = torch.randn( |
| 68 | (e, 2 * n, k), device=current_platform.device_type, dtype=torch.float16 |
| 69 | ).div_(100) |
| 70 | w2 = torch.randn( |
| 71 | (e, k, n), device=current_platform.device_type, dtype=torch.float16 |
| 72 | ).div_(100) |
| 73 | |
| 74 | n_b_scales = 2 * n |
| 75 | k_b_scales = k |
| 76 | w1_q = torch.empty_like(w1, dtype=dtype) |
| 77 | w2_q = torch.empty_like(w2, dtype=dtype) |
| 78 | w1_scale = torch.empty((e, n_b_scales, 1), device="cuda", dtype=torch.float32) |
| 79 | w2_scale = torch.empty((e, k_b_scales, 1), device="cuda", dtype=torch.float32) |
| 80 | for expert in range(e): |
| 81 | w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant( |
| 82 | w1[expert], use_per_token_if_dynamic=True |
| 83 | ) |
| 84 | w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant( |
| 85 | w2[expert], use_per_token_if_dynamic=True |
| 86 | ) |
| 87 | return w1_q, w2_q, w1_scale, w2_scale |
| 88 | |
| 89 | |
| 90 | @dataclasses.dataclass |
| 91 | class TestConfig: |
| 92 | dtype: torch.dtype |
| 93 | topk: int |
| 94 | m: int |
| 95 | k: int |
| 96 | n: int |
| 97 | num_experts: int |
| 98 | |
| 99 | |
| 100 | @dataclasses.dataclass |
| 101 | class TestTensors: |
| 102 | rank_tokens: torch.Tensor # all ranks make this many tokens |
| 103 | rank_token_scales: torch.Tensor | None |
| 104 | topk: torch.Tensor |
| 105 | topk_weights: torch.Tensor |
| 106 | config: TestConfig |
| 107 | |
| 108 | @staticmethod |
| 109 | def make(config: TestConfig, low_latency_mode: bool) -> "TestTensors": |
| 110 | # TODO (varun) - check that float16 works ? |
| 111 | assert config.dtype in [torch.bfloat16, current_platform.fp8_dtype()] |
| 112 | token_dtype = ( |
| 113 | torch.bfloat16 |
| 114 | if config.dtype == current_platform.fp8_dtype() |
| 115 | else config.dtype |
| 116 | ) |
| 117 | rank_tokens = ( |
| 118 | torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10 |
| 119 | ) |
| 120 | rank_token_scales = None |
| 121 | |
| 122 | topk = torch.randint( |
| 123 | low=0, high=config.num_experts, size=(config.m, config.topk), device="cuda" |
| 124 | ).to(dtype=torch.int64) |
| 125 | topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda") |
| 126 | return TestTensors( |
| 127 | rank_tokens=rank_tokens, |
| 128 | rank_token_scales=rank_token_scales, |
| 129 | topk=topk, |
| 130 | topk_weights=topk_weights, |
| 131 | config=config, |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | def make_modular_kernel( |
| 136 | pg: ProcessGroup, |
| 137 | pgi: ProcessGroupInfo, |
| 138 | low_latency_mode: bool, |
| 139 | hidden_size: int, |
| 140 | dp_size: int, |
| 141 | num_experts: int, |
| 142 | num_local_experts: int, |
| 143 | q_dtype: torch.dtype | None, |
| 144 | use_fp8_dispatch: bool, |
| 145 | quant_config: FusedMoEQuantConfig, |
| 146 | ) -> FusedMoEKernel: |
| 147 | ht_args: DeepEPHTArgs | None = None |
| 148 | ll_args: DeepEPLLArgs | None = None |
| 149 | |
| 150 | if low_latency_mode: |
| 151 | ll_args = DeepEPLLArgs( |
| 152 | max_tokens_per_rank=MAX_TOKENS_PER_RANK, |
| 153 | hidden_size=hidden_size, |
| 154 | num_experts=num_experts, |
| 155 | use_fp8_dispatch=use_fp8_dispatch, |
| 156 | ) |
| 157 | else: |
| 158 | assert not use_fp8_dispatch, ( |
| 159 | "FP8 Dispatch is valid only for low-latency kernels" |
| 160 | ) |
| 161 | ht_args = DeepEPHTArgs(num_local_experts=num_local_experts) |
| 162 | |
| 163 | a2a: DeepEPHTPrepareAndFinalize | DeepEPLLPrepareAndFinalize = make_deepep_a2a( |
| 164 | pg=pg, |
| 165 | pgi=pgi, |
| 166 | dp_size=dp_size, |
| 167 | q_dtype=q_dtype, |
| 168 | block_shape=None, |
| 169 | deepep_ht_args=ht_args, |
| 170 | deepep_ll_args=ll_args, |
| 171 | ) |
| 172 | |
| 173 | num_dispatchers = pgi.world_size // dp_size |
| 174 | |
| 175 | moe_config = make_dummy_moe_config() |
| 176 | |
| 177 | if low_latency_mode: |
| 178 | assert not quant_config.per_act_token_quant, "not supported in ll mode" |
| 179 | fused_experts = BatchedTritonExperts( |
| 180 | max_num_tokens=MAX_TOKENS_PER_RANK, |
| 181 | num_dispatchers=num_dispatchers, |
| 182 | moe_config=moe_config, |
| 183 | quant_config=quant_config, |
| 184 | ) |
| 185 | else: |
| 186 | fused_experts = TritonExperts( |
| 187 | moe_config=moe_config, |
| 188 | quant_config=quant_config, |
| 189 | ) |
| 190 | |
| 191 | mk = FusedMoEKernel( |
| 192 | prepare_finalize=a2a, |
| 193 | fused_experts=fused_experts, |
| 194 | ) |
| 195 | return mk |
| 196 | |
| 197 | |
| 198 | def deep_ep_moe_impl( |
| 199 | pg: ProcessGroup, |
| 200 | pgi: ProcessGroupInfo, |
| 201 | low_latency_mode: bool, |
| 202 | dp_size: int, |
| 203 | test_tensors: TestTensors, |
| 204 | w1: torch.Tensor, |
| 205 | w2: torch.Tensor, |
| 206 | w1_scale: torch.Tensor | None, |
| 207 | w2_scale: torch.Tensor | None, |
| 208 | num_experts: int, |
| 209 | use_fp8_dispatch: bool, |
| 210 | per_act_token_quant: bool, |
| 211 | ) -> torch.Tensor: |
| 212 | num_local_experts = w1.size(0) |
| 213 | |
| 214 | def build_expert_map(): |
| 215 | num_local_experts = w1.size(0) |
| 216 | expert_map = torch.full((num_experts,), fill_value=-1, dtype=torch.int32) |
| 217 | s = pgi.rank * num_local_experts |
| 218 | e = s + num_local_experts |
| 219 | expert_map[s:e] = torch.tensor(list(range(num_local_experts))) |
| 220 | device = torch.accelerator.current_device_index() |
| 221 | return expert_map.to(device=device, dtype=torch.int32) |
| 222 | |
| 223 | hidden_size = test_tensors.rank_tokens.size(1) |
| 224 | is_quantized = w1.dtype == current_platform.fp8_dtype() |
| 225 | q_dtype = None |
| 226 | if is_quantized: |
| 227 | q_dtype = current_platform.fp8_dtype() |
| 228 | |
| 229 | out_hidden_states = torch.empty_like(test_tensors.rank_tokens) |
| 230 | total_num_tokens = test_tensors.rank_tokens.size(0) |
| 231 | |
| 232 | quant_config = FusedMoEQuantConfig.make( |
| 233 | q_dtype, |
| 234 | w1_scale=w1_scale, |
| 235 | w2_scale=w2_scale, |
| 236 | per_act_token_quant=per_act_token_quant, |
| 237 | a1_scale=test_tensors.rank_token_scales, |
| 238 | ) |
| 239 | |
| 240 | # Build the kernel (and its DeepEP buffer) once and reuse it across chunks. |
| 241 | # Re-creating it per chunk re-inits rocSHMEM, which only allows one |
| 242 | # allocation per process on ROCm. The buffer is sized by max_tokens_per_rank |
| 243 | # so it is valid for every chunk (mirrors production's cached all2all handle). |
| 244 | mk: FusedMoEKernel = make_modular_kernel( |
| 245 | pg, |
| 246 | pgi, |
| 247 | low_latency_mode, |
| 248 | hidden_size, |
| 249 | dp_size, |
| 250 | num_experts, |
| 251 | num_local_experts, |
| 252 | q_dtype, |
| 253 | use_fp8_dispatch, |
| 254 | quant_config, |
| 255 | ) |
| 256 | |
| 257 | def process_chunk(chunk_start, chunk_end, skip_result_store=False): |
| 258 | rank_tokens_chunk = test_tensors.rank_tokens[chunk_start:chunk_end] |
| 259 | topk_weights_chunk = test_tensors.topk_weights[chunk_start:chunk_end] |
| 260 | topk_chunk = test_tensors.topk[chunk_start:chunk_end] |
| 261 | |
| 262 | if low_latency_mode: |
| 263 | # Reusing one buffer leaves it dirty; the low-latency kernels need |
| 264 | # the zero-initialized regions reset before each dispatch. |
| 265 | mk.prepare_finalize.buffer.clean_low_latency_buffer( |
| 266 | MAX_TOKENS_PER_RANK, |
| 267 | hidden_size, |
| 268 | num_experts, |
| 269 | ) |
| 270 | |
| 271 | out = mk.apply( |
| 272 | hidden_states=rank_tokens_chunk, |
| 273 | w1=w1, |
| 274 | w2=w2, |
| 275 | topk_weights=topk_weights_chunk, |
| 276 | topk_ids=topk_chunk, |
| 277 | activation=MoEActivation.SILU, |
| 278 | global_num_experts=num_experts, |
| 279 | expert_map=build_expert_map(), |
| 280 | apply_router_weight_on_input=False, |
| 281 | ) |
| 282 | |
| 283 | if not skip_result_store: |
| 284 | out_hidden_states[chunk_start:chunk_end, :].copy_(out, non_blocking=True) |
| 285 | |
| 286 | max_num_tokens_per_dp = ( |
| 287 | MAX_TOKENS_PER_RANK if low_latency_mode else total_num_tokens |
| 288 | ) |
| 289 | |
| 290 | for chunk_start_ in range(0, total_num_tokens, max_num_tokens_per_dp): |
| 291 | chunk_start = chunk_start_ |
| 292 | chunk_end = min(chunk_start + max_num_tokens_per_dp, total_num_tokens) |
| 293 | # clamp start and end |
| 294 | chunk_start = min(chunk_start, total_num_tokens - 1) |
| 295 | chunk_end = min(chunk_end, total_num_tokens) |
| 296 | |
| 297 | process_chunk( |
| 298 | chunk_start, chunk_end, skip_result_store=chunk_start_ >= total_num_tokens |
| 299 | ) |
| 300 | |
| 301 | return out_hidden_states |
| 302 | |
| 303 | |
| 304 | def torch_moe_impl( |
| 305 | test_tensors: TestTensors, |
| 306 | w1: torch.Tensor, |
| 307 | w2: torch.Tensor, |
| 308 | w1_scale: torch.Tensor | None, |
| 309 | w2_scale: torch.Tensor | None, |
| 310 | using_fp8_dispatch: bool, |
| 311 | per_act_token_quant: bool, |
| 312 | ): |
| 313 | a, topk_ids, topk_weights = ( |
| 314 | test_tensors.rank_tokens, |
| 315 | test_tensors.topk, |
| 316 | test_tensors.topk_weights, |
| 317 | ) |
| 318 | if using_fp8_dispatch: |
| 319 | # The DeepEP implementation is requested to dispatch using FP8. |
| 320 | # For numerical stability for testing, emulate the fp8 dispatch by |
| 321 | # blockwise quant and de-quant. |
| 322 | assert not per_act_token_quant |
| 323 | a = test_tensors.rank_tokens |
| 324 | aq, aq_scale = per_token_group_quant_fp8(a, 128, use_ue8m0=False) |
| 325 | a = ( |
| 326 | (aq.view(-1, 128).to(torch.float32) * aq_scale.view(-1, 1)) |
| 327 | .view(a.shape) |
| 328 | .to(a.dtype) |
| 329 | ) |
| 330 | |
| 331 | is_quantized = w1.dtype == current_platform.fp8_dtype() |
| 332 | a_dtype = a.dtype |
| 333 | if is_quantized: |
| 334 | w1 = w1.to(dtype=torch.float32) * w1_scale |
| 335 | w2 = w2.to(dtype=torch.float32) * w2_scale |
| 336 | a = a.to(dtype=torch.float32) |
| 337 | |
| 338 | m, _ = a.shape |
| 339 | topk = topk_ids.size(1) |
| 340 | out = torch.zeros_like(a) |
| 341 | |
| 342 | for i in range(m): |
| 343 | a_i = a[i] |
| 344 | o_i = out[i] |
| 345 | for j in range(topk): |
| 346 | e = topk_ids[i][j] |
| 347 | e_w = topk_weights[i][j] |
| 348 | w1_e = w1[e] |
| 349 | w2_e = w2[e] |
| 350 | o_i += ( |
| 351 | SiluAndMul()(a_i @ w1_e.transpose(0, 1)) @ w2_e.transpose(0, 1) |
| 352 | ) * e_w |
| 353 | |
| 354 | if is_quantized: |
| 355 | out = out.to(dtype=a_dtype) |
| 356 | |
| 357 | return out |
| 358 | |
| 359 | |
| 360 | def assert_deepep_close( |
| 361 | expected: torch.Tensor, |
| 362 | actual: torch.Tensor, |
| 363 | k: int, |
| 364 | use_fp8_dispatch: bool, |
| 365 | ) -> None: |
| 366 | torch.testing.assert_close( |
| 367 | expected, |
| 368 | actual, |
| 369 | atol=6e-2, |
| 370 | rtol=6e-2, |
| 371 | ) |
| 372 | |
| 373 | |
| 374 | def _deep_ep_moe( |
| 375 | pgi: ProcessGroupInfo, |
| 376 | low_latency_mode: bool, |
| 377 | dp_size: int, |
| 378 | config: TestConfig, |
| 379 | w1: torch.Tensor, |
| 380 | w2: torch.Tensor, |
| 381 | w1_scale: torch.Tensor | None, |
| 382 | w2_scale: torch.Tensor | None, |
| 383 | use_fp8_dispatch: bool, |
| 384 | per_act_token_quant: bool, |
| 385 | ): |
| 386 | # Set seed in worker process for deterministic tensor generation. |
| 387 | set_random_seed(7) |
| 388 | |
| 389 | device = torch.device(f"cuda:{pgi.local_rank}") |
| 390 | init_workspace_manager(device) |
| 391 | |
| 392 | if not low_latency_mode: |
| 393 | assert not use_fp8_dispatch, ( |
| 394 | "FP8 dispatch interface is available only in low-latency mode" |
| 395 | ) |
| 396 | |
| 397 | is_quantized = w1.dtype == current_platform.fp8_dtype() |
| 398 | device_idx = torch.accelerator.current_device_index() |
| 399 | w1 = w1.to(device=device_idx) |
| 400 | w2 = w2.to(device=device_idx) |
| 401 | if is_quantized: |
| 402 | assert w1_scale is not None and w2_scale is not None |
| 403 | w1_scale = w1_scale.to(device=device_idx) |
| 404 | w2_scale = w2_scale.to(device=device_idx) |
| 405 | |
| 406 | if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP: |
| 407 | pg = torch.distributed.split_group( |
| 408 | split_ranks=[list(range(pgi.world_size))], |
| 409 | group_desc="deepep_test", |
| 410 | ) |
| 411 | else: |
| 412 | pg = torch.distributed.new_group(list(range(pgi.world_size))) |
| 413 | test_tensors = TestTensors.make(config, low_latency_mode) |
| 414 | |
| 415 | with set_current_vllm_config(VllmConfig()): |
| 416 | # Reference |
| 417 | torch_combined = torch_moe_impl( |
| 418 | test_tensors, |
| 419 | w1, |
| 420 | w2, |
| 421 | w1_scale, |
| 422 | w2_scale, |
| 423 | use_fp8_dispatch, |
| 424 | per_act_token_quant, |
| 425 | ) |
| 426 | |
| 427 | # Splice experts for this rank. |
| 428 | num_local_experts = config.num_experts // pgi.world_size |
| 429 | e_start = num_local_experts * pgi.rank |
| 430 | e_end = e_start + num_local_experts |
| 431 | w1_ep = w1[e_start:e_end] |
| 432 | w2_ep = w2[e_start:e_end] |
| 433 | |
| 434 | w1_scale_ep, w2_scale_ep = None, None |
| 435 | if is_quantized: |
| 436 | w1_scale_ep = w1_scale[e_start:e_end] # type: ignore |
| 437 | w2_scale_ep = w2_scale[e_start:e_end] # type: ignore |
| 438 | deepep_combined = deep_ep_moe_impl( |
| 439 | pg, |
| 440 | pgi, |
| 441 | low_latency_mode, |
| 442 | dp_size, |
| 443 | test_tensors, |
| 444 | w1_ep, |
| 445 | w2_ep, |
| 446 | w1_scale_ep, |
| 447 | w2_scale_ep, |
| 448 | config.num_experts, |
| 449 | use_fp8_dispatch, |
| 450 | per_act_token_quant, |
| 451 | ) |
| 452 | |
| 453 | assert_deepep_close(torch_combined, deepep_combined, config.k, use_fp8_dispatch) |
| 454 | |
| 455 | |
| 456 | MNKs = [ |
| 457 | (1, 128, 128), |
| 458 | (2, 128, 512), |
| 459 | (3, 1024, 2048), |
| 460 | (32, 128, 1024), |
| 461 | (45, 512, 2048), |
| 462 | (64, 1024, 1024), |
| 463 | (222, 1024, 2048), |
| 464 | ] |
| 465 | |
| 466 | DTYPES = [torch.bfloat16, current_platform.fp8_dtype()] |
| 467 | |
| 468 | |
| 469 | @pytest.mark.parametrize("dtype", DTYPES) |
| 470 | @pytest.mark.parametrize("m,n,k", MNKs) |
| 471 | @pytest.mark.parametrize("num_experts", [32]) |
| 472 | @pytest.mark.parametrize("topk", [6]) |
| 473 | @pytest.mark.parametrize("world_dp_size", [(2, 1)]) |
| 474 | @pytest.mark.parametrize("per_act_token_quant", [False, True]) |
| 475 | @multi_gpu_test(num_gpus=2) |
| 476 | @requires_deep_ep |
| 477 | def test_deep_ep_moe( |
| 478 | dtype: torch.dtype, |
| 479 | m: int, |
| 480 | n: int, |
| 481 | k: int, |
| 482 | num_experts: int, |
| 483 | topk: int, |
| 484 | world_dp_size: tuple[int, int], |
| 485 | per_act_token_quant: bool, |
| 486 | workspace_init, |
| 487 | ): |
| 488 | low_latency_mode = False |
| 489 | use_fp8_dispatch = False |
| 490 | |
| 491 | set_random_seed(7) |
| 492 | world_size, dp_size = world_dp_size |
| 493 | config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts) |
| 494 | |
| 495 | w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype) |
| 496 | |
| 497 | parallel_launch( |
| 498 | world_size, |
| 499 | _deep_ep_moe, |
| 500 | low_latency_mode, |
| 501 | dp_size, |
| 502 | config, |
| 503 | w1, |
| 504 | w2, |
| 505 | w1_scale, |
| 506 | w2_scale, |
| 507 | use_fp8_dispatch, |
| 508 | per_act_token_quant, |
| 509 | ) |
| 510 | |
| 511 | |
| 512 | MNKs = [ |
| 513 | (1, 128, 2560), |
| 514 | (2, 128, 2560), |
| 515 | (3, 1024, 2560), |
| 516 | (32, 128, 2560), |
| 517 | (45, 512, 2560), |
| 518 | (64, 1024, 2560), |
| 519 | (222, 1024, 2560), |
| 520 | ] |
| 521 | DTYPES = [current_platform.fp8_dtype(), torch.bfloat16] |
| 522 | USE_FP8_DISPATCH = [True, False] |
| 523 | |
| 524 | |
| 525 | @pytest.mark.parametrize("dtype", DTYPES) |
| 526 | @pytest.mark.parametrize("m,n,k", MNKs) |
| 527 | @pytest.mark.parametrize("num_experts", [32]) |
| 528 | @pytest.mark.parametrize("topk", [6]) |
| 529 | @pytest.mark.parametrize("world_dp_size", [(2, 1)]) |
| 530 | @pytest.mark.parametrize("use_fp8_dispatch", USE_FP8_DISPATCH) |
| 531 | @multi_gpu_test(num_gpus=2) |
| 532 | @requires_deep_ep |
| 533 | def test_low_latency_deep_ep_moe( |
| 534 | dtype: torch.dtype, |
| 535 | m: int, |
| 536 | n: int, |
| 537 | k: int, |
| 538 | num_experts: int, |
| 539 | topk: int, |
| 540 | world_dp_size: tuple[int, int], |
| 541 | use_fp8_dispatch: bool, |
| 542 | workspace_init, |
| 543 | ): |
| 544 | low_latency_mode = True |
| 545 | |
| 546 | if low_latency_mode and k not in DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES: |
| 547 | pytest.skip( |
| 548 | f"Skipping test as hidden size {k} is not in list of supported " |
| 549 | f"hidden sizes {DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES}" |
| 550 | ) |
| 551 | |
| 552 | set_random_seed(7) |
| 553 | world_size, dp_size = world_dp_size |
| 554 | config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts) |
| 555 | |
| 556 | w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype) |
| 557 | |
| 558 | parallel_launch( |
| 559 | world_size, |
| 560 | _deep_ep_moe, |
| 561 | low_latency_mode, |
| 562 | dp_size, |
| 563 | config, |
| 564 | w1, |
| 565 | w2, |
| 566 | w1_scale, |
| 567 | w2_scale, |
| 568 | use_fp8_dispatch, |
| 569 | False, |
| 570 | ) |