[JIT sgl-kernel] Jit support per tensor quant (#15709)

python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py

@@ -0,0 +1,135 @@
+import itertools
+import os
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.testing
+
+from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+
+try:
+    from vllm import _custom_ops as ops
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    ops = None
+    VLLM_AVAILABLE = False
+
+try:
+    from sglang.srt.utils import is_hip
+
+    _is_hip = is_hip()
+except ImportError:
+    _is_hip = False
+
+IS_CI = (
+    os.getenv("CI", "false").lower() == "true"
+    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
+)
+
+fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+
+
+def vllm_scaled_fp8_quant(
+    input: torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if not VLLM_AVAILABLE:
+        return sglang_scaled_fp8_quant(input, scale)
+    return ops.scaled_fp8_quant(input, scale)
+
+
+def sglang_scaled_fp8_quant(
+    input: torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    fp8_type_: torch.dtype = torch.float8_e4m3fn
+    output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
+    is_static = True
+    if scale is None:
+        scale = torch.zeros(1, device=input.device, dtype=torch.float32)
+        is_static = False
+    per_tensor_quant_fp8(input, output, scale, is_static)
+
+    return output, scale
+
+
+def calculate_diff(batch_size: int, seq_len: int):
+    device = torch.device("cuda")
+    x = torch.rand((batch_size, seq_len), dtype=torch.float16, device=device)
+
+    if not VLLM_AVAILABLE:
+        print("vLLM not available, skipping comparison")
+        return
+
+    vllm_out, vllm_scale = vllm_scaled_fp8_quant(x)
+    sglang_out, sglang_scale = sglang_scaled_fp8_quant(x)
+
+    scale_diff = torch.abs(vllm_scale - sglang_scale).item()
+    output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
+
+    if torch.allclose(
+        vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    ) and torch.allclose(vllm_scale, sglang_scale, rtol=1e-3, atol=1e-5):
+        print("All implementations match")
+    else:
+        print("Implementations differ")
+
+
+if IS_CI:
+    batch_size_range = [16]
+    seq_len_range = [64]
+else:
+    batch_size_range = [16, 32, 64, 128]
+    seq_len_range = [64, 128, 256, 512, 1024, 2048]
+
+configs = list(itertools.product(batch_size_range, seq_len_range))
+
+
+if VLLM_AVAILABLE:
+    line_vals = ["vllm", "sglang"]
+    line_names = ["VLLM", "SGL Kernel"]
+    styles = [("blue", "-"), ("green", "-")]
+else:
+    line_vals = ["sglang"]
+    line_names = ["SGL Kernel"]
+    styles = [("green", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=line_vals,
+        line_names=line_names,
+        styles=styles,
+        ylabel="us",
+        plot_name="per-tensor-quant-fp8-performance",
+        args={},
+    )
+)
+def benchmark(batch_size, seq_len, provider):
+    dtype = torch.float16
+    device = torch.device("cuda")
+
+    x = torch.randn(batch_size * seq_len, 4096, device=device, dtype=dtype)
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "vllm":
+        fn = lambda: vllm_scaled_fp8_quant(x.clone())
+    elif provider == "sglang":
+        fn = lambda: sglang_scaled_fp8_quant(x.clone())
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    calculate_diff(batch_size=4, seq_len=4096)
+    benchmark.run(print_data=True)

python/sglang/jit_kernel/csrc/add_constant.cuh

@@ -53,8 +53,6 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
       static_cast<int32_t*>(dst.data_ptr()),
       static_cast<int32_t*>(src.data_ptr()),
       num_elements);
-  // You can also manually check the last CUDA error code via:
-  // RuntimeDeviceCheck();
 }
 
 }  // namespace

python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh

@@ -0,0 +1,160 @@
+#include <sgl_kernel/fp8_utils.cuh>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/utils.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cub/block/block_reduce.cuh>
+#include <flashinfer/vec_dtypes.cuh>
+
+namespace {
+
+using device::atomicMaxFloat;
+using device::blockReduceMax;
+using device::FP8_E4M3_MAX;
+
+template <typename T>
+__global__ void
+per_tensor_absmax_kernel(const T* __restrict__ input, float* __restrict__ output_s, const int64_t num_elements) {
+  float max_value = 0.0f;
+  unsigned int tid = threadIdx.x;
+  unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int grid_size = blockDim.x * gridDim.x;
+
+  constexpr uint32_t vec_size = 16 / sizeof(T);
+  using vec_t = flashinfer::vec_t<T, vec_size>;
+
+  const int32_t num_vec_elems = num_elements / vec_size;
+
+  for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
+    vec_t input_vec;
+    input_vec.cast_load(input + i * vec_size);
+
+#pragma unroll
+    for (uint32_t j = 0; j < vec_size; ++j) {
+      float val = static_cast<float>(input_vec[j]);
+      max_value = fmaxf(max_value, fabsf(val));
+    }
+  }
+
+  const int32_t remaining_start = num_vec_elems * vec_size;
+  for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
+    float val = static_cast<float>(input[idx]);
+    max_value = fmaxf(max_value, fabsf(val));
+  }
+
+  max_value = blockReduceMax(max_value);
+
+  if (tid == 0) {
+    atomicMaxFloat(output_s, max_value / FP8_E4M3_MAX);
+  }
+}
+
+template <typename T, typename DST_DTYPE>
+__global__ void per_tensor_quant_fp8_kernel(
+    const T* __restrict__ input,
+    DST_DTYPE* __restrict__ output,
+    const float* __restrict__ scale,
+    const int64_t num_elements) {
+  const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int grid_size = blockDim.x * gridDim.x;
+  const float scale_val = 1.0f / (*scale);
+
+  const uint32_t VEC_SIZE = 16;
+  using vec_t = flashinfer::vec_t<T, VEC_SIZE>;
+
+  const int32_t num_vec_elems = num_elements / VEC_SIZE;
+
+  for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
+    vec_t input_vec;
+    input_vec.cast_load(input + i * VEC_SIZE);
+
+    DST_DTYPE output_arr[VEC_SIZE];
+#pragma unroll
+    for (uint32_t j = 0; j < VEC_SIZE; ++j) {
+      float val = fmax(fmin(static_cast<float>(input_vec[j]) * scale_val, FP8_E4M3_MAX), -FP8_E4M3_MAX);
+#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
+      output_arr[j] = static_cast<DST_DTYPE>(val);
+#else
+      output_arr[j] = c10::Float8_e4m3fnuz(
+          __hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
+          c10::Float8_e4m3fnuz::from_bits());
+#endif
+    }
+    *(uint4*)(output + i * VEC_SIZE) = *(uint4*)output_arr;
+  }
+
+  const int32_t remaining_start = num_vec_elems * VEC_SIZE;
+  for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
+    float val = fmax(-FP8_E4M3_MAX, fmin(static_cast<float>(input[idx]) * scale_val, FP8_E4M3_MAX));
+#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
+    output[idx] = static_cast<DST_DTYPE>(val);
+#else
+    output[idx] = c10::Float8_e4m3fnuz(
+        __hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
+        c10::Float8_e4m3fnuz::from_bits());
+#endif
+  }
+}
+
+constexpr size_t kBlockSize = 256;
+
+template <bool kIsStatic>
+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;
+
+  TensorMatcher({num_tokens, hidden_dim})  //
+      .with_dtype<float, __half, __nv_bfloat16>(input_dtype)
+      .with_device<kDLCUDA>(device_)
+      .verify(input);
+
+  TensorMatcher({num_tokens, hidden_dim})  //
+      .with_dtype<__nv_fp8_e4m3>()
+      .with_device<kDLCUDA>(device_)
+      .verify(output_q);
+
+  TensorMatcher({1})  //
+      .with_dtype<float>()
+      .with_device<kDLCUDA>(device_)
+      .verify(output_s);
+
+  const size_t total_elements = num_tokens.unwrap() * hidden_dim.unwrap();
+  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 = [&]<typename T>() {
+    if constexpr (!kIsStatic) {
+      LaunchKernel(num_blocks, kBlockSize, device)(
+          per_tensor_absmax_kernel<T>,
+          static_cast<const T*>(input.data_ptr()),
+          static_cast<float*>(output_s.data_ptr()),
+          static_cast<int64_t>(total_elements));
+    }
+
+    LaunchKernel(num_blocks, kBlockSize, device)(
+        per_tensor_quant_fp8_kernel<T, __nv_fp8_e4m3>,
+        static_cast<const T*>(input.data_ptr()),
+        static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
+        static_cast<const float*>(output_s.data_ptr()),
+        static_cast<int64_t>(total_elements));
+  };
+
+  const DLDataType dtype = input_dtype.unwrap();
+  if (dtype.code == kDLFloat && dtype.bits == 32) {
+    launch_kernels.template operator()<float>();
+  } else if (dtype.code == kDLBfloat && dtype.bits == 16) {
+    launch_kernels.template operator()<__nv_bfloat16>();
+  } else if (dtype.code == kDLFloat && dtype.bits == 16) {
+    launch_kernels.template operator()<__half>();
+  }
+}
+
+}  // namespace

