@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 18f6bf5a21

Fork 0

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 3, 19:00 UTC

[MoE Refactor] Add sequence parallel tests to test_moe_layer.py (#41299)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed May 12, 2026 at 21:52 UTC 18f6bf5a214fd51d0cb08fcf1b59edc5629e7c19

4 files changed +190 -73

tests/kernels/moe/test_moe_layer.py

+175 -64

@@ -31,12 +31,13 @@ from vllm.config import (

31 VllmConfig,

32 set_current_vllm_config,

33 )

34 -from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator

35 -from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace

36 -from vllm.distributed.parallel_state import (

34 +from vllm.distributed import (

35 get_ep_group,

36 get_eplb_group,

37 + tensor_model_parallel_all_gather,

38 )

39 +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator

40 +from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace

41 from vllm.forward_context import set_forward_context

42 from vllm.model_executor.layers.fused_moe import FusedMoE, fused_experts

43 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

@@ -49,6 +50,7 @@ from vllm.model_executor.layers.quantization.modelopt import (

50 ModelOptFp8Config,

51 ModelOptNvFp4Config,

52 )

53 +from vllm.model_executor.models.utils import sequence_parallel_chunk

54 from vllm.platforms import current_platform

55 from vllm.utils.flashinfer import (

56 has_flashinfer_nvlink_one_sided,

@@ -81,6 +83,9 @@ PARALLEL_COMBOS = [

83 [1, 4, False],

84 [2, 1, True],

85 [4, 1, True],

86 + # This combination indicates sequence parallel.

87 + # See ParallelConfig.use_sequence_parallel.

88 + [2, 2, True],

89 ]

91 # TODO: should this even be set manually? let oracles handle this

@@ -112,24 +117,24 @@ QUANT_METHODS = [

117 # Which quantization methods each backend supports.

118 # fmt: off

119 BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {

115 - "allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},

116 - "mori": {None, "fp8", "modelopt_fp8"},

117 - "flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"},

118 - "flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"},

119 - "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"},

120 + "allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501

121 + "mori": {None, "fp8", "modelopt_fp8"},

122 + "flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501

123 + "flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501

124 + "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501

125 "deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501

121 - "nixl_ep": {None, "fp8", "modelopt_fp8"},

126 + "nixl_ep": {None, "fp8_blocked", "modelopt_fp8"},

127 }

128

124 -# Map from backend -> (DP/EP support, DP support, TP support)

125 -BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool]] = {

126 - "allgather_reducescatter": (True, True, True),

127 - "mori": (True, False, False),

128 - "flashinfer_nvlink_two_sided": (False, True, False),

129 - "flashinfer_nvlink_one_sided": (False, True, False),

130 - "deepep_low_latency": (True, False, False),

131 - "deepep_high_throughput": (True, False, False),

132 - "nixl_ep": (True, False, False),

129 +# Map from backend -> (DP/EP support, DP support, TP support, SP support)

130 +BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool, bool]] = {

131 + "allgather_reducescatter": (True, True, True, True),

132 + "mori": (True, False, False, True),

133 + "flashinfer_nvlink_two_sided": (False, True, False, False),

134 + "flashinfer_nvlink_one_sided": (False, True, False, False),

135 + "deepep_low_latency": (True, False, False, True),

136 + "deepep_high_throughput": (True, False, False, True),

137 + "nixl_ep": (True, False, False, True),

138 }

139 # fmt: on

140

@@ -163,6 +168,45 @@ def override_normalize_e4m3fn_to_e4m3fnuz():

168 vllm.model_executor.layers.quantization.utils.w8a8_utils.normalize_e4m3fn_to_e4m3fnuz = mock_normalize_e4m3fn_to_e4m3fnuz # noqa: E501

169

170

171 +def sp_wrapper(

172 + fn: Callable | FusedMoE, is_sequence_parallel: bool | None = None

173 +) -> Callable:

