chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

EP Support for Piecewise Cuda Graph (#14164)

Browse Source 
Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>
This commit is contained in:
Yuwei An
2025-12-19 09:59:27 -08:00
committed by GitHub
parent 05eb0bcc61
commit 9d0347b33a
11 changed files with 224 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/compilation/backend.py
View File

@@ -26,16 +26,6 @@ from sglang.srt.utils.common import is_npu, rank0_log
logger = logging.getLogger(__name__)
SPLIT_OPS = [
"sglang.unified_attention_with_output",
"sglang.gdn_with_output",
]
def add_split_ops(ops):
SPLIT_OPS.extend(ops)
def make_compiler(config: CompilationConfig):
if config.compiler == "eager":
return EagerAdapter()
@@ -433,7 +423,7 @@ class SGLangBackend:
self.split_gm, self.piecewise_graphs = split_graph(
graph,
SPLIT_OPS,
self.compile_config.split_ops,
)
from torch._dynamo.utils import lazy_format_graph_code

7
python/sglang/srt/compilation/compilation_config.py
View File

@@ -15,6 +15,13 @@ class CompilationConfig:
self.capture_sizes = capture_sizes
self.compiler = compiler
self.enable_debug_mode = enable_debug_mode
self.split_ops = [
"sglang.unified_attention_with_output",
"sglang.gdn_with_output",
]
def add_split_op(self, op: str):
self.split_ops.append(op)
def add_traced_file(self, file_path: str):
self.traced_files.add(file_path)

11
python/sglang/srt/compilation/piecewise_context_manager.py
View File

@@ -26,6 +26,8 @@ class ForwardContext:
def __init__(self):
self.forward_batch = None
self.attention_layer = None
self.quant_config = None
self.moe_layers = None
def set_forward_batch(self, forward_batch: ForwardBatch):
self.forward_batch = forward_batch
@@ -36,6 +38,9 @@ class ForwardContext:
def set_quant_config(self, quant_config: Any):
self.quant_config = quant_config
def set_moe_layers(self, layers: List[Any]):
self.moe_layers = layers
_forward_context: Optional[ForwardContext] = None
@@ -48,13 +53,17 @@ def get_forward_context() -> Optional[ForwardContext]:
@contextmanager
def set_forward_context(
forward_batch: ForwardBatch, attention_layers: List[Any], quant_config: Any
forward_batch: ForwardBatch,
attention_layers: List[Any],
quant_config: Any,
moe_layers: List[Any],
):
global _forward_context
_forward_context = ForwardContext()
_forward_context.set_forward_batch(forward_batch)
_forward_context.set_attention_layers(attention_layers)
_forward_context.set_quant_config(quant_config)
_forward_context.set_moe_layers(moe_layers)
try:
yield
finally:

1
python/sglang/srt/layers/communicator.py
View File

@@ -545,6 +545,7 @@ class LayerCommunicator:
return True
return False
# NOTE: This function will cause torch recompilation
def should_fuse_mlp_allreduce_with_next_layer(
self, forward_batch: ForwardBatch
) -> bool:

25
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Union
import torch
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A16Int4DynamicMoEMethod,
@@ -20,7 +21,7 @@ from sglang.srt.layers.moe.token_dispatcher.deepep import (
DeepEPLLCombineInput,
DeepEPNormalCombineInput,
)
from sglang.srt.layers.moe.topk import TopKOutput
from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
@@ -90,7 +91,6 @@ class DeepEPMoE(FusedMoE):
routed_scaling_factor=routed_scaling_factor,
**kwargs,
)
if _use_aiter or _is_npu:
self.deprecate_flag = False
elif deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and isinstance(
@@ -152,9 +152,28 @@ class DeepEPMoE(FusedMoE):
hidden_states: torch.Tensor,
topk_output: TopKOutput,
):
if is_in_piecewise_cuda_graph():
assert TopKOutputChecker.format_is_standard(
topk_output
), "Only standard topk output is supported for piecewise cuda graph"
return torch.ops.sglang.moe_forward_piecewise_cuda_graph_impl(
hidden_states,
topk_output.topk_weights,
topk_output.topk_ids,
topk_output.router_logits,
self.layer_id,
)
else:
return self.forward_impl(hidden_states, topk_output)
def forward_impl(
self,
hidden_states: torch.Tensor,
topk_output: TopKOutput,
):
if self.deprecate_flag:
return super().forward(
return super().forward_impl(
hidden_states,
topk_output,
)

86
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -8,6 +8,10 @@ import torch
from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher
from sglang.srt.compilation.piecewise_context_manager import (
get_forward_context,
is_in_piecewise_cuda_graph,
)
from sglang.srt.distributed import (
get_moe_expert_parallel_rank,
get_moe_expert_parallel_world_size,
@@ -37,7 +41,7 @@ from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardDispatcher,
StandardDispatchOutput,
)
from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker
from sglang.srt.layers.moe.topk import StandardTopKOutput, TopKOutput, TopKOutputChecker
from sglang.srt.layers.moe.utils import RoutingMethodType
from sglang.srt.layers.quantization.base_config import (
FusedMoEMethodBase,
@@ -57,6 +61,7 @@ from sglang.srt.utils import (
next_power_of_2,
round_up,
)
from sglang.srt.utils.common import direct_register_custom_op
if is_flashinfer_available():
from flashinfer import fp4_quantize
@@ -882,6 +887,21 @@ class FusedMoE(torch.nn.Module):
)
def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
if is_in_piecewise_cuda_graph():
assert TopKOutputChecker.format_is_standard(
topk_output
), "Only standard topk output is supported for piecewise cuda graph"
return torch.ops.sglang.moe_forward_piecewise_cuda_graph_impl(
hidden_states,
topk_output.topk_weights,
topk_output.topk_ids,
topk_output.router_logits,
self.layer_id,
)
else:
return self.forward_impl(hidden_states, topk_output)
def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
origin_hidden_states_dim = hidden_states.shape[-1]
assert self.quant_method is not None
@@ -1039,6 +1059,21 @@ class FlashInferFusedMoE(FusedMoE):
super().__init__(*args, **kwargs)
def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
if is_in_piecewise_cuda_graph():
assert TopKOutputChecker.format_is_standard(
topk_output
), "Only standard topk output is supported for piecewise cuda graph"
return torch.ops.sglang.moe_forward_piecewise_cuda_graph_impl(
hidden_states,
topk_output.topk_weights,
topk_output.topk_ids,
topk_output.router_logits,
self.layer_id,
)
else:
return self.forward_impl(hidden_states, topk_output)
def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
assert (
self.moe_runner_config.activation == "silu"
), "Only silu is supported for flashinfer trtllm moe"
@@ -1150,6 +1185,21 @@ class FlashInferFP4MoE(FusedMoE):
return hs_fp4, hs_sf
def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
if is_in_piecewise_cuda_graph():
assert TopKOutputChecker.format_is_standard(
topk_output
), "Only standard topk output is supported for piecewise cuda graph"
return torch.ops.sglang.moe_forward_piecewise_cuda_graph_impl(
hidden_states,
topk_output.topk_weights,
topk_output.topk_ids,
topk_output.router_logits,
self.layer_id,
)
else:
return self.forward_impl(hidden_states, topk_output)
def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
"""Forward pass using FP4 TRTLLM kernel.
Args:
@@ -1234,3 +1284,37 @@ class FlashInferFP4MoE(FusedMoE):
)[0]
return result
def moe_forward_piecewise_cuda_graph_impl(
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
router_logits: torch.Tensor,
layer_id: int,
) -> torch.Tensor:
# only standard topk output is supported for piecewise cuda graph
topk_output = StandardTopKOutput(
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits
)
forward_context = get_forward_context()
moe_layer = forward_context.moe_layers[layer_id]
return moe_layer.forward_impl(hidden_states, topk_output)
def moe_forward_piecewise_cuda_graph_impl_fake(
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
router_logits: torch.Tensor,
layer_id: int,
) -> torch.Tensor:
return torch.empty_like(hidden_states)
direct_register_custom_op(
op_name="moe_forward_piecewise_cuda_graph_impl",
op_func=moe_forward_piecewise_cuda_graph_impl,
mutates_args=[],
fake_impl=moe_forward_piecewise_cuda_graph_impl_fake,
)

19
python/sglang/srt/layers/moe/topk.py
View File

@@ -17,7 +17,7 @@ from __future__ import annotations
import logging
import math
from dataclasses import dataclass
from enum import Enum, auto
from enum import IntEnum, auto
from typing import (
TYPE_CHECKING,
Callable,
@@ -120,33 +120,24 @@ class TopKOutputChecker:
@staticmethod
def format_is_standard(topk_output: TopKOutput) -> TypeGuard[StandardTopKOutput]:
return topk_output.format.is_standard()
return isinstance(topk_output, StandardTopKOutput)
@staticmethod
def format_is_triton_kernels(
topk_output: TopKOutput,
) -> TypeGuard[TritonKernelTopKOutput]:
return topk_output.format.is_triton_kernels()
return isinstance(topk_output, TritonKernelTopKOutput)
@staticmethod
def format_is_bypassed(topk_output: TopKOutput) -> TypeGuard[BypassedTopKOutput]:
return topk_output.format.is_bypassed()
return isinstance(topk_output, BypassedTopKOutput)
class TopKOutputFormat(Enum):
class TopKOutputFormat(IntEnum):
STANDARD = auto()
TRITON_KERNEL = auto()
BYPASSED = auto()
def is_standard(self) -> bool:
return self == TopKOutputFormat.STANDARD
def is_triton_kernels(self) -> bool:
return self == TopKOutputFormat.TRITON_KERNEL
def is_bypassed(self) -> bool:
return self == TopKOutputFormat.BYPASSED
@runtime_checkable
class TopKOutput(Protocol):

24
python/sglang/srt/model_executor/model_runner.py
View File

@@ -97,6 +97,7 @@ from sglang.srt.layers.dp_attention import (
set_is_extend_in_batch,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
from sglang.srt.layers.sampler import create_sampler
@@ -1723,6 +1724,13 @@ class ModelRunner:
"Disable piecewise CUDA graph because piecewise_cuda_graph does not support PP",
)
return False
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
# TODO(yuwei): fix the compilation errors for MOE A2A backend
log_info_on_rank0(
logger,
"Disable piecewise CUDA graph due to existing compilation errors",
)
return False
return True
def init_memory_pool(
@@ -2623,9 +2631,10 @@ class ModelRunner:
):
return
# Collect attention layers from the model
self.attention_layers = []
# Collect attention layers and moe layers from the model
self.model.model = resolve_language_model(self.model)
self.attention_layers = []
self.moe_layers = []
for layer in self.model.model.layers:
if hasattr(layer, "self_attn"):
if hasattr(layer.self_attn, "attn"):
@@ -2643,6 +2652,17 @@ class ModelRunner:
if hasattr(layer.attention, "attn"):
self.attention_layers.append(layer.attention.attn)
moe_block = None
if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
moe_block = layer.mlp.experts
if hasattr(layer, "block_sparse_moe") and hasattr(
layer.block_sparse_moe, "experts"
):
moe_block = layer.block_sparse_moe.experts
if hasattr(layer, "moe") and hasattr(layer.moe, "experts"):
moe_block = layer.moe.experts
self.moe_layers.append(moe_block)
if len(self.attention_layers) < self.model_config.num_hidden_layers:
# TODO(yuwei): support Non-Standard GQA
log_info_on_rank0(

49
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
View File

@@ -44,6 +44,7 @@ from sglang.srt.layers.dp_attention import (
set_is_extend_in_batch,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
@@ -195,6 +196,11 @@ class PiecewiseCudaGraphRunner:
self.model_runner.server_args.piecewise_cuda_graph_compiler,
self.model_runner.server_args.enable_torch_compile_debug_mode,
)
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
self.compile_config.add_split_op(
"sglang.moe_forward_piecewise_cuda_graph_impl"
)
self.quant_config = getattr(self.model_runner.model, "quant_config", None)
# Batch sizes to capture
@@ -243,6 +249,7 @@ class PiecewiseCudaGraphRunner:
)
self.attention_layers = self.model_runner.attention_layers
self.moe_layers = self.model_runner.moe_layers
if get_global_graph_memory_pool() is None:
set_global_graph_memory_pool(self.device_module.graph_pool_handle())
@@ -260,10 +267,9 @@ class PiecewiseCudaGraphRunner:
compile_config=self.compile_config,
graph_pool=get_global_graph_memory_pool(),
)
with set_compiled(True):
self.warmup_torch_compile()
with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc):
with set_compiled(True):
self.warmup_torch_compile()
# Capture
try:
self.capture()
@@ -277,28 +283,24 @@ class PiecewiseCudaGraphRunner:
def warmup_torch_compile(self):
"""Warmup the model with a simple forward pass before CUDA graph capture."""
num_tokens = 2
input_ids = self.input_ids[:num_tokens]
input_embeds = self.input_embeds[:num_tokens] if self.is_multimodal else None
out_cache_loc = self.out_cache_loc[:num_tokens]
out_cache_loc_swa = (
self.out_cache_loc_swa[:num_tokens]
if self.out_cache_loc_swa is not None
else None
)
positions = self.positions[:num_tokens]
mrope_positions = (
self.mrope_positions[:, :num_tokens] if self.is_multimodal else None
)
with torch.device(self.device):
forward_batch = ForwardBatch(
forward_mode=ForwardMode.EXTEND,
batch_size=1,
input_ids=(torch.randint(0, 100, (num_tokens,), device=self.device)),
input_embeds=(
torch.randn(
num_tokens,
self.model_runner.model_config.hidden_size,
dtype=self.model_runner.dtype,
device=self.device,
)
if self.is_multimodal
else None
),
input_ids=input_ids,
input_embeds=input_embeds,
req_pool_indices=torch.arange(1, device=self.device),
seq_lens=torch.tensor([num_tokens], device=self.device),
next_token_logits_buffer=None,
@@ -319,14 +321,12 @@ class PiecewiseCudaGraphRunner:
extend_prefix_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"),
positions=torch.arange(num_tokens, device=self.device),
positions=positions,
global_num_tokens_gpu=None,
global_num_tokens_for_logprob_gpu=None,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=None,
mrope_positions=(
self.mrope_positions[:, :num_tokens] if self.is_multimodal else None
),
mrope_positions=mrope_positions,
spec_algorithm=None,
spec_info=None,
capture_hidden_mode=CaptureHiddenMode.NULL,
@@ -338,7 +338,7 @@ class PiecewiseCudaGraphRunner:
# Attention backend
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
with set_forward_context(
forward_batch, self.attention_layers, self.quant_config
forward_batch, self.attention_layers, self.quant_config, self.moe_layers
), disable_ca_comm(self.model_runner.tp_group):
_ = self.model_runner.model.forward(
forward_batch.input_ids,
@@ -483,7 +483,7 @@ class PiecewiseCudaGraphRunner:
kwargs = {}
with set_forward_context(
forward_batch, self.attention_layers, self.quant_config
forward_batch, self.attention_layers, self.quant_config, self.moe_layers
):
self.model_runner.model.forward(
forward_batch.input_ids,
@@ -608,7 +608,10 @@ class PiecewiseCudaGraphRunner:
static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
# Replay
with set_forward_context(
static_forward_batch, self.attention_layers, self.quant_config
static_forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
):
with set_compiled(True):
output = self.model_runner.model.forward(