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

direct register custom op for mm_fp4 (#13699)

Browse Source 
Co-authored-by: Brayden Zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
b8zhong
2025-12-11 22:33:17 -08:00
committed by GitHub
co-authored by Brayden Zhong
parent a7fa31ff17
commit dab31e4cc7
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/modelopt_quant.py
+9 -2
View File
@@ -45,6 +45,7 @@ from sglang.srt.layers.quantization.utils import (
) )
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.utils.common import ( from sglang.srt.utils.common import (
direct_register_custom_op,
get_bool_env_var, get_bool_env_var,
is_cuda, is_cuda,
is_sm120_supported, is_sm120_supported,
@@ -100,7 +101,6 @@ except ImportError:
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@torch.library.custom_op("sglang::fp4_gemm", mutates_args=())
def _sglang_fp4_gemm( def _sglang_fp4_gemm(
input: torch.Tensor, input: torch.Tensor,
weight: torch.Tensor, weight: torch.Tensor,
@@ -119,7 +119,6 @@ def _sglang_fp4_gemm(
return fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype) return fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype)
@torch.library.register_fake("sglang::fp4_gemm")
def _sglang_fp4_gemm_fake( def _sglang_fp4_gemm_fake(
input, input,
weight, weight,
@@ -134,6 +133,14 @@ def _sglang_fp4_gemm_fake(
return input.new_empty((M, N), dtype=out_dtype) return input.new_empty((M, N), dtype=out_dtype)
direct_register_custom_op(
op_name="fp4_gemm",
op_func=_sglang_fp4_gemm,
mutates_args=[],
fake_impl=_sglang_fp4_gemm_fake,
)
if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None): if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None):
@register_fake_if_exists("sgl_kernel::scaled_fp4_quant") @register_fake_if_exists("sgl_kernel::scaled_fp4_quant")