A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 21:15 UTC
[MoE Refactor] Add more MoE layer tests (#39349)
Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
bnellnm committed
Apr 21, 2026 at 18:12 UTC
9db4650e5e4c726eb5ae29330cd55e796567469c
3 files changed
+97
-43
tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+1
-1
@@ -77,8 +77,8 @@ def _worker_parallel_launch(
77
*args: Any,
78
) -> None:
79
rank = node_rank * world_local_size + local_rank
80
- torch.accelerator.set_device_index(local_rank)
80
device = torch.device("cuda", local_rank)
81
+ torch.accelerator.set_device_index(device)
82
torch.distributed.init_process_group(
83
backend="cpu:gloo,cuda:nccl",
84
init_method=init_method,
tests/kernels/moe/test_moe_layer.py
+95
-40
@@ -65,8 +65,8 @@ fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype
65
66
SHAPE_COMBOS = [
67
(1, 128, 256),
68
- (32, 1024, 512),
69
- (222, 2048, 2048),
68
+ (32, 512, 512),
69
+ (222, 1024, 2048),
70
]
71
MAX_M = max([x[0] for x in SHAPE_COMBOS])
72
@@ -95,7 +95,7 @@ if has_flashinfer_nvlink_one_sided():
95
BACKENDS += ["flashinfer_nvlink_one_sided"]
96
97
if has_deep_ep():
98
- BACKENDS += ["deepep_low_latency", "deepep_high_throughput"]
98
+ BACKENDS += ["deepep_high_throughput", "deepep_low_latency"]
99
100
if has_nixl_ep():
101
BACKENDS += ["nixl_ep"]
@@ -103,6 +103,7 @@ if has_nixl_ep():
103
QUANT_METHODS = [
104
None,
105
"fp8",
106
+ "fp8_blocked",
107
"modelopt_fp8",
108
"modelopt_fp4",
109
]
@@ -114,10 +115,21 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
115
"mori": {None, "fp8", "modelopt_fp8"},
116
"flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"},
117
"flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"},
117
- "deepep_low_latency": {None, "modelopt_fp8", "modelopt_fp4"},
118
- "deepep_high_throughput": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
118
+ "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"},
119
+ "deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
120
"nixl_ep": {None, "fp8", "modelopt_fp8"},
121
}
122
+
123
+# Map from backend -> (DP/EP support, DP support, TP support)
124
+BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool]] = {
125
+ "allgather_reducescatter": (True, True, True),
126
+ "mori": (True, False, False),
127
+ "flashinfer_nvlink_two_sided": (False, True, False),
128
+ "flashinfer_nvlink_one_sided": (False, True, False),
129
+ "deepep_low_latency": (True, False, False),
130
+ "deepep_high_throughput": (True, False, False),
131
+ "nixl_ep": (True, False, False),
132
+}
133
# fmt: on
134
135
# Which quantization methods support EPLB.
@@ -424,27 +436,35 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
436
f"Skipping unsupported K {config.k} in {config.backend} w/o EP.",
437
)
438
427
- if config.enable_eplb and config.ep_size == 1:
428
- return False, "EPLB requires EP."
439
+ if config.backend is not None:
440
+ supports_ep_dp, supports_dp, supports_tp = BACKEND_EP_DP_TP_SUPPORT[
441
+ config.backend
442
+ ]
443
430
- if config.enable_eplb and config.quantization not in EPLB_SUPPORTED_QUANTS:
431
- return False, f"EPLB not supported with {config.quantization} quantization."
444
+ if config.tp_size > 1 and not supports_tp:
445
+ return False, f"{config.backend} does not support TP."
446
433
- if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS:
434
- return False, f"EPLB not supported with {config.backend}."
447
+ if config.dp_size > 1 and config.ep_size == 1 and not supports_dp:
448
+ return False, f"{config.backend} does not support DP."
449
436
- if (
437
- config.backend is not None
438
- and config.backend.startswith("flashinfer_nvlink")
439
- and config.ep_size > 1
440
- ):
441
- return False, "flashinfer_nvlink EP not yet supported."
450
+ if config.dp_size > 1 and config.ep_size > 1 and not supports_ep_dp:
451
+ return False, f"{config.backend} does not support EP/DP."
452
+ else:
453
+ if config.tp_size > 1 or config.ep_size > 1 or config.dp_size > 1:
454
+ return False, "An all2all backend is required for parallelism."
455
443
- if config.enable_eplb and config.num_experts % config.dp_size != 0:
444
- return False, "EPLB requires num_experts divisible by ep_size"
456
+ if config.enable_eplb:
457
+ if config.ep_size == 1:
458
+ return False, "EPLB requires EP."
459
446
- if config.enable_eplb and config.ep_size == 1:
447
- return False, "EPLB only works with EP+DP"
460
+ if config.quantization not in EPLB_SUPPORTED_QUANTS:
461
+ return False, f"EPLB not supported with {config.quantization} quantization."
462
+
463
+ if config.backend not in EPLB_SUPPORTED_BACKENDS:
464
+ return False, f"EPLB not supported with {config.backend}."
465
+
466
+ if config.num_experts % config.dp_size != 0:
467
+ return False, "EPLB requires num_experts divisible by ep_size"
468
469
# Disable fp4 tests until flashinfer is updated or the Dockerfile is
470
# modified to install cublasLt.h. See #39525.
@@ -507,27 +527,48 @@ class QuantizedWeights:
527
def _quantize_fp8_halves(
528
w1: torch.Tensor,
529
w2: torch.Tensor,
530
+ block_shape: list[int] | None = None,
531
) -> QuantizedWeights:
532
"""Quantize w13 gate/up halves separately to FP8, producing per-shard scales."""
533
half = w1.shape[1] // 2
534
w1q_a, w1s_a, _ = moe_quantize_weights(
514
- w1[:, :half, :], None, fp8_dtype, False, None
535
+ w1[:, :half, :],
536
+ None,
537
+ fp8_dtype,
538
+ False,
539
+ block_shape,
540
)
541
w1q_b, w1s_b, _ = moe_quantize_weights(
517
- w1[:, half:, :], None, fp8_dtype, False, None
542
+ w1[:, half:, :],
543
+ None,
544
+ fp8_dtype,
545
+ False,
546
+ block_shape,
547
)
548
assert w1s_a is not None and w1s_b is not None
549
521
- w2q, w2s, _ = moe_quantize_weights(w2, None, fp8_dtype, False, None)
550
+ w2q, w2s, _ = moe_quantize_weights(w2, None, fp8_dtype, False, block_shape)
551
assert w2s is not None
552
553
+ if block_shape is not None:
554
+ # Blocked quantization: scales have shape (E, n_tiles, k_tiles)
555
+ # Concatenate gate and up scales along the n_tiles dimension (dim=1)
556
+ # to match the concatenation of gate and up weights
557
+ w13_weight_scale = torch.cat([w1s_a, w1s_b], dim=1)
558
+ # w2 scales keep their blocked shape (E, k_tiles, n_tiles)
559
+ w2_weight_scale = w2s
560
+ else:
561
+ # Non-blocked quantization: scales have shape (E, 1, 1)
562
+ # Each w1s_x is (E, 1, 1) -> reshape to (E, 1), cat to (E, 2)
563
+ w13_weight_scale = torch.cat([w1s_a.view(-1, 1), w1s_b.view(-1, 1)], dim=1)
564
+ # w2s is (E, 1, 1) -> reshape to (E,)
565
+ w2_weight_scale = w2s.view(-1)
566
+
567
return QuantizedWeights(
568
w13_weight=torch.cat([w1q_a, w1q_b], dim=1),
569
w2_weight=w2q,
527
- # Each w1s_x is (E, 1, 1) -> reshape to (E, 1), cat to (E, 2)
528
- w13_weight_scale=torch.cat([w1s_a.view(-1, 1), w1s_b.view(-1, 1)], dim=1),
529
- # w2s is (E, 1, 1) -> reshape to (E,)
530
- w2_weight_scale=w2s.view(-1),
570
+ w13_weight_scale=w13_weight_scale,
571
+ w2_weight_scale=w2_weight_scale,
572
)
573
574
@@ -536,7 +577,7 @@ def quantization_to_quant_dtype(
577
) -> torch.dtype | str | None:
578
if quantization is None:
579
return None
539
- elif quantization in ["fp8", "modelopt_fp8"]:
580
+ elif quantization in ["fp8", "fp8_blocked", "modelopt_fp8"]:
581
return fp8_dtype
582
elif quantization in ["modelopt_fp4"]:
583
return "nvfp4"
@@ -558,6 +599,12 @@ def make_quant_config(
599
if quantization == "fp8":
600
return Fp8Config(True), _quantize_fp8_halves(w1, w2)
601
602
+ if quantization == "fp8_blocked":
603
+ block_shape = [128, 128]
604
+ return Fp8Config(True, weight_block_size=block_shape), _quantize_fp8_halves(
605
+ w1, w2, block_shape
606
+ )
607
+
608
if quantization == "modelopt_fp8":
609
qw = _quantize_fp8_halves(w1, w2)
610
# why?
@@ -896,11 +943,13 @@ def make_fused_moe_layer(
943
**kwargs,
944
)
945
946
+ weight_scale_name = getattr(layer.quant_method, "weight_scale_name", "weight_scale")
947
+
948
for name, value in [
949
("w13_weight", qw.w13_weight),
950
("w2_weight", qw.w2_weight),
902
- ("w13_weight_scale", qw.w13_weight_scale),
903
- ("w2_weight_scale", qw.w2_weight_scale),
951
+ (f"w13_{weight_scale_name}", qw.w13_weight_scale),
952
+ (f"w2_{weight_scale_name}", qw.w2_weight_scale),
953
("w13_weight_scale_2", qw.w13_weight_scale_2),
954
("w2_weight_scale_2", qw.w2_weight_scale_2),
955
("w13_input_scale", qw.w13_input_scale),
@@ -922,7 +971,7 @@ def make_fake_moe_layer(
971
top_k: int,
972
global_num_experts: int,
973
in_dtype: torch.dtype,
925
- quant_dtype: torch.dtype | None,
974
+ quantization: str | None,
975
renormalize: bool = False,
976
shared_experts_config: SharedExpertsConfig | None = None,
977
use_grouped_topk: bool = False,
@@ -948,6 +997,7 @@ def make_fake_moe_layer(
997
dp_size: int = 1,
998
ep_size: int = 1,
999
) -> Callable:
1000
+ quant_dtype = None
1001
activation = MoEActivation.from_str(activation)
1002
1003
router = create_fused_moe_router(
@@ -1139,7 +1189,6 @@ def _test_body_eplb(
1189
routed_output_transform=routed_output_transform,
1190
)
1191
1142
- # Necessary?
1192
if eplb_moe_layer._expert_map is not None:
1193
eplb_moe_layer._expert_map = eplb_moe_layer._expert_map.to(device)
1194
@@ -1267,6 +1316,7 @@ def _run_one_config(
1316
gate = test_data.gate
1317
routed_input_transform = test_data.routed_input_transform
1318
routed_output_transform = test_data.routed_output_transform
1319
+ activation = "silu"
1320
1321
baseline_layer = make_fake_moe_layer(
1322
w1=w1,
@@ -1274,7 +1324,7 @@ def _run_one_config(
1324
top_k=top_k,
1325
global_num_experts=num_experts,
1326
in_dtype=in_dtype,
1277
- quant_dtype=None, # quantization_to_quant_dtype(quantization),
1327
+ quantization=quantization,
1328
renormalize=False,
1329
shared_experts_config=shared_experts_config,
1330
gate=gate,
@@ -1284,6 +1334,7 @@ def _run_one_config(
1334
tp_size=tp_size,
1335
ep_size=ep_size,
1336
dp_size=dp_size,
1337
+ activation=activation,
1338
)
1339
1340
baseline_output = baseline_layer(hidden_states, router_logits)
@@ -1328,9 +1379,9 @@ def _run_one_config(
1379
gate=gate,
1380
routed_input_transform=routed_input_transform,
1381
routed_output_transform=routed_output_transform,
1382
+ activation=activation,
1383
)
1384
1333
- # Necessary?
1385
if moe_layer._expert_map is not None:
1386
moe_layer._expert_map = moe_layer._expert_map.to(device)
1387
@@ -1377,13 +1428,17 @@ def _run_one_config(
1428
atol, rtol = 7.6e-2, 7.6e-2
1429
else:
1430
atol, rtol = 3.5e-2, 3.5e-2
1380
- elif quantization in ("fp8", "modelopt_fp8"):
1431
+ elif quantization in ("fp8", "fp8_blocked", "modelopt_fp8"):
1432
+ atol, rtol = 6e-2, 6e-2
1433
+ elif quantization == "modelopt_fp4":
1434
if k >= 2048:
1382
- atol, rtol = 7.6e-2, 7.6e-2
1435
+ atol = rtol = 1e-1 + (k * 1e-4)
1436
else:
1384
- atol, rtol = 6e-2, 6e-2
1385
- elif quantization == "modelopt_fp4":
1386
- atol = rtol = 1e-1 + k * 5e-4
1437
+ atol = rtol = 1e-1
1438
+
1439
+ if backend == "allgather_reducescatter" and tp_size > 1:
1440
+ atol += 2e-1
1441
+ rtol += 2e-1
1442
else:
1443
atol, rtol = 6e-2, 6e-2
1444
vllm/model_executor/layers/fused_moe/config.py
+1
-2
@@ -990,8 +990,7 @@ class FusedMoEParallelConfig:
990
991
@property
992
def use_batched_activation_format(self):
993
- # TODO(bnell): nixl also uses batched format
994
- return self.use_deepep_ll_kernels
993
+ return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels
994
995
@property
996
def use_ag_rs_all2all_kernels(self):