From 0fee6bc6323e2e7aaaedfbc932289b8f7627945c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Mon, 5 Jan 2026 10:15:00 +0800 Subject: [PATCH] [JIT kernel] Apply jit per_tensor_quant_fp8 kernel (#15836) --- .../csrc/gemm/per_tensor_quant_fp8.cuh | 61 ++++++++++++------- .../sglang/jit_kernel/per_tensor_quant_fp8.py | 17 ++++-- .../tests/test_per_tensor_quant_fp8.py | 16 +++++ .../sglang/srt/layers/moe/cutlass_w4a8_moe.py | 16 ++--- .../srt/layers/quantization/fp8_kernel.py | 10 +-- 5 files changed, 79 insertions(+), 41 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh b/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh index d32171a04..3a0d677f8 100644 --- a/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh +++ b/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh @@ -104,31 +104,51 @@ template void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) { using namespace host; - SymbolicSize num_tokens = {"num_tokens"}; - SymbolicSize hidden_dim = {"hidden_dim"}; - SymbolicDevice device_; - SymbolicDType input_dtype; + const DLDevice device = input.device(); + RuntimeCheck(device.device_type == kDLCUDA, "input must be on CUDA"); + RuntimeCheck(input.is_contiguous(), "input must be contiguous"); - TensorMatcher({num_tokens, hidden_dim}) // - .with_dtype(input_dtype) - .with_device(device_) - .verify(input); + const int64_t ndim = input.dim(); + RuntimeCheck(ndim >= 1, "input.ndim must be >= 1, but got ", ndim); - TensorMatcher({num_tokens, hidden_dim}) // - .with_dtype<__nv_fp8_e4m3>() - .with_device(device_) - .verify(output_q); + RuntimeCheck(output_q.device() == device, "output_q must be on the same device as input"); + RuntimeCheck(output_q.is_contiguous(), "output_q must be contiguous"); + RuntimeCheck(output_q.dim() == ndim, "output_q.ndim must match input.ndim"); + for (int64_t i = 0; i < ndim; ++i) { + RuntimeCheck( + output_q.size(i) == input.size(i), + "output_q.shape mismatch at dim ", + i, + ": expected ", + input.size(i), + " but got ", + output_q.size(i)); + } TensorMatcher({1}) // .with_dtype() - .with_device(device_) + .with_device() .verify(output_s); + RuntimeCheck(output_s.device() == device, "output_s must be on the same device as input"); - const size_t total_elements = num_tokens.unwrap() * hidden_dim.unwrap(); + const DLDataType in_dtype = input.dtype(); + const bool in_ok = (in_dtype.code == kDLFloat && in_dtype.bits == 32) || + (in_dtype.code == kDLFloat && in_dtype.bits == 16) || + (in_dtype.code == kDLBfloat && in_dtype.bits == 16); + RuntimeCheck(in_ok, "input dtype must be fp32/fp16/bf16, but got ", in_dtype); + + const DLDataType out_dtype = output_q.dtype(); + RuntimeCheck( + out_dtype.code == kDLFloat8_e4m3fn && out_dtype.bits == 8, + "output_q dtype must be fp8_e4m3fn, but got ", + out_dtype); + + size_t total_elements = 1; + for (const auto s : input.shape()) { + RuntimeCheck(s > 0, "Input tensor must be non-empty"); + total_elements *= static_cast(s); + } const size_t num_blocks = std::min((total_elements + kBlockSize - 1) / kBlockSize, size_t(1024)); - const DLDevice device = device_.unwrap(); - - RuntimeCheck(total_elements > 0, "Input tensor must be non-empty"); auto launch_kernels = [&]() { if constexpr (!kIsStatic) { @@ -147,12 +167,11 @@ void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView outpu static_cast(total_elements)); }; - const DLDataType dtype = input_dtype.unwrap(); - if (dtype.code == kDLFloat && dtype.bits == 32) { + if (in_dtype.code == kDLFloat && in_dtype.bits == 32) { launch_kernels.template operator()(); - } else if (dtype.code == kDLBfloat && dtype.bits == 16) { + } else if (in_dtype.code == kDLBfloat && in_dtype.bits == 16) { launch_kernels.template operator()<__nv_bfloat16>(); - } else if (dtype.code == kDLFloat && dtype.bits == 16) { + } else if (in_dtype.code == kDLFloat && in_dtype.bits == 16) { launch_kernels.template operator()<__half>(); } } diff --git a/python/sglang/jit_kernel/per_tensor_quant_fp8.py b/python/sglang/jit_kernel/per_tensor_quant_fp8.py index 9a63357ca..bee07a207 100644 --- a/python/sglang/jit_kernel/per_tensor_quant_fp8.py +++ b/python/sglang/jit_kernel/per_tensor_quant_fp8.py @@ -1,6 +1,5 @@ from __future__ import annotations -import functools import os from typing import TYPE_CHECKING @@ -8,13 +7,14 @@ import flashinfer import torch from torch.utils.cpp_extension import CUDA_HOME -from sglang.jit_kernel.utils import load_jit, make_cpp_args +from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from tvm_ffi.module import Module -@functools.cache +@cache_once def _jit_per_tensor_quant_fp8_module(is_static: bool) -> Module: args = make_cpp_args(is_static) @@ -32,6 +32,10 @@ def _jit_per_tensor_quant_fp8_module(is_static: bool) -> Module: ) +@register_custom_op( + op_name="per_tensor_quant_fp8", + mutates_args=["output_q", "output_s"], +) def per_tensor_quant_fp8( input: torch.Tensor, output_q: torch.Tensor, @@ -44,8 +48,13 @@ def per_tensor_quant_fp8( Args: input: Input tensor to quantize (float, half, or bfloat16) output_q: Output quantized tensor (fp8_e4m3) - output_s: Output scale tensor (float scalar) + output_s: Output scale tensor (float scalar or 1D tensor with 1 element) is_static: If True, assumes scale is pre-computed and skips absmax computation """ + # Ensure output_s has shape [1] instead of being a 0D scalar + # The JIT kernel expects a 1D tensor + if output_s.ndim == 0: + output_s = output_s.reshape(1) + module = _jit_per_tensor_quant_fp8_module(is_static) module.per_tensor_quant_fp8(input, output_q, output_s) diff --git a/python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py b/python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py index 76ed08681..23357c57e 100644 --- a/python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py +++ b/python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py @@ -57,6 +57,22 @@ def test_jit_per_tensor_quant_compare_implementations( sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3 ) + +@pytest.mark.parametrize("shape", [(4, 8, 64), (2, 16, 128)]) +def test_jit_per_tensor_quant_supports_3d(shape): + device = torch.device("cuda") + x = torch.rand(shape, dtype=torch.bfloat16, device=device) + out = torch.empty_like(x, device=x.device, dtype=fp8_type_) + scale = torch.zeros(1, device=x.device, dtype=torch.float32) + + per_tensor_quant_fp8(x, out, scale, is_static=False) + + x_2d = x.flatten(0, -2) + out_ref_2d = torch_scaled_fp8_quant(x_2d, scale) + out_ref = out_ref_2d.reshape(shape) + + torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-3, atol=1e-3) + scale = torch.rand(1, dtype=torch.float32, device=device) sglang_out, sglang_scale = sglang_scaled_fp8_quant(x, scale) torch_out = torch_scaled_fp8_quant(x, scale) diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index 0e17d5cc7..0bf89b2d0 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py @@ -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) diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index df231f1a0..27f6c86a8 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -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