174 + """Wrapper to handle sequence parallelism chunking and gathering.

175 +

176 + For SP with EP:

177 + - The TP group is created with the original tensor_parallel_size (e.g., 2)

178 + - get_tp_group() has the correct world_size for SP operations

179 + - sequence_parallel_chunk() uses get_tensor_model_parallel_world_size()

180 + - tensor_model_parallel_all_gather() uses get_tp_group()

181 + - Both should work correctly even when EP is enabled

182 + """

183 + if isinstance(fn, FusedMoE):

184 + assert is_sequence_parallel is None

185 + is_sequence_parallel = fn.is_sequence_parallel

186 + else:

187 + assert is_sequence_parallel is not None

188 +

189 + if is_sequence_parallel:

190 +

191 + def wrapper(

192 + hidden_states: torch.Tensor,

193 + router_logits: torch.Tensor,

194 + ) -> torch.Tensor:

195 + # Split sequence across TP ranks

196 + # Both hidden_states and router_logits have [num_tokens, ...] shape

197 + hidden_states = sequence_parallel_chunk(hidden_states)

198 + router_logits = sequence_parallel_chunk(router_logits)

199 + # Run MoE on local chunk

200 + result = fn(hidden_states, router_logits)

201 + # Gather results from all TP ranks

202 + result = tensor_model_parallel_all_gather(result, 0)

203 + # Remove any padding added by SP.

204 + return result[: hidden_states.shape[0]]

205 +

206 + return wrapper

207 + return fn

208 +

209 +

210 def maybe_roundup_layer_hidden_size(

211 hidden_size: int,

212 act_dtype: torch.dtype,

@@ -272,6 +316,15 @@ class MoETestConfig:

316 dp_size: int = 1

317 tp_size: int = 1

318

319 + @property

320 + def is_sequence_parallel(self) -> bool:

321 + # Sequence parallelism: EP enabled + TP dimension used for sequence splitting

322 + # In test config: ep_size represents total expert parallel size

323 + # tp_size represents the original TP dimension (becomes sp_size in FusedMoE)

324 + # dp_size represents data parallel size

325 + # For SP: we need EP enabled (ep_size > 1) and sequence splitting (tp_size > 1)

326 + return self.ep_size > 1 and self.tp_size > 1

327 +

328 # TODO: add more error messages

329 def id(self) -> str:

330 def proc(s: str) -> str:

@@ -404,11 +457,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

457 "leads to large differences.",

458 )

459

407 - # gate requires shared_experts (use_overlapped mode)

408 - # TODO: also not sure this is true

409 - if config.use_gate and not config.use_shared_experts:

410 - return False, "gate requires shared_experts (use_overlapped mode)"

411 -

460 # Skip modelopt_fp4 if not on B100+ (compute capability 10.0+)

461 if (

462 config.quantization == "modelopt_fp4"

@@ -445,7 +493,7 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

493 )

494

495 if config.backend == "nixl_ep":

448 - from vllm.model_executor.layers.fused_moe.nixl_ep_prepare_finalize import ( # noqa: E501

496 + from vllm.model_executor.layers.fused_moe.prepare_finalize.nixl_ep import ( # noqa: E501

497 NixlEPPrepareAndFinalize,

498 )

499

@@ -456,11 +504,11 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

504 )

505

506 if config.backend is not None:

459 - supports_ep_dp, supports_dp, supports_tp = BACKEND_EP_DP_TP_SUPPORT[

460 - config.backend

461 - ]

507 + supports_ep_dp, supports_dp, supports_tp, supports_sp = (

508 + BACKEND_EP_DP_TP_SUPPORT[config.backend]

509 + )

510

463 - if config.tp_size > 1 and not supports_tp:

511 + if config.tp_size > 1 and not supports_tp and not config.is_sequence_parallel:

512 return False, f"{config.backend} does not support TP."

513

514 if config.dp_size > 1 and config.ep_size == 1 and not supports_dp:

@@ -468,10 +516,34 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

516

517 if config.dp_size > 1 and config.ep_size > 1 and not supports_ep_dp:

518 return False, f"{config.backend} does not support EP/DP."

519 +

520 + if config.is_sequence_parallel and not supports_sp:

521 + return False, f"{config.backend} does not support SP."

522 else:

