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

[PCG] GPT OSS Triton Kernel Support (#18405)

Browse Source 
Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>
This commit is contained in:
Yuwei An
2026-02-11 17:23:55 -08:00
committed by GitHub
parent f06ab17a73
commit 2bd8363486
6 changed files with 68 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -56,6 +56,7 @@ class ForwardContext:
self.attention_layer = None
self.quant_config = None
self.moe_layers = None
self.moe_fusions = None
def set_forward_batch(self, forward_batch: ForwardBatch):
self.forward_batch = forward_batch
@@ -69,6 +70,9 @@ class ForwardContext:
def set_moe_layers(self, layers: List[Any]):
self.moe_layers = layers
def set_moe_fusions(self, fusions: List[Any]):
self.moe_fusions = fusions
_forward_context: Optional[ForwardContext] = None
@@ -85,6 +89,7 @@ def set_forward_context(
attention_layers: List[Any],
quant_config: Any,
moe_layers: List[Any],
moe_fusions: List[Any],
):
global _forward_context
_forward_context = ForwardContext()
@@ -92,6 +97,7 @@ def set_forward_context(
_forward_context.set_attention_layers(attention_layers)
_forward_context.set_quant_config(quant_config)
_forward_context.set_moe_layers(moe_layers)
_forward_context.set_moe_fusions(moe_fusions)
try:
yield
finally:

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

@@ -957,16 +957,17 @@ 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 moe_forward_piecewise_cuda_graph_impl(
hidden_states,
topk_output.topk_weights,
topk_output.topk_ids,
topk_output.router_logits,
self.layer_id,
)
if not TopKOutputChecker.format_is_standard(topk_output):
# Make sure there is torch lib op registration for the whole moe layer
return self.forward_impl(hidden_states, topk_output)
else:
return 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)
@@ -1129,16 +1130,17 @@ class FlashInferFusedMoE(FusedMoE):
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 moe_forward_piecewise_cuda_graph_impl(
hidden_states,
topk_output.topk_weights,
topk_output.topk_ids,
topk_output.router_logits,
self.layer_id,
)
if not TopKOutputChecker.format_is_standard(topk_output):
# Make sure there is torch lib op registration for the whole moe layer
return self.forward_impl(hidden_states, topk_output)
else:
return 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)

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

@@ -2126,6 +2126,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
language_model = getattr(self.model, "language_model", self.model)
self.attention_layers = []
self.moe_layers = []
self.moe_fusions = []
for layer in language_model.model.layers:
if hasattr(layer, "self_attn"):
if hasattr(layer.self_attn, "attn"):
@@ -2144,15 +2145,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.attention_layers.append(layer.attention.attn)
moe_block = None
moe_fusion = None
if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
moe_block = layer.mlp.experts
moe_fusion = layer.mlp
if hasattr(layer, "block_sparse_moe") and hasattr(
layer.block_sparse_moe, "experts"
):
moe_block = layer.block_sparse_moe.experts
moe_fusion = layer.block_sparse_moe
if hasattr(layer, "moe") and hasattr(layer.moe, "experts"):
moe_block = layer.moe.experts
moe_fusion = layer.moe
self.moe_layers.append(moe_block)
self.moe_fusions.append(moe_fusion)
if len(self.attention_layers) < self.model_config.num_hidden_layers:
# TODO(yuwei): support Non-Standard GQA

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

@@ -233,6 +233,7 @@ class PiecewiseCudaGraphRunner:
self.attention_layers = self.model_runner.attention_layers
self.moe_layers = self.model_runner.moe_layers
self.moe_fusions = self.model_runner.moe_fusions
if get_global_graph_memory_pool() is None:
set_global_graph_memory_pool(self.device_module.graph_pool_handle())
@@ -358,7 +359,11 @@ class PiecewiseCudaGraphRunner:
set_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len())
set_is_extend_in_batch(False)
with set_forward_context(
forward_batch, self.attention_layers, self.quant_config, self.moe_layers
forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
):
_ = self.model_runner.model.forward(
forward_batch.input_ids,
@@ -520,7 +525,11 @@ class PiecewiseCudaGraphRunner:
kwargs = {}
with set_forward_context(
forward_batch, self.attention_layers, self.quant_config, self.moe_layers
forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
):
self.model_runner.model.forward(
forward_batch.input_ids,
@@ -684,6 +693,7 @@ class PiecewiseCudaGraphRunner:
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
):
with set_compiled(True):
output = self.model_runner.model.forward(

25
python/sglang/srt/models/gpt_oss.py
View File

@@ -25,6 +25,10 @@ import torch
from torch import nn
from transformers import PretrainedConfig
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,
@@ -72,6 +76,7 @@ from sglang.srt.models.utils import (
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import LazyValue, add_prefix, is_cuda, is_npu, make_layers
from sglang.srt.utils.custom_op import register_custom_op
_is_cuda = is_cuda()
_is_npu = is_npu()
@@ -183,10 +188,12 @@ class GptOssSparseMoeBlock(nn.Module):
should_allreduce_fusion: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
router_logits, _ = self.router(hidden_states)
topk_output = self.topk(hidden_states, router_logits)
final_hidden_states = self.experts(hidden_states, topk_output)
if is_in_piecewise_cuda_graph():
final_hidden_states = moe_impl(self.layer_id, hidden_states)
else:
router_logits, _ = self.router(hidden_states)
topk_output = self.topk(hidden_states, router_logits)
final_hidden_states = self.experts(hidden_states, topk_output)
if self.tp_size > 1 and not should_allreduce_fusion:
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
@@ -195,6 +202,16 @@ class GptOssSparseMoeBlock(nn.Module):
return ans
@register_custom_op(out_shape="hidden_states")
def moe_impl(layer_id: int, hidden_states: torch.Tensor) -> torch.Tensor:
forward_context = get_forward_context()
moe_fusion = forward_context.moe_fusions[layer_id]
router_logits, _ = moe_fusion.router(hidden_states)
topk_output = moe_fusion.topk(hidden_states, router_logits)
final_hidden_states = moe_fusion.experts(hidden_states, topk_output)
return final_hidden_states
class GptOssAttention(nn.Module):
def __init__(
self,

7
python/sglang/srt/server_args.py
View File

@@ -1363,12 +1363,7 @@ class ServerArgs:
self.dtype = "bfloat16"
if self.moe_runner_backend == "auto":
if self.enable_piecewise_cuda_graph:
self.moe_runner_backend = "auto"
logger.warning(
"Enable piecewise CUDA graph, enabling auto MOE kernel."
)
elif is_blackwell_supported() and is_mxfp4_quant_format:
if is_blackwell_supported() and is_mxfp4_quant_format:
self.moe_runner_backend = "flashinfer_mxfp4"
logger.warning(
"Detected SM100 and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel."