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

[RL] Support mxfp8 DeepSeek V3 (#21280)

Browse Source
This commit is contained in:
Ziang Li
2026-04-03 21:57:45 -07:00
committed by GitHub
parent 68f4c52d34
commit 990c7590b8
3 changed files with 105 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+12 -7
View File
@@ -291,13 +291,18 @@ class FusedMoE(torch.nn.Module):
self.quant_method.create_moe_runner(self, self.moe_runner_config)
self.dispatcher = create_moe_dispatcher(self.moe_runner_config)
self.should_fuse_routed_scaling_factor_in_topk = isinstance(
self.quant_method, ModelOptNvFp4FusedMoEMethod
) or (
isinstance(self.quant_method, Fp8MoEMethod)
and (
get_moe_runner_backend().is_cutlass()
or get_moe_runner_backend().is_flashinfer_trtllm_routed()
self.should_fuse_routed_scaling_factor_in_topk = (
isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod)
or (
isinstance(self.quant_method, Fp8MoEMethod)
and (
get_moe_runner_backend().is_cutlass()
or get_moe_runner_backend().is_flashinfer_trtllm_routed()
)
)
or (
isinstance(self.quant_method, UnquantizedFusedMoEMethod)
and get_moe_runner_backend().is_flashinfer_trtllm_routed()
)
)
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+86 -38
View File
@@ -47,6 +47,10 @@ elif is_cuda_alike():
else:
fp4_quantize = None
_flashinfer_trtllm_shuffle_row_indices_cache_mxfp8: dict[
tuple, dict[str, torch.Tensor]
] = {}
def align_fp8_moe_weights_for_flashinfer_trtllm(
layer: Module, swap_w13_halves: bool = False
@@ -126,10 +130,13 @@ def align_fp8_moe_weights_for_flashinfer_trtllm(
def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
"""Prepare MXFP8 MoE weights/scales for FlashInfer TRT-LLM kernels."""
from flashinfer import (
reorder_rows_for_gated_act_gemm,
shuffle_matrix_a,
shuffle_matrix_sf_a,
from flashinfer import block_scale_interleave
from flashinfer.fused_moe.core import (
get_reorder_rows_for_gated_act_gemm_row_indices,
)
from flashinfer.utils import (
get_shuffle_matrix_a_row_indices,
get_shuffle_matrix_sf_a_row_indices,
)
w13_weight = cast(torch.Tensor, layer.w13_weight).contiguous()
@@ -144,52 +151,93 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
_, hidden_size, _ = w2_weight.shape
epilogue_tile_m = 128
w13_interleaved = [
reorder_rows_for_gated_act_gemm(w13_weight[i]) for i in range(num_experts)
]
w13_scale_interleaved = [
reorder_rows_for_gated_act_gemm(w13_scale[i]) for i in range(num_experts)
]
# Reuse precomputed row-index transforms whenever shape/device are unchanged.
w13_weight_u8 = w13_weight.view(torch.uint8)
w2_weight_u8 = w2_weight.view(torch.uint8)
cache_key = (
two_n,
hidden_size,
w2_weight.shape[-1],
w13_scale.shape[-1],
w2_scale.shape[-1],
epilogue_tile_m,
(w13_weight.device.type, w13_weight.device.index),
(w2_weight.device.type, w2_weight.device.index),
(w13_scale.device.type, w13_scale.device.index),
(w2_scale.device.type, w2_scale.device.index),
)
cache = _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8.get(cache_key)
if cache is None:
reorder_row_indices = get_reorder_rows_for_gated_act_gemm_row_indices(
w13_weight_u8[0]
).to(w13_weight.device)
w13_shuffle_row_indices = get_shuffle_matrix_a_row_indices(
w13_weight_u8[0], epilogue_tile_m
).to(w13_weight.device)
w2_shuffle_row_indices = get_shuffle_matrix_a_row_indices(
w2_weight_u8[0], epilogue_tile_m
).to(w2_weight.device)
w13_scale_shuffle_row_indices = get_shuffle_matrix_sf_a_row_indices(
w13_scale[0].reshape(two_n, -1), epilogue_tile_m
).to(w13_scale.device)
w2_scale_shuffle_row_indices = get_shuffle_matrix_sf_a_row_indices(
w2_scale[0].reshape(hidden_size, -1), epilogue_tile_m
).to(w2_scale.device)
cache = {
"reorder_row_indices": reorder_row_indices,
"w13_shuffle_row_indices": w13_shuffle_row_indices,
"w2_shuffle_row_indices": w2_shuffle_row_indices,
"w13_scale_shuffle_row_indices": w13_scale_shuffle_row_indices,
"w2_scale_shuffle_row_indices": w2_scale_shuffle_row_indices,
}
_flashinfer_trtllm_shuffle_row_indices_cache_mxfp8[cache_key] = cache
w13_shuffled = [
shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m)
for i in range(num_experts)
]
w2_shuffled = [
shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m)
for i in range(num_experts)
]
w13_scale_shuffled = [
shuffle_matrix_sf_a(
w13_scale_interleaved[i].view(torch.uint8).reshape(two_n, -1),
epilogue_tile_m,
reorder_row_indices = cache["reorder_row_indices"]
w13_shuffle_row_indices = cache["w13_shuffle_row_indices"]
w2_shuffle_row_indices = cache["w2_shuffle_row_indices"]
w13_scale_shuffle_row_indices = cache["w13_scale_shuffle_row_indices"]
w2_scale_shuffle_row_indices = cache["w2_scale_shuffle_row_indices"]
w13_shuffled_u8 = torch.empty_like(w13_weight_u8)
w2_shuffled_u8 = torch.empty_like(w2_weight_u8)
w13_scale_shuffled = torch.empty_like(w13_scale)
w2_scale_shuffled = torch.empty_like(w2_scale)
for i in range(num_experts):
w13_interleaved_u8 = w13_weight_u8[i].index_select(0, reorder_row_indices)
w13_scale_interleaved = w13_scale[i].index_select(0, reorder_row_indices)
w13_shuffled_u8[i].copy_(
w13_interleaved_u8.index_select(0, w13_shuffle_row_indices)
)
for i in range(num_experts)
]
w2_scale_shuffled = [
shuffle_matrix_sf_a(
w2_scale[i].view(torch.uint8).reshape(hidden_size, -1),
epilogue_tile_m,
w2_shuffled_u8[i].copy_(w2_weight_u8[i].index_select(0, w2_shuffle_row_indices))
w13_scale_linear = w13_scale_interleaved.reshape(two_n, -1)
w13_scale_shuffled[i].copy_(
block_scale_interleave(
w13_scale_linear.index_select(0, w13_scale_shuffle_row_indices)
).reshape_as(w13_scale_shuffled[i])
)
w2_scale_linear = w2_scale[i].reshape(hidden_size, -1)
w2_scale_shuffled[i].copy_(
block_scale_interleave(
w2_scale_linear.index_select(0, w2_scale_shuffle_row_indices)
).reshape_as(w2_scale_shuffled[i])
)
for i in range(num_experts)
]
# Keep parameter identities stable for CUDA graph capture reuse.
copy_or_rebind_param(
layer, "w13_weight", torch.stack(w13_shuffled).view(torch.float8_e4m3fn)
)
copy_or_rebind_param(
layer, "w2_weight", torch.stack(w2_shuffled).view(torch.float8_e4m3fn)
)
copy_or_rebind_param(layer, "w13_weight", w13_shuffled_u8.view(torch.float8_e4m3fn))
copy_or_rebind_param(layer, "w2_weight", w2_shuffled_u8.view(torch.float8_e4m3fn))
copy_or_rebind_param(
layer,
"w13_weight_scale_inv",
torch.stack(w13_scale_shuffled).reshape_as(w13_scale).contiguous(),
w13_scale_shuffled.contiguous(),
)
copy_or_rebind_param(
layer,
"w2_weight_scale_inv",
torch.stack(w2_scale_shuffled).reshape_as(w2_scale).contiguous(),
w2_scale_shuffled.contiguous(),
)
layer.w13_weight_scale_inv.format_ue8m0 = True
layer.w2_weight_scale_inv.format_ue8m0 = True
python/sglang/srt/layers/quantization/fp8.py
+7
View File
@@ -93,6 +93,7 @@ if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
from sglang.srt.layers.moe.topk import TopKOutput
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
from sglang.srt.models.utils import WeightsMapper
_is_hip = is_hip()
_is_cuda = is_cuda()
@@ -241,6 +242,12 @@ class Fp8Config(QuantizationConfig):
def get_scaled_act_names(self) -> List[str]:
return []
def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"):
if self.ignored_layers:
self.ignored_layers = list(
dict.fromkeys(hf_to_sglang_mapper.apply_list(self.ignored_layers))
)
class Fp8LinearMethod(LinearMethodBase):
"""Linear method for FP8.