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

Replace [silu_and_mul_]scaled_fp4_group_quant by Flashinfer equivalent (#12376)

Browse Source
This commit is contained in:
Shu Wang
2025-11-13 02:26:00 -06:00
committed by GitHub
parent c2d69e8b56
commit 6664083522
6 changed files with 42 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/layers/moe/flashinfer_cutedsl_moe.py
View File

@@ -1,11 +1,11 @@
from typing import Optional
import torch
from flashinfer.cute_dsl.blockscaled_gemm import grouped_gemm_nt_masked
from sgl_kernel.gemm import (
scaled_fp4_grouped_quant,
silu_and_mul_scaled_fp4_grouped_quant,
from flashinfer import (
scaled_fp4_grouped_quantize,
silu_and_mul_scaled_nvfp4_experts_quantize,
)
from flashinfer.cute_dsl.blockscaled_gemm import grouped_gemm_nt_masked
def get_cute_dtype(input: torch.Tensor) -> str:
@@ -97,10 +97,10 @@ def flashinfer_cutedsl_moe_masked(
num_experts,
), f"input_global_scale must be (l,), got {input_global_scale.shape}"
a_q, a_q_sf = scaled_fp4_grouped_quant(
a_q, a_q_sf = scaled_fp4_grouped_quantize(
hidden_states[0],
input_global_scale,
masked_m,
input_global_scale,
)
assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}"
@@ -148,10 +148,10 @@ def flashinfer_cutedsl_moe_masked(
) # in logical [m, n, l]
# SILU and quantization
diq, diq_sf = silu_and_mul_scaled_fp4_grouped_quant(
diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize(
gateup_output.permute(2, 0, 1),
a2_global_scale,
masked_m,
a2_global_scale,
)
if down_start_event is not None: