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

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 13, 08:15 UTC

[MoE][Kernel][Perf] Improve Shared Expert Stream Overlap (#28406)

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

Alexander Matveev committed Nov 12, 2025 at 18:37 UTC 69d0e903137109dd9a1e0f88bc0af20db7ca70fb

3 files changed +81 -36

vllm/envs.py

@@ -222,6 +222,7 @@ if TYPE_CHECKING:

222 VLLM_USE_FBGEMM: bool = False

223 VLLM_GC_DEBUG: str = ""

224 VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False

225 + VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256

226 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"

227 VLLM_FLAT_LOGPROBS: bool = False

228

@@ -1476,6 +1477,13 @@ environment_variables: dict[str, Callable[[], Any]] = {

1477 "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool(

1478 int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0"))

1479 ),

1480 + # Limits when we run shared_experts in a separate stream.

1481 + # We found out that for large batch sizes, the separate stream

1482 + # execution is not beneficial (most likely because of the input clone)

1483 + # TODO(alexm-redhat): Tune to be more dynamic based on GPU type

1484 + "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int(

1485 + int(os.getenv("VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256))

1486 + ),

1487 # Format for saving torch.compile cache artifacts

1488 # - "binary": saves as binary file

1489 # Safe for multiple vllm serve processes accessing the same torch compile cache.

vllm/model_executor/layers/fused_moe/layer.py

+49 -36

@@ -48,7 +48,11 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

48 )

49 from vllm.platforms import current_platform

50 from vllm.utils.math_utils import cdiv, round_up

51 -from vllm.utils.torch_utils import current_stream, direct_register_custom_op

51 +from vllm.utils.torch_utils import (

52 + aux_stream,

53 + current_stream,

54 + direct_register_custom_op,

55 +)

56 from vllm.v1.worker.ubatching import dbo_current_ubatch_id

58 if current_platform.is_cuda_alike():

@@ -331,7 +335,11 @@ class FusedMoE(CustomOp):

335 logger.info_once("Disabling MoE shared_experts cuda stream")

336 self.shared_experts_stream = None

337 else:

334 - self.shared_experts_stream = torch.cuda.Stream()

338 + # TODO(rob): enable shared expert overlap with non-cuda.

339 + # aux_stream() returns None on non-cuda platforms.

340 + self.shared_experts_stream = aux_stream()

341 + if self.shared_experts_stream is not None:

342 + logger.info_once("Enabled separate cuda stream for MoE shared_experts")

343

344 if params_dtype is None:

345 params_dtype = torch.get_default_dtype()

@@ -1606,7 +1614,9 @@ class FusedMoE(CustomOp):

1614 if has_separate_shared_experts:

1615 assert not isinstance(final_hidden_states, tuple)

1616 assert self.shared_experts is not None

1617 +

1618 shared_output = self.shared_experts(staged_hidden_states)

1619 +

1620 final_hidden_states = (

1621 shared_output,

1622 final_hidden_states,

@@ -1684,13 +1694,34 @@ class FusedMoE(CustomOp):

1694

1695 use_chunked_impl = self.use_dp_chunking

1696

1687 - if (

1697 + use_shared_experts_stream = (

1698 has_separate_shared_experts

1699 and not use_chunked_impl

1700 and self.shared_experts_stream is not None

1691 - ):

1692 - # Start the separate shared experts stream here since we want

1693 - # to run in parallel with the router/gate (next op below)

1701 + and (

1702 + hidden_states.shape[0]

1703 + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

1704 + )

1705 + )

1706 +

1707 + if use_shared_experts_stream:

1708 + assert self.shared_experts_stream is not None

1709 +

1710 + # Clone BEFORE switching streams to avoid race condition

1711 + # where routed_expert kernel may mutate hidden_states.

1712 + hidden_states_clone = hidden_states.clone()

1713 +

1714 + # Record that the clone will be used by shared_experts_stream

1715 + # to avoid gc issue from deallocation of hidden_states_clone

1716 + # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501

1717 + # NOTE: We dont need shared_output.record_stream(current_stream())

1718 + # because we synch the streams before using shared_output.

1719 + hidden_states_clone.record_stream(self.shared_experts_stream)

1720 +

1721 + # Mark sync start point for the separate shared experts

1722 + # stream here since we want to run in parallel with the

1723 + # router/gate (next op below)

1724 + assert self.shared_experts_stream is not None

1725 self.shared_experts_stream.wait_stream(current_stream())

1726

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

@@ -1709,33 +1740,6 @@ class FusedMoE(CustomOp):

1740 self.quant_method, FusedMoEModularMethod

1741 )

1742

1712 - # If there are shared experts but we are not using a modular kernel, the

1713 - # shared experts must be called here

1714 - if has_separate_shared_experts:

1715 - assert self.shared_experts is not None

1716 -

1717 - if self.shared_experts_stream is not None:

1718 - # Clone BEFORE switching streams to avoid race condition

1719 - # where routed_expert kernel may mutate hidden_states.

1720 - hidden_states_clone = hidden_states.clone()

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

1722 -

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

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

1725 - shared_output = self.shared_experts(hidden_states_clone)

1726 -

1727 - # Record that the clone will be used by shared_experts_stream

1728 - # to avoid gc issue from deallocation of hidden_states_clone

1729 - # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501

1730 - # NOTE: we dont need shared_output.record_stream(current_stream())

1731 - # because we synch the streams before using shared_output.

1732 - hidden_states_clone.record_stream(self.shared_experts_stream)

1733 -

1734 - else:

1735 - shared_output = self.shared_experts(hidden_states)

1736 - else:

1737 - shared_output = None

1738 -

1743 ctx = get_forward_context()

1744 sp_ctx = (

1745 ctx.dp_metadata.sp_local_sizes(self.sp_size)

@@ -1776,12 +1780,21 @@ class FusedMoE(CustomOp):

1780 )

1781

1782 if has_separate_shared_experts:

1779 - assert not isinstance(final_hidden_states, tuple)

1783 assert self.shared_experts is not None

1784

1782 - # Wait for the parallel shared experts stream to finish here

1783 - if self.shared_experts_stream is not None:

1785 + if use_shared_experts_stream:

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

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

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

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

1790 + # graph replay later.

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

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

1793 + # conflict with the main stream

1794 + shared_output = self.shared_experts(hidden_states_clone)

1795 current_stream().wait_stream(self.shared_experts_stream)

1796 + else:

1797 + shared_output = self.shared_experts(hidden_states)

1798

1799 final_hidden_states = (

1800 shared_output,

vllm/utils/torch_utils.py

+24

@@ -409,6 +409,30 @@ def current_stream() -> torch.cuda.Stream:

409 return _current_stream_tls.value

410

411

412 +# Global auxilary stream for running operations in background streams.

413 +# We have single global auxilary stream to avoid an explosion of streams

414 +# for every layer (and make profiling look sane).

415 +#

416 +# aux_stream() is currently used for:

417 +# - MoE shared_expert overlap with router

418 +_aux_stream: torch.cuda.Stream | None = None

419 +

420 +

421 +def aux_stream() -> torch.cuda.Stream | None:

422 + """

423 + Ensures aux_stream is initialized only once

424 + """

425 + global _aux_stream

426 +

427 + from vllm.platforms import current_platform

428 +

429 + # TODO: validate this works properly on ROCm platform.

430 + if _aux_stream is None and current_platform.is_cuda():

431 + _aux_stream = torch.cuda.Stream()

432 +

433 + return _aux_stream

434 +

435 +

436 @lru_cache(maxsize=8)

437 def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:

438 # Note: cuda_visible_devices is not used, but we keep it as an argument for