[Bugfix] vLLM crashes at startup when DeepEP v2 is used with `--enforce-eager` wiht TRTLLM Bf16 (#51824)
Sage Moore committed
Aug 19, 2026 at 17:18 UTC
5bf0dbd6fd7b0e9d0aa744e483037563111ce887
3 files changed
+170
-111
tests/kernels/moe/parallel_utils.py
+1
@@ -267,6 +267,7 @@ def make_deepep_v2_a2a(
267
hidden=v2_args.hidden_size,
268
num_topk=v2_args.num_topk,
269
use_fp8_dispatch=v2_args.use_fp8_dispatch,
270
+ allow_hybrid_mode=False,
271
explicitly_destroy=True,
272
)
273
return DeepEPV2PrepareAndFinalize(
tests/kernels/moe/test_deepep_v2_moe.py
+168
-110
@@ -11,6 +11,7 @@ import pytest
11
import torch.distributed
12
from torch.distributed import ProcessGroup
13
14
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
15
from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights
16
from tests.kernels.utils import torch_experts
17
from vllm.config import VllmConfig, set_current_vllm_config
@@ -21,6 +22,8 @@ from vllm.model_executor.layers.fused_moe.config import (
22
FusedMoEQuantConfig,
23
)
24
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
25
+from vllm.platforms import current_platform
26
+from vllm.utils.flashinfer import has_flashinfer
27
from vllm.utils.import_utils import has_deep_ep_v2
28
from vllm.utils.torch_utils import set_random_seed
29
from vllm.v1.worker.workspace import init_workspace_manager
@@ -356,18 +359,136 @@ def test_deep_ep_v2_moe(
359
)
360
361
359
-def _deep_ep_v2_moe_cudagraph(
362
+EXPERTS_BACKENDS = [
363
+ "flashinfer_trtllm",
364
+ "flashinfer_cutlass",
365
+ "trtllm_fp8",
366
+]
367
+
368
+
369
+def _make_experts(
370
+ experts_backend: str,
371
+ config: TestConfig,
372
+ moe_config,
373
+ num_local_experts: int,
374
+ rank: int,
375
+ w1_bf16: torch.Tensor,
376
+ w2_bf16: torch.Tensor,
377
+ test_tensors: TestTensors,
378
+):
379
+ e_start = num_local_experts * rank
380
+ e_end = e_start + num_local_experts
381
+
382
+ if experts_backend != "trtllm_fp8":
383
+ from vllm.model_executor.layers.fused_moe.config import (
384
+ FUSED_MOE_UNQUANTIZED_CONFIG,
385
+ )
386
+ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
387
+ backend_to_kernel_cls,
388
+ convert_to_unquantized_kernel_format,
389
+ map_unquantized_backend,
390
+ )
391
+
392
+ torch_combined = torch_experts(
393
+ test_tensors.rank_tokens,
394
+ w1_bf16,
395
+ w2_bf16,
396
+ test_tensors.topk_weights,
397
+ test_tensors.topk,
398
+ )
399
+ backend = map_unquantized_backend(experts_backend)
400
+ w1_ep, w2_ep = convert_to_unquantized_kernel_format(
401
+ backend,
402
+ moe_config,
403
+ w1_bf16[e_start:e_end],
404
+ w2_bf16[e_start:e_end],
405
+ )
406
+ experts_cls = next(
407
+ cls
408
+ for cls in backend_to_kernel_cls(backend)
409
+ if issubclass(cls, mk.FusedMoEExpertsModular)
410
+ )
411
+ fused_experts = experts_cls(
412
+ moe_config=moe_config,
413
+ quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
414
+ )
415
+ return fused_experts, w1_ep, w2_ep, torch_combined, 1e-1, 2e-1
416
+
417
+ from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves
418
+ from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
419
+ TrtLlmFp8ExpertsModular,
420
+ )
421
+ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
422
+ Fp8MoeBackend,
423
+ convert_to_fp8_moe_kernel_format,
424
+ )
425
+
426
+ w1_ref = w1_bf16.to(torch.float8_e4m3fn).to(torch.bfloat16)
427
+ w2_ref = w2_bf16.to(torch.float8_e4m3fn).to(torch.bfloat16)
428
+
429
+ block_shape = [128, 128]
430
+ qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape)
431
+ assert qw.w13_weight_scale is not None
432
+ assert qw.w2_weight_scale is not None
433
+
434
+ reference_topk_weights = test_tensors.topk_weights.to(torch.bfloat16).to(
435
+ torch.float32
436
+ )
437
+ torch_combined = torch_experts(
438
+ test_tensors.rank_tokens,
439
+ qw.w13_weight,
440
+ qw.w2_weight,
441
+ reference_topk_weights,
442
+ test_tensors.topk,
443
+ w1_scale=qw.w13_weight_scale,
444
+ w2_scale=qw.w2_weight_scale,
445
+ quant_dtype=torch.float8_e4m3fn,
446
+ block_shape=block_shape,
447
+ )
448
+
449
+ class _MockLayer:
450
+ weight_block_size = block_shape
451
+
452
+ class moe_config:
453
+ is_act_and_mul = True
454
+ intermediate_size_per_partition = config.n
455
+
456
+ class activation:
457
+ is_gated = True
458
+
459
+ w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = convert_to_fp8_moe_kernel_format(
460
+ fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM,
461
+ layer=_MockLayer(),
462
+ w13=qw.w13_weight[e_start:e_end],
463
+ w2=qw.w2_weight[e_start:e_end],
464
+ w13_scale=qw.w13_weight_scale[e_start:e_end],
465
+ w2_scale=qw.w2_weight_scale[e_start:e_end],
466
+ w13_input_scale=None,
467
+ w2_input_scale=None,
468
+ )
469
+
470
+ fused_experts = TrtLlmFp8ExpertsModular(
471
+ moe_config=moe_config,
472
+ quant_config=FusedMoEQuantConfig.make(
473
+ torch.float8_e4m3fn,
474
+ block_shape=block_shape,
475
+ w1_scale=w1_scale_ep,
476
+ w2_scale=w2_scale_ep,
477
+ ),
478
+ )
479
+ return fused_experts, w1_ep, w2_ep, torch_combined, 6e-2, 6e-2
480
+
481
+
482
+def _deep_ep_v2_moe_backends(
483
pgi: ProcessGroupInfo,
484
dp_size: int,
485
config: TestConfig,
363
- w1: torch.Tensor,
364
- w2: torch.Tensor,
365
- w1_scale: torch.Tensor | None,
366
- w2_scale: torch.Tensor | None,
486
+ use_cudagraph: bool,
487
+ experts_backend: str,
488
):
368
- """Worker function: verify DeepEP v2 + TrtLLM FP8 with do_expand=False."""
489
import tempfile
490
491
+ from vllm.config import KernelConfig
492
from vllm.distributed import (
493
init_distributed_environment,
494
initialize_model_parallel,
@@ -402,19 +523,10 @@ def _deep_ep_v2_moe_cudagraph(
523
torch.distributed.broadcast(w1_bf16, src=0, group=pg)
524
torch.distributed.broadcast(w2_bf16, src=0, group=pg)
525
405
- # Round-trip through FP8 before constructing the reference and kernel weights.
406
- w1_fp8 = w1_bf16.to(torch.float8_e4m3fn)
407
- w2_fp8 = w2_bf16.to(torch.float8_e4m3fn)
408
- w1_ref = w1_fp8.to(torch.bfloat16)
409
- w2_ref = w2_fp8.to(torch.bfloat16)
410
-
411
- from vllm.config import KernelConfig
412
-
526
vllm_cfg = VllmConfig()
527
vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm")
528
529
with set_current_vllm_config(vllm_cfg):
417
- # Initialize vLLM parallel state (needed by MoERunner layer)
530
temp_file = tempfile.mktemp()
531
init_distributed_environment(
532
world_size=pgi.world_size,
@@ -424,77 +536,7 @@ def _deep_ep_v2_moe_cudagraph(
536
backend="nccl",
537
)
538
initialize_model_parallel(tensor_model_parallel_size=1)
427
- # Mirror production weight processing: quantize, EP-slice, then
428
- # convert to the TrtLLM BlockMajorK format.
429
- from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves
430
- from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
431
- TrtLlmFp8ExpertsModular,
432
- )
433
- from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
434
- Fp8MoeBackend,
435
- convert_to_fp8_moe_kernel_format,
436
- )
539
438
- block_shape = [128, 128]
439
- qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape)
440
- assert qw.w13_weight_scale is not None
441
- assert qw.w2_weight_scale is not None
442
-
443
- # Reference MoE using the same blockwise FP8 quantization scheme as
444
- # the production kernel. torch_experts quantizes activations before
445
- # both GEMMs and dequantizes the operands for the reference matmuls.
446
- reference_topk_weights = test_tensors.topk_weights.to(torch.bfloat16).to(
447
- torch.float32
448
- )
449
- torch_combined = torch_experts(
450
- test_tensors.rank_tokens,
451
- qw.w13_weight,
452
- qw.w2_weight,
453
- reference_topk_weights,
454
- test_tensors.topk,
455
- w1_scale=qw.w13_weight_scale,
456
- w2_scale=qw.w2_weight_scale,
457
- quant_dtype=torch.float8_e4m3fn,
458
- block_shape=block_shape,
459
- )
460
-
461
- # EP-slice before format conversion
462
- e_start = num_local_experts * pgi.rank
463
- e_end = e_start + num_local_experts
464
- w1_ep = qw.w13_weight[e_start:e_end]
465
- w2_ep = qw.w2_weight[e_start:e_end]
466
- w1_scale_ep = qw.w13_weight_scale[e_start:e_end]
467
- w2_scale_ep = qw.w2_weight_scale[e_start:e_end]
468
-
469
- # Convert to TrtLLM format (W31 swap + BlockMajorK shuffle)
470
- class _MockLayer:
471
- weight_block_size = block_shape
472
-
473
- class moe_config:
474
- is_act_and_mul = True
475
- intermediate_size_per_partition = config.n
476
-
477
- class activation:
478
- is_gated = True
479
-
480
- w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = convert_to_fp8_moe_kernel_format(
481
- fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM,
482
- layer=_MockLayer(),
483
- w13=w1_ep,
484
- w2=w2_ep,
485
- w13_scale=w1_scale_ep,
486
- w2_scale=w2_scale_ep,
487
- w13_input_scale=None,
488
- w2_input_scale=None,
489
- )
490
-
491
- # Build TrtLLM expert with correct EP params
492
- quant_config = FusedMoEQuantConfig.make(
493
- torch.float8_e4m3fn,
494
- block_shape=block_shape,
495
- w1_scale=w1_scale_ep,
496
- w2_scale=w2_scale_ep,
497
- )
540
moe_config = make_dummy_moe_config(
541
num_experts=config.num_experts,
542
num_local_experts=num_local_experts,
@@ -502,20 +544,33 @@ def _deep_ep_v2_moe_cudagraph(
544
hidden_dim=hidden_size,
545
intermediate_size=config.n,
546
)
505
- moe_parallel_config = dataclasses.replace(
506
- moe_config.moe_parallel_config,
507
- ep_size=pgi.world_size,
508
- ep_rank=pgi.rank,
509
- use_ep=True,
510
- all2all_backend="deepep_v2",
511
- )
547
moe_config = dataclasses.replace(
548
moe_config,
514
- moe_parallel_config=moe_parallel_config,
549
+ moe_parallel_config=dataclasses.replace(
550
+ moe_config.moe_parallel_config,
551
+ ep_size=pgi.world_size,
552
+ ep_rank=pgi.rank,
553
+ use_ep=True,
554
+ all2all_backend="deepep_v2",
555
+ ),
556
)
516
- fused_experts = TrtLlmFp8ExpertsModular(
517
- moe_config=moe_config,
518
- quant_config=quant_config,
557
+
558
+ (
559
+ fused_experts,
560
+ w1_ep,
561
+ w2_ep,
562
+ torch_combined,
563
+ atol,
564
+ rtol,
565
+ ) = _make_experts(
566
+ experts_backend,
567
+ config,
568
+ moe_config,
569
+ num_local_experts,
570
+ pgi.rank,
571
+ w1_bf16,
572
+ w2_bf16,
573
+ test_tensors,
574
)
575
576
v2_args = DeepEPV2Args(
@@ -531,7 +586,7 @@ def _deep_ep_v2_moe_cudagraph(
586
pgi=pgi,
587
dp_size=dp_size,
588
v2_args=v2_args,
534
- use_cudagraph=True,
589
+ use_cudagraph=use_cudagraph,
590
)
591
mk_kernel = FusedMoEKernel(
592
prepare_finalize=a2a,
@@ -552,33 +607,38 @@ def _deep_ep_v2_moe_cudagraph(
607
apply_router_weight_on_input=False,
608
)
609
555
- torch.testing.assert_close(
556
- torch_combined,
557
- out,
558
- atol=6e-2,
559
- rtol=6e-2,
560
- )
610
+ torch.testing.assert_close(torch_combined, out, atol=atol, rtol=rtol)
611
612
613
@pytest.mark.parametrize("m,n,k", [(32, 256, 1024)])
614
@pytest.mark.parametrize("num_experts", [32])
615
@pytest.mark.parametrize("topk", [6])
616
@pytest.mark.parametrize("world_dp_size", [(2, 1)])
617
+@pytest.mark.parametrize("experts_backend", EXPERTS_BACKENDS)
618
+@pytest.mark.parametrize("use_cudagraph", [True, False])
619
@multi_gpu_test(num_gpus=2)
620
@requires_deep_ep_v2
569
-def test_deep_ep_v2_moe_cudagraph(
621
+@pytest.mark.skipif(
622
+ not has_flashinfer() or not current_platform.has_device_capability(100),
623
+ reason="Requires FlashInfer TRTLLM fused MoE (SM100)",
624
+)
625
+def test_deep_ep_v2_moe_backends(
626
m: int,
627
n: int,
628
k: int,
629
num_experts: int,
630
topk: int,
631
world_dp_size: tuple[int, int],
632
+ experts_backend: str,
633
+ use_cudagraph: bool,
634
workspace_init,
635
):
636
set_random_seed(7)
637
world_size, dp_size = world_dp_size
638
config = TestConfig(
581
- dtype=torch.float8_e4m3fn,
639
+ dtype=torch.float8_e4m3fn
640
+ if experts_backend == "trtllm_fp8"
641
+ else torch.bfloat16,
642
topk=topk,
643
m=m,
644
k=k,
@@ -588,11 +648,9 @@ def test_deep_ep_v2_moe_cudagraph(
648
649
parallel_launch(
650
world_size,
591
- _deep_ep_v2_moe_cudagraph,
651
+ _deep_ep_v2_moe_backends,
652
dp_size,
653
config,
594
- None, # weights created inside worker
595
- None,
596
- None,
597
- None,
654
+ use_cudagraph,
655
+ experts_backend,
656
)
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+1
-1
@@ -207,7 +207,7 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular)
207
gemm1_weights=w1,
208
gemm2_weights=w2,
209
num_experts=global_num_experts,
210
- top_k=self.topk,
210
+ top_k=topk_ids.size(1),
211
n_group=None,
212
topk_group=None,
213
intermediate_size=self.intermediate_size_per_partition,