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

Mistral Large 3 NVFP4 TRTLLM MoE support (#15049)

Browse Source
This commit is contained in:
elvischenv
2025-12-18 11:11:42 +08:00
committed by GitHub
parent 9e7656be80
commit 9970ee34e8
7 changed files with 340 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -548,8 +548,9 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]):
quant_config is None
or quant_config.get_name() == "fp8"
or quant_config.get_name() == "modelopt_fp8"
or quant_config.get_name() == "compressed_tensors"
):
# FlashInferFusedMoE support bf16 and fp8
# FlashInferFusedMoE support bf16, fp8 and compressed_tensors
return FlashInferFusedMoE
if get_moe_runner_backend().is_flashinfer_cutlass():

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

@@ -1093,7 +1093,6 @@ class FlashInferFusedMoE(FusedMoE):
else:
# FP8 Matrix multiply.
final_hidden_states = self.quant_method.apply_with_router_logits(
layer=self,
dispatch_output=StandardDispatchOutput(

214
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py
View File

@@ -11,10 +11,15 @@ import torch
from compressed_tensors import CompressionFormat
from compressed_tensors.quantization import QuantizationStrategy
from sglang.srt.distributed import get_tensor_model_parallel_world_size
from sglang.srt.distributed import get_tensor_model_parallel_world_size, get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
from sglang.srt.layers.moe.utils import get_moe_runner_backend
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
from sglang.srt.layers.quantization.compressed_tensors.schemes import (
WNA16_SUPPORTED_BITS,
@@ -29,10 +34,18 @@ from sglang.srt.layers.quantization.marlin_utils import marlin_moe_permute_scale
from sglang.srt.layers.quantization.utils import (
all_close_1d,
per_tensor_dequantize,
prepare_static_weights_for_trtllm_fp4_moe,
reorder_w1w3_to_w3w1,
replace_parameter,
swizzle_blockscale,
)
from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, set_weight_attrs
from sglang.srt.utils import (
get_bool_env_var,
is_cuda,
is_hip,
next_power_of_2,
set_weight_attrs,
)
if TYPE_CHECKING:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
@@ -115,6 +128,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
)
self.quant_config = quant_config
self.group_size = 16
self.use_flashinfer_trtllm = get_moe_runner_backend().is_flashinfer_trtllm()
def create_weights(
self,
@@ -127,7 +141,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
):
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
layer.num_experts = num_experts
layer.params_dtype = params_dtype
w13_weight = torch.nn.Parameter(
@@ -240,6 +253,13 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
)
delattr(layer, "w2_weight_packed")
if self.use_flashinfer_trtllm:
w, s = reorder_w1w3_to_w3w1(
layer.w13_weight.data, layer.w13_weight_scale.data, dim=-2
)
layer.w13_weight = torch.nn.Parameter(w, requires_grad=False)
layer.w13_weight_scale = torch.nn.Parameter(s, requires_grad=False)
if not torch.allclose(
layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1]
):
@@ -258,9 +278,16 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
)
# w13
w13_input_global_scale = layer.w13_input_global_scale.min(dim=1).values.to(
torch.float32
)
if self.use_flashinfer_trtllm:
w13_input_global_scale = (
layer.w13_input_global_scale.min()
.to(torch.float32)
.expand(layer.num_local_experts)
)
else:
w13_input_global_scale = layer.w13_input_global_scale.min(dim=1).values.to(
torch.float32
)
layer.g1_alphas = torch.nn.Parameter(
((1 / w13_input_global_scale) * layer.w13_weight_scale_2),
requires_grad=False,
@@ -271,7 +298,14 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
)
# w2
w2_input_global_scale = layer.w2_input_global_scale
if self.use_flashinfer_trtllm:
w2_input_global_scale = (
layer.w2_input_global_scale.min()
.to(torch.float32)
.expand(layer.num_local_experts)
)
else:
w2_input_global_scale = layer.w2_input_global_scale
layer.g2_alphas = torch.nn.Parameter(
((1 / w2_input_global_scale) * layer.w2_weight_scale_2).to(torch.float32),
@@ -282,22 +316,66 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
(w2_input_global_scale), requires_grad=False
)
# swizzle weight scales
layer.w13_weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.w13_weight_scale), requires_grad=False
)
# TensorRT-LLM specific processing
if self.use_flashinfer_trtllm:
# Prepare static weights for TRT-LLM kernel
(
gemm1_weights_fp4_shuffled,
gemm1_scales_fp4_shuffled,
gemm2_weights_fp4_shuffled,
gemm2_scales_fp4_shuffled,
) = prepare_static_weights_for_trtllm_fp4_moe(
layer.w13_weight,
layer.w2_weight,
layer.w13_weight_scale,
layer.w2_weight_scale,
layer.w2_weight.size(-2), # hidden_size
layer.w13_weight.size(-2) // 2, # intermediate_size
layer.w13_weight.size(0), # num_experts
)
logger.debug("Finished shuffling weights for TRT-LLM MOE")
layer.w2_weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.w2_weight_scale), requires_grad=False
)
layer.gemm1_weights_fp4_shuffled = torch.nn.Parameter(
gemm1_weights_fp4_shuffled, requires_grad=False
)
layer.gemm2_weights_fp4_shuffled = torch.nn.Parameter(
gemm2_weights_fp4_shuffled, requires_grad=False
)
layer.gemm1_scales_fp4_shuffled = torch.nn.Parameter(
gemm1_scales_fp4_shuffled, requires_grad=False
)
layer.gemm2_scales_fp4_shuffled = torch.nn.Parameter(
gemm2_scales_fp4_shuffled, requires_grad=False
)
layer.cutlass_moe_params = CutlassMoEParams(
CutlassMoEType.BlockscaledFP4,
layer.w13_weight.device,
num_experts=layer.num_experts,
intermediate_size_per_partition=layer.w2_weight.shape[2] * 2,
hidden_size=layer.w13_weight.shape[2] * 2,
)
# Additional parameter needed for TRT-LLM
layer.g1_scale_c = torch.nn.Parameter(
(layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32),
requires_grad=False,
)
# Clean up weights that won't be used by TRT-LLM
del layer.w2_weight
del layer.w2_weight_scale
del layer.w13_weight
del layer.w13_weight_scale
else:
# swizzle weight scales
layer.w13_weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.w13_weight_scale), requires_grad=False
)
layer.w2_weight_scale = torch.nn.Parameter(
swizzle_blockscale(layer.w2_weight_scale), requires_grad=False
)
layer.cutlass_moe_params = CutlassMoEParams(
CutlassMoEType.BlockscaledFP4,
layer.w13_weight.device,
num_experts=layer.num_experts,
intermediate_size_per_partition=layer.w2_weight.shape[2] * 2,
hidden_size=layer.w13_weight.shape[2] * 2,
)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
@@ -336,6 +414,100 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
return StandardCombineInput(hidden_states=output)
def apply_with_router_logits(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> torch.Tensor:
assert self.use_flashinfer_trtllm
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
from flashinfer import fp4_quantize, trtllm_fp4_block_scale_moe
from sglang.srt.layers.moe.utils import RoutingMethodType
router_logits = topk_output.router_logits
topk_config = topk_output.topk_config
# Quantize input hidden states using fp4_quantize
hs_fp4_bytes, hs_sf_bytes = fp4_quantize(
x,
layer.w13_input_scale_quant,
self.group_size, # sf_vec_size
False, # use_ue8m0
False, # is_sf_swizzled_layout
)
hs_fp4 = hs_fp4_bytes.reshape(x.shape[0], x.shape[1] // 2)
hs_scale = hs_sf_bytes.view(torch.float8_e4m3fn).reshape(-1)
correction_bias = (
None
if topk_config.correction_bias is None
else topk_config.correction_bias.to(x.dtype)
)
assert layer.routing_method_type is not None
# DeepSeekV3 style routing requires float32 router logits
if layer.routing_method_type == RoutingMethodType.DeepSeekV3:
router_logits = router_logits.to(torch.float32)
routed_scaling_factor = self.moe_runner_config.routed_scaling_factor
routed_scaling_factor = (
routed_scaling_factor if routed_scaling_factor is not None else 1.0
)
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
num_tokens = hs_fp4.shape[0]
hidden_size = (
hs_fp4.shape[-1] * 2
if hs_fp4.dtype == torch.uint8
else hs_fp4.shape[-1]
)
symm_output = torch.empty(
num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device
)
return trtllm_fp4_block_scale_moe(
routing_logits=router_logits,
routing_bias=correction_bias,
hidden_states=hs_fp4,
hidden_states_scale=hs_scale,
gemm1_weights=layer.gemm1_weights_fp4_shuffled,
gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.view(
torch.float8_e4m3fn
),
gemm1_bias=None,
gemm1_alpha=None,
gemm1_beta=None,
gemm1_clamp_limit=None,
gemm2_weights=layer.gemm2_weights_fp4_shuffled,
gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.view(
torch.float8_e4m3fn
),
gemm2_bias=None,
output1_scale_scalar=layer.g1_scale_c,
output1_scale_gate_scalar=layer.g1_alphas,
output2_scale_scalar=layer.g2_alphas,
num_experts=layer.num_experts,
top_k=topk_config.top_k,
n_group=topk_config.num_expert_group,
topk_group=topk_config.topk_group,
intermediate_size=layer.intermediate_size_per_partition,
local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
local_num_experts=layer.num_local_experts,
routed_scaling_factor=routed_scaling_factor,
tile_tokens_dim=None,
routing_method_type=layer.routing_method_type,
do_finalize=True,
tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]),
output=symm_output,
)[0]
class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

