Refactor FP4 quantization and remove deprecated JIT kernels (#30448)

Co-authored-by: root <root@sgl-b300-inference.datacrunch.io>
This commit is contained in:
Brayden Zhong
2026-07-14 09:22:07 +08:00
committed by GitHub
co-authored by root
parent 0c01971eeb
commit 9756f768a6
38 changed files with 120 additions and 5705 deletions
@@ -43,11 +43,11 @@ try:
from sgl_kernel import fused_add_rmsnorm as SGL_FUSED_ADD_RMS_NORM
from sgl_kernel import rmsnorm as SGL_RMS_NORM
from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as SGL_SCALED_FP4_QUANT
from sglang.srt.layers.quantization.fp4_utils import fp4_quantize as SGL_FP4_QUANT
except Exception: # pragma: no cover - fallback on non-supported platforms
SGL_FUSED_ADD_RMS_NORM = None
SGL_RMS_NORM = None
SGL_SCALED_FP4_QUANT = None
SGL_FP4_QUANT = None
FP8_DTYPE = SGLANG_FP8_DTYPE
@@ -386,9 +386,9 @@ def standard_allreduce_rmsnorm_fp4_quant(
residual_out = allreduce_out
# Finally FP4 quantization
if SGL_SCALED_FP4_QUANT is None:
raise RuntimeError("scaled_fp4_quant is not available on this platform")
quant_res, output_scale_res = SGL_SCALED_FP4_QUANT(quant_input, input_global_scale)
if SGL_FP4_QUANT is None:
raise RuntimeError("fp4_quantize is not available on this platform")
quant_res, output_scale_res = SGL_FP4_QUANT(quant_input, input_global_scale)
if residual is not None:
return quant_res, residual_out, output_scale_res
else:
@@ -464,9 +464,9 @@ def standard_allreduce_rmsnorm_fp4_quant_native(
residual_out = allreduce_out
# Apply FP4 quantization (still using fused CUDA op as there's no native FP4)
if SGL_SCALED_FP4_QUANT is None:
raise RuntimeError("scaled_fp4_quant is not available on this platform")
quant_res, output_scale_res = SGL_SCALED_FP4_QUANT(quant_input, input_global_scale)
if SGL_FP4_QUANT is None:
raise RuntimeError("fp4_quantize is not available on this platform")
quant_res, output_scale_res = SGL_FP4_QUANT(quant_input, input_global_scale)
if residual is not None:
return quant_res, residual_out, output_scale_res