523 if config.tp_size > 1 or config.ep_size > 1 or config.dp_size > 1:

524 return False, "An all2all backend is required for parallelism."

525

526 + # Sequence parallelism specific validations

527 + if config.is_sequence_parallel:

528 + if config.ep_size == 1:

529 + return False, "Sequence parallelism requires EP to be enabled (ep_size > 1)"

530 +

531 + if config.tp_size == 1:

532 + return (

533 + False,

534 + "Sequence parallelism requires tp_size > 1 for sequence splitting",

535 + )

536 +

537 + # SP is essentially EP + sequence splitting

538 + # Verify the relationship: ep_size should equal dp_size * tp_size

539 + # (when pcp_size=1).

540 + expected_ep_size = config.dp_size * config.tp_size

541 + if config.ep_size != expected_ep_size:

542 + return False, (

543 + f"For sequence parallelism: ep_size ({config.ep_size}) should equal "

544 + f"dp_size * tp_size ({expected_ep_size})"

545 + )

546 +

547 if config.enable_eplb:

548 if config.ep_size == 1:

549 return False, "EPLB requires EP."

@@ -485,14 +557,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

557 if config.num_experts % config.dp_size != 0:

558 return False, "EPLB requires num_experts divisible by ep_size"

559

488 - # Disable fp4 tests until flashinfer is updated or the Dockerfile is

489 - # modified to install cublasLt.h. See #39525.

490 - if (

491 - config.quantization == "modelopt_fp4"

492 - and current_platform.is_device_capability_family(100)

493 - ):

494 - return False, "Temporarily skip until #39525 is resolved"

495 -

560 return True, None

561

562

@@ -751,7 +815,8 @@ def create_shared_experts_from_config(

815 in_dtype: torch.dtype,

816 tp_size: int = 1,

817 tp_rank: int = 0,

754 - device: torch.device | str | None = None,

818 + is_sequence_parallel: bool = False,

819 + device: torch.device | str | None = "cuda",

820 ) -> TestMLP | None:

821 """Create TestMLP for shared experts from config.

822

@@ -772,7 +837,7 @@ def create_shared_experts_from_config(

837 s_w2 = shared_experts_config.w2

838

839 # Apply TP chunking if needed

775 - if tp_size > 1:

840 + if tp_size > 1 and not is_sequence_parallel:

841 s_w1 = tp_chunk_gate_up(s_w1, tp_rank, tp_size, dim=1, device=device)

842 s_w2 = chunk_by_rank(s_w2, tp_rank, tp_size, dim=0, device=device)

843 else:

@@ -920,6 +985,7 @@ def make_fused_moe_layer(

985 routed_input_transform: torch.nn.Module | None = None,

986 routed_output_transform: torch.nn.Module | None = None,

987 pcp_size: int | None = 1,

988 + is_sequence_parallel: bool = False,

989 ) -> FusedMoE:

990 quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts)

991

@@ -959,6 +1025,7 @@ def make_fused_moe_layer(

1025 enable_eplb=enable_eplb,

1026 num_redundant_experts=num_redundant_experts,

1027 has_bias=has_bias,

1028 + is_sequence_parallel=is_sequence_parallel,

1029 **kwargs,

1030 )

1031

@@ -1014,6 +1081,7 @@ def make_fake_moe_layer(

1081 tp_size: int = 1,

1082 dp_size: int = 1,

1083 ep_size: int = 1,

1084 + is_sequence_parallel: bool = False,

1085 ) -> Callable:

1086 quant_dtype = None

1087 activation = MoEActivation.from_str(activation)

@@ -1043,7 +1111,8 @@ def make_fake_moe_layer(

1111 w2_s = None

1112

1113 shared_experts = create_shared_experts_from_config(

1046 - shared_experts_config, in_dtype, 1, 0, "cuda"

1114 + shared_experts_config,

1115 + in_dtype,

1116 )

1117

1118 quant_config = FusedMoEQuantConfig.make(

@@ -1114,7 +1183,7 @@ def make_fake_moe_layer(

1183

1184

1185 def _test_body_regular(

1117 - moe_layer: Callable,

1186 + moe_layer: FusedMoE,

1187 hidden_states: torch.Tensor,

1188 router_logits: torch.Tensor,

1189 vllm_config: VllmConfig,

@@ -1131,7 +1200,7 @@ def _test_body_regular(

1200 num_tokens=num_tokens,

1201 num_tokens_across_dp=num_tokens_across_dp,

1202 ):

1134 - output = moe_layer(hidden_states, router_logits)

1203 + output = sp_wrapper(moe_layer)(hidden_states, router_logits)

1204

1205 return baseline_output, output

1206

@@ -1164,6 +1233,8 @@ def _test_body_eplb(

1233 ) -> tuple[torch.Tensor, torch.Tensor]:

1234 device = torch.accelerator.current_accelerator()

1235

1236 + is_sequence_parallel = moe_layer.is_sequence_parallel

1237 +

1238 """EPLB test body: compare output before and after expert weight rearrangement."""

1239 # Get "before" output with original weight arrangement

1240 with set_forward_context(

@@ -1172,7 +1243,7 @@ def _test_body_eplb(

1243 num_tokens=num_tokens,

1244 num_tokens_across_dp=num_tokens_across_dp,

1245 ):

1175 - output_before = moe_layer(hidden_states, router_logits)

1246 + output_before = sp_wrapper(moe_layer)(hidden_states, router_logits)

1247

1248 # Create a fresh FusedMoE layer with enable_eplb=True

1249 # Delete the original layer's registration so the constructor can

@@ -1203,6 +1274,7 @@ def _test_body_eplb(

1274 gate=gate,

1275 routed_input_transform=routed_input_transform,

1276 routed_output_transform=routed_output_transform,

1277 + is_sequence_parallel=is_sequence_parallel,

1278 )

1279

1280 if eplb_moe_layer._expert_map is not None:

@@ -1264,7 +1336,7 @@ def _test_body_eplb(

1336 num_tokens=num_tokens,

1337 num_tokens_across_dp=num_tokens_across_dp,

1338 ):

1267 - output_after = eplb_moe_layer(hidden_states, router_logits)

1339 + output_after = sp_wrapper(eplb_moe_layer)(hidden_states, router_logits)

1340

1341 return output_before, output_after

1342

@@ -1272,11 +1344,12 @@ def _test_body_eplb(

1344 # TODO: make this take a MoETestConfig

1345 def _run_one_config(

1346 vllm_config: VllmConfig,

1275 - ep_size: int,

1276 - dp_size: int,

1277 - tp_size: int,

1278 - dp_rank: int,

1279 - tp_rank: int,

1347 + ep_size: int, # Expert parallel size (total across all ranks)

1348 + dp_size: int, # Data parallel size (number of DP groups)

1349 + tp_size: int, # Tensor parallel size OR sequence parallel size (when use_ep=True)

1350 + dp_rank: int, # Current rank in data parallel dimension

1351 + tp_rank: int, # Current rank in tensor/sequence parallel dimension

1352 + is_sequence_parallel: bool, # Whether to use sequence parallelism

1353 m: int,

1354 n: int,

1355 k: int,

@@ -1290,15 +1363,22 @@ def _run_one_config(

1363 use_routed_input_transform: bool,

1364 **kwargs,

1365 ) -> None:

1293 - set_random_seed(7)

1294 -

1366 """Generic test loop that sets up environment and delegates to test_body_fn.

1367

1297 - This function is called directly by test_moe_layer and test_moe_layer_eplb

1298 - via parallel_launch_with_config, passing either _test_body_regular or

1299 - _test_body_eplb as the test_body_fn parameter.

1368 + Parameter Interpretation:

1369 + - When is_sequence_parallel=False (standard TP or EP):

1370 + * ep_size: Number of expert parallel ranks (or 1 if no EP)

1371 + * tp_size: Number of tensor parallel ranks (or 1 if no TP)

1372 + * Weights are chunked by ep_size (experts) and tp_size (tensors)

1373 +

1374 + - When is_sequence_parallel=True (EP + sequence splitting):

