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

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, 06:45 UTC

[Bugfix][MoE] FlashInfer one-sided: workspace union across heterogeneous layers (#42976)

Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com>

tomeras91 committed May 19, 2026 at 21:43 UTC f54721bcc3e072d71b0e09c0b0bd6d692eb06161

2 files changed +142 -27

tests/distributed/test_mnnvl_alltoall.py

+84

@@ -342,6 +342,90 @@ def test_one_sided_manager_lifecycle(world_size):

342 )

343

344

345 +# ---------------------------------------------------------------------------

346 +# Test 2b: One-sided manager grows workspace across heterogeneous MoE layers

347 +# ---------------------------------------------------------------------------

348 +#

349 +# Models with heterogeneous MoE quantization — most notably a quantized base

350 +# MoE combined with an unquantized MTP head — can call initialize() multiple

351 +# times with different per-token dispatch payload sizes. The shared workspace

352 +# must grow to the union and the MoeAlltoAll must be rebuilt; otherwise a

353 +# later layer's combine call overruns the workspace sized for the first

354 +# layer's smaller payload and trips FlashInfer's combinePayloadOffset assert.

355 +# ---------------------------------------------------------------------------

356 +

357 +

358 +def _one_sided_workspace_grow_worker(rank, world_size):

359 + from vllm.distributed.device_communicators.all2all import (

360 + FlashInferNVLinkOneSidedManager,

361 + )

362 + from vllm.distributed.parallel_state import get_dp_group

363 +

364 + cpu_group = get_dp_group().cpu_group

365 + manager = FlashInferNVLinkOneSidedManager(cpu_group)

366 +

367 + base_kwargs = dict(

368 + max_num_tokens=1024,

369 + top_k=2,

370 + num_experts=world_size * 8,

371 + hidden_size=4096,

372 + )

373 + nvfp4_kwargs = dict(

374 + dispatch_dtype_bytes_per_elem=0,

375 + dispatch_scale_bytes_per_token=base_kwargs["hidden_size"] // 16,

376 + )

377 + bf16_kwargs = dict(

378 + dispatch_dtype_bytes_per_elem=2,

379 + dispatch_scale_bytes_per_token=0,

380 + )

381 +

382 + # First init: NVFP4-like (hidden_bytes = hidden // 2 + hidden // 16).

383 + manager.initialize(**base_kwargs, **nvfp4_kwargs)

384 + assert manager.initialized

385 + nvfp4_workspace_size = manager.workspace_size

386 + nvfp4_moe_alltoall = manager.moe_alltoall

387 +

388 + torch.distributed.barrier()

389 +

390 + # Second init: bf16-like (hidden_bytes = hidden * 2). Models the case of

391 + # a quantized base MoE followed by an unquantized MoE layer (e.g. an MTP

392 + # head). Per-token dispatch payload is ~4x larger, so the union workspace

393 + # must grow and MoeAlltoAll must be rebuilt.

394 + manager.initialize(**base_kwargs, **bf16_kwargs)

395 + assert manager.initialized

396 + assert manager.workspace_size > nvfp4_workspace_size

397 + assert manager.moe_alltoall is not nvfp4_moe_alltoall

398 + bf16_workspace_size = manager.workspace_size

399 + bf16_moe_alltoall = manager.moe_alltoall

400 +

401 + torch.distributed.barrier()

402 +

403 + # Third init: back to NVFP4-like shape. Existing workspace already covers

404 + # it, so initialize() must no-op — no shrink, no rebuild.

405 + manager.initialize(**base_kwargs, **nvfp4_kwargs)

406 + assert manager.initialized

407 + assert manager.workspace_size == bf16_workspace_size

408 + assert manager.moe_alltoall is bf16_moe_alltoall

409 +

410 + torch.distributed.barrier()

411 + manager.cleanup()

412 +

413 +

414 +@requires_multi_gpu

415 +@requires_one_sided

416 +@requires_ptrace

417 +@pytest.mark.parametrize("world_size", [2])

418 +def test_one_sided_manager_workspace_grow(world_size):

419 + """A later initialize() with a larger per-token payload must grow the

420 + workspace and rebuild MoeAlltoAll; a later initialize() with a smaller

421 + payload must no-op."""

422 + _spawn_workers(

423 + _one_sided_workspace_grow_worker,

424 + world_size,

425 + dp_size=world_size,

426 + )

427 +

428 +

429 # ---------------------------------------------------------------------------

430 # Test 3: AgRs dispatch/combine with value validation

431 # ---------------------------------------------------------------------------

vllm/distributed/device_communicators/all2all.py

+58 -27

@@ -571,6 +571,10 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):

571 self.initialized = False

572 self.moe_alltoall: MoeAlltoAll | None = None

573 self.mapping = None

574 + self.workspace_size = 0

575 + self.max_num_tokens = 0

576 + self.top_k = 0

577 + self.num_experts = 0

578

