9279c59a0e · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Callable

4 from contextlib import nullcontext

5 from typing import TYPE_CHECKING

6

@@ -82,9 +83,22 @@ def _moe_forward(

83 layer = get_layer_from_name(_resolve_layer_name(layer_name))

84 # TODO(bnell): this can be removed after MK migration is complete.

85 layer.ensure_moe_quant_config_init()

85 - return layer.runner.forward_impl(

86 - layer, hidden_states, router_logits, shared_experts_input

87 - )

86 + runner = layer.runner

87 + with runner._sequence_parallel_context():

88 + if runner.use_dp_chunking:

89 + return runner.forward_impl_chunked(

90 + layer,

91 + hidden_states,

92 + router_logits,

93 + shared_experts_input,

94 + )

95 + else:

96 + return runner.forward_impl(

97 + layer,

98 + hidden_states,

99 + router_logits,

100 + shared_experts_input,

101 + )

102

103

104 def _moe_forward_fake(

@@ -105,9 +119,22 @@ def _moe_forward_shared(

119 layer = get_layer_from_name(_resolve_layer_name(layer_name))

120 # TODO(bnell): this can be removed after MK migration is complete.

121 layer.ensure_moe_quant_config_init()

108 - return layer.runner.forward_impl(

109 - layer, hidden_states, router_logits, shared_experts_input

110 - )

122 + runner = layer.runner

123 + with runner._sequence_parallel_context():

124 + if runner.use_dp_chunking:

125 + return runner.forward_impl_chunked(

126 + layer,

127 + hidden_states,

128 + router_logits,

129 + shared_experts_input,

130 + )

131 + else:

132 + return runner.forward_impl(

133 + layer,

134 + hidden_states,

135 + router_logits,

136 + shared_experts_input,

137 + )

138

139

140 def _moe_forward_shared_fake(

@@ -191,10 +218,17 @@ class DefaultMoERunner(MoERunner):

218 self.reduce_results = reduce_results

219 self.enable_dbo = enable_dbo

220

221 + # Chunked all2all staging tensor

222 + # TODO(bnell) rename these?

223 + self.batched_hidden_states: torch.Tensor | None = None

224 + self.batched_router_logits: torch.Tensor | None = None

225 + self._maybe_init_dp_chunking()

226 +

227 # Allow disabling of the separate shared experts stream for

228 # debug purposes.

229 # TODO: Remove this after more extensive testings with TP/DP

230 # and other execution modes

231 + self.use_shared_experts_stream = False

232 if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM:

233 logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local")

234 self.shared_experts_stream = None

@@ -210,23 +244,20 @@ class DefaultMoERunner(MoERunner):

244 # Needed for string -> FusedMoE layer lookup in custom ops.

245 self.layer_name = layer.layer_name

246

247 + self.moe_forward = self._select_forward(layer)

248 +

249 + def _select_forward(self, layer: torch.nn.Module) -> Callable:

250 if current_platform.is_tpu() or current_platform.is_cpu():

251 # TODO: Once the OOM issue for the TPU backend is resolved, we

252 # will switch to using the moe_forward custom op.

253 # Note: CPU doesn't require wrapped forward_impl.

217 - if self.shared_experts is None:

218 - self.moe_forward = _moe_forward

219 - else:

220 - self.moe_forward = _moe_forward_shared

221 - else:

222 - if self.shared_experts is None:

223 - self.moe_forward = torch.ops.vllm.moe_forward

224 - else:

225 - self.moe_forward = torch.ops.vllm.moe_forward_shared

254 + return _moe_forward if self.shared_experts is None else _moe_forward_shared

255

227 - # Chunked all2all staging tensor

228 - self.batched_hidden_states: torch.Tensor | None = None

229 - self.batched_router_logits: torch.Tensor | None = None

256 + return (

257 + torch.ops.vllm.moe_forward

258 + if self.shared_experts is None

259 + else torch.ops.vllm.moe_forward_shared

260 + )

261

262 @property

263 def use_dp_chunking(self) -> bool:

@@ -241,22 +272,8 @@ class DefaultMoERunner(MoERunner):

272 self,

273 hidden_states: torch.Tensor,

274 shared_input: torch.Tensor | None,

244 - has_separate_shared_experts: bool,

245 - use_chunked_impl: bool,

246 - ) -> tuple[bool, torch.Tensor | None]:

247 - use_shared_experts_stream = (

248 - current_platform.is_cuda()

249 - and has_separate_shared_experts

250 - and not use_chunked_impl

251 - and self.shared_experts_stream is not None

252 - and (

253 - hidden_states.shape[0]

254 - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

255 - )

256 - )

257 -

258 - shared_experts_input: torch.Tensor | None = None

259 - if use_shared_experts_stream:

275 + ):

276 + if self.use_shared_experts_stream:

277 assert self.shared_experts_stream is not None

278 assert self.moe_config.disable_inplace

279

@@ -278,12 +295,11 @@ class DefaultMoERunner(MoERunner):

295 assert self.shared_experts_stream is not None

296 self.shared_experts_stream.wait_stream(current_stream())

297

281 - return use_shared_experts_stream, shared_experts_input

282 -

283 - def ensure_dp_chunking_init(self):

284 - if not self.use_dp_chunking or self.batched_hidden_states is not None:

298 + def _maybe_init_dp_chunking(self):

299 + if not self.use_dp_chunking:

300 return

301

302 + assert self.batched_hidden_states is None

303 states_shape: tuple[int, ...]

304 logits_shape: tuple[int, ...]

305

@@ -309,6 +325,38 @@ class DefaultMoERunner(MoERunner):

325 device=device,

326 )

327

328 + @property

329 + def has_separate_shared_experts(self) -> bool:

330 + return (

331 + not self.quant_method.mk_owns_shared_expert

332 + and self.shared_experts is not None

333 + )

334 +

335 + def _apply_shared_experts(

336 + self,

337 + hidden_states: torch.Tensor,

338 + allow_streaming: bool = False,

339 + ) -> torch.Tensor | None:

340 + shared_output: torch.Tensor | None = None

341 + if self.has_separate_shared_experts:

342 + assert self.shared_experts is not None

343 +

344 + if self.use_shared_experts_stream and allow_streaming:

345 + # Run shared experts in parallel on a separate stream

346 + # NOTE: We start the separate stream here and mark the

347 + # sync end point immediately after it is done. This is

348 + # important to avoid excessive stream allocations by the cuda

349 + # graph replay later.

350 + with torch.cuda.stream(self.shared_experts_stream):

351 + # Note that hidden_states clone() is necessary here to avoid

352 + # conflict with the main stream

353 + shared_output = self.shared_experts(hidden_states)

354 + current_stream().wait_stream(self.shared_experts_stream)

355 + else:

356 + shared_output = self.shared_experts(hidden_states)

357 +

358 + return shared_output

359 +

360 def must_reduce_shared_expert_outputs(self) -> bool:

361 """

362 The shared_experts are typically computed using the RowParallelLinear

@@ -322,7 +370,6 @@ class DefaultMoERunner(MoERunner):

370 Therefore it is required that we reduce the shared_experts output

371 early.

372 """

325 - assert self.quant_method is not None

373 return (

374 self.quant_method.moe_kernel is not None

375 and self.quant_method.moe_kernel.output_is_reduced()

@@ -357,7 +404,7 @@ class DefaultMoERunner(MoERunner):

404 return result

405 return hidden_states

406

360 - def _reduce_output(

407 + def _maybe_reduce_output(

408 self,

409 states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

410 trunc_sizes: list[int],

@@ -397,23 +444,16 @@ class DefaultMoERunner(MoERunner):

444 return "from_forward_context"

445 return self.layer_name

446

400 - def forward(

447 + def _maybe_pad_hidden_states(

448 self,

449 + original_hidden_states: torch.Tensor | None,

450 hidden_states: torch.Tensor,

403 - router_logits: torch.Tensor,

404 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

405 - # For latent MoE: save ORIGINAL hidden_states before transform

406 - # (shared_experts need original dimension, routed experts use transformed)

407 - if self.shared_experts is not None:

408 - original_hidden_states = hidden_states

409 - original_hidden_dim = hidden_states.shape[-1]

410 - else:

411 - original_hidden_states = None

412 -

413 - # Apply transform for routed experts (e.g., latent projection for latent MoE)

414 - hidden_states = self.apply_routed_input_transform(hidden_states)

415 -

416 - # This is the dimension after transform (for routed expert output slicing)

451 + ) -> tuple[torch.Tensor, list[int]]:

452 + original_hidden_dim = (

453 + original_hidden_states.shape[-1]

454 + if original_hidden_states is not None

455 + else 0

456 + )

457 transformed_hidden_dim = hidden_states.shape[-1]

458 if (

459 not self.quant_method.skip_forward_padding

@@ -426,134 +466,235 @@ class DefaultMoERunner(MoERunner):

466 value=0.0,

467 )

468

429 - fused_output = self.moe_forward(

430 - hidden_states,

431 - router_logits,

432 - original_hidden_states,

433 - self._encode_layer_name(),

434 - )

435 -

469 if self.shared_experts is not None:

470 orig_hidden_dims = [original_hidden_dim, transformed_hidden_dim]

471 else:

472 orig_hidden_dims = [transformed_hidden_dim]

473

441 - return self._reduce_output(fused_output, orig_hidden_dims)

474 + return hidden_states, orig_hidden_dims

475

443 - def forward_impl_chunked(

476 + def _apply_quant_method(

477 self,

478 layer: torch.nn.Module,

446 - full_hidden_states: torch.Tensor,

447 - full_router_logits: torch.Tensor,

448 - full_shared_input: torch.Tensor | None,

449 - has_separate_shared_experts: bool,

450 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

479 + hidden_states: torch.Tensor,

480 + router_logits: torch.Tensor,

481 + shared_input: torch.Tensor | None,

482 + run_shared_experts_before: bool = True,

483 + ) -> tuple[torch.Tensor | None, torch.Tensor]:

484 + shared_input = shared_input if shared_input is not None else hidden_states

485 + shared_output: torch.Tensor | None = None

486 +

487 + # Run this before quant_method to avoid inplace issues.

488 + if run_shared_experts_before:

489 + shared_output = self._apply_shared_experts(shared_input, False)

490 +

491 + if self.quant_method.is_monolithic:

492 + result = self.quant_method.apply_monolithic(

493 + layer=layer,

494 + x=hidden_states,

495 + router_logits=router_logits,

496 + )

497 + else:

498 + topk_weights, topk_ids = self.router.select_experts(

499 + hidden_states=hidden_states,

500 + router_logits=router_logits,

501 + )

502 +

503 + result = self.quant_method.apply(

504 + layer=layer,

505 + x=hidden_states,

506 + topk_weights=topk_weights,

507 + topk_ids=topk_ids,

508 + shared_experts_input=shared_input,

509 + )

510 +

511 + if isinstance(result, tuple):

512 + assert shared_output is None

513 + shared_output, hidden_states = result

514 + else:

515 + hidden_states = result

516 +

517 + if not run_shared_experts_before and self.has_separate_shared_experts:

518 + assert shared_output is None

519 + shared_output = self._apply_shared_experts(shared_input, True)

520 +

521 + return shared_output, hidden_states

522 +

523 + def _sequence_parallel_context(self):

524 + ctx = get_forward_context()

525 + return (

526 + ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)

527 + if ctx.dp_metadata

528 + else nullcontext()

529 + )

530 +

531 + def _allocate_dp_chunking_outputs(

532 + self,

533 + hidden_states: torch.Tensor,

534 + router_logits: torch.Tensor,

535 + ) -> tuple[torch.Tensor | None, torch.Tensor]:

536 + assert self.use_dp_chunking

537 +

538 + # Assert the inputs are of the proper type and shape.

539 assert self.batched_hidden_states is not None

540 assert self.batched_router_logits is not None

453 - assert self.batched_hidden_states.dtype == full_hidden_states.dtype, (

454 - f"{self.batched_hidden_states.dtype} == {full_hidden_states.dtype}"

541 +

542 + assert self.batched_hidden_states.dtype == hidden_states.dtype, (

543 + f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}"

544 )

456 - assert self.batched_router_logits.dtype == full_router_logits.dtype, (

457 - f"{self.batched_router_logits.dtype} == {full_router_logits.dtype}"

545 + assert self.batched_router_logits.dtype == router_logits.dtype, (

546 + f"{self.batched_router_logits.dtype} == {router_logits.dtype}"

547 )

459 - # Check size compatibility.

460 - assert self.batched_hidden_states.size(-1) == full_hidden_states.size(-1)

461 - assert self.batched_router_logits.size(-1) == full_router_logits.size(-1)

548

463 - # TODO(bnell): Fix shared_expert_inputs w/chunking.

464 - # assert shared_input is None, (

465 - # "Routed input transform is not currently supported with DP chunking."

466 - # )

549 + # Check size compatibility.

550 + assert self.batched_hidden_states.size(-1) == hidden_states.size(-1)

551 + assert self.batched_router_logits.size(-1) == router_logits.size(-1)

552

468 - full_fused_final_hidden_states = torch.empty_like(full_hidden_states)

553 + final_fused_hidden_states = torch.empty_like(hidden_states)

554 if self.shared_experts is not None:

470 - full_shared_final_hidden_states = torch.empty_like(full_hidden_states)

471 -

472 - def process_chunk(chunk_start, chunk_end, skip_result_store=False):

473 - chunk_size = chunk_end - chunk_start

474 - hidden_states = full_hidden_states[chunk_start:chunk_end, :]

475 - router_logits = full_router_logits[chunk_start:chunk_end, :]

476 - shared_input = (

477 - full_shared_input[chunk_start:chunk_end, :]

478 - if full_shared_input is not None

479 - else None

480 - )

555 + final_shared_hidden_states = torch.empty_like(hidden_states)

556 + else:

557 + final_shared_hidden_states = None

558

482 - assert self.batched_hidden_states is not None

483 - assert self.batched_router_logits is not None

484 - # This is only true when DBO has been enabled in the config.

485 - # Both tensors will have an outer dimension for the ubatch id

486 - if self.batched_hidden_states.dim() == 3:

487 - assert self.batched_router_logits.dim() == 3

488 - batch_buffer_idx = dbo_current_ubatch_id()

489 - batched_hidden_states = self.batched_hidden_states[batch_buffer_idx, :]

490 - batched_router_logits = self.batched_router_logits[batch_buffer_idx, :]

491 - else:

492 - batched_hidden_states = self.batched_hidden_states

493 - batched_router_logits = self.batched_router_logits

559 + return final_shared_hidden_states, final_fused_hidden_states

560 +

561 + def _maybe_gate(

562 + self,

563 + hidden_states: torch.Tensor,

564 + router_logits: torch.Tensor,

565 + ) -> torch.Tensor:

566 + # If router/gate provided, then apply it here.

567 + # (Note: This code runs only when "overlapped mode" is on to allow

568 + # parallel execution of shared experts with the FusedMoE via

569 + # separate cuda stream)

570 + if self.gate is not None:

571 + router_logits, _ = self.gate(hidden_states)

572 + return router_logits

573 +

574 + @property

575 + def do_naive_dispatch_combine(self) -> bool:

576 + return (

577 + self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk

578 + )

579

495 - assert (

496 - batched_hidden_states.size(0) # type: ignore

497 - >= chunk_size

580 + def _maybe_dispatch(

581 + self,

582 + layer: torch.nn.Module,

583 + hidden_states: torch.Tensor,

584 + router_logits: torch.Tensor,

585 + ) -> tuple[torch.Tensor, torch.Tensor]:

586 + # For naive dispatch/combine Dp/Ep, dispatch the hidden states and

587 + # router logits to all experts.

588 + # NOTE: this will be removed once all kernels are migrated into the

589 + # MoEKernel framework.

590 + if self.do_naive_dispatch_combine:

591 + hidden_states, router_logits = get_ep_group().dispatch_router_logits(

592 + hidden_states,

593 + router_logits,

594 + self.moe_config.is_sequence_parallel,

595 )

499 - assert (

500 - batched_router_logits.size(0) # type: ignore

501 - >= chunk_size

596 +

597 + # NOTE: Similar with DP, PCP also needs dispatch and combine. For

598 + # simplicity, AgRsAll2All was added separately for PCP here. Maybe

599 + # we should modify All2AllManager abstraction to better support PCP.

600 + if self.moe_config.pcp_size > 1:

601 + hidden_states = get_pcp_group().all_gather(

602 + hidden_states,

603 + dim=0,

604 )

503 - staged_hidden_states = batched_hidden_states[:chunk_size, :] # type: ignore

504 - staged_router_logits = batched_router_logits[:chunk_size, :] # type: ignore

505 - staged_hidden_states.copy_(hidden_states, non_blocking=True)

506 - staged_router_logits.copy_(router_logits, non_blocking=True)

605 + router_logits = get_pcp_group().all_gather(

606 + router_logits,

607 + dim=0,

608 + )

609 +

610 + return hidden_states, router_logits

611

508 - shared_input = (

509 - shared_input if shared_input is not None else staged_hidden_states

612 + def _maybe_combine(

613 + self,

614 + shared_output: torch.Tensor | None,

615 + hidden_states: torch.Tensor,

616 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:

617 + if self.do_naive_dispatch_combine:

618 + hidden_states = get_ep_group().combine(

619 + hidden_states, self.moe_config.is_sequence_parallel

620 )

621

512 - # Matrix multiply.

513 - if self.quant_method.is_monolithic:

514 - assert has_separate_shared_experts or self.shared_experts is None

515 - final_hidden_states = self.quant_method.apply_monolithic(

516 - layer=layer,

517 - x=staged_hidden_states,

518 - router_logits=staged_router_logits,

519 - )

520 - else:

521 - topk_weights, topk_ids = self.router.select_experts(

522 - hidden_states=staged_hidden_states,

523 - router_logits=staged_router_logits,

524 - )

622 + if self.moe_config.pcp_size > 1:

623 + hidden_states = get_pcp_group().reduce_scatter(

624 + hidden_states,

625 + dim=0,

626 + )

627 + # need RS for shared_output?

628

526 - final_hidden_states = self.quant_method.apply(

527 - layer=layer,

528 - x=staged_hidden_states,

529 - topk_weights=topk_weights,

530 - topk_ids=topk_ids,

531 - shared_experts_input=shared_input,

532 - )

629 + if self.shared_experts is not None:

630 + assert shared_output is not None

631 + return shared_output, hidden_states

632 + else:

633 + return hidden_states

634

534 - if has_separate_shared_experts:

535 - assert not isinstance(final_hidden_states, tuple)

536 - assert self.shared_experts is not None

635 + def forward(

636 + self,

637 + hidden_states: torch.Tensor,

638 + router_logits: torch.Tensor,

639 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

640 + # For latent MoE: save ORIGINAL hidden_states before transform

641 + # (shared_experts need original dimension, routed experts use transformed)

642 + if self.shared_experts is not None:

643 + original_hidden_states = hidden_states

644 + else:

645 + original_hidden_states = None

646

538 - shared_output = self.shared_experts(shared_input)

647 + # Apply transform for routed experts (e.g., latent projection for latent MoE)

648 + hidden_states = self.apply_routed_input_transform(hidden_states)

649

540 - final_hidden_states = (

541 - shared_output,

542 - final_hidden_states,

543 - )

650 + hidden_states, og_hidden_dims = self._maybe_pad_hidden_states(

651 + original_hidden_states,

652 + hidden_states,

653 + )

654

545 - if not skip_result_store:

546 - if self.shared_experts is None:

547 - full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_(

548 - final_hidden_states, non_blocking=True

549 - )

550 - else:

551 - full_shared_final_hidden_states[chunk_start:chunk_end, :].copy_(

552 - final_hidden_states[0], non_blocking=True

553 - )

554 - full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_(

555 - final_hidden_states[1], non_blocking=True

556 - )

655 + fused_output = self.moe_forward(

656 + hidden_states,

657 + router_logits,

658 + original_hidden_states,

659 + self._encode_layer_name(),

660 + )

661 +

662 + return self._maybe_reduce_output(fused_output, og_hidden_dims)

663 +

664 + def _slice_and_copy_input(

665 + self,

666 + out_slice: torch.Tensor,

667 + orig: torch.Tensor | None,

668 + start: int,

669 + end: int,

670 + ) -> torch.Tensor:

671 + assert orig is not None

672 + slice_size = end - start

673 + orig_slice = orig[start:end, :]

674 + if self.enable_dbo:

675 + assert out_slice.dim() == 3

676 + batch_buffer_idx = dbo_current_ubatch_id()

677 + out_slice = out_slice[batch_buffer_idx, :]

678 +

679 + assert out_slice.size(0) >= slice_size

680 + out_slice = out_slice[:slice_size, :]

681 + out_slice.copy_(orig_slice, non_blocking=True)

682 + return out_slice

683 +

684 + def forward_impl_chunked(

685 + self,

686 + layer: torch.nn.Module,

687 + hidden_states: torch.Tensor,

688 + router_logits: torch.Tensor,

689 + shared_input: torch.Tensor | None,

690 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

691 + # Gate overlap not supported when chunking is enabled. Run the

692 + # gate first.

693 + router_logits = self._maybe_gate(hidden_states, router_logits)

694 +

695 + final_shared_hidden_states, final_fused_hidden_states = (

696 + self._allocate_dp_chunking_outputs(hidden_states, router_logits)

697 + )

698

699 ctx = get_forward_context()

700 # flashinfer_cutlass_kernels can handle: optional DP + TP/EP

@@ -567,7 +708,7 @@ class DefaultMoERunner(MoERunner):

708 max_tokens_across_dispatchers, self.moe_config.sp_size

709 )

710

570 - num_tokens = full_hidden_states.size(0)

711 + num_tokens = hidden_states.size(0)

712 for chunk_idx, chunk_start_ in enumerate(

713 range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank)

714 ):

@@ -578,17 +719,55 @@ class DefaultMoERunner(MoERunner):

719 # clamp start and end

720 chunk_start = min(chunk_start, num_tokens - 1)

721 chunk_end = min(chunk_end, num_tokens)

581 - with ctx.dp_metadata.chunked_sizes(

722 + chunk_sizes = ctx.dp_metadata.chunked_sizes(

723 self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx

583 - ):

584 - process_chunk(

585 - chunk_start, chunk_end, skip_result_store=chunk_start_ >= num_tokens

724 + )

725 + with chunk_sizes:

726 + hidden_states_chunk = self._slice_and_copy_input(

727 + self.batched_hidden_states,

728 + hidden_states,

729 + chunk_start,

730 + chunk_end,

731 + )

732 +

733 + router_logits_chunk = self._slice_and_copy_input(

734 + self.batched_router_logits,

735 + router_logits,

736 + chunk_start,

737 + chunk_end,

738 + )

739 +

740 + shared_input_chunk = (

741 + shared_input[chunk_start:chunk_end, :]

742 + if shared_input is not None

743 + else None

744 + )

745 +

746 + shared_output_chunk, hidden_states_chunk = self._apply_quant_method(

747 + layer=layer,

748 + hidden_states=hidden_states_chunk,

749 + router_logits=router_logits_chunk,

750 + shared_input=shared_input_chunk,

751 )

752

753 + # Store outputs

754 + # TODO(bnell): document when chunk_start >= num_tokens

755 + if chunk_start < num_tokens:

756 + final_fused_hidden_states[chunk_start:chunk_end, :].copy_(

757 + hidden_states_chunk, non_blocking=True

758 + )

759 + if self.shared_experts is not None:

760 + assert shared_output_chunk is not None

761 + assert final_shared_hidden_states is not None

762 + final_shared_hidden_states[chunk_start:chunk_end, :].copy_(

763 + shared_output_chunk, non_blocking=True

764 + )

765 +

766 if self.shared_experts is None:

589 - return full_fused_final_hidden_states

767 + return final_fused_hidden_states

768 else:

591 - return (full_shared_final_hidden_states, full_fused_final_hidden_states)

769 + assert final_shared_hidden_states is not None

770 + return (final_shared_hidden_states, final_fused_hidden_states)

771

772 def forward_impl(

773 self,

@@ -597,148 +776,51 @@ class DefaultMoERunner(MoERunner):

776 router_logits: torch.Tensor,

777 shared_input: torch.Tensor | None,

778 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

600 - assert self.quant_method is not None

601 -

602 - self.ensure_dp_chunking_init()

603 -

604 - has_separate_shared_experts = (

605 - not self.quant_method.mk_owns_shared_expert

606 - and self.shared_experts is not None

779 + self.use_shared_experts_stream = (

780 + current_platform.is_cuda()

781 + and self.has_separate_shared_experts

782 + and not self.use_dp_chunking

783 + and self.shared_experts_stream is not None

784 + and (

785 + hidden_states.shape[0]

786 + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

787 + )

788 )

789

609 - use_chunked_impl = self.use_dp_chunking

790 + # Check if we need to run shared experts before matrix multiply because

791 + # matrix multiply may modify the hidden_states.

792 + run_shared_experts_before = (

793 + self.has_separate_shared_experts and not self.use_shared_experts_stream

794 + )

795

611 - use_shared_experts_stream, shared_experts_input = (

796 + # The shared experts stream must be set up before calling the gate so they

797 + # can be overlapped.

798 + if not run_shared_experts_before:

799 self._maybe_setup_shared_experts_stream(

800 hidden_states,

801 shared_input,

615 - has_separate_shared_experts,

616 - use_chunked_impl,

802 )

618 - )

803

620 - # If router/gate provided, then apply it here.

621 - # (Note: This code runs only when "overlapped mode" is on to allow

622 - # parallel execution of shared experts with the FusedMoE via

623 - # separate cuda stream)

624 - if self.gate is not None:

625 - router_logits, _ = self.gate(hidden_states)

626 -

627 - if use_chunked_impl:

628 - return self.forward_impl_chunked(

629 - layer,

630 - hidden_states,

631 - router_logits,

632 - shared_input,

633 - has_separate_shared_experts,

634 - )

804 + router_logits = self._maybe_gate(hidden_states, router_logits)

805

636 - # NOTE(rob): once we finish migrating all the quant methods to use

637 - # MKs, we can remove the naive dispatch/combine path from here.

638 - do_naive_dispatch_combine = (

639 - self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk

806 + # TODO(bnell): parts of the dispatch/combine steps will go away once

807 + # #32567 lands and the remaining kernels are made MKs. The PCP

808 + # code will probably remain

809 + hidden_states, router_logits = self._maybe_dispatch(

810 + layer,

811 + hidden_states,

812 + router_logits,

813 )

814

642 - ctx = get_forward_context()

643 - sp_ctx = (

644 - ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)

645 - if ctx.dp_metadata

646 - else nullcontext()

815 + shared_output, hidden_states = self._apply_quant_method(

816 + layer=layer,

817 + hidden_states=hidden_states,

818 + router_logits=router_logits,

819 + shared_input=shared_input,

820 + run_shared_experts_before=run_shared_experts_before,

821 )

822

649 - with sp_ctx:

650 - # Run shared experts before matrix multiply.

651 - # because matrix multiply maybe modify the hidden_states.

652 - if has_separate_shared_experts and not use_shared_experts_stream:

653 - assert self.shared_experts is not None

654 - shared_input = (

655 - shared_input if shared_input is not None else hidden_states

656 - )

657 - shared_output = self.shared_experts(shared_input)

658 -

659 - # For naive dispatch/combine Dp/Ep, dispatch the hidden states and

660 - # router logits to all experts.

661 - # NOTE: this will be removed once all kernels are migrated into the

662 - # MoEKernel framework.

663 - if do_naive_dispatch_combine:

664 - hidden_states, router_logits = get_ep_group().dispatch_router_logits(

665 - hidden_states,

666 - router_logits,

667 - self.moe_config.is_sequence_parallel,

668 - )

669 -

670 - # NOTE: Similar with DP, PCP also needs dispatch and combine. For

671 - # simplicity, AgRsAll2All was added separately for PCP here. Maybe

672 - # we should modify All2AllManager abstract to better support PCP.

673 - if self.moe_config.pcp_size > 1:

674 - hidden_states = get_pcp_group().all_gather(

675 - hidden_states,

676 - dim=0,

677 - )

678 - router_logits = get_pcp_group().all_gather(

679 - router_logits,

680 - dim=0,

681 - )

682 -

683 - # Matrix multiply.

684 - if self.quant_method.is_monolithic:

685 - final_hidden_states = self.quant_method.apply_monolithic(

686 - layer=layer,

687 - x=hidden_states,

688 - router_logits=router_logits,

689 - )

690 - else:

691 - topk_weights, topk_ids = self.router.select_experts(

692 - hidden_states=hidden_states,

693 - router_logits=router_logits,

694 - )

695 -

696 - final_hidden_states = self.quant_method.apply(

697 - layer=layer,

698 - x=hidden_states,

699 - topk_weights=topk_weights,

700 - topk_ids=topk_ids,

701 - shared_experts_input=shared_input,

702 - )

703 -

704 - if has_separate_shared_experts:

705 - assert self.shared_experts is not None

706 -

707 - if use_shared_experts_stream:

708 - # Run shared experts in parallel on a separate stream

709 - # NOTE: We start the separate stream here and mark the

710 - # sync end point immediately after it is done. This is

711 - # important to avoid excessive stream allocations by the cuda

712 - # graph replay later.

713 - with torch.cuda.stream(self.shared_experts_stream):

714 - # Note that hidden_states clone() is necessary here to avoid

715 - # conflict with the main stream

716 - shared_output = self.shared_experts(shared_experts_input)

717 - current_stream().wait_stream(self.shared_experts_stream)

718 -

719 - final_hidden_states = (

720 - shared_output,

721 - final_hidden_states,

722 - )

723 -

724 - def combine_output(states: torch.Tensor) -> torch.Tensor:

725 - if do_naive_dispatch_combine:

726 - states = get_ep_group().combine(

727 - states, self.moe_config.is_sequence_parallel

728 - )

729 -

730 - if self.moe_config.pcp_size > 1:

731 - states = get_pcp_group().reduce_scatter(

732 - states,

733 - dim=0,

734 - )

735 -

736 - return states

737 -

738 - if self.shared_experts is not None:

739 - return (

740 - final_hidden_states[0],

741 - combine_output(final_hidden_states[1]),

742 - )

743 - else:

744 - return combine_output(final_hidden_states)

823 + return self._maybe_combine(

824 + shared_output,

825 + hidden_states,

826 + )

[MoE Refactor] DefaultMoERunner simplifcation (#33049)