diff --git a/python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py b/python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py new file mode 100644 index 000000000..870057a22 --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/csrc/add_constant.cuh b/python/sglang/jit_kernel/csrc/add_constant.cuh index 33f37d119..d4e5bed67 100644 --- a/python/sglang/jit_kernel/csrc/add_constant.cuh +++ b/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(dst.data_ptr()), static_cast(src.data_ptr()), num_elements); - // You can also manually check the last CUDA error code via: - // RuntimeDeviceCheck(); } } // namespace 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 new file mode 100644 index 000000000..d32171a04 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh @@ -0,0 +1,160 @@ +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace { + +using device::atomicMaxFloat; +using device::blockReduceMax; +using device::FP8_E4M3_MAX; + +template +__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; + + 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(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(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 +__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; + + 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(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(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(input[idx]) * scale_val, FP8_E4M3_MAX)); +#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3) + output[idx] = static_cast(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 +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(input_dtype) + .with_device(device_) + .verify(input); + + TensorMatcher({num_tokens, hidden_dim}) // + .with_dtype<__nv_fp8_e4m3>() + .with_device(device_) + .verify(output_q); + + TensorMatcher({1}) // + .with_dtype() + .with_device(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 = [&]() { + if constexpr (!kIsStatic) { + LaunchKernel(num_blocks, kBlockSize, device)( + per_tensor_absmax_kernel, + static_cast(input.data_ptr()), + static_cast(output_s.data_ptr()), + static_cast(total_elements)); + } + + LaunchKernel(num_blocks, kBlockSize, device)( + per_tensor_quant_fp8_kernel, + static_cast(input.data_ptr()), + static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + static_cast(total_elements)); + }; + + const DLDataType dtype = input_dtype.unwrap(); + if (dtype.code == kDLFloat && dtype.bits == 32) { + launch_kernels.template operator()(); + } 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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/fp8_utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/fp8_utils.cuh new file mode 100644 index 000000000..3989b9443 --- /dev/null +++ b/python/sglang/jit_kernel/include/sgl_kernel/fp8_utils.cuh @@ -0,0 +1,13 @@ +#pragma once + +#ifdef __CUDACC__ +#include +#include +#include +#endif + +namespace device { + +inline constexpr float FP8_E4M3_MAX = 448.0f; + +} // namespace device diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h index 9eb7bc6d5..b577d8f48 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h +++ b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h @@ -23,6 +23,7 @@ #ifdef __CUDACC__ #include #include +#include #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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index 2f1051288..14a90aa9d 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/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 { diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.h b/python/sglang/jit_kernel/include/sgl_kernel/utils.h index f76c35bfd..3f59349a5 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.h +++ b/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 +#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 +#endif +#else // no __CUDACC__ #include #endif @@ -33,7 +39,6 @@ #include #include #include -#include #include #include diff --git a/python/sglang/jit_kernel/per_tensor_quant_fp8.py b/python/sglang/jit_kernel/per_tensor_quant_fp8.py new file mode 100644 index 000000000..9a63357ca --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/test_add_constant.py b/python/sglang/jit_kernel/tests/test_add_constant.py similarity index 100% rename from python/sglang/jit_kernel/test_add_constant.py rename to python/sglang/jit_kernel/tests/test_add_constant.py 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 new file mode 100644 index 000000000..1c1c2044c --- /dev/null +++ b/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__]) diff --git a/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py b/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py index ead9e9aa1..d44b017a0 100644 --- a/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py +++ b/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py @@ -90,14 +90,24 @@ else: 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=["vllm", "sglang"], - line_names=["VLLM", "SGL Kernel"], - styles=[("blue", "-"), ("green", "-")], + line_vals=line_vals, + line_names=line_names, + styles=styles, ylabel="us", plot_name="per-tensor-quant-fp8-performance", args={},