127
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -42,6 +42,7 @@ from sglang.srt.layers.quantization.utils import (
convert_to_channelwise,
is_layer_skipped,
per_tensor_dequantize,
prepare_static_weights_for_trtllm_fp4_moe,
requantize_with_max_scale,
swizzle_blockscale,
)
@@ -1398,130 +1399,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
w2_input_scale._sglang_require_global_experts = True
layer.register_parameter("w2_input_scale", w2_input_scale)
def prepare_static_weights_for_kernel(
self,
# args_dequant,
# args,
gemm1_weights,
gemm2_weights,
gemm1_scales_linear_fp4_bytes,
gemm2_scales_linear_fp4_bytes,
hidden_size,
intermediate_size,
num_experts,
):
from flashinfer import nvfp4_block_scale_interleave
from flashinfer.fused_moe.core import (
_maybe_get_cached_w3_w1_permute_indices,
get_w2_permute_indices_with_cache,
)
"""Prepare quantized weights for kernel (done offline with weights)."""
epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
# Convert quantized weights to proper formats
gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(
num_experts, 2 * intermediate_size, hidden_size // 2
) # packed fp4
gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(
torch.float8_e4m3fn
).reshape(
num_experts, 2 * intermediate_size, hidden_size // 16
) # fp8 scaling factors
gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(
num_experts, hidden_size, intermediate_size // 2
) # packed fp4
gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view(
torch.float8_e4m3fn
).reshape(
num_experts, hidden_size, intermediate_size // 16
) # fp8 scaling factors
gemm1_weights_fp4_shuffled = []
gemm1_scales_fp4_shuffled = []
gemm2_weights_fp4_shuffled = []
gemm2_scales_fp4_shuffled = []
for i in range(num_experts):
# Calculate the permute indices for the following:
# 1. Reorder rows of W1 and scales for fused gated activation
# 2. Shuffle weights and scaling factors for transposed mma output
# for both w3_w1 and w2 weights and scale factors
permute_indices = _maybe_get_cached_w3_w1_permute_indices(
self._cache_permute_indices,
gemm1_weights_fp4[i].view(torch.uint8),
epilogue_tile_m,
)
gemm1_weights_fp4_shuffled.append(
gemm1_weights_fp4[i]
.view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)]
.contiguous()
)
permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(
self._cache_permute_indices,
gemm1_scales_linear_fp4[i].view(torch.uint8),
epilogue_tile_m,
num_elts_per_sf=16,
)
gemm1_scales_fp4_shuffled.append(
nvfp4_block_scale_interleave(
gemm1_scales_linear_fp4[i]
.view(torch.uint8)[
permute_sf_indices.to(gemm1_scales_linear_fp4.device)
]
.contiguous()
)
)
permute_indices = get_w2_permute_indices_with_cache(
self._cache_permute_indices,
gemm2_weights_fp4[i].view(torch.uint8),
epilogue_tile_m,
)
gemm2_weights_fp4_shuffled.append(
gemm2_weights_fp4[i]
.view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)]
.contiguous()
)
permute_sf_indices = get_w2_permute_indices_with_cache(
self._cache_permute_indices,
gemm2_scales_linear_fp4[i].view(torch.uint8),
epilogue_tile_m,
num_elts_per_sf=16,
)
gemm2_scales_fp4_shuffled.append(
nvfp4_block_scale_interleave(
gemm2_scales_linear_fp4[i]
.view(torch.uint8)[
permute_sf_indices.to(gemm2_scales_linear_fp4.device)
]
.contiguous()
)
)
# Stack weights for all experts
gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)
gemm1_scales_fp4_shuffled = (
torch.stack(gemm1_scales_fp4_shuffled)
.view(torch.float8_e4m3fn)
.reshape(num_experts, 2 * intermediate_size, hidden_size // 16)
)
gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)
gemm2_scales_fp4_shuffled = (
torch.stack(gemm2_scales_fp4_shuffled)
.view(torch.float8_e4m3fn)
.reshape(num_experts, hidden_size, intermediate_size // 16)
)
return (
gemm1_weights_fp4_shuffled,
gemm1_scales_fp4_shuffled,
gemm2_weights_fp4_shuffled,
gemm2_scales_fp4_shuffled,
)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Process FP4 MoE weights after loading from serialized checkpoint.
@@ -1633,7 +1510,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
gemm1_scales_fp4_shuffled,
gemm2_weights_fp4_shuffled,
gemm2_scales_fp4_shuffled,
) = self.prepare_static_weights_for_kernel(
) = prepare_static_weights_for_trtllm_fp4_moe(
layer.w13_weight,
layer.w2_weight,
layer.w13_weight_scale,

140
python/sglang/srt/layers/quantization/utils.py
View File

@@ -592,3 +592,143 @@ def swizzle_blockscale(scale: torch.Tensor):
if scale_ndim == 2
else swizzled_scale.reshape(B, M_padded, K_padded)
)
def reorder_w1w3_to_w3w1(
weight: torch.Tensor, scale: torch.Tensor, dim: int = -2
) -> tuple[torch.Tensor, torch.Tensor]:
"""Re-order the concatenated `[w1, w3]` tensors to `[w3, w1]`"""
size = weight.size(dim)
assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
half = size // 2
w1, w3 = weight.split(half, dim=dim)
s1, s3 = scale.split(half, dim=dim)
return (
torch.cat([w3, w1], dim=dim).contiguous(),
torch.cat([s3, s1], dim=dim).contiguous(),
)
def prepare_static_weights_for_trtllm_fp4_moe(
gemm1_weights,
gemm2_weights,
gemm1_scales_linear_fp4_bytes,
gemm2_scales_linear_fp4_bytes,
hidden_size,
intermediate_size,
num_experts,
):
from flashinfer import nvfp4_block_scale_interleave
from flashinfer.fused_moe.core import (
_maybe_get_cached_w3_w1_permute_indices,
get_w2_permute_indices_with_cache,
)
"""Prepare quantized weights for kernel (done offline with weights)."""
_cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
# Convert quantized weights to proper formats
gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(
num_experts, 2 * intermediate_size, hidden_size // 2
) # packed fp4
gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(
torch.float8_e4m3fn
).reshape(
num_experts, 2 * intermediate_size, hidden_size // 16
) # fp8 scaling factors
gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(
num_experts, hidden_size, intermediate_size // 2
) # packed fp4
gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view(
torch.float8_e4m3fn
).reshape(
num_experts, hidden_size, intermediate_size // 16
) # fp8 scaling factors
gemm1_weights_fp4_shuffled = []
gemm1_scales_fp4_shuffled = []
gemm2_weights_fp4_shuffled = []
gemm2_scales_fp4_shuffled = []
for i in range(num_experts):
# Calculate the permute indices for the following:
# 1. Reorder rows of W1 and scales for fused gated activation
# 2. Shuffle weights and scaling factors for transposed mma output
# for both w3_w1 and w2 weights and scale factors
permute_indices = _maybe_get_cached_w3_w1_permute_indices(
_cache_permute_indices,
gemm1_weights_fp4[i].view(torch.uint8),
epilogue_tile_m,
)
gemm1_weights_fp4_shuffled.append(
gemm1_weights_fp4[i]
.view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)]
.contiguous()
)
permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(
_cache_permute_indices,
gemm1_scales_linear_fp4[i].view(torch.uint8),
epilogue_tile_m,
num_elts_per_sf=16,
)
gemm1_scales_fp4_shuffled.append(
nvfp4_block_scale_interleave(
gemm1_scales_linear_fp4[i]
.view(torch.uint8)[
permute_sf_indices.to(gemm1_scales_linear_fp4.device)
]
.contiguous()
)
)
permute_indices = get_w2_permute_indices_with_cache(
_cache_permute_indices,
gemm2_weights_fp4[i].view(torch.uint8),
epilogue_tile_m,
)
gemm2_weights_fp4_shuffled.append(
gemm2_weights_fp4[i]
.view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)]
.contiguous()
)
permute_sf_indices = get_w2_permute_indices_with_cache(
_cache_permute_indices,
gemm2_scales_linear_fp4[i].view(torch.uint8),
epilogue_tile_m,
num_elts_per_sf=16,
)
gemm2_scales_fp4_shuffled.append(
nvfp4_block_scale_interleave(
gemm2_scales_linear_fp4[i]
.view(torch.uint8)[
permute_sf_indices.to(gemm2_scales_linear_fp4.device)
]
.contiguous()
)
)
# Stack weights for all experts
gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)
gemm1_scales_fp4_shuffled = (
torch.stack(gemm1_scales_fp4_shuffled)
.view(torch.float8_e4m3fn)
.reshape(num_experts, 2 * intermediate_size, hidden_size // 16)
)
gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)
gemm2_scales_fp4_shuffled = (
torch.stack(gemm2_scales_fp4_shuffled)
.view(torch.float8_e4m3fn)
.reshape(num_experts, hidden_size, intermediate_size // 16)
)
return (
gemm1_weights_fp4_shuffled,
gemm1_scales_fp4_shuffled,
gemm2_weights_fp4_shuffled,
gemm2_scales_fp4_shuffled,
)

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

@@ -1790,8 +1790,9 @@ class ServerArgs:
"modelopt_fp4",
"fp8",
"modelopt_fp8",
"compressed-tensors",
None,
], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', or bfloat16 (None)."
], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', 'compressed-tensors', or bfloat16 (None)."
self.disable_shared_experts_fusion = True
logger.warning(
"FlashInfer TRTLLM MoE is enabled. --disable-shared-experts-fusion is automatically set."

3
python/sglang/srt/utils/mistral_utils.py
View File

@@ -81,8 +81,6 @@ def adapt_config_dict(
config_dict = _remap_mistral_vision_args(config_dict)
if is_audio:
config_dict = _remap_mistral_audio_args(config_dict)
if is_eagle:
config_dict["routing_method_type"] = 1 # RoutingMethodType.Renormalize
config = PretrainedConfig.from_dict(config_dict)
@@ -234,6 +232,7 @@ def _remap_moe_args(config: dict) -> dict:
config["topk_method"] = None
config["scoring_func"] = "softmax"
config["routing_method_type"] = 1 # RoutingMethodType.Renormalize
return config