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

Add piecewise cuda graph for Qwen3-Next FP8 flashinfer_trtllm moe backend (#18184)

Browse Source
This commit is contained in:
Xiaowei Wang
2026-03-15 04:03:31 +08:00
committed by GitHub
parent 3e643967e6
commit 574dbe23b2
3 changed files with 578 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

274
python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py Normal file
View File

@@ -0,0 +1,274 @@
from typing import Optional
import torch
from sglang.srt.utils.custom_op import register_custom_op
def _fake_fp8_block_scale_moe(
routing_logits: torch.Tensor,
routing_bias: Optional[torch.Tensor],
hidden_states: torch.Tensor,
hidden_states_scale: torch.Tensor,
gemm1_weights: torch.Tensor,
gemm1_weights_scale: torch.Tensor,
gemm2_weights: torch.Tensor,
gemm2_weights_scale: torch.Tensor,
num_experts: int,
top_k: int,
n_group: Optional[int],
topk_group: Optional[int],
intermediate_size: int,
local_expert_offset: int,
local_num_experts: int,
routed_scaling_factor: Optional[float],
routing_method_type: int = 0,
use_shuffled_weight: bool = False,
weight_layout: int = 0,
enable_pdl: Optional[bool] = None,
tune_max_num_tokens: int = 8192,
fp8_quantization_type: Optional[int] = None,
) -> torch.Tensor:
return torch.empty(
hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
)
@register_custom_op(fake_impl=_fake_fp8_block_scale_moe)
def trtllm_fp8_block_scale_moe_wrapper(
routing_logits: torch.Tensor,
routing_bias: Optional[torch.Tensor],
hidden_states: torch.Tensor,
hidden_states_scale: torch.Tensor,
gemm1_weights: torch.Tensor,
gemm1_weights_scale: torch.Tensor,
gemm2_weights: torch.Tensor,
gemm2_weights_scale: torch.Tensor,
num_experts: int,
top_k: int,
n_group: Optional[int],
topk_group: Optional[int],
intermediate_size: int,
local_expert_offset: int,
local_num_experts: int,
routed_scaling_factor: Optional[float],
routing_method_type: int = 0,
use_shuffled_weight: bool = False,
weight_layout: int = 0,
enable_pdl: Optional[bool] = None,
tune_max_num_tokens: int = 8192,
fp8_quantization_type: Optional[int] = None,
) -> torch.Tensor:
try:
from flashinfer.fused_moe import trtllm_fp8_block_scale_moe
except ImportError as e:
raise ImportError(
"Can't import trtllm_fp8_block_scale_moe from flashinfer. "
"Please check flashinfer version."
) from e
kwargs = {
"routing_logits": routing_logits,
"routing_bias": routing_bias,
"hidden_states": hidden_states,
"hidden_states_scale": hidden_states_scale,
"gemm1_weights": gemm1_weights,
"gemm1_weights_scale": gemm1_weights_scale,
"gemm2_weights": gemm2_weights,
"gemm2_weights_scale": gemm2_weights_scale,
"num_experts": num_experts,
"top_k": top_k,
"n_group": n_group,
"topk_group": topk_group,
"intermediate_size": intermediate_size,
"local_expert_offset": local_expert_offset,
"local_num_experts": local_num_experts,
"routed_scaling_factor": routed_scaling_factor,
"routing_method_type": routing_method_type,
"use_shuffled_weight": use_shuffled_weight,
"weight_layout": weight_layout,
"enable_pdl": enable_pdl,
"tune_max_num_tokens": tune_max_num_tokens,
}
if fp8_quantization_type is not None:
from flashinfer.fused_moe import Fp8QuantizationType
kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
return trtllm_fp8_block_scale_moe(**kwargs)
def _fake_fp8_block_scale_routed_moe(
topk_ids: torch.Tensor,
routing_bias: Optional[torch.Tensor],
hidden_states: torch.Tensor,
hidden_states_scale: torch.Tensor,
gemm1_weights: torch.Tensor,
gemm1_weights_scale: torch.Tensor,
gemm2_weights: torch.Tensor,
gemm2_weights_scale: torch.Tensor,
num_experts: int,
top_k: int,
n_group: Optional[int],
topk_group: Optional[int],
intermediate_size: int,
local_expert_offset: int,
local_num_experts: int,
routed_scaling_factor: Optional[float],
routing_method_type: int = 0,
use_shuffled_weight: bool = False,
weight_layout: int = 0,
enable_pdl: Optional[bool] = None,
tune_max_num_tokens: int = 8192,
fp8_quantization_type: Optional[int] = None,
) -> torch.Tensor:
return torch.empty(
hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
)
@register_custom_op(fake_impl=_fake_fp8_block_scale_routed_moe)
def trtllm_fp8_block_scale_routed_moe_wrapper(
topk_ids: torch.Tensor,
routing_bias: Optional[torch.Tensor],
hidden_states: torch.Tensor,
hidden_states_scale: torch.Tensor,
gemm1_weights: torch.Tensor,
gemm1_weights_scale: torch.Tensor,
gemm2_weights: torch.Tensor,
gemm2_weights_scale: torch.Tensor,
num_experts: int,
top_k: int,
n_group: Optional[int],
topk_group: Optional[int],
intermediate_size: int,
local_expert_offset: int,
local_num_experts: int,
routed_scaling_factor: Optional[float],
routing_method_type: int = 0,
use_shuffled_weight: bool = False,
weight_layout: int = 0,
enable_pdl: Optional[bool] = None,
tune_max_num_tokens: int = 8192,
fp8_quantization_type: Optional[int] = None,
) -> torch.Tensor:
try:
from flashinfer.fused_moe import trtllm_fp8_block_scale_routed_moe
except ImportError as e:
raise ImportError(
"Can't import trtllm_fp8_block_scale_routed_moe from flashinfer. "
"Please check flashinfer version."
) from e
kwargs = {
"topk_ids": topk_ids,
"routing_bias": routing_bias,
"hidden_states": hidden_states,
"hidden_states_scale": hidden_states_scale,
"gemm1_weights": gemm1_weights,
"gemm1_weights_scale": gemm1_weights_scale,
"gemm2_weights": gemm2_weights,
"gemm2_weights_scale": gemm2_weights_scale,
"num_experts": num_experts,
"top_k": top_k,
"n_group": n_group,
"topk_group": topk_group,
"intermediate_size": intermediate_size,
"local_expert_offset": local_expert_offset,
"local_num_experts": local_num_experts,
"routed_scaling_factor": routed_scaling_factor,
"routing_method_type": routing_method_type,
"use_shuffled_weight": use_shuffled_weight,
"weight_layout": weight_layout,
"enable_pdl": enable_pdl,
"tune_max_num_tokens": tune_max_num_tokens,
}
if fp8_quantization_type is not None:
from flashinfer.fused_moe import Fp8QuantizationType
kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
return trtllm_fp8_block_scale_routed_moe(**kwargs)
def _fake_fp8_per_tensor_scale_moe(
routing_logits: torch.Tensor,
routing_bias: Optional[torch.Tensor],
hidden_states: torch.Tensor,
gemm1_weights: torch.Tensor,
output1_scales_scalar: torch.Tensor,
output1_scales_gate_scalar: torch.Tensor,
gemm2_weights: torch.Tensor,
output2_scales_scalar: torch.Tensor,
num_experts: int,
top_k: int,
n_group: Optional[int],
topk_group: Optional[int],
intermediate_size: int,
local_expert_offset: int,
local_num_experts: int,
routed_scaling_factor: Optional[float],
use_routing_scales_on_input: bool,
routing_method_type: int = 0,
enable_pdl: Optional[bool] = None,
tune_max_num_tokens: int = 8192,
) -> torch.Tensor:
return torch.empty(
hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
)
@register_custom_op(fake_impl=_fake_fp8_per_tensor_scale_moe)
def trtllm_fp8_per_tensor_scale_moe_wrapper(
routing_logits: torch.Tensor,
routing_bias: Optional[torch.Tensor],
hidden_states: torch.Tensor,
gemm1_weights: torch.Tensor,
output1_scales_scalar: torch.Tensor,
output1_scales_gate_scalar: torch.Tensor,
gemm2_weights: torch.Tensor,
output2_scales_scalar: torch.Tensor,
num_experts: int,
top_k: int,
n_group: Optional[int],
topk_group: Optional[int],
intermediate_size: int,
local_expert_offset: int,
local_num_experts: int,
routed_scaling_factor: Optional[float],
use_routing_scales_on_input: bool,
routing_method_type: int = 0,
enable_pdl: Optional[bool] = None,
tune_max_num_tokens: int = 8192,
) -> torch.Tensor:
# lazy import
try:
from flashinfer.fused_moe import trtllm_fp8_per_tensor_scale_moe
except ImportError as e:
raise ImportError(
"Can't import trtllm_fp8_per_tensor_scale_moe from flashinfer. "
"Please check flashinfer version."
) from e
kwargs = {
"routing_logits": routing_logits,
"routing_bias": routing_bias,
"hidden_states": hidden_states,
"gemm1_weights": gemm1_weights,
"output1_scales_scalar": output1_scales_scalar,
"output1_scales_gate_scalar": output1_scales_gate_scalar,
"gemm2_weights": gemm2_weights,
"output2_scales_scalar": output2_scales_scalar,
"num_experts": num_experts,
"top_k": top_k,
"n_group": n_group,
"topk_group": topk_group,
"intermediate_size": intermediate_size,
"local_expert_offset": local_expert_offset,
"local_num_experts": local_num_experts,
"routed_scaling_factor": routed_scaling_factor,
"use_routing_scales_on_input": use_routing_scales_on_input,
"routing_method_type": routing_method_type,
"enable_pdl": enable_pdl,
"tune_max_num_tokens": tune_max_num_tokens,
}
return trtllm_fp8_per_tensor_scale_moe(**kwargs)

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

@@ -40,6 +40,7 @@ from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher
from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardDispatcher,
StandardDispatchOutput,
)
from sglang.srt.layers.moe.topk import (
BypassedTopKOutput,
@@ -967,10 +968,7 @@ class FusedMoE(torch.nn.Module):
def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
if is_in_piecewise_cuda_graph():
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:
if TopKOutputChecker.format_is_standard(topk_output):
return moe_forward_piecewise_cuda_graph_impl(
hidden_states,
topk_output.topk_weights,
@@ -978,6 +976,20 @@ class FusedMoE(torch.nn.Module):
topk_output.router_logits,
self.layer_id,
)
elif TopKOutputChecker.format_is_bypassed(topk_output):
return fused_moe_bypassed_piecewise_cuda_graph_impl(
hidden_states,
topk_output.router_logits,
topk_output.topk_config.top_k,
topk_output.topk_config.topk_group,
topk_output.topk_config.num_expert_group,
topk_output.topk_config.correction_bias,
topk_output.topk_config.renormalize,
self.layer_id,
)
else:
# Make sure there is torch lib op registration for the whole moe layer
return self.forward_impl(hidden_states, topk_output)
else:
return self.forward_impl(hidden_states, topk_output)
@@ -1134,6 +1146,116 @@ class FusedMoE(torch.nn.Module):
self.meta_overlap_args = None
class FlashInferFusedMoE(FusedMoE):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
assert TopKOutputChecker.format_is_bypassed(
topk_output
), "Only bypassed topk output is supported for flashinfer trtllm moe"
if is_in_piecewise_cuda_graph():
return flashinfer_bf16_moe_forward_piecewise_cuda_graph_impl(
hidden_states,
topk_output.router_logits,
topk_output.topk_config.top_k,
topk_output.topk_config.topk_group,
topk_output.topk_config.num_expert_group,
topk_output.topk_config.correction_bias,
topk_output.topk_config.renormalize,
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"
assert self.quant_method is not None
assert (
topk_output.topk_config.renormalize
), "Renormalize is required for flashinfer trtllm moe"
assert (
self.num_fused_shared_experts == 0
), "Fused shared experts are not supported for flashinfer trtllm moe"
assert (
self.moe_runner_config.is_gated
), "Only gated MoEs are supported for flashinfer trtllm moe"
router_logits = topk_output.router_logits
topk_config = topk_output.topk_config
correction_bias = topk_config.correction_bias
routed_scaling_factor = self.moe_runner_config.routed_scaling_factor
if isinstance(self.quant_method, UnquantizedFusedMoEMethod):
# lazy import
try:
from flashinfer.fused_moe import trtllm_bf16_moe
except ImportError as e:
raise ImportError(
"Can't import trtllm_bf16_moe from flashinfer. "
"Please check flashinfer version to use bf16 with flashinfer_trtllm backend."
) from e
# Allocate output inside symmetric memory context
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
# TODO: Now trtllm_bf16_moe doesn't support inplace output,
# we can move this out when it support that.
symm_output = torch.empty(
hidden_states.shape[0],
hidden_states.shape[1],
dtype=hidden_states.dtype,
device=hidden_states.device,
)
# Move kernel call outside context manager to avoid graph breaks
# during torch.compile for piecewise cuda graph
moe_result = trtllm_bf16_moe(
routing_logits=router_logits,
routing_bias=correction_bias,
hidden_states=hidden_states,
gemm1_weights=self.w13_weight,
gemm2_weights=self.w2_weight,
num_experts=self.num_experts,
top_k=topk_config.top_k,
n_group=topk_config.num_expert_group,
topk_group=topk_config.topk_group,
intermediate_size=self.intermediate_size_per_partition,
local_expert_offset=self.moe_ep_rank * self.num_local_experts,
local_num_experts=self.num_local_experts,
routing_method_type=self.routing_method_type,
tune_max_num_tokens=next_power_of_2(hidden_states.shape[0]),
)
# Copy result to symmetric memory output
symm_output.copy_(moe_result)
final_hidden_states = symm_output
else:
final_hidden_states = self.quant_method.apply(
layer=self,
dispatch_output=StandardDispatchOutput(
hidden_states=hidden_states,
hidden_states_scale=None,
topk_output=topk_output,
),
).hidden_states
# NOTE for symmetric memory tagging:
# We do not create the context in this function.
# Instead, we create the context and tagging inside each FusedMoEMethodBase
# This can allow fine-grained tagging.
if self.reduce_results and (self.moe_tp_size > 1 or self.moe_ep_size > 1):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states
class FlashInferFP4MoE(FusedMoE):
"""FP4 TRTLLM MoE implementation using FlashInfer."""
@@ -1300,6 +1422,60 @@ def moe_forward_piecewise_cuda_graph_impl(
return moe_layer.forward_impl(hidden_states, topk_output)
@register_custom_op(out_shape="hidden_states")
def fused_moe_bypassed_piecewise_cuda_graph_impl(
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
top_k: int,
topk_group: Optional[int],
num_expert_group: Optional[int],
correction_bias: Optional[torch.Tensor],
renormalize: bool,
layer_id: int,
) -> torch.Tensor:
topk_output = BypassedTopKOutput(
hidden_states=hidden_states,
router_logits=router_logits,
topk_config=TopKConfig(
top_k=top_k,
topk_group=topk_group,
num_expert_group=num_expert_group,
correction_bias=correction_bias,
renormalize=renormalize,
),
)
forward_context = get_forward_context()
moe_layer = forward_context.moe_layers[layer_id]
return moe_layer.forward_impl(hidden_states, topk_output)
@register_custom_op(out_shape="hidden_states")
def flashinfer_bf16_moe_forward_piecewise_cuda_graph_impl(
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
top_k: int,
topk_group: Optional[int],
num_expert_group: Optional[int],
correction_bias: Optional[torch.Tensor],
renormalize: bool,
layer_id: int,
) -> torch.Tensor:
topk_output = BypassedTopKOutput(
hidden_states=hidden_states,
router_logits=router_logits,
topk_config=TopKConfig(
top_k=top_k,
topk_group=topk_group,
num_expert_group=num_expert_group,
correction_bias=correction_bias,
renormalize=renormalize,
),
)
forward_context = get_forward_context()
moe_layer = forward_context.moe_layers[layer_id]
return moe_layer.forward_impl(hidden_states, topk_output)
@register_custom_op(out_shape="hidden_states")
def flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl(
hidden_states: torch.Tensor,

237
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
View File

@@ -7,6 +7,8 @@ import torch
from torch.nn import Module
from torch.nn.parameter import Parameter
# Import to register custom ops for torch.compile compatibility
import sglang.srt.layers.moe.flashinfer_trtllm_moe # noqa: F401
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
@@ -298,12 +300,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
runner_config: MoeRunnerConfig,
use_routed_topk: bool = False,
) -> StandardCombineInput:
from flashinfer.fused_moe import (
Fp8QuantizationType,
trtllm_fp8_block_scale_moe,
trtllm_fp8_block_scale_routed_moe,
trtllm_fp8_per_tensor_scale_moe,
)
from flashinfer.fused_moe import Fp8QuantizationType
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker
@@ -354,92 +351,95 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
)
a_sf_t = a_sf.t().contiguous()
# Allocate output inside symmetric memory context
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
if use_routed_topk:
assert (
runner_config.top_k is not None
), "runner_config.top_k is required for flashinfer_trtllm_routed."
assert TopKOutputChecker.format_is_standard(topk_output)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
topk_ids=topk_output.topk_ids,
topk_weights=topk_output.topk_weights,
)
symm_output = torch.empty(
hidden_states.shape[0],
hidden_states.shape[1],
dtype=torch.bfloat16,
device=hidden_states.device,
)
output = trtllm_fp8_block_scale_routed_moe(
topk_ids=packed_topk_ids,
routing_bias=None,
hidden_states=a_q,
hidden_states_scale=a_sf_t,
gemm1_weights=quant_info.w13_weight,
gemm1_weights_scale=quant_info.w13_weight_scale_inv,
gemm2_weights=quant_info.w2_weight,
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
num_experts=quant_info.global_num_experts,
top_k=runner_config.top_k,
n_group=None,
topk_group=None,
intermediate_size=quant_info.intermediate_size,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
routed_scaling_factor=(
runner_config.routed_scaling_factor
if runner_config.routed_scaling_factor is not None
else 1.0
),
routing_method_type=(
RoutingMethodType.TopK
if routing_method_type == RoutingMethodType.DeepSeekV3
else routing_method_type
),
use_shuffled_weight=use_shuffled_weight,
weight_layout=0,
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
fp8_quantization_type=fp8_quantization_type,
)
else:
assert TopKOutputChecker.format_is_bypassed(topk_output)
# Move kernel call outside context manager to avoid graph breaks
# during torch.compile for piecewise cuda graph.
# Use custom op wrapper for torch.compile compatibility.
if use_routed_topk:
assert (
runner_config.top_k is not None
), "runner_config.top_k is required for flashinfer_trtllm_routed."
assert TopKOutputChecker.format_is_standard(topk_output)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
topk_ids=topk_output.topk_ids,
topk_weights=topk_output.topk_weights,
)
# FIXME: there is a bug in the trtllm_fp8_block_scale_moe.
# It ignored the `output` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325
# so we put the whole function under the ``use_symmetric_memory`` context manager.
# If the bug is fixed, we can only put the output tensor allocation under the context manager.
output = trtllm_fp8_block_scale_moe(
routing_logits=(
router_logits.to(torch.float32)
if routing_method_type == RoutingMethodType.DeepSeekV3
else router_logits
),
routing_bias=correction_bias,
hidden_states=a_q,
hidden_states_scale=a_sf_t,
gemm1_weights=quant_info.w13_weight,
gemm1_weights_scale=quant_info.w13_weight_scale_inv,
gemm2_weights=quant_info.w2_weight,
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
num_experts=quant_info.global_num_experts,
top_k=topk_config.top_k,
n_group=(
topk_config.num_expert_group
if topk_config.num_expert_group
else 0
),
topk_group=topk_config.topk_group if topk_config.topk_group else 0,
intermediate_size=quant_info.intermediate_size,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
routed_scaling_factor=(
runner_config.routed_scaling_factor
if runner_config.routed_scaling_factor is not None
else 1.0
),
routing_method_type=routing_method_type,
use_shuffled_weight=use_shuffled_weight,
weight_layout=0,
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
fp8_quantization_type=fp8_quantization_type,
)
output = torch.ops.sglang.trtllm_fp8_block_scale_routed_moe_wrapper(
topk_ids=packed_topk_ids,
routing_bias=None,
hidden_states=a_q,
hidden_states_scale=a_sf_t,
gemm1_weights=quant_info.w13_weight,
gemm1_weights_scale=quant_info.w13_weight_scale_inv,
gemm2_weights=quant_info.w2_weight,
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
num_experts=quant_info.global_num_experts,
top_k=runner_config.top_k,
n_group=None,
topk_group=None,
intermediate_size=quant_info.intermediate_size,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
routed_scaling_factor=(
runner_config.routed_scaling_factor
if runner_config.routed_scaling_factor is not None
else 1.0
),
routing_method_type=(
RoutingMethodType.TopK
if routing_method_type == RoutingMethodType.DeepSeekV3
else routing_method_type
),
use_shuffled_weight=use_shuffled_weight,
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
fp8_quantization_type=int(fp8_quantization_type),
)
else:
assert TopKOutputChecker.format_is_bypassed(topk_output)
output = torch.ops.sglang.trtllm_fp8_block_scale_moe_wrapper(
routing_logits=(
router_logits.to(torch.float32)
if routing_method_type == RoutingMethodType.DeepSeekV3
else router_logits
),
routing_bias=correction_bias,
hidden_states=a_q,
hidden_states_scale=a_sf_t,
gemm1_weights=quant_info.w13_weight,
gemm1_weights_scale=quant_info.w13_weight_scale_inv,
gemm2_weights=quant_info.w2_weight,
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
num_experts=quant_info.global_num_experts,
top_k=topk_config.top_k,
n_group=topk_config.num_expert_group,
topk_group=topk_config.topk_group,
intermediate_size=quant_info.intermediate_size,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
routed_scaling_factor=(
runner_config.routed_scaling_factor
if runner_config.routed_scaling_factor is not None
else 1.0
),
routing_method_type=routing_method_type,
use_shuffled_weight=use_shuffled_weight,
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
fp8_quantization_type=int(fp8_quantization_type),
)
symm_output.copy_(output)
output = symm_output
else:
assert quant_info.w13_input_scale is not None
assert quant_info.output1_scales_scalar is not None
@@ -451,37 +451,48 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
None if correction_bias is None else correction_bias.to(torch.bfloat16)
)
# Allocate output inside symmetric memory context
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
output = trtllm_fp8_per_tensor_scale_moe(
routing_logits=router_logits.to(torch.bfloat16),
routing_bias=routing_bias_cast,
hidden_states=a_q,
gemm1_weights=quant_info.w13_weight,
output1_scales_scalar=quant_info.output1_scales_scalar,
output1_scales_gate_scalar=quant_info.output1_scales_gate_scalar,
gemm2_weights=quant_info.w2_weight,
output2_scales_scalar=quant_info.output2_scales_scalar,
num_experts=quant_info.global_num_experts,
top_k=topk_config.top_k,
n_group=(
topk_config.num_expert_group if topk_config.num_expert_group else 0
),
topk_group=topk_config.topk_group if topk_config.topk_group else 0,
intermediate_size=quant_info.intermediate_size,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
routed_scaling_factor=(
runner_config.routed_scaling_factor
if runner_config.routed_scaling_factor is not None
else 1.0
),
use_routing_scales_on_input=quant_info.use_routing_scales_on_input,
routing_method_type=routing_method_type,
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
symm_output = torch.empty(
hidden_states.shape[0],
hidden_states.shape[1],
dtype=torch.bfloat16,
device=hidden_states.device,
)
# Move kernel call outside context manager to avoid graph breaks
# during torch.compile for piecewise cuda graph.
# Use custom op wrapper for torch.compile compatibility.
output = torch.ops.sglang.trtllm_fp8_per_tensor_scale_moe(
routing_logits=router_logits.to(torch.bfloat16),
routing_bias=routing_bias_cast,
hidden_states=a_q,
gemm1_weights=quant_info.w13_weight,
output1_scales_scalar=quant_info.output1_scales_scalar,
output1_scales_gate_scalar=quant_info.output1_scales_gate_scalar,
gemm2_weights=quant_info.w2_weight,
output2_scales_scalar=quant_info.output2_scales_scalar,
num_experts=quant_info.global_num_experts,
top_k=topk_config.top_k,
n_group=topk_config.num_expert_group,
topk_group=topk_config.topk_group,
intermediate_size=int(quant_info.w2_weight.shape[2]),
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
routed_scaling_factor=(
runner_config.routed_scaling_factor
if runner_config.routed_scaling_factor is not None
else 1.0
),
use_routing_scales_on_input=False,
routing_method_type=routing_method_type,
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
)
symm_output.copy_(output)
output = symm_output
return StandardCombineInput(hidden_states=output)