1375 + * ep_size: Number of expert parallel ranks (equals dp_size * tp_size)

1376 + * tp_size: Number of ranks to split sequence across (becomes sp_size in FusedMoE)

1377 + * Weights are chunked by ep_size (experts) but NOT by tp_size

1378 + * Input sequences are chunked by tp_size (via sp_wrapper)

1379 """

1301 - world_size = tp_size * dp_size

1380 + set_random_seed(7)

1381 +

1382 use_ep = ep_size > 1

1383

1384 assert vllm_config.parallel_config.enable_expert_parallel == use_ep

@@ -1334,6 +1414,8 @@ def _run_one_config(

1414 routed_output_transform = test_data.routed_output_transform

1415 activation = "silu"

1416

1417 + # Create baseline layer with FULL weights (no EP chunking)

1418 + # Baseline represents the expected output using full model

1419 baseline_layer = make_fake_moe_layer(

1420 w1=w1,

1421 w2=w2,

@@ -1351,26 +1433,45 @@ def _run_one_config(

1433 ep_size=ep_size,

1434 dp_size=dp_size,

1435 activation=activation,

1436 + is_sequence_parallel=is_sequence_parallel,

1437 )

1438

1356 - baseline_output = baseline_layer(hidden_states, router_logits)

1439 + with set_current_vllm_config(vllm_config):

1440 + # Compute baseline output with SP wrapper if needed

1441 + # sp_wrapper handles sequence chunking/gathering for SP

1442 + baseline_output = sp_wrapper(baseline_layer, is_sequence_parallel)(

1443 + hidden_states, router_logits

1444 + )

1445

1446 del baseline_layer

1447 torch.accelerator.empty_cache()

1448

1449 with set_current_vllm_config(vllm_config):

1362 - # Chunk weights for EP/TP (after baseline is created)

1450 + # Chunk weights for EP BEFORE creating FusedMoE

1451 + # FusedMoE uses EP-chunked weights and handles reductions internally

1452 if ep_size > 1:

1364 - w1 = chunk_by_rank(w1, dp_rank, dp_size, dim=0, device=device)

1365 - w2 = chunk_by_rank(w2, dp_rank, dp_size, dim=0, device=device)

1366 -

1367 - if tp_size > 1:

1453 + # Split experts across ranks (dimension 0 is the expert dimension)

1454 + # When EP is enabled, use EP group rank and ep_size for chunking

1455 + ep_rank = get_ep_group().rank_in_group

1456 + w1 = chunk_by_rank(w1, ep_rank, ep_size, dim=0, device=device)

1457 + w2 = chunk_by_rank(w2, ep_rank, ep_size, dim=0, device=device)

1458 +

1459 + # Chunk weights for TP (only if NOT doing sequence parallelism)

1460 + # Sequence parallelism splits tokens/sequences, not weight tensors

1461 + if tp_size > 1 and not is_sequence_parallel:

1462 w1 = tp_chunk_gate_up(w1, tp_rank, tp_size, dim=1, device=device)

1463 w2 = chunk_by_rank(w2, tp_rank, tp_size, dim=2, device=device)

1464

1465 # Setup shared experts if needed

1466 + # In SP mode, shared experts should NOT be TP-chunked (same as routed experts)

1467 + # tp_size is used for sequence splitting, not weight splitting

1468 shared_experts = create_shared_experts_from_config(

1373 - shared_experts_config, in_dtype, tp_size, tp_rank, device

1469 + shared_experts_config,

1470 + in_dtype,

1471 + tp_size,

1472 + tp_rank,

1473 + is_sequence_parallel,

1474 + device,

1475 )

1476

1477 # Determine hidden size for MoE layer

@@ -1396,14 +1497,17 @@ def _run_one_config(

1497 routed_input_transform=routed_input_transform,

1498 routed_output_transform=routed_output_transform,

1499 activation=activation,

1500 + is_sequence_parallel=is_sequence_parallel,

1501 )

1502

1503 if moe_layer._expert_map is not None:

1504 moe_layer._expert_map = moe_layer._expert_map.to(device)

1505

1506 num_tokens = m

1507 + # num_tokens_across_dp should have one entry per DP group, not per total rank

1508 + # When EP is enabled, dp_size represents the number of DP groups

1509 num_tokens_across_dp = torch.tensor(

1406 - [num_tokens] * world_size,

1510 + [num_tokens] * dp_size,

1511 device=device,

1512 dtype=torch.int,

1513 )

@@ -1445,7 +1549,7 @@ def _run_one_config(

1549 else:

1550 atol, rtol = 3.5e-2, 3.5e-2

1551 elif quantization in ("fp8", "fp8_blocked", "modelopt_fp8"):

1448 - atol, rtol = 6e-2, 6e-2

1552 + atol, rtol = 6.5e-2, 6.5e-2

1553 elif quantization == "modelopt_fp4":

1554 if k >= 2048:

1555 atol = rtol = 1e-1 + (k * 1e-4)

@@ -1529,6 +1633,7 @@ def test_moe_layer_no_parallel(

1633 test_config.tp_size,

1634 0,

1635 0,

1636 + False,

1637 test_config.m,

1638 test_config.n,

1639 test_config.k,

@@ -1589,6 +1694,7 @@ def _parallel_worker(

1694 test_config.tp_size,

1695 dp_rank,

1696 tp_rank,

1697 + test_config.is_sequence_parallel,

1698 test_config.m,

1699 test_config.n,

1700 test_config.k,

@@ -1673,7 +1779,12 @@ def test_moe_layer(

1779 """

