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

[JIT kernel] Apply jit per_tensor_quant_fp8 kernel (#15836)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-01-05 10:15:00 +08:00
committed by GitHub
parent 0ff3747ca1
commit 0fee6bc632
5 changed files with 79 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,13 +3,9 @@
from typing import Optional
import torch
from sgl_kernel import (
cutlass_w4a8_moe_mm,
get_cutlass_w4a8_moe_mm_data,
sgl_per_tensor_quant_fp8,
silu_and_mul,
)
from sgl_kernel import cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_data, silu_and_mul
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
from sglang.srt.distributed import get_moe_expert_parallel_world_size
from sglang.srt.layers.moe.ep_moe.kernels import (
cutlass_w4_run_moe_ep_preproess,
@@ -321,9 +317,7 @@ def cutlass_w4a8_moe_deepep_normal(
gateup_input = torch.empty(
gateup_input_pre_reorder.shape, dtype=torch.float8_e4m3fn, device=device
)
sgl_per_tensor_quant_fp8(
gateup_input_pre_reorder, gateup_input, a1_scale.float(), True
)
per_tensor_quant_fp8(gateup_input_pre_reorder, gateup_input, a1_scale.float(), True)
del gateup_input_pre_reorder
local_topk_ids = topk_ids_
local_topk_ids = (
@@ -367,7 +361,7 @@ def cutlass_w4a8_moe_deepep_normal(
intermediate_q = torch.empty(
intermediate.shape, dtype=torch.float8_e4m3fn, device=device
)
sgl_per_tensor_quant_fp8(intermediate, intermediate_q, a2_scale.float(), True)
per_tensor_quant_fp8(intermediate, intermediate_q, a2_scale.float(), True)
cutlass_w4a8_moe_mm(
c2,
@@ -495,7 +489,7 @@ def cutlass_w4a8_moe_deepep_ll(
)
gateup_input = torch.empty(a.shape, dtype=torch.float8_e4m3fn, device=device)
sgl_per_tensor_quant_fp8(a, gateup_input, a1_scale.float(), True)
per_tensor_quant_fp8(a, gateup_input, a1_scale.float(), True)
c1 = torch.empty((num_experts, m, n * 2), device=device, dtype=torch.bfloat16)
c2 = torch.empty((num_experts, m, k), device=device, dtype=torch.bfloat16)

10
python/sglang/srt/layers/quantization/fp8_kernel.py
View File

@@ -42,7 +42,11 @@ _is_cpu = is_cpu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _is_cuda:
from sgl_kernel import sgl_per_tensor_quant_fp8, sgl_per_token_quant_fp8
from sgl_kernel import sgl_per_token_quant_fp8
from sglang.jit_kernel.per_tensor_quant_fp8 import (
per_tensor_quant_fp8 as sgl_per_tensor_quant_fp8,
)
# Temporary
try:
@@ -1864,7 +1868,3 @@ if _is_cuda:
@torch.library.register_fake("sgl_kernel::sgl_per_token_quant_fp8")
def _(input, output_q, output_s):
return
@torch.library.register_fake("sgl_kernel::sgl_per_tensor_quant_fp8")
def _sgl_per_tensor_quant_fp8(input, output_q, output_s, is_static):
return