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

[MoE Refactor] Migrate SM100 trtllm-gen mxfp4 MoE onto MoeRunner (#32405)

Browse Source
This commit is contained in:
Beihao Zhou
2026-09-05 13:48:10 +00:00
committed by GitHub
parent eda10c3678
commit 1e6f18bfeb
6 changed files with 709 additions and 303 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
+33 -7
View File
@@ -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
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+236
View File
@@ -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."""
python/sglang/srt/layers/quantization/mxfp4.py
+52 -250
View File
@@ -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)