1780 num_gpus = current_platform.device_count()

1781 world_size = tp_size * dp_size

1676 - ep_size = 1 if not use_ep else world_size # or dp_size?

1782 + # When use_ep=True: FusedMoEParallelConfig flattens tp_size across dp ranks

1783 + # Result: ep_size = dp_size * pcp_size * tp_size

1784 + # Since pcp_size=1 in these tests: ep_size = dp_size * tp_size = world_size

1785 + # When use_ep=False: no expert parallelism, ep_size = 1

1786 + ep_size = 1 if not use_ep else world_size

1787 +

1788 assert world_size > 1

1789

1790 # Check if enough GPUs available

vllm/distributed/device_communicators/all2all.py

@@ -335,6 +335,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):

335 _lock = threading.Lock()

336

337 def __init__(self, cpu_group, tcp_store_group=None):

338 + assert tcp_store_group is not None

339 super().__init__(cpu_group, tcp_store_group)

340

341 self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS

vllm/distributed/device_communicators/cuda_communicator.py

+1 -1

@@ -270,7 +270,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

270 input_tensor = input_.movedim(0, dim).contiguous()

271

272 if sizes is not None:

273 - assert len(sizes) == world_size

273 + assert len(sizes) == world_size, f"{len(sizes)} == {world_size}"

274 assert input_tensor.shape[0] == sum(sizes)

275 chunk_size = sizes[self.rank_in_group]

276 else:

vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_two_sided.py

+13 -8

@@ -185,16 +185,21 @@ def flashinfer_alltoall_dispatch(

185 ep_size,

186 )

187

188 - x_sf = MnnvlMoe.mnnvl_moe_alltoallv(

189 - x_sf,

190 - alltoall_info,

191 - all2all_manager.workspace_tensor, # type: ignore[attr-defined]

192 - ep_rank,

193 - ep_size,

194 - )

188 + if x_sf is not None:

189 + x_sf = MnnvlMoe.mnnvl_moe_alltoallv(

190 + x_sf,

191 + alltoall_info,

192 + all2all_manager.workspace_tensor, # type: ignore[attr-defined]

193 + ep_rank,

194 + ep_size,

195 + )

196

197 # Swizzle after the A2A if MoE kernel expects swizzled scales.

197 - if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled:

198 + if (

199 + x_sf is not None

200 + and quant_config.quant_dtype == "nvfp4"

201 + and quant_config.is_scale_swizzled

202 + ):

203 if x_sf.element_size() == 1:

204 x_sf = x_sf.view(torch.uint8)

205 x_sf = nvfp4_block_scale_interleave(x_sf)