python/sglang/jit_kernel/include/sgl_kernel/fp8_utils.cuh

@@ -0,0 +1,13 @@
+#pragma once
+
+#ifdef __CUDACC__
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#endif
+
+namespace device {
+
+inline constexpr float FP8_E4M3_MAX = 448.0f;
+
+}  // namespace device

python/sglang/jit_kernel/include/sgl_kernel/tensor.h

@@ -23,6 +23,7 @@
 #ifdef __CUDACC__
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
+#include <cuda_fp8.h>
 #endif
 
 namespace host {
@@ -65,6 +66,10 @@ template <>
 struct dtype_trait<__nv_bfloat16> {
   inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1};
 };
+template <>
+struct dtype_trait<__nv_fp8_e4m3> {
+  inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat8_e4m3fn, .bits = 8, .lanes = 1};
+};
 #endif
 
 template <DLDeviceType Code>

python/sglang/jit_kernel/include/sgl_kernel/utils.cuh

@@ -12,6 +12,49 @@
 namespace device {
 
 inline constexpr auto kWarpThreads = 32u;
+inline constexpr auto kFullMask = 0xffffffffu;
+
+__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
+#ifndef USE_ROCM
+  float old;
+  old = (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
+                     : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value)));
+  return old;
+#else
+  int* addr_as_i = (int*)addr;
+  int old = *addr_as_i, assumed;
+  do {
+    assumed = old;
+    old = atomicCAS(addr_as_i, assumed, __float_as_int(fmaxf(value, __int_as_float(assumed))));
+  } while (assumed != old);
+  return __int_as_float(old);
+#endif
+}
+
+__device__ __forceinline__ float warpReduceMax(float value) {
+  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 16));
+  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 8));
+  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 4));
+  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 2));
+  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 1));
+  return value;
+}
+
+__device__ __forceinline__ float blockReduceMax(float value) {
+  static __shared__ float warpLevelMaxs[kWarpThreads];
+  const int laneId = threadIdx.x % kWarpThreads;
+  const int warpId = threadIdx.x / kWarpThreads;
+
+  value = warpReduceMax(value);
+
+  if (laneId == 0) warpLevelMaxs[warpId] = value;
+  __syncthreads();
+
+  value = (threadIdx.x < blockDim.x / kWarpThreads) ? warpLevelMaxs[laneId] : 0;
+  if (warpId == 0) value = warpReduceMax(value);
+
+  return value;
+}
 
 namespace pointer {
 

python/sglang/jit_kernel/include/sgl_kernel/utils.h

@@ -2,6 +2,9 @@
 
 // ref: https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3
 #ifdef __CUDACC__
+#include <cuda.h>
+#if CUDA_VERSION <= 12010
+
 #pragma push_macro("__cpp_consteval")
 #pragma push_macro("_NODISCARD")
 #pragma push_macro("__builtin_LINE")
@@ -23,7 +26,10 @@
 #undef consteval
 #pragma pop_macro("__cpp_consteval")
 #pragma pop_macro("_NODISCARD")
-#else
+#else  // __CUDACC__ && CUDA_VERSION > 12010
+#include <source_location>
+#endif
+#else  // no __CUDACC__
 #include <source_location>
 #endif
 
@@ -33,7 +39,6 @@
 #include <cstddef>
 #include <ostream>
 #include <ranges>
-#include <source_location>
 #include <sstream>
 #include <utility>
 

python/sglang/jit_kernel/per_tensor_quant_fp8.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import functools
+import os
+from typing import TYPE_CHECKING
+
+import flashinfer
+import torch
+from torch.utils.cpp_extension import CUDA_HOME
+
+from sglang.jit_kernel.utils import load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@functools.cache
+def _jit_per_tensor_quant_fp8_module(is_static: bool) -> Module:
+    args = make_cpp_args(is_static)
+
+    flashinfer_include = os.path.join(
+        os.path.dirname(flashinfer.__file__), "data", "include"
+    )
+    cub_include = os.path.join(CUDA_HOME, "include")
+
+    return load_jit(
+        "per_tensor_quant_fp8",
+        *args,
+        cuda_files=["gemm/per_tensor_quant_fp8.cuh"],
+        cuda_wrappers=[("per_tensor_quant_fp8", f"per_tensor_quant_fp8<{args}>")],
+        extra_include_paths=[flashinfer_include, cub_include],
+    )
+
+
+def per_tensor_quant_fp8(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+    is_static: bool = False,
+) -> None:
+    """
+    Per-tensor quantization to FP8 format.
+
+    Args:
+        input: Input tensor to quantize (float, half, or bfloat16)
+        output_q: Output quantized tensor (fp8_e4m3)
+        output_s: Output scale tensor (float scalar)
+        is_static: If True, assumes scale is pre-computed and skips absmax computation
+    """
+    module = _jit_per_tensor_quant_fp8_module(is_static)
+    module.per_tensor_quant_fp8(input, output_q, output_s)

python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py

@@ -0,0 +1,70 @@
+import itertools
+from typing import Optional, Tuple
+
+import pytest
+import torch
+
+from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+
+try:
+    from sglang.srt.utils import is_hip
+
+    _is_hip = is_hip()
+except ImportError:
+    _is_hip = False
+
+fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+
+
+def sglang_scaled_fp8_quant(
+    input: torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    fp8_type_: torch.dtype = torch.float8_e4m3fn
+    output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
+    is_static = True
+    if scale is None:
+        scale = torch.zeros(1, device=input.device, dtype=torch.float32)
+        is_static = False
+    per_tensor_quant_fp8(input, output, scale, is_static)
+
+    return output, scale
+
+
+def torch_scaled_fp8_quant(tensor, inv_scale):
+    finfo = torch.finfo(torch.float8_e4m3fn)
+    scale = inv_scale.reciprocal()
+    qweight = (tensor.to(torch.float32) * scale).clamp(min=finfo.min, max=finfo.max)
+    qweight = qweight.to(torch.float8_e4m3fn)
+    return qweight
+
+
+@pytest.mark.parametrize(
+    "num_tokens,hidden_dim",
+    list(itertools.product([128, 256, 512], [512, 2048, 4096])),
+)
+def test_per_tensor_quant_compare_implementations(
+    num_tokens: int,
+    hidden_dim: int,
+):
+    device = torch.device("cuda")
+    x = torch.rand((num_tokens, hidden_dim), dtype=torch.float16, device=device)
+
+    sglang_out, sglang_scale = sglang_scaled_fp8_quant(x)
+    torch_out = torch_scaled_fp8_quant(x, sglang_scale)
+
+    torch.testing.assert_close(
+        sglang_out.float(), torch_out.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)
+
+    torch.testing.assert_close(
+        sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])