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

[Fix] Support Kimi-K3 ModelOpt mixed NVFP4/FP8 checkpoint (#35077)

Browse Source
This commit is contained in:
YAMY
2026-08-19 08:13:45 -07:00
committed by GitHub
parent 41c018a9ec
commit 5f12839591
6 changed files with 205 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+13 -2
View File
@@ -517,6 +517,7 @@ def _compute_g1_scale_c(
g1_alphas: torch.Tensor,
g1_alphas_up: torch.Tensor,
is_gated: bool,
activation: Optional[str] = None,
) -> torch.Tensor:
"""TRT-LLM GEMM1-output scale for the up (w3) half.
@@ -526,6 +527,11 @@ def _compute_g1_scale_c(
scale passes g1_alphas as g1_alphas_up and recovers the single-scale value;
non-gated (Relu2) has no gate half, so it is just 1/a2_scale per expert.
"""
if activation == "situ":
# SiTU consumes both GEMM1 scales before tanh; scale_c carries only
# the GEMM2 input requantization factor.
num_experts = g1_alphas.shape[0]
return w2_input_scale_quant.to(torch.float32).expand(num_experts).contiguous()
if is_gated:
return (w2_input_scale_quant * g1_alphas_up).to(torch.float32)
num_experts = g1_alphas.shape[0]
@@ -596,7 +602,11 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
g1_alphas = cast(torch.Tensor, layer.g1_alphas)
g1_alphas_up = cast(torch.Tensor, getattr(layer, "g1_alphas_up", g1_alphas))
g1_scale_c = _compute_g1_scale_c(
w2_input_scale_quant, g1_alphas, g1_alphas_up, layer.moe_runner_config.is_gated
w2_input_scale_quant,
g1_alphas,
g1_alphas_up,
layer.moe_runner_config.is_gated,
activation=layer.moe_runner_config.activation,
)
copy_or_rebind_param(layer, "g1_scale_c", g1_scale_c)
@@ -612,6 +622,7 @@ def get_activation_type(activation: str, is_gated: bool = True) -> int:
_ACTIVATION_STR_TO_TYPE = {
"silu": ActivationType.Swiglu,
"gelu": ActivationType.Geglu,
"situ": ActivationType.Situ,
}
else:
_ACTIVATION_STR_TO_TYPE = {
@@ -956,7 +967,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.layers.moe.utils import RoutingMethodType
_SUPPORTED_FP4_ACTIVATIONS = {"silu", "relu2", "gelu"}
_SUPPORTED_FP4_ACTIVATIONS = {"silu", "relu2", "gelu", "situ"}
assert runner_config.activation in _SUPPORTED_FP4_ACTIVATIONS, (
f"Only {_SUPPORTED_FP4_ACTIVATIONS} are supported for FP4 MoE, "
f"got '{runner_config.activation}'."
python/sglang/srt/layers/quantization/modelopt_quant.py
+30 -11
View File
@@ -776,6 +776,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
packed_modules_mapping: Optional[Dict[str, List[str]]],
quantized_layers: Dict[str, Dict[str, Any]],
fp8_config: ModelOptFp8Config,
fp8_pb_wo_config: Fp8Config,
nvfp4_config: ModelOptFp4Config,
nvfp4a16_config: ModelOptFp4Config,
mxfp8_config: Fp8Config,
@@ -783,6 +784,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping)
self.quantized_layers = quantized_layers
self.fp8_config = fp8_config
self.fp8_pb_wo_config = fp8_pb_wo_config
self.mxfp8_config = mxfp8_config
self.nvfp4_config = nvfp4_config
self.nvfp4a16_config = nvfp4a16_config
@@ -867,6 +869,12 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
exclude_modules=[],
packed_modules_mapping=packed_modules_mapping,
)
fp8_pb_wo_config = Fp8Config(
is_checkpoint_fp8_serialized=True,
activation_scheme="dynamic",
weight_block_size=[128, 128],
packed_modules_mapping=packed_modules_mapping,
)
mxfp8_config = Fp8Config(
is_checkpoint_fp8_serialized=True,
activation_scheme="dynamic",
@@ -896,6 +904,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
packed_modules_mapping=packed_modules_mapping,
quantized_layers=quantized_layers,
fp8_config=fp8_config,
fp8_pb_wo_config=fp8_pb_wo_config,
mxfp8_config=mxfp8_config,
nvfp4_config=nvfp4_config,
nvfp4a16_config=nvfp4a16_config,
@@ -977,6 +986,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
return UnquantizedLinearMethod()
if quant_algo == "FP8":
return ModelOptFp8LinearMethod(self.fp8_config)
if quant_algo == "FP8_PB_WO":
return Fp8LinearMethod(self.fp8_pb_wo_config)
if quant_algo == "MXFP8":
return Fp8LinearMethod(self.mxfp8_config)
if quant_algo == "NVFP4":
@@ -2511,9 +2522,12 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
)
if layer.moe_runner_config.is_gated and self.enable_flashinfer_trtllm_moe:
runner_config = layer.moe_runner_config
is_situ = runner_config.activation == "situ"
gemm1_clamp_limit = (
layer.moe_runner_config.gemm1_clamp_limit
or layer.moe_runner_config.swiglu_limit
None
if is_situ
else (runner_config.gemm1_clamp_limit or runner_config.swiglu_limit)
)
if gemm1_clamp_limit is not None:
copy_or_rebind_param(
@@ -2522,21 +2536,26 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
(gemm1_clamp_limit / layer.g1_alphas).to(torch.float32),
)
if layer.moe_runner_config.gemm1_alpha is not None:
if runner_config.gemm1_alpha is not None:
copy_or_rebind_param(
layer,
"gemm1_alpha",
torch.full_like(
layer.g1_alphas,
layer.moe_runner_config.gemm1_alpha,
runner_config.gemm1_alpha,
dtype=torch.float32,
),
)
copy_or_rebind_param(
layer,
"gemm1_beta",
(1.0 / layer.g1_alphas).to(torch.float32),
gemm1_beta = (
torch.full_like(
layer.g1_alphas,
runner_config.gemm1_clamp_limit,
dtype=torch.float32,
)
if is_situ
else (1.0 / layer.g1_alphas).to(torch.float32)
)
copy_or_rebind_param(layer, "gemm1_beta", gemm1_beta)
# TODO: for flashinfer always do MOE_NVFP4_DISPATCH
use_dispatch_fp4 = not self.quant_config.use_per_token_activation and (
@@ -2802,9 +2821,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
self, "_moe_runner_backend", get_moe_runner_backend()
)
assert (
activation in _SUPPORTED_ACT_STRS
), f"{activation=} not in supported {_SUPPORTED_ACT_STRS}"
assert activation in _SUPPORTED_ACT_STRS or (
activation == "situ" and moe_runner_backend.is_flashinfer_trtllm()
), f"{activation=} is unsupported by {moe_runner_backend}"
moe_runner_config = self.moe_runner_config
if moe_runner_backend.is_marlin():
python/sglang/srt/models/kimi_k3.py
+82 -23
View File
@@ -73,6 +73,7 @@ from sglang.srt.layers.moe.utils import (
get_moe_runner_backend,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.fp8_utils import block_quant_dequant
from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
from sglang.srt.layers.vocab_parallel_embedding import (
@@ -142,6 +143,35 @@ def _cdiv(a: int, b: int) -> int:
return (a + b - 1) // b
def _uses_modelopt_fp8_pb_wo(
quant_config: Optional[QuantizationConfig], prefix: str
) -> bool:
resolver = getattr(quant_config, "_resolve_quant_algo", None)
return resolver is not None and resolver(prefix) == "FP8_PB_WO"
def _maybe_map_fp8_pb_scale_name(name: str, params_dict: dict) -> str:
"""Map ModelOpt FP8_PB_WO scale keys to SGLang block-FP8 params."""
if name.endswith(".weight_scale"):
candidate = name.removesuffix(".weight_scale") + ".weight_scale_inv"
if candidate in params_dict:
return candidate
return name
def _get_k3_dense_weight(module: nn.Module) -> torch.Tensor:
"""Return a dense weight with serialized block-FP8 scales applied."""
weight = module.weight.data
if not hasattr(module, "weight_scale_inv"):
return weight
return block_quant_dequant(
weight,
module.weight_scale_inv,
module.quant_method.weight_block_size,
module.params_dtype,
)
def _k3_bf16_gemm(
x: torch.Tensor,
weight: torch.Tensor,
@@ -457,17 +487,17 @@ class KimiK3MoE(nn.Module):
quant_config=quant_config,
routed_scaling_factor=self.routed_scaling_factor,
apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
# flashinfer_mxfp4 + situ consumes precomputed routing
# (PackedPrecomputed): keep the radix router in the TopK layer
# and hand its ids/weights to the MoE op. Other quantized paths
# keep the runner-resolved format (marlin -> standard anyway,
# bypassed only for the public logits-routing path).
# TRT-LLM cannot consume fused-front's row-strided router logits;
# keep K3's FP32 router and pass precomputed top-k instead.
output_format=(
TopKOutputFormat.STANDARD
if quant_config is None
or (
config.hidden_act == "situ"
and get_moe_runner_backend().is_flashinfer_mxfp4()
and (
get_moe_runner_backend().is_flashinfer_mxfp4()
or get_moe_runner_backend().is_flashinfer_trtllm()
)
)
# mega pre-dispatch consumes raw topk_ids/topk_weights
or get_moe_a2a_backend().is_megamoe()
@@ -1385,11 +1415,12 @@ class KimiK3DeltaAttention(nn.Module):
self.use_full_rank_gate = config.linear_attn_config.get(
"use_full_rank_gate", False
)
self._bfa_uses_block_fp8 = self.use_full_rank_gate and _uses_modelopt_fp8_pb_wo(
quant_config, f"{prefix}.b_proj"
)
# The fused path hardcodes tp_size sharding, so require attn_tp == tp.
# For the full-rank gate (K3) the checkpoint quantizes only the MoE
# experts; attention linears resolve to UnquantizedLinearMethod, so a
# non-None quant_config is fine for the merged projection.
# Full-rank K3 also fuses mixed block-FP8 attention projections.
self.do_fuse_qkvbfg = self.attn_tp_size == self.tp_size and (
quant_config is None or self.use_full_rank_gate
)
@@ -1425,6 +1456,8 @@ class KimiK3DeltaAttention(nn.Module):
tp_rank=self.attn_tp_rank,
tp_size=self.attn_tp_size,
prefix=f"{prefix}.b_proj",
# TP8 shards K3's 96 beta rows below the 128-row FP8 block.
skip_block_quant_check=self._bfa_uses_block_fp8,
)
self.f_a_proj = ReplicatedLinear(
self.hidden_size,
@@ -1445,6 +1478,7 @@ class KimiK3DeltaAttention(nn.Module):
# Merged [f_a | b] weight, built after weight loading by
# _merge_bfa_weights().
self._bfa_w: Optional[torch.Tensor] = None
self._bfa_f_b_w: Optional[torch.Tensor] = None
elif self.do_fuse_qkvbfg:
self.qkvb_sizes = [
projection_size,
@@ -1661,15 +1695,24 @@ class KimiK3DeltaAttention(nn.Module):
and the width is padded to a multiple of 8 so every fused-output row
stays 16-byte aligned for vectorized consumers (tiny-GEMM on f_b).
Called once from load_weights (after all weights are loaded, before
cuda graph capture)."""
Called once after weight loading. Block-FP8 inputs are dequantized into
the BF16 tiny-GEMM buffers here."""
if not self.use_full_rank_gate:
return
if _is_npu:
return
self._bfa_w, sizes = _merge_weights_as_views(
[self.f_a_proj, self.b_proj], pad_rows_to=8
)
mods = [self.f_a_proj, self.b_proj]
if self._bfa_uses_block_fp8:
weights = [_get_k3_dense_weight(mod) for mod in mods]
sizes = [weight.shape[0] for weight in weights]
pad = (-sum(sizes)) % 8
if pad:
weights.append(weights[0].new_zeros((pad, weights[0].shape[1])))
self._bfa_w = torch.cat(weights, dim=0).contiguous()
self._bfa_f_b_w = _get_k3_dense_weight(self.f_b_proj).contiguous()
else:
self._bfa_w, sizes = _merge_weights_as_views(mods, pad_rows_to=8)
self._bfa_f_b_w = self.f_b_proj.weight
self._bfa_fa_size, self._bfa_b_size = sizes
def _prepare_fused_decode(self) -> None:
@@ -1748,7 +1791,7 @@ class KimiK3DeltaAttention(nn.Module):
alt.wait_stream(cur)
with torch.cuda.stream(alt):
bfa = gemm(hidden_states, w)
forget_gate = gemm(bfa[..., :n_fa], self.f_b_proj.weight)
forget_gate = gemm(bfa[..., :n_fa], self._bfa_f_b_w)
beta = bfa[..., n_fa : n_fa + n_b]
fused_states, _ = self.fused_qkvg_proj(hidden_states)
qkv, g_proj_states = torch.split(
@@ -1760,7 +1803,7 @@ class KimiK3DeltaAttention(nn.Module):
fused_states, _ = self.fused_qkvg_proj(hidden_states)
qkv, g_proj_states = torch.split(fused_states, self.split_sizes, dim=-1)
bfa = gemm(hidden_states, w)
forget_gate = gemm(bfa[..., :n_fa], self.f_b_proj.weight)
forget_gate = gemm(bfa[..., :n_fa], self._bfa_f_b_w)
beta = bfa[..., n_fa : n_fa + n_b]
else:
fused_states, _ = self.fused_qkvg_proj(hidden_states)
@@ -2883,6 +2926,8 @@ class KimiK3LinearForCausalLM(nn.Module):
for args in weights:
name, loaded_weight = args[:2]
kwargs = args[2] if len(args) > 2 else {}
if name.endswith(".weight_scale") and loaded_weight.ndim == 4:
loaded_weight = loaded_weight[:, 0, :, 0]
layer_id = get_layer_id(name)
if layer_id is not None and (
@@ -2904,13 +2949,20 @@ class KimiK3LinearForCausalLM(nn.Module):
# MLA: fuse q_a_proj + kv_a_proj_with_mqa → fused_qkv_a_proj_with_mqa
if ".q_a_proj." in name or ".kv_a_proj_with_mqa." in name:
is_q_a = ".q_a_proj." in name
fused_name = name.replace(".q_a_proj.", ".fused_qkv_a_proj_with_mqa.")
fused_name = fused_name.replace(
".kv_a_proj_with_mqa.", ".fused_qkv_a_proj_with_mqa."
)
fused_name = _maybe_map_fp8_pb_scale_name(fused_name, params_dict)
if fused_name in params_dict:
param = params_dict[fused_name]
if ".q_a_proj." in name:
if fused_name.endswith(".weight_scale_inv"):
offset = 0 if is_q_a else _cdiv(self.config.q_lora_rank, 128)
param.data[offset : offset + loaded_weight.shape[0]].copy_(
loaded_weight
)
elif is_q_a:
param.data[: loaded_weight.shape[0]].copy_(loaded_weight)
else:
q_lora_rank = self.config.q_lora_rank or 0
@@ -2947,6 +2999,7 @@ class KimiK3LinearForCausalLM(nn.Module):
name = name.replace(weight_name, param_name)
if name.endswith(".bias") and name not in params_dict:
continue
name = _maybe_map_fp8_pb_scale_name(name, params_dict)
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
@@ -2983,13 +3036,18 @@ class KimiK3LinearForCausalLM(nn.Module):
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
continue
name = _maybe_map_fp8_pb_scale_name(name, params_dict)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight, **kwargs)
if name.endswith(".b_proj.weight_scale_inv"):
# All TP ranks share K3's single beta output-scale block.
param.data.copy_(loaded_weight)
else:
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight, **kwargs)
loaded_params.add(name)
self.post_load_weights()
@@ -3007,7 +3065,8 @@ class KimiK3LinearForCausalLM(nn.Module):
if isinstance(layer, PPMissingLayer):
continue
self_attn = layer.self_attn
w_kc, w_vc = self_attn.kv_b_proj.weight.unflatten(
kv_b_weight = _get_k3_dense_weight(self_attn.kv_b_proj)
w_kc, w_vc = kv_b_weight.unflatten(
0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)
@@ -3065,7 +3124,7 @@ class KimiK3LinearForCausalLM(nn.Module):
if precompile_k3_recompute_w_u_kernel(
num_heads=layer.self_attn.local_num_heads,
dtype=layer.self_attn.o_proj.weight.dtype,
dtype=layer.self_attn.o_proj.params_dtype,
device=layer.self_attn.dt_bias.device,
):
rank0_log("Precompiled the Kimi-K3 KDA prefill kernel.")
test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_scales.py
+20
View File
@@ -291,6 +291,26 @@ class TestG1ScaleC(CustomTestCase):
self.assertTrue(g1_scale_c.is_contiguous())
torch.testing.assert_close(g1_scale_c, torch.full((num_experts,), 20.0))
def test_situ_keeps_both_dequant_scales_inside_activation(self):
# SiTU keeps both GEMM1 scales before tanh; scale_c contains only the
# GEMM2 input requantization.
num_experts = GATED_CONFIGS[0][1]
w2_input_scale_quant = torch.tensor(20.0)
gate = _global_scales(num_experts, 1, seed=8)
up = _global_scales(num_experts, 1, seed=9)
g1_scale_c = _compute_g1_scale_c(
w2_input_scale_quant,
gate,
up,
is_gated=True,
activation="situ",
)
self.assertEqual(g1_scale_c.shape, (num_experts,))
self.assertTrue(g1_scale_c.is_contiguous())
torch.testing.assert_close(g1_scale_c, torch.full((num_experts,), 20.0))
def test_scale_c_is_float32(self):
# Lower-precision inputs are upcast to fp32 for the kernel.
num_experts = GATED_CONFIGS[0][1]
test/registered/unit/model_loader/test_modelopt_loader.py
+22
View File
@@ -696,6 +696,28 @@ class TestModelOptFp4LoaderSelection(CustomTestCase):
class TestModelOptMixedPrecisionConfig(CustomTestCase):
def test_fp8_pb_wo_dispatches_to_native_block_fp8(self):
quant_config = ModelOptMixedPrecisionConfig.from_config(
{
"quant_algo": "MIXED_PRECISION",
"quantized_layers": {
"model.layers.0.self_attn.q_proj": {"quant_algo": "FP8_PB_WO"},
},
"packed_modules_mapping": {},
}
)
# Type dispatch only needs a LinearBase instance; skip GPU weight setup.
linear = ReplicatedLinear.__new__(ReplicatedLinear)
method = quant_config.get_quant_method(
linear, "model.layers.0.self_attn.q_proj"
)
self.assertIsInstance(method, Fp8LinearMethod)
self.assertEqual(method.quant_config.weight_block_size, [128, 128])
self.assertTrue(method.quant_config.is_checkpoint_fp8_serialized)
self.assertEqual(method.quant_config.activation_scheme, "dynamic")
def test_incomplete_inline_config_falls_back_to_hf_quant_config_file(self):
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
test/registered/unit/models/test_kimi_k3_bfa_overlap.py
+38 -2
View File
@@ -8,7 +8,10 @@ from unittest.mock import patch
import torch
from sglang.srt.models.kimi_k3 import KimiK3DeltaAttention
from sglang.srt.models.kimi_k3 import (
KimiK3DeltaAttention,
_get_k3_dense_weight,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -39,9 +42,9 @@ def _make_owner(with_stream: bool):
owner = SimpleNamespace(
use_full_rank_gate=True,
_bfa_w=_randn(_BFA_W_ROWS, _H).contiguous(),
_bfa_f_b_w=_randn(1536, _N_FA).contiguous(),
_bfa_fa_size=_N_FA,
_bfa_b_size=_N_B,
f_b_proj=SimpleNamespace(weight=_randn(1536, _N_FA).contiguous()),
fused_qkvg_proj=fused_qkvg_proj,
split_sizes=[3 * 1536, 1536],
_bfa_alt_stream=torch.cuda.Stream() if with_stream else None,
@@ -97,6 +100,39 @@ class TestKimiK3BfaOverlap(CustomTestCase):
for got, ref in zip(overlap, serial):
self.assertTrue(torch.equal(got, ref))
def test_block_fp8_weight_is_dequantized_for_tiny_gemm(self):
module = SimpleNamespace(
weight=torch.nn.Parameter(
torch.ones((130, 129), device="cuda", dtype=torch.float8_e4m3fn),
requires_grad=False,
),
weight_scale_inv=torch.nn.Parameter(
torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda"),
requires_grad=False,
),
quant_method=SimpleNamespace(weight_block_size=[128, 128]),
params_dtype=torch.bfloat16,
)
weight = _get_k3_dense_weight(module)
self.assertEqual(weight.dtype, torch.bfloat16)
torch.testing.assert_close(
weight[[0, 0, 128, 128], [0, 128, 0, 128]].float(),
torch.tensor([1.0, 2.0, 3.0, 4.0], device="cuda"),
)
def test_per_tensor_fp8_weight_is_not_block_dequantized(self):
weight = torch.nn.Parameter(
torch.ones((2, 2), device="cuda", dtype=torch.float8_e4m3fn),
requires_grad=False,
)
module = SimpleNamespace(
weight=weight, weight_scale=torch.ones(1, device="cuda")
)
self.assertEqual(_get_k3_dense_weight(module).data_ptr(), weight.data_ptr())
if __name__ == "__main__":
unittest.main()