[MoE Refactor] Migrate SM100 trtllm-gen mxfp4 MoE onto MoeRunner (#32405)
This commit is contained in:
Beihao Zhou
@@ -1,9 +1,9 @@
|
||||
"""FlashInfer CUTLASS MoE fused funcs.
|
||||
|
||||
This module owns the FlashInfer ``cutlass_fused_moe`` calls used by the
|
||||
unquantized, ModelOpt FP8, ModelOpt NVFP4, and MXFP4 MoE paths.
|
||||
Quantization methods prepare a small quant_info payload and route through
|
||||
``MoeRunner``.
|
||||
unquantized, ModelOpt FP8, ModelOpt NVFP4, and CUTLASS MXFP4 MoE paths, plus
|
||||
the shared ``flashinfer_mxfp4`` dispatcher. Quantization methods prepare a
|
||||
small quant_info payload and route through ``MoeRunner``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -312,14 +312,40 @@ def fused_experts_none_to_flashinfer_mxfp4(
|
||||
quant_info: MoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
) -> StandardCombineInput:
|
||||
"""Run the FlashInfer CUTLASS MXFP4 fused experts."""
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
"""Dispatch flashinfer_mxfp4 by quant-info type.
|
||||
|
||||
assert isinstance(quant_info, FlashInferCutlassMxfp4MoeQuantInfo), (
|
||||
Both mxfp4 paths register under this single ``("none", "flashinfer_mxfp4")``
|
||||
key but call different kernels.
|
||||
"""
|
||||
if isinstance(quant_info, FlashInferCutlassMxfp4MoeQuantInfo):
|
||||
return _fused_experts_flashinfer_mxfp4_cutlass(
|
||||
dispatch_output, quant_info, runner_config
|
||||
)
|
||||
|
||||
# Keep one fused-op registration for the shared backend while loading the
|
||||
# TRT-LLM implementation only when its quant-info type is dispatched.
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
FlashInferTrtllmGenMxfp4MoeQuantInfo,
|
||||
_fused_experts_flashinfer_mxfp4_sm100_trtllm_gen,
|
||||
)
|
||||
|
||||
if isinstance(quant_info, FlashInferTrtllmGenMxfp4MoeQuantInfo):
|
||||
return _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen(
|
||||
dispatch_output, quant_info, runner_config
|
||||
)
|
||||
raise TypeError(
|
||||
f"Unexpected quant_info type for flashinfer_mxfp4: {type(quant_info)}"
|
||||
)
|
||||
|
||||
|
||||
def _fused_experts_flashinfer_mxfp4_cutlass(
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
quant_info: FlashInferCutlassMxfp4MoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
) -> StandardCombineInput:
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
|
||||
flashinfer_cutlass_fused_moe, ActivationType = _flashinfer_cutlass_fused_moe()
|
||||
|
||||
x = dispatch_output.hidden_states
|
||||
|
||||
@@ -958,6 +958,242 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlashInferTrtllmGenMxfp4MoeQuantInfo(MoeQuantInfo):
|
||||
"""Payload for the SM100 (Blackwell) trtllm-gen MXFP4 MoE path."""
|
||||
|
||||
# Packed MXFP4 weights: uint8, e2m1 x 2.
|
||||
w13_weight: torch.Tensor
|
||||
w2_weight: torch.Tensor
|
||||
w13_weight_scale: torch.Tensor
|
||||
w2_weight_scale: torch.Tensor
|
||||
|
||||
# fp32 per expert. GPT-OSS sets these.
|
||||
w13_weight_bias: torch.Tensor
|
||||
w2_weight_bias: torch.Tensor
|
||||
gemm1_alpha: torch.Tensor
|
||||
gemm1_beta: torch.Tensor
|
||||
gemm1_clamp_limit: torch.Tensor
|
||||
|
||||
global_num_experts: int
|
||||
local_expert_offset: int
|
||||
local_num_experts: int
|
||||
intermediate_size_per_partition: int
|
||||
|
||||
hidden_size: int
|
||||
flashinfer_mxfp4_moe_precision: str
|
||||
routing_bias: Optional[torch.Tensor] = None
|
||||
|
||||
|
||||
def _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen(
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
quant_info: FlashInferTrtllmGenMxfp4MoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
) -> StandardCombineInput:
|
||||
"""SM100 (Blackwell) trtllm-gen MXFP4 fused experts."""
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
|
||||
x = dispatch_output.hidden_states
|
||||
topk_output = dispatch_output.topk_output
|
||||
|
||||
origin_hidden_states_dim = x.shape[-1]
|
||||
prepared_packed_topk = None
|
||||
if quant_info.flashinfer_mxfp4_moe_precision == "bf16":
|
||||
assert x.dtype == torch.bfloat16
|
||||
x_quant = x
|
||||
x_scale = None
|
||||
if quant_info.hidden_size != origin_hidden_states_dim:
|
||||
x_quant = torch.nn.functional.pad(
|
||||
x_quant,
|
||||
(0, quant_info.hidden_size - origin_hidden_states_dim),
|
||||
mode="constant",
|
||||
value=0.0,
|
||||
)
|
||||
elif quant_info.flashinfer_mxfp4_moe_precision == "default":
|
||||
# Deferred import: mxfp4 dispatches into this module, and the helper
|
||||
# stays there so its registered unit test keeps its import path.
|
||||
from sglang.srt.layers.quantization.mxfp4 import (
|
||||
_prepare_flashinfer_mxfp8_activations,
|
||||
)
|
||||
|
||||
x, prepared_packed_topk, x_quant, x_scale = (
|
||||
_prepare_flashinfer_mxfp8_activations(x, quant_info.hidden_size)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported flashinfer_mxfp4_moe_precision: "
|
||||
f"{quant_info.flashinfer_mxfp4_moe_precision}"
|
||||
)
|
||||
|
||||
assert x_quant.shape[-1] == quant_info.hidden_size
|
||||
is_standard = TopKOutputChecker.format_is_standard(topk_output)
|
||||
assert is_standard or TopKOutputChecker.format_is_bypassed(topk_output), (
|
||||
f"unsupported topk format: {topk_output.format}"
|
||||
)
|
||||
if is_standard:
|
||||
assert runner_config.activation == "situ", (
|
||||
"standard topk output only wired for the situ path"
|
||||
)
|
||||
top_k = topk_output.topk_ids.shape[1]
|
||||
router_logits = None
|
||||
else:
|
||||
top_k = topk_output.topk_config.top_k
|
||||
router_logits = topk_output.router_logits
|
||||
|
||||
num_tokens = x_quant.shape[0]
|
||||
from sglang.srt.layers import zero_copy_context
|
||||
|
||||
symm_output = zero_copy_context.get_moe_output_spec(
|
||||
torch.Size((num_tokens, origin_hidden_states_dim)),
|
||||
torch.bfloat16,
|
||||
x_quant.device,
|
||||
)
|
||||
if symm_output is None:
|
||||
with use_symmetric_memory(
|
||||
get_tp_group(), disabled=not is_allocation_symmetric()
|
||||
):
|
||||
symm_output = torch.empty(
|
||||
num_tokens,
|
||||
origin_hidden_states_dim,
|
||||
dtype=torch.bfloat16,
|
||||
device=x_quant.device,
|
||||
)
|
||||
|
||||
if runner_config.activation == "situ":
|
||||
from flashinfer import trtllm_fp4_block_scale_moe
|
||||
from flashinfer.fused_moe import trtllm_fp4_block_scale_routed_moe
|
||||
from flashinfer.tllm_enums import ActivationType, RoutingMethodType
|
||||
|
||||
if is_standard:
|
||||
if prepared_packed_topk is not None:
|
||||
packed_topk = prepared_packed_topk
|
||||
else:
|
||||
packed_topk = PackTopkIds.execute(
|
||||
topk_output.topk_ids, topk_output.topk_weights
|
||||
)
|
||||
|
||||
defer_finalize = _deferred_finalize_enabled.get()
|
||||
result = trtllm_fp4_block_scale_routed_moe(
|
||||
topk_ids=packed_topk,
|
||||
routing_bias=None,
|
||||
hidden_states=x_quant,
|
||||
hidden_states_scale=x_scale,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
gemm1_weights_scale=quant_info.w13_weight_scale,
|
||||
gemm1_bias=None,
|
||||
gemm1_alpha=quant_info.gemm1_alpha,
|
||||
gemm1_beta=quant_info.gemm1_clamp_limit,
|
||||
gemm1_clamp_limit=None,
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
gemm2_weights_scale=quant_info.w2_weight_scale,
|
||||
gemm2_bias=None,
|
||||
output1_scale_scalar=None,
|
||||
output1_scale_gate_scalar=None,
|
||||
output2_scale_scalar=None,
|
||||
num_experts=quant_info.global_num_experts,
|
||||
top_k=packed_topk.shape[1],
|
||||
n_group=None,
|
||||
topk_group=None,
|
||||
intermediate_size=quant_info.intermediate_size_per_partition,
|
||||
local_expert_offset=quant_info.local_expert_offset,
|
||||
local_num_experts=quant_info.local_num_experts,
|
||||
routed_scaling_factor=None,
|
||||
routing_method_type=RoutingMethodType.TopK.value,
|
||||
activation_type=ActivationType.Situ.value,
|
||||
tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
|
||||
output=symm_output,
|
||||
do_finalize=not defer_finalize,
|
||||
enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]),
|
||||
)
|
||||
if defer_finalize:
|
||||
gemm2_out, topk_weights, expanded_idx = result
|
||||
result = FlashInferTrtllmDeferredFinalizeOutput(
|
||||
gemm2_out=gemm2_out,
|
||||
expert_weights=topk_weights,
|
||||
expanded_idx_to_permuted_idx=expanded_idx,
|
||||
top_k=packed_topk.shape[1],
|
||||
)
|
||||
return StandardCombineInput(hidden_states=result)
|
||||
# The finalized kernel writes to its explicit output argument. Do
|
||||
# not propagate the FFI return tensor: some SiTU runner versions
|
||||
# return a distinct wrapper/allocation even though symm_output
|
||||
# contains the published result. Returning the destination makes
|
||||
# the pointer contract explicit for K3's zero-copy latent buffer.
|
||||
return StandardCombineInput(hidden_states=symm_output)
|
||||
|
||||
trtllm_fp4_block_scale_moe(
|
||||
routing_logits=router_logits.to(torch.bfloat16).contiguous(),
|
||||
routing_bias=quant_info.routing_bias,
|
||||
hidden_states=x_quant,
|
||||
hidden_states_scale=x_scale,
|
||||
gemm1_weights=quant_info.w13_weight,
|
||||
gemm1_weights_scale=quant_info.w13_weight_scale,
|
||||
gemm1_bias=None,
|
||||
gemm1_alpha=quant_info.gemm1_alpha,
|
||||
gemm1_beta=quant_info.gemm1_clamp_limit,
|
||||
gemm1_clamp_limit=None,
|
||||
gemm2_weights=quant_info.w2_weight,
|
||||
gemm2_weights_scale=quant_info.w2_weight_scale,
|
||||
gemm2_bias=None,
|
||||
output1_scale_scalar=None,
|
||||
output1_scale_gate_scalar=None,
|
||||
output2_scale_scalar=None,
|
||||
num_experts=quant_info.global_num_experts,
|
||||
top_k=top_k,
|
||||
n_group=topk_output.topk_config.num_expert_group,
|
||||
topk_group=topk_output.topk_config.topk_group,
|
||||
intermediate_size=quant_info.intermediate_size_per_partition,
|
||||
routed_scaling_factor=(
|
||||
topk_output.topk_config.routed_scaling_factor or 1.0
|
||||
),
|
||||
routing_method_type=RoutingMethodType.DeepSeekV3.value,
|
||||
activation_type=ActivationType.Situ.value,
|
||||
norm_topk_prob=topk_output.topk_config.renormalize,
|
||||
local_expert_offset=quant_info.local_expert_offset,
|
||||
local_num_experts=quant_info.local_num_experts,
|
||||
tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
|
||||
output=symm_output,
|
||||
enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]),
|
||||
)
|
||||
return StandardCombineInput(hidden_states=symm_output)
|
||||
|
||||
from flashinfer import trtllm_fp4_block_scale_moe
|
||||
|
||||
trtllm_gen_output = trtllm_fp4_block_scale_moe(
|
||||
router_logits.to(torch.bfloat16),
|
||||
None, # routing_bias
|
||||
x_quant,
|
||||
x_scale,
|
||||
quant_info.w13_weight, # uint8 (e2m1 x 2)
|
||||
quant_info.w13_weight_scale, # uint8 (e4m3 x 2)
|
||||
quant_info.w13_weight_bias, # fp32 per expert per channel
|
||||
quant_info.gemm1_alpha, # fp32 per expert
|
||||
quant_info.gemm1_beta, # fp32 per expert
|
||||
quant_info.gemm1_clamp_limit, # fp32 per expert
|
||||
quant_info.w2_weight, # uint8 (e2m1 x 2)
|
||||
quant_info.w2_weight_scale, # ue8m0
|
||||
quant_info.w2_weight_bias, # fp32 per expert per channel
|
||||
None, # output1_scale_scalar
|
||||
None, # output1_scale_gate_scalar
|
||||
None, # output2_scale_scalar
|
||||
quant_info.global_num_experts,
|
||||
top_k,
|
||||
None, # n_group # TODO: support n_group
|
||||
None, # topk_group # TODO: support topk_group
|
||||
quant_info.intermediate_size_per_partition, # padded to multiple of 128
|
||||
quant_info.local_expert_offset,
|
||||
quant_info.local_num_experts,
|
||||
None, # routed_scaling_factor
|
||||
1, # routing_method_type, renormalize
|
||||
True, # do finalize
|
||||
tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
|
||||
output=symm_output,
|
||||
enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]),
|
||||
)[0]
|
||||
return StandardCombineInput(hidden_states=trtllm_gen_output)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlashInferTrtllmFp4MoeQuantInfo(MoeQuantInfo):
|
||||
"""Quantization payload consumed by FlashInfer TRT-LLM FP4 MoE kernels."""
|
||||
|
||||
@@ -29,17 +29,11 @@ from torch.nn.parameter import Parameter
|
||||
# cutlass_fused_moe. Its C++ logger reads TLLM_LOG_LEVEL on first kernel launch;
|
||||
# setdefault preserves any explicit user override.
|
||||
os.environ.setdefault("TLLM_LOG_LEVEL", "INFO")
|
||||
from sglang.srt.distributed import get_tp_group
|
||||
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
|
||||
use_symmetric_memory,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers import zero_copy_context
|
||||
from sglang.srt.layers.amx_utils import (
|
||||
CPUQuantMethod,
|
||||
_amx_process_weight_after_loading,
|
||||
)
|
||||
from sglang.srt.layers.dp_attention import is_allocation_symmetric
|
||||
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
|
||||
from sglang.srt.layers.moe.utils import get_moe_a2a_backend, get_moe_runner_backend
|
||||
@@ -62,7 +56,6 @@ from sglang.srt.utils import (
|
||||
is_hip,
|
||||
is_triton_kernels_available,
|
||||
is_xpu,
|
||||
next_power_of_2,
|
||||
round_up,
|
||||
set_weight_attrs,
|
||||
use_intel_amx_backend,
|
||||
@@ -127,13 +120,10 @@ def _prepare_flashinfer_mxfp8_activations(
|
||||
if is_flashinfer_available():
|
||||
from flashinfer import (
|
||||
nvfp4_block_scale_interleave,
|
||||
trtllm_fp4_block_scale_moe,
|
||||
)
|
||||
from flashinfer.fused_moe import (
|
||||
trtllm_fp4_block_scale_routed_moe,
|
||||
from flashinfer.fused_moe.core import (
|
||||
get_w2_permute_indices_with_cache,
|
||||
)
|
||||
from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache
|
||||
from flashinfer.tllm_enums import ActivationType, RoutingMethodType
|
||||
|
||||
# SM90 mixed-input helpers landed in FlashInfer #3084 (post-0.6.10). Older
|
||||
# versions don't ship them; gate at import so unrelated code paths still load.
|
||||
@@ -729,6 +719,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
torch.tensor([_limit] * E, dtype=torch.float32).cuda(),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer._situ_routing_bias_bf16 = None
|
||||
sf_block_size = 32 # mxfp4 block size
|
||||
|
||||
assert (
|
||||
@@ -1415,6 +1406,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
):
|
||||
self.moe_runner_config = moe_runner_config
|
||||
moe_runner_backend = get_moe_runner_backend()
|
||||
if self.use_mega_moe and moe_runner_backend.is_flashinfer_mxfp4():
|
||||
# MegaMoE uses this method for weight preparation only and calls
|
||||
# DeepGEMM directly instead of dispatching through a MoeRunner.
|
||||
return
|
||||
|
||||
if moe_runner_backend.is_auto():
|
||||
# Must match apply() priority: _use_aiter before use_triton_kernels.
|
||||
if _use_aiter and get_moe_a2a_backend().supports_aiter():
|
||||
@@ -1442,16 +1438,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
elif moe_runner_backend.is_flashinfer_mxfp4() and self._fi_kernel in (
|
||||
"cutlass_sm90",
|
||||
"cutlass_sm120",
|
||||
"trtllm_sm100",
|
||||
):
|
||||
# Register the fused func at runner construction so the FusedOpPool
|
||||
# lookup at `MoeRunner.__init__` finds it.
|
||||
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401
|
||||
|
||||
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
|
||||
else:
|
||||
# Legacy bypass path (e.g. SM100 trtllm-gen under flashinfer_mxfp4)
|
||||
# routes through `apply` without a MoeRunner. TODO(cwan): migrate.
|
||||
pass
|
||||
raise NotImplementedError(
|
||||
f"Mxfp4MoEMethod has no MoeRunner for backend={moe_runner_backend} "
|
||||
f"/ _fi_kernel={self._fi_kernel}."
|
||||
)
|
||||
|
||||
def _apply_sm90_cutlass(self, layer, dispatch_output):
|
||||
"""SM90 MXFP4 x BF16/FP8 MoE via FlashInfer's mixed-input kernels.
|
||||
@@ -1534,6 +1528,44 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
)
|
||||
return self.runner.run(dispatch_output, quant_info)
|
||||
|
||||
def _apply_sm100_trtllm_gen(self, layer, dispatch_output):
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
FlashInferTrtllmGenMxfp4MoeQuantInfo,
|
||||
)
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
|
||||
routing_bias = layer._situ_routing_bias_bf16
|
||||
topk_output = dispatch_output.topk_output
|
||||
if (
|
||||
self.moe_runner_config.activation == "situ"
|
||||
and routing_bias is None
|
||||
and TopKOutputChecker.format_is_bypassed(topk_output)
|
||||
):
|
||||
correction_bias = topk_output.topk_config.correction_bias
|
||||
if correction_bias is not None:
|
||||
routing_bias = correction_bias.to(torch.bfloat16)
|
||||
layer._situ_routing_bias_bf16 = routing_bias
|
||||
|
||||
quant_info = FlashInferTrtllmGenMxfp4MoeQuantInfo(
|
||||
w13_weight=layer.w13_weight,
|
||||
w2_weight=layer.w2_weight,
|
||||
w13_weight_scale=layer.w13_weight_scale,
|
||||
w2_weight_scale=layer.w2_weight_scale,
|
||||
w13_weight_bias=layer.w13_weight_bias,
|
||||
w2_weight_bias=layer.w2_weight_bias,
|
||||
gemm1_alpha=layer.gemm1_alpha,
|
||||
gemm1_beta=layer.gemm1_beta,
|
||||
gemm1_clamp_limit=layer.gemm1_clamp_limit,
|
||||
global_num_experts=layer.num_experts,
|
||||
local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
|
||||
local_num_experts=layer.num_local_experts,
|
||||
intermediate_size_per_partition=self.intermediate_size_per_partition,
|
||||
hidden_size=self.hidden_size,
|
||||
flashinfer_mxfp4_moe_precision=self.flashinfer_mxfp4_moe_precision,
|
||||
routing_bias=routing_bias,
|
||||
)
|
||||
return self.runner.run(dispatch_output, quant_info)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
@@ -1649,237 +1681,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
if self._fi_kernel == "cutlass_sm120":
|
||||
return self._apply_sm120_cutlass(layer, dispatch_output)
|
||||
if self.use_flashinfer:
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
trtllm_moe_enable_pdl,
|
||||
)
|
||||
|
||||
# When bf16 mode is enabled, we don't need to quantize the input,
|
||||
# TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations,
|
||||
# which can theoretically improve performance
|
||||
origin_hidden_states_dim = x.shape[-1]
|
||||
# Filled by the staged K3 route+pack+quant fusion below; the pack
|
||||
# site further down falls back to PackTopkIds when it is None.
|
||||
prepared_packed_topk = None
|
||||
if self.flashinfer_mxfp4_moe_precision == "bf16":
|
||||
assert x.dtype == torch.bfloat16
|
||||
x_quant = x
|
||||
x_scale = None
|
||||
|
||||
# May be fused later if this code branch is frequently needed
|
||||
if self.hidden_size != origin_hidden_states_dim:
|
||||
x_quant = torch.nn.functional.pad(
|
||||
x_quant,
|
||||
(0, self.hidden_size - origin_hidden_states_dim),
|
||||
mode="constant",
|
||||
value=0.0,
|
||||
)
|
||||
elif self.flashinfer_mxfp4_moe_precision == "default":
|
||||
x, prepared_packed_topk, x_quant, x_scale = (
|
||||
_prepare_flashinfer_mxfp8_activations(x, self.hidden_size)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
assert x_quant.shape[-1] == self.hidden_size
|
||||
is_standard = TopKOutputChecker.format_is_standard(topk_output)
|
||||
# The situ path accepts precomputed (standard) routing; the
|
||||
# public path below is logits-only.
|
||||
assert is_standard or TopKOutputChecker.format_is_bypassed(topk_output), (
|
||||
f"unsupported topk format: {topk_output.format}"
|
||||
)
|
||||
if is_standard:
|
||||
assert self.moe_runner_config.activation == "situ", (
|
||||
"standard topk output only wired for the situ path"
|
||||
)
|
||||
top_k = topk_output.topk_ids.shape[1]
|
||||
router_logits = None
|
||||
else:
|
||||
top_k = topk_output.topk_config.top_k
|
||||
router_logits = topk_output.router_logits
|
||||
|
||||
num_tokens = x_quant.shape[0]
|
||||
hidden_size = origin_hidden_states_dim
|
||||
# The K3 fused-front path publishes its [latent | shared] buffer
|
||||
# slice as the output destination (zero_copy_context); writing
|
||||
# the finalize output there directly skips this allocation and
|
||||
# the copy_ back in _forward_fused. The slice lives in the same
|
||||
# symmetric buffer the caller all-reduces.
|
||||
symm_output = zero_copy_context.get_moe_output_spec(
|
||||
torch.Size((num_tokens, hidden_size)),
|
||||
torch.bfloat16,
|
||||
x_quant.device,
|
||||
)
|
||||
if symm_output is None:
|
||||
with use_symmetric_memory(
|
||||
get_tp_group(), disabled=not is_allocation_symmetric()
|
||||
):
|
||||
symm_output = torch.empty(
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype=torch.bfloat16,
|
||||
device=x_quant.device,
|
||||
)
|
||||
|
||||
if self.moe_runner_config.activation == "situ":
|
||||
# FlashInfer 0.6.17+ ships the SiTU TRT-LLM-gen kernels.
|
||||
# Routing must be noaux_tc (sigmoid + correction bias,
|
||||
# DeepSeekV3), not the renormalize-softmax default below.
|
||||
# EP is cubin-internal: each rank computes its local expert slice
|
||||
# [offset, +num_local) and the caller all-reduces. ep=1 -> TP path.
|
||||
local_expert_offset = layer.moe_ep_rank * layer.num_local_experts
|
||||
if TopKOutputChecker.format_is_standard(topk_output):
|
||||
# Precomputed routing (radix router upstream): skip the
|
||||
# in-op routing kernels entirely. At small T the in-op
|
||||
# single-CTA routing costs ~22 us/layer vs ~6 us for the
|
||||
# external radix router.
|
||||
if prepared_packed_topk is not None:
|
||||
packed_topk = prepared_packed_topk
|
||||
else:
|
||||
from sglang.kernels.ops.moe.pack_topk_ids import PackTopkIds
|
||||
|
||||
packed_topk = PackTopkIds.execute(
|
||||
topk_output.topk_ids, topk_output.topk_weights
|
||||
)
|
||||
# Deferred finalize (K3 forward_deferred_finalize): return
|
||||
# the finalize inputs instead of the finalized output.
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
_deferred_finalize_enabled,
|
||||
)
|
||||
|
||||
defer_finalize = _deferred_finalize_enabled.get()
|
||||
result = trtllm_fp4_block_scale_routed_moe(
|
||||
topk_ids=packed_topk,
|
||||
routing_bias=None,
|
||||
hidden_states=x_quant,
|
||||
hidden_states_scale=x_scale,
|
||||
gemm1_weights=layer.w13_weight,
|
||||
gemm1_weights_scale=layer.w13_weight_scale,
|
||||
gemm1_bias=None,
|
||||
gemm1_alpha=layer.gemm1_alpha,
|
||||
# SiTU beta is the linear-half tanh clip; K3 stores it
|
||||
# in gemm1_clamp_limit.
|
||||
gemm1_beta=layer.gemm1_clamp_limit,
|
||||
gemm1_clamp_limit=None,
|
||||
gemm2_weights=layer.w2_weight,
|
||||
gemm2_weights_scale=layer.w2_weight_scale,
|
||||
gemm2_bias=None,
|
||||
output1_scale_scalar=None,
|
||||
output1_scale_gate_scalar=None,
|
||||
output2_scale_scalar=None,
|
||||
num_experts=layer.num_experts,
|
||||
top_k=packed_topk.shape[1],
|
||||
n_group=None,
|
||||
topk_group=None,
|
||||
intermediate_size=self.intermediate_size_per_partition,
|
||||
local_expert_offset=local_expert_offset,
|
||||
local_num_experts=layer.num_local_experts,
|
||||
routed_scaling_factor=None,
|
||||
routing_method_type=RoutingMethodType.TopK.value,
|
||||
activation_type=ActivationType.Situ.value,
|
||||
tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
|
||||
output=symm_output,
|
||||
do_finalize=not defer_finalize,
|
||||
enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]),
|
||||
)
|
||||
if defer_finalize:
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
FlashInferTrtllmDeferredFinalizeOutput,
|
||||
)
|
||||
|
||||
gemm2_out, topk_weights, expanded_idx = result
|
||||
result = FlashInferTrtllmDeferredFinalizeOutput(
|
||||
gemm2_out=gemm2_out,
|
||||
expert_weights=topk_weights,
|
||||
expanded_idx_to_permuted_idx=expanded_idx,
|
||||
top_k=packed_topk.shape[1],
|
||||
)
|
||||
return StandardCombineInput(hidden_states=result)
|
||||
# The finalized kernel writes to its explicit output
|
||||
# argument. Do not propagate the FFI return tensor: some
|
||||
# SiTU runner versions return a distinct wrapper/allocation
|
||||
# even though symm_output contains the published result.
|
||||
# Returning the destination makes the pointer contract
|
||||
# explicit for K3's zero-copy latent buffer.
|
||||
return StandardCombineInput(hidden_states=symm_output)
|
||||
|
||||
# Bypassed topk: route from logits inside the op.
|
||||
correction_bias = topk_output.topk_config.correction_bias
|
||||
bias_bf16 = getattr(layer, "_situ_routing_bias_bf16", None)
|
||||
if bias_bf16 is None and correction_bias is not None:
|
||||
bias_bf16 = correction_bias.to(torch.bfloat16)
|
||||
layer._situ_routing_bias_bf16 = bias_bf16
|
||||
trtllm_fp4_block_scale_moe(
|
||||
# router_logits is a row-strided slice of the K3 fused
|
||||
# front GEMM output; the FFI reads it as dense.
|
||||
routing_logits=router_logits.to(torch.bfloat16).contiguous(),
|
||||
routing_bias=bias_bf16,
|
||||
hidden_states=x_quant,
|
||||
hidden_states_scale=x_scale,
|
||||
gemm1_weights=layer.w13_weight,
|
||||
gemm1_weights_scale=layer.w13_weight_scale,
|
||||
gemm1_bias=None,
|
||||
gemm1_alpha=layer.gemm1_alpha,
|
||||
# SiTU beta is the linear-half tanh clip; K3 stores it in
|
||||
# gemm1_clamp_limit (situ_linear_beta).
|
||||
gemm1_beta=layer.gemm1_clamp_limit,
|
||||
gemm1_clamp_limit=None,
|
||||
gemm2_weights=layer.w2_weight,
|
||||
gemm2_weights_scale=layer.w2_weight_scale,
|
||||
gemm2_bias=None,
|
||||
output1_scale_scalar=None,
|
||||
output1_scale_gate_scalar=None,
|
||||
output2_scale_scalar=None,
|
||||
num_experts=layer.num_experts,
|
||||
top_k=top_k,
|
||||
n_group=topk_output.topk_config.num_expert_group,
|
||||
topk_group=topk_output.topk_config.topk_group,
|
||||
intermediate_size=self.intermediate_size_per_partition,
|
||||
routed_scaling_factor=(
|
||||
topk_output.topk_config.routed_scaling_factor or 1.0
|
||||
),
|
||||
routing_method_type=RoutingMethodType.DeepSeekV3.value,
|
||||
activation_type=ActivationType.Situ.value,
|
||||
norm_topk_prob=topk_output.topk_config.renormalize,
|
||||
local_expert_offset=local_expert_offset,
|
||||
local_num_experts=layer.num_local_experts,
|
||||
tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
|
||||
output=symm_output,
|
||||
enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]),
|
||||
)
|
||||
return StandardCombineInput(hidden_states=symm_output)
|
||||
|
||||
trtllm_gen_output = trtllm_fp4_block_scale_moe(
|
||||
router_logits.to(torch.bfloat16),
|
||||
None, # routing_bias
|
||||
x_quant,
|
||||
x_scale,
|
||||
layer.w13_weight, # uint8 (e2m1 x 2)
|
||||
layer.w13_weight_scale, # uint8 (e4m3 x 2)
|
||||
layer.w13_weight_bias, # fp32 per expert per channel
|
||||
layer.gemm1_alpha, # fp32 per expert
|
||||
layer.gemm1_beta, # fp32 per expert
|
||||
layer.gemm1_clamp_limit, # fp32 per expert
|
||||
layer.w2_weight, # uint8 (e2m1 x 2)
|
||||
layer.w2_weight_scale, # ue8m0
|
||||
layer.w2_weight_bias, # fp32 per expert per channel
|
||||
None, # output1_scale_scalar
|
||||
None, # output1_scale_gate_scalar
|
||||
None, # output2_scale_scalar
|
||||
layer.num_experts,
|
||||
top_k,
|
||||
None, # n_group # TODO: support n_group
|
||||
None, # topk_group # TODO: support topk_group
|
||||
self.intermediate_size_per_partition, # padded to multiple of 256
|
||||
layer.moe_ep_rank * layer.num_local_experts, # local_expert_offset
|
||||
layer.num_local_experts, # local num experts
|
||||
None, # routed_scaling_factor
|
||||
1, # routing_method_type, renormalize
|
||||
True, # do finalize
|
||||
tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
|
||||
output=symm_output,
|
||||
enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]),
|
||||
)[0]
|
||||
return StandardCombineInput(hidden_states=trtllm_gen_output)
|
||||
return self._apply_sm100_trtllm_gen(layer, dispatch_output)
|
||||
if _use_aiter:
|
||||
return self._apply_aiter(layer, dispatch_output)
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
@@ -16,9 +16,19 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires
|
||||
|
||||
|
||||
def test_situ_routed_moe_returns_published_output_buffer():
|
||||
from sglang.srt.layers.moe import route_quant_handoff
|
||||
# Import mxfp4 before flashinfer_trtllm to avoid the pre-existing
|
||||
# compressed_tensors circular import.
|
||||
# isort: off
|
||||
from sglang.srt.layers.quantization import mxfp4 as mxfp4_module
|
||||
from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
|
||||
from sglang.srt.layers.moe.moe_runner import (
|
||||
flashinfer_trtllm as flashinfer_trtllm_module,
|
||||
)
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
FlashInferTrtllmGenMxfp4MoeQuantInfo,
|
||||
_fused_experts_flashinfer_mxfp4_sm100_trtllm_gen,
|
||||
)
|
||||
# isort: on
|
||||
|
||||
tokens, hidden, top_k = 3, 128, 2
|
||||
x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda")
|
||||
@@ -38,27 +48,23 @@ def test_situ_routed_moe_returns_published_output_buffer():
|
||||
topk_output=topk_output,
|
||||
)
|
||||
|
||||
method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod)
|
||||
method.use_deep_gemm = False
|
||||
method.use_marlin = False
|
||||
method.use_flashinfer = True
|
||||
method._fi_kernel = None
|
||||
method.flashinfer_mxfp4_moe_precision = "default"
|
||||
method.hidden_size = hidden
|
||||
method.intermediate_size_per_partition = 128
|
||||
method.moe_runner_config = SimpleNamespace(activation="situ")
|
||||
|
||||
dummy = torch.empty(1, dtype=torch.uint8, device="cuda")
|
||||
layer = SimpleNamespace(
|
||||
moe_ep_rank=0,
|
||||
num_local_experts=1,
|
||||
num_experts=1,
|
||||
quant_info = FlashInferTrtllmGenMxfp4MoeQuantInfo(
|
||||
w13_weight=dummy,
|
||||
w13_weight_scale=dummy,
|
||||
gemm1_alpha=None,
|
||||
gemm1_clamp_limit=None,
|
||||
w2_weight=dummy,
|
||||
w13_weight_scale=dummy,
|
||||
w2_weight_scale=dummy,
|
||||
w13_weight_bias=dummy,
|
||||
w2_weight_bias=dummy,
|
||||
gemm1_alpha=dummy,
|
||||
gemm1_beta=dummy,
|
||||
gemm1_clamp_limit=dummy,
|
||||
global_num_experts=1,
|
||||
local_expert_offset=0,
|
||||
local_num_experts=1,
|
||||
intermediate_size_per_partition=128,
|
||||
hidden_size=hidden,
|
||||
flashinfer_mxfp4_moe_precision="default",
|
||||
)
|
||||
expected = (
|
||||
torch.arange(tokens * hidden, dtype=torch.float32, device="cuda")
|
||||
@@ -74,33 +80,40 @@ def test_situ_routed_moe_returns_published_output_buffer():
|
||||
returned_ptr = ffi_result.data_ptr()
|
||||
return ffi_result
|
||||
|
||||
flashinfer = ModuleType("flashinfer")
|
||||
flashinfer.__path__ = []
|
||||
flashinfer.trtllm_fp4_block_scale_moe = None
|
||||
fused_moe = ModuleType("flashinfer.fused_moe")
|
||||
fused_moe.trtllm_fp4_block_scale_routed_moe = fake_routed_moe
|
||||
tllm_enums = ModuleType("flashinfer.tllm_enums")
|
||||
tllm_enums.RoutingMethodType = SimpleNamespace(TopK=SimpleNamespace(value=0))
|
||||
tllm_enums.ActivationType = SimpleNamespace(Situ=SimpleNamespace(value=0))
|
||||
|
||||
latent = torch.empty_like(x)
|
||||
with (
|
||||
patch.object(
|
||||
route_quant_handoff,
|
||||
"take",
|
||||
return_value=(packed_topk, x_quant, x_scale),
|
||||
mxfp4_module,
|
||||
"_prepare_flashinfer_mxfp8_activations",
|
||||
return_value=(x, packed_topk, x_quant, x_scale),
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.layers.quantization.mxfp4.trtllm_fp4_block_scale_routed_moe",
|
||||
side_effect=fake_routed_moe,
|
||||
create=True,
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"flashinfer": flashinfer,
|
||||
"flashinfer.fused_moe": fused_moe,
|
||||
"flashinfer.tllm_enums": tllm_enums,
|
||||
},
|
||||
),
|
||||
patch.object(
|
||||
mxfp4_module,
|
||||
"RoutingMethodType",
|
||||
SimpleNamespace(TopK=SimpleNamespace(value=0)),
|
||||
create=True,
|
||||
),
|
||||
patch.object(
|
||||
mxfp4_module,
|
||||
"ActivationType",
|
||||
SimpleNamespace(Situ=SimpleNamespace(value=0)),
|
||||
create=True,
|
||||
flashinfer_trtllm_module, "trtllm_moe_enable_pdl", return_value=False
|
||||
),
|
||||
zero_copy_context.set_moe_output(latent),
|
||||
):
|
||||
combine_input = method.apply(layer, dispatch_output)
|
||||
combine_input = _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen(
|
||||
dispatch_output,
|
||||
quant_info,
|
||||
MoeRunnerConfig(activation="situ"),
|
||||
)
|
||||
|
||||
assert returned_ptr is not None
|
||||
assert returned_ptr != latent.data_ptr()
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
"""Focused SM100 trtllm-gen MXFP4 MoE regression test.
|
||||
|
||||
``Mxfp4MoEMethod.apply`` (SM100 branch, via the unified MoeRunner) must feed
|
||||
``trtllm_fp4_block_scale_moe`` the same args a direct kernel call does, so the
|
||||
two outputs stay bit-exact.
|
||||
|
||||
Fixtures are raw checkpoint-order MXFP4, converted by the production
|
||||
``process_weights_after_loading`` so the kernel sees the interleaved weights and
|
||||
float8_e4m3fn block scales it gets in a real run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import nullcontext
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from flashinfer import trtllm_fp4_block_scale_moe
|
||||
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=120, stage="base-b", runner_config="4-gpu-b200")
|
||||
|
||||
if not is_sm100_supported():
|
||||
pytest.skip(
|
||||
reason="trtllm-gen MXFP4 requires SM100 (Blackwell).",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
GROUP_SIZE = 32 # MXFP4 block size
|
||||
|
||||
|
||||
class _MockLayer:
|
||||
"""Hand-built ``FusedMoE`` stand-in (avoids distributed init)."""
|
||||
|
||||
|
||||
def _make_random_mxfp4(num_experts, hidden, inter, seed=0):
|
||||
g = torch.Generator(device="cuda").manual_seed(seed)
|
||||
w13 = torch.randint(
|
||||
0,
|
||||
256,
|
||||
(num_experts, 2 * inter, hidden // 2),
|
||||
dtype=torch.uint8,
|
||||
device="cuda",
|
||||
generator=g,
|
||||
)
|
||||
w2 = torch.randint(
|
||||
0,
|
||||
256,
|
||||
(num_experts, hidden, inter // 2),
|
||||
dtype=torch.uint8,
|
||||
device="cuda",
|
||||
generator=g,
|
||||
)
|
||||
# E8M0 scales centered around 127 (= 2^0); narrow band keeps dequant values
|
||||
# sane so the SwiGLU clamp doesn't dominate.
|
||||
w13_s = torch.randint(
|
||||
125,
|
||||
130,
|
||||
(num_experts, 2 * inter, hidden // GROUP_SIZE),
|
||||
dtype=torch.uint8,
|
||||
device="cuda",
|
||||
generator=g,
|
||||
)
|
||||
w2_s = torch.randint(
|
||||
125,
|
||||
130,
|
||||
(num_experts, hidden, inter // GROUP_SIZE),
|
||||
dtype=torch.uint8,
|
||||
device="cuda",
|
||||
generator=g,
|
||||
)
|
||||
w13_b = (
|
||||
torch.randn(
|
||||
num_experts, 2 * inter, dtype=torch.float32, device="cuda", generator=g
|
||||
)
|
||||
* 0.01
|
||||
)
|
||||
w2_b = (
|
||||
torch.randn(
|
||||
num_experts, hidden, dtype=torch.float32, device="cuda", generator=g
|
||||
)
|
||||
* 0.01
|
||||
)
|
||||
return w13, w2, w13_s, w2_s, w13_b, w2_b
|
||||
|
||||
|
||||
def _build_mock_layer(num_experts, hidden, inter, fixtures):
|
||||
"""Raw checkpoint-order weights; ``process_weights_after_loading`` converts
|
||||
them in place and seeds the SwiGLU scalars, so nothing is pre-applied here."""
|
||||
w13, w2, w13_s, w2_s, w13_b, w2_b = fixtures
|
||||
layer = _MockLayer()
|
||||
layer.w13_weight = torch.nn.Parameter(w13.clone(), requires_grad=False)
|
||||
layer.w2_weight = torch.nn.Parameter(w2.clone(), requires_grad=False)
|
||||
layer.w13_weight_scale = torch.nn.Parameter(w13_s.clone(), requires_grad=False)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(w2_s.clone(), requires_grad=False)
|
||||
layer.w13_weight_bias = torch.nn.Parameter(w13_b.clone(), requires_grad=False)
|
||||
layer.w2_weight_bias = torch.nn.Parameter(w2_b.clone(), requires_grad=False)
|
||||
layer.num_experts = num_experts
|
||||
layer.num_local_experts = num_experts # tests run with EP size = 1
|
||||
layer.moe_ep_rank = 0
|
||||
return layer
|
||||
|
||||
|
||||
def _build_method(num_experts, hidden, inter, precision):
|
||||
from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
|
||||
|
||||
method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod)
|
||||
method._fi_kernel = "trtllm_sm100"
|
||||
method.use_flashinfer = True
|
||||
method.use_marlin = False
|
||||
method.use_deep_gemm = False
|
||||
method.use_mega_moe = False
|
||||
method.num_experts = num_experts
|
||||
method.hidden_size = hidden
|
||||
method.intermediate_size_per_partition = inter
|
||||
method.flashinfer_mxfp4_moe_precision = precision
|
||||
method.runner = _build_flashinfer_mxfp4_runner(num_experts, hidden, inter)
|
||||
method.moe_runner_config = method.runner.config
|
||||
return method
|
||||
|
||||
|
||||
def _build_flashinfer_mxfp4_runner(num_experts, hidden, inter):
|
||||
# Bypass create_moe_runner (needs a live server arg context); the fused func
|
||||
# only reads dispatch_output / quant_info, so a minimal config suffices.
|
||||
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
|
||||
from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
||||
|
||||
cfg = MoeRunnerConfig(
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_experts,
|
||||
hidden_size=hidden,
|
||||
intermediate_size_per_partition=inter,
|
||||
top_k=None,
|
||||
activation="silu",
|
||||
is_gated=True,
|
||||
)
|
||||
return MoeRunner(MoeRunnerBackend.FLASHINFER_MXFP4, cfg)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_mega_moe", [False, True])
|
||||
def test_create_moe_runner_handles_flashinfer_for_megamoe(monkeypatch, use_mega_moe):
|
||||
import sglang.srt.layers.quantization.mxfp4 as mxfp4_mod
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
||||
|
||||
runner = object()
|
||||
|
||||
def build_runner(backend, config):
|
||||
assert not use_mega_moe
|
||||
assert backend == MoeRunnerBackend.FLASHINFER_MXFP4
|
||||
assert config is runner_config
|
||||
return runner
|
||||
|
||||
monkeypatch.setattr(
|
||||
mxfp4_mod,
|
||||
"get_moe_runner_backend",
|
||||
lambda: MoeRunnerBackend.FLASHINFER_MXFP4,
|
||||
)
|
||||
monkeypatch.setattr(mxfp4_mod, "MoeRunner", build_runner)
|
||||
|
||||
method = mxfp4_mod.Mxfp4MoEMethod.__new__(mxfp4_mod.Mxfp4MoEMethod)
|
||||
method._fi_kernel = "trtllm_sm100"
|
||||
method.use_mega_moe = use_mega_moe
|
||||
runner_config = MoeRunnerConfig()
|
||||
|
||||
method.create_moe_runner(object(), runner_config)
|
||||
|
||||
assert method.moe_runner_config is runner_config
|
||||
if use_mega_moe:
|
||||
# FusedMoEMethodBase declares ``runner: MoeRunner | None = None``, so the
|
||||
# early return leaves the class default rather than no attribute at all.
|
||||
assert method.runner is None
|
||||
else:
|
||||
assert method.runner is runner
|
||||
|
||||
|
||||
class _MockDispatchOutput:
|
||||
# SM100 keeps BYPASSED topk (kernel routes from router_logits), so the
|
||||
# dispatch output must carry a real BypassedTopKOutput.
|
||||
def __init__(self, hidden_states, router_logits, top_k):
|
||||
from sglang.srt.layers.moe.topk import BypassedTopKOutput, TopKConfig
|
||||
|
||||
self.hidden_states = hidden_states
|
||||
self.topk_output = BypassedTopKOutput(
|
||||
hidden_states=hidden_states,
|
||||
router_logits=router_logits,
|
||||
topk_config=TopKConfig(top_k=top_k, renormalize=True),
|
||||
)
|
||||
|
||||
|
||||
def _quant_input(x, precision, hidden_size):
|
||||
# Mirror the SM100 helper's input-quant branch so the reference feeds the
|
||||
# kernel the same x_quant / x_scale the SGLang path does.
|
||||
origin = x.shape[-1]
|
||||
if precision == "bf16":
|
||||
x_quant = x
|
||||
x_scale = None
|
||||
if hidden_size != origin:
|
||||
x_quant = torch.nn.functional.pad(
|
||||
x_quant, (0, hidden_size - origin), mode="constant", value=0.0
|
||||
)
|
||||
elif precision == "default":
|
||||
if x.shape[-1] == hidden_size:
|
||||
if x.dim() > 2:
|
||||
x = x.view(-1, x.shape[-1])
|
||||
from sglang.kernels.ops.quantization.per_token_group_quant import (
|
||||
per_token_group_quant,
|
||||
)
|
||||
|
||||
x_quant, x_scale = per_token_group_quant(x, group_size=32, scale_ue8m0=True)
|
||||
x_scale = x_scale.view(torch.float8_e4m3fn)
|
||||
else:
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
flashinfer_mxfp8_quantize,
|
||||
)
|
||||
|
||||
x_quant, x_scale = flashinfer_mxfp8_quantize(
|
||||
x, False, alignment=hidden_size
|
||||
)
|
||||
x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1)
|
||||
else:
|
||||
raise AssertionError(precision)
|
||||
return x_quant, x_scale
|
||||
|
||||
|
||||
def _ref_trtllm(x, layer, method, precision, top_k, router_logits):
|
||||
# Direct kernel call mirroring the SM100 helper's arg list.
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
|
||||
trtllm_moe_enable_pdl,
|
||||
)
|
||||
from sglang.srt.utils.common import next_power_of_2
|
||||
|
||||
x_quant, x_scale = _quant_input(x, precision, method.hidden_size)
|
||||
# zeros, not empty: the output is compared bit-exact, so any row the kernel
|
||||
# leaves unwritten must not carry allocator garbage.
|
||||
out = torch.zeros(
|
||||
x_quant.shape[0], x.shape[-1], dtype=torch.bfloat16, device=x_quant.device
|
||||
)
|
||||
return trtllm_fp4_block_scale_moe(
|
||||
router_logits.to(torch.bfloat16),
|
||||
None,
|
||||
x_quant,
|
||||
x_scale,
|
||||
layer.w13_weight,
|
||||
layer.w13_weight_scale,
|
||||
layer.w13_weight_bias,
|
||||
layer.gemm1_alpha,
|
||||
layer.gemm1_beta,
|
||||
layer.gemm1_clamp_limit,
|
||||
layer.w2_weight,
|
||||
layer.w2_weight_scale,
|
||||
layer.w2_weight_bias,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
layer.num_experts,
|
||||
top_k,
|
||||
None,
|
||||
None,
|
||||
method.intermediate_size_per_partition,
|
||||
layer.moe_ep_rank * layer.num_local_experts,
|
||||
layer.num_local_experts,
|
||||
None,
|
||||
1,
|
||||
True,
|
||||
tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
|
||||
output=out,
|
||||
enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]),
|
||||
)[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("precision", ["default", "bf16"])
|
||||
@pytest.mark.parametrize(
|
||||
"tokens,num_experts,hidden,inter,top_k",
|
||||
[
|
||||
(4, 4, 256, 256, 2),
|
||||
(16, 8, 512, 512, 2),
|
||||
(32, 8, 1024, 1024, 4),
|
||||
],
|
||||
)
|
||||
def test_apply_trtllm_gen_matches_flashinfer_direct(
|
||||
tokens, num_experts, hidden, inter, top_k, precision, monkeypatch
|
||||
):
|
||||
"""``Mxfp4MoEMethod.apply`` (SM100 branch) must produce the same output as a
|
||||
direct ``trtllm_fp4_block_scale_moe`` call fed the same inputs.
|
||||
|
||||
Turns red if apply mis-wires a ``FlashInferTrtllmGenMxfp4MoeQuantInfo`` field
|
||||
into the kernel (e.g. swapping ``local_expert_offset`` / ``local_num_experts``
|
||||
or dropping the bf16-vs-default input-quant branch)."""
|
||||
method = _build_method(num_experts, hidden, inter, precision)
|
||||
|
||||
import sglang.srt.layers.moe.moe_runner.flashinfer_trtllm as fi_trtllm_mod
|
||||
|
||||
# Bypass symmetric-memory / TP-group in the fused-func module, where the
|
||||
# kernel call now lives.
|
||||
monkeypatch.setattr(
|
||||
fi_trtllm_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
|
||||
)
|
||||
monkeypatch.setattr(fi_trtllm_mod, "is_allocation_symmetric", lambda: False)
|
||||
monkeypatch.setattr(fi_trtllm_mod, "get_tp_group", lambda: None)
|
||||
|
||||
fixtures = _make_random_mxfp4(num_experts, hidden, inter)
|
||||
x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1
|
||||
g = torch.Generator(device="cuda").manual_seed(1234)
|
||||
router_logits = torch.randn(
|
||||
tokens, num_experts, dtype=torch.float32, device="cuda", generator=g
|
||||
)
|
||||
layer = _build_mock_layer(num_experts, hidden, inter, fixtures)
|
||||
layer.moe_runner_config = method.moe_runner_config
|
||||
|
||||
# Convert via the production path so the fixtures can't drift from it.
|
||||
method.process_weights_after_loading(layer)
|
||||
|
||||
# ---- FlashInfer-direct reference ----
|
||||
out_ref = _ref_trtllm(x, layer, method, precision, top_k, router_logits)
|
||||
|
||||
# ---- SGLang path (same x + router_logits) ----
|
||||
out_sglang = method.apply(
|
||||
layer, _MockDispatchOutput(x.clone(), router_logits, top_k)
|
||||
).hidden_states
|
||||
|
||||
assert torch.equal(out_sglang, out_ref), (
|
||||
f"SGLang vs FlashInfer-direct mismatch (precision={precision}); "
|
||||
f"max abs diff = {(out_sglang.float() - out_ref.float()).abs().max().item():.4g}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -345,15 +345,9 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
|
||||
here we just verify that ``apply`` calls the kernel with the right
|
||||
arguments (incl. input padding + output trim)."""
|
||||
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as fi_cutlass_mod
|
||||
import sglang.srt.layers.quantization.mxfp4 as mxfp4_mod
|
||||
|
||||
# Bypass symmetric-memory / TP-group in both the legacy quant_method and
|
||||
# the new fused-func module (where the kernel call now lives).
|
||||
monkeypatch.setattr(
|
||||
mxfp4_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
|
||||
)
|
||||
monkeypatch.setattr(mxfp4_mod, "is_allocation_symmetric", lambda: False)
|
||||
monkeypatch.setattr(mxfp4_mod, "get_tp_group", lambda: None)
|
||||
# Bypass symmetric-memory / TP-group in the fused-func module, which is where
|
||||
# the kernel call lives.
|
||||
monkeypatch.setattr(
|
||||
fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user