579 def initialize(

580 self,

@@ -581,9 +585,54 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):

585 dispatch_dtype_bytes_per_elem: int = 0,

586 dispatch_scale_bytes_per_token: int = 0,

587 ):

584 - """Initialize the MoeAlltoAll workspace."""

588 + """Initialize (or grow) the MoeAlltoAll workspace."""

589 + if dispatch_dtype_bytes_per_elem == 0:

590 + hidden_bytes = hidden_size // 2

591 + else:

592 + hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem

593 + total_dispatch_payload_size_per_token = (

594 + hidden_bytes

595 + + dispatch_scale_bytes_per_token

596 + + top_k * 4 # int32 topks ids

597 + + top_k * 4 # float32 topk weights

598 + )

599 + combine_payload_size_per_token = hidden_size * 2 # bf16 hidden states

600 + needed_workspace_size = moe_a2a_get_workspace_size_per_rank(

601 + ep_size=self.world_size,

602 + max_num_tokens=max_num_tokens,

603 + total_dispatch_payload_size_per_token=total_dispatch_payload_size_per_token,

604 + combine_payload_size_per_token=combine_payload_size_per_token,

605 + )

606 + # workspace_size and max_num_tokens are kernel-side max-bounds, so

607 + # heterogeneous MoE layers (e.g. NVFP4 base + bf16 MTP head) only

608 + # need the shared workspace grown to the union. top_k and num_experts

609 + # must match across layers: top_k is a strict-equality assert at

610 + # dispatch (FlashInfer csrc/trtllm_moe_alltoall.cu), and num_experts

611 + # feeds the expert-to-rank routing math, so any mismatch would crash

612 + # or silently corrupt routing. All ranks see the same MoE layers in

613 + # the same order with identical shapes, so the skip / rebuild

614 + # branches are taken consistently across ranks.

615 if self.initialized:

586 - return

616 + assert top_k == self.top_k, (

617 + "FlashInfer one-sided MoeAlltoAll does not support "

618 + f"heterogeneous top_k across MoE layers (got {top_k}, "

619 + f"was built with {self.top_k})"

620 + )

621 + assert num_experts == self.num_experts, (

622 + "FlashInfer one-sided MoeAlltoAll does not support "

623 + f"heterogeneous num_experts across MoE layers (got "

624 + f"{num_experts}, was built with {self.num_experts})"

625 + )

626 + if (

627 + needed_workspace_size <= self.workspace_size

628 + and max_num_tokens <= self.max_num_tokens

629 + ):

630 + return

631 +

632 + self.workspace_size = max(self.workspace_size, needed_workspace_size)

633 + self.max_num_tokens = max(self.max_num_tokens, max_num_tokens)

634 + self.top_k = top_k

635 + self.num_experts = num_experts

636

637 self.cleanup()

638 gpus_per_node = torch.accelerator.device_count()

@@ -610,38 +659,17 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):

659 ep_config = MnnvlConfig(

660 comm_backend=CustomCommunicator(self.cpu_group),

661 )

613 - if dispatch_dtype_bytes_per_elem == 0:

614 - hidden_bytes = hidden_size // 2

615 - else:

616 - hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem

617 - total_dispatch_payload_size_per_token = (

618 - hidden_bytes

619 - + dispatch_scale_bytes_per_token

620 - + top_k * 4 # int32 topks ids

621 - + top_k * 4 # float32 topk weights

622 - )

623 - combine_payload_size_per_token = hidden_size * 2 # bf16 hidden states

624 - self.workspace_size = moe_a2a_get_workspace_size_per_rank(

625 - ep_size=self.world_size,

626 - max_num_tokens=max_num_tokens,

627 - total_dispatch_payload_size_per_token=total_dispatch_payload_size_per_token,

628 - combine_payload_size_per_token=combine_payload_size_per_token,

629 - )

662

663 self.moe_alltoall = MoeAlltoAll(

664 mapping=self.mapping,

633 - max_num_tokens=max_num_tokens,

634 - top_k=top_k,

635 - num_experts=num_experts,

665 + max_num_tokens=self.max_num_tokens,

666 + top_k=self.top_k,

667 + num_experts=self.num_experts,

668 workspace_size_per_rank=self.workspace_size,

669 mnnvl_config=ep_config,

670 )

671

672 self.gpus_per_node = gpus_per_node

641 - self.max_num_tokens = max_num_tokens

642 - self.top_k = top_k

643 - self.num_experts = num_experts

644 - self.hidden_size = hidden_size

673 self.initialized = True

674

675 logger.info(

@@ -649,7 +677,10 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):

677 self.rank,

678 self.world_size,

679 )

652 - dist.barrier()

680 + # Scope barrier to the EP group: with PP, different EP groups can

681 + # rebuild a different number of times if their MoE layers have

682 + # different shape sequences, so a world-level barrier would deadlock.

683 + dist.barrier(group=self.cpu_group)

684

685 def get_handle(self, kwargs):

686 return self