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

[jit kernel] Support per_token_group_quant_8bit jit kernel (#18905)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-02-20 21:01:05 +08:00
committed by GitHub
parent 38a69652e6
commit 7d953440ec
6 changed files with 811 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

290
python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py Normal file
View File

@@ -0,0 +1,290 @@
import itertools
import os
from typing import Any, Dict, List
import torch
import triton
from sgl_kernel.test_utils import create_per_token_group_quant_test_data
from sglang.jit_kernel.benchmark.utils import (
get_benchmark_range,
)
from sglang.jit_kernel.per_token_group_quant_8bit import (
per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
)
from sglang.srt.layers.quantization.fp8_kernel import (
create_per_token_group_quant_fp8_output_scale,
)
from sglang.srt.layers.quantization.fp8_kernel import (
per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
)
from sglang.srt.utils import is_hip
from sglang.srt.utils.bench_utils import bench_kineto
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
_is_hip = is_hip()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
NUM_TESTS = 300 if IS_CI else 30
GROUP_SIZE_RANGE = [128]
DST_DTYPE_RANGE = [fp8_type_]
# ---- GEMM-like branch (num_ranks=None) ----
NUM_TOKENS_RANGE_GEMM = get_benchmark_range(
full_range=[1, 4, 16, 64, 256, 768, 2048, 8192, 16384],
ci_range=[768],
)
HIDDEN_DIM_RANGE_GEMM = [1536, 7168, 16384]
NUM_RANKS_RANGE_GEMM = [None]
FLAGS_GEMM_FULL: List[Dict[str, Any]] = [
dict(
column_major_scales=False,
scale_tma_aligned=False,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=False,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
]
FLAGS_GEMM_CI: List[Dict[str, Any]] = [
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
]
FLAGS_RANGE_GEMM = get_benchmark_range(
full_range=FLAGS_GEMM_FULL, ci_range=FLAGS_GEMM_CI
)
CONFIGS_GEMM = list(
itertools.product(
NUM_TOKENS_RANGE_GEMM,
HIDDEN_DIM_RANGE_GEMM,
GROUP_SIZE_RANGE,
NUM_RANKS_RANGE_GEMM,
DST_DTYPE_RANGE,
FLAGS_RANGE_GEMM,
)
)
# ---- MoE-like / multi-rank branch (hidden_dim=2048, num_ranks in {8,16,32,48}) ----
NUM_TOKENS_RANGE_MOE = get_benchmark_range(
full_range=[1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
ci_range=[768 * 8],
)
HIDDEN_DIM_RANGE_MOE = [2048]
NUM_RANKS_RANGE_MOE = get_benchmark_range(
full_range=[8, 16, 32, 48],
ci_range=[48],
)
FLAGS_MOE: List[Dict[str, Any]] = [
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="balanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="imbalanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="extreme",
),
]
FLAGS_RANGE_MOE = get_benchmark_range(full_range=FLAGS_MOE, ci_range=FLAGS_MOE)
CONFIGS_MOE = list(
itertools.product(
NUM_TOKENS_RANGE_MOE,
HIDDEN_DIM_RANGE_MOE,
GROUP_SIZE_RANGE,
NUM_RANKS_RANGE_MOE,
DST_DTYPE_RANGE,
FLAGS_RANGE_MOE,
)
)
# ---- Final configs ----
CONFIGS = CONFIGS_GEMM + CONFIGS_MOE
LINE_VALS = ["triton", "sglang"]
LINE_NAMES = ["Triton (Inaccurate)", "SGL Kernel"]
STYLES = [("blue", "-"), ("green", "-")]
def _flatten_to_2d(t: torch.Tensor) -> torch.Tensor:
"""Reshape a tensor with 3+ dims to 2D by merging all leading dims."""
if t.ndim <= 2:
return t
return t.reshape(-1, t.shape[-1])
def _make_sglang_bench_fn(
x: torch.Tensor,
group_size: int,
dst_dtype: torch.dtype,
flags: dict,
):
"""
Adapter that pre-allocates output tensors and returns a zero-arg callable
matching the JIT kernel's signature.
The JIT kernel does not support fuse_silu_and_mul, so when enabled we
pre-compute silu+mul on the input. bench_kineto only times the kernel
matching the given name, so the pre-processing is not included.
The JIT kernel expects 2D tensors, so any higher-dimensional inputs
(e.g. from masked_layout_mode) are flattened to 2D.
"""
fuse_silu_and_mul = flags.get("fuse_silu_and_mul", False)
column_major_scales = flags.get("column_major_scales", False)
scale_tma_aligned = flags.get("scale_tma_aligned", False)
scale_ue8m0 = flags.get("scale_ue8m0", False)
# JIT kernel does not support fuse_silu_and_mul; pre-compute it
if fuse_silu_and_mul:
half = x.shape[-1] // 2
x_input = torch.nn.functional.silu(x[..., :half]) * x[..., half:]
else:
x_input = x
# JIT kernel expects 2D (num_tokens, hidden_dim); flatten if needed
x_input = _flatten_to_2d(x_input.contiguous())
out_shape = x_input.shape
output_q = torch.empty(out_shape, device=x.device, dtype=dst_dtype)
fp8_max = torch.finfo(dst_dtype).max
fp8_min = -fp8_max
output_s = create_per_token_group_quant_fp8_output_scale(
x_shape=out_shape,
device=x.device,
group_size=group_size,
column_major_scales=column_major_scales,
scale_tma_aligned=scale_tma_aligned,
scale_ue8m0=scale_ue8m0,
)
def _run():
sglang_per_token_group_quant_8bit(
input=x_input,
output_q=output_q,
output_s=output_s,
group_size=group_size,
eps=1e-10,
fp8_min=fp8_min,
fp8_max=fp8_max,
scale_ue8m0=scale_ue8m0,
)
return _run
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=[
"num_tokens",
"hidden_dim",
"group_size",
"num_ranks",
"dst_dtype",
"flags",
],
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
# Triton has multi kernels and we only report the time for the core one
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="per-token-group-quant-8bit-performance",
args={},
)
)
def benchmark(
num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags, provider
):
print(
f"Testing: {num_tokens=} {hidden_dim=} {group_size=} {num_ranks=} {dst_dtype=} {flags=} {provider=}"
)
x, masked_m = create_per_token_group_quant_test_data(
num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags
)
if provider == "triton":
fn = triton_per_token_group_quant_8bit
kernel_names = "_per_token_group_quant_8bit|_silu_and_mul_post_quant_kernel"
bench_fn = lambda: fn(
x=x,
masked_m=masked_m,
group_size=group_size,
dst_dtype=dst_dtype,
**{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
)
elif provider == "sglang":
kernel_names = "per_token_group_quant_8bit_kernel"
bench_fn = _make_sglang_bench_fn(
x=x,
group_size=group_size,
dst_dtype=dst_dtype,
flags=flags,
)
else:
raise ValueError(f"Unknown provider: {provider}")
time_s = bench_kineto(bench_fn, kernel_names=kernel_names, num_tests=NUM_TESTS)
return time_s * 1e6
if __name__ == "__main__":
benchmark.run(print_data=True)

218
python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh Normal file
View File

@@ -0,0 +1,218 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/atomic.cuh>
#include <sgl_kernel/cta.cuh>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <cstddef>
#include <cstdint>
namespace {
constexpr int kThreadsPerGroup = 16;
__device__ __forceinline__ float GroupReduceMax(float val, const int tid) {
unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;
val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
val = fmaxf(val, __shfl_xor_sync(mask, val, 4));
val = fmaxf(val, __shfl_xor_sync(mask, val, 2));
val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
return val;
}
template <bool kScaleUE8M0>
using scale_packed_t_t = std::conditional_t<kScaleUE8M0, uint32_t, float>;
template <bool kScaleUE8M0>
using scale_element_t_t = std::conditional_t<kScaleUE8M0, uint8_t, float>;
template <typename T, typename DST_DTYPE, bool kIsColumnMajor, bool kScaleUE8M0>
__global__ void per_token_group_quant_8bit_kernel(
const T* __restrict__ input,
DST_DTYPE* __restrict__ output_q,
scale_packed_t_t<kScaleUE8M0>* __restrict__ output_s,
const int group_size,
const int num_groups,
const int groups_per_block,
const float eps,
const float min_8bit,
const float max_8bit,
const int num_groups_per_row = 0,
const int scale_stride = 0) {
using namespace device;
namespace math = device::math;
(void)num_groups;
const int local_group_id = static_cast<int>(threadIdx.x / kThreadsPerGroup);
const int lane_id = threadIdx.x % kThreadsPerGroup;
const int64_t block_group_id = blockIdx.x * groups_per_block;
const int64_t global_group_id = block_group_id + local_group_id;
const int64_t block_group_offset = global_group_id * group_size;
float local_absmax = eps;
using scale_packed_t = scale_packed_t_t<kScaleUE8M0>;
using scale_element_t = scale_element_t_t<kScaleUE8M0>;
static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
const T* group_input = input + block_group_offset;
DST_DTYPE* group_output = static_cast<DST_DTYPE*>(output_q) + block_group_offset;
scale_element_t* scale_output = nullptr;
if constexpr (kIsColumnMajor) {
constexpr int kElemsPerPack = static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
const int row_idx = global_group_id / num_groups_per_row;
const int col_idx_unpacked = global_group_id % num_groups_per_row;
const int col_idx = col_idx_unpacked / kElemsPerPack;
const int pack_idx = col_idx_unpacked % kElemsPerPack;
scale_output = reinterpret_cast<scale_element_t*>(output_s) +
(col_idx * scale_stride * kElemsPerPack + row_idx * kElemsPerPack + pack_idx);
} else {
static_assert(!kScaleUE8M0);
scale_output = output_s + global_group_id;
}
constexpr uint32_t kVecSize = 16 / sizeof(T);
using vec_t = AlignedVector<T, kVecSize>;
const auto gmem_in = tile::Memory<vec_t>::thread();
const int32_t num_vec_elems = group_size / kVecSize;
for (int32_t i = lane_id; i < num_vec_elems; i += kThreadsPerGroup) {
const vec_t input_vec = gmem_in.load(group_input, i);
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
const float val = static_cast<float>(input_vec[j]);
local_absmax = math::max(local_absmax, math::abs(val));
}
}
local_absmax = GroupReduceMax(local_absmax, lane_id);
float y_s = local_absmax / max_8bit;
if constexpr (kScaleUE8M0) {
y_s = exp2f(ceilf(log2f(math::max(y_s, 1e-10f))));
}
scale_element_t y_s_quant;
if constexpr (kScaleUE8M0) {
y_s_quant = static_cast<uint8_t>(((int)log2f(y_s)) + 127);
} else {
y_s_quant = y_s;
}
if (lane_id == 0) {
*scale_output = y_s_quant;
}
for (int32_t i = lane_id; i < num_vec_elems; i += kThreadsPerGroup) {
const vec_t input_vec = gmem_in.load(group_input, i);
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
const float val = static_cast<float>(input_vec[j]);
const float q_val = math::min(math::max(val / y_s, min_8bit), max_8bit);
group_output[i * kVecSize + j] = DST_DTYPE(q_val);
}
}
}
inline int compute_groups_per_block(int64_t num_groups) {
if (num_groups % 16 == 0) return 16;
if (num_groups % 8 == 0) return 8;
if (num_groups % 4 == 0) return 4;
if (num_groups % 2 == 0) return 2;
return 1;
}
template <typename DType, typename OutType>
void per_token_group_quant_8bit(
tvm::ffi::TensorView input,
tvm::ffi::TensorView output_q,
tvm::ffi::TensorView output_s,
int64_t group_size,
double eps,
double min_8bit,
double max_8bit,
bool scale_ue8m0) {
using namespace host;
auto device = SymbolicDevice{};
auto M = SymbolicSize{"num_tokens"};
auto K = SymbolicSize{"hidden_dim"};
device.set_options<kDLCUDA>();
TensorMatcher({M, K}).with_dtype<DType>().with_device(device).verify(input);
TensorMatcher({M, K}).with_dtype<OutType>().with_device(device).verify(output_q);
const auto num_tokens = M.unwrap();
const auto hidden_dim = K.unwrap();
const int64_t num_groups_per_row = hidden_dim / group_size;
const int64_t num_groups = num_tokens * num_groups_per_row;
const int groups_per_block = compute_groups_per_block(num_groups);
const int num_blocks = num_groups / groups_per_block;
const int num_threads = groups_per_block * kThreadsPerGroup;
const bool is_column_major = output_s.stride(0) < output_s.stride(1);
const int scale_stride = output_s.stride(1);
const float feps = static_cast<float>(eps);
const float fmin8 = static_cast<float>(min_8bit);
const float fmax8 = static_cast<float>(max_8bit);
if (is_column_major) {
if (scale_ue8m0) {
LaunchKernel(num_blocks, num_threads, input.device())(
per_token_group_quant_8bit_kernel<DType, OutType, true, true>,
static_cast<const DType*>(input.data_ptr()),
static_cast<OutType*>(output_q.data_ptr()),
static_cast<uint32_t*>(output_s.data_ptr()),
static_cast<int>(group_size),
static_cast<int>(num_groups),
static_cast<int>(groups_per_block),
feps,
fmin8,
fmax8,
static_cast<int>(num_groups_per_row),
scale_stride);
} else {
LaunchKernel(num_blocks, num_threads, input.device())(
per_token_group_quant_8bit_kernel<DType, OutType, true, false>,
static_cast<const DType*>(input.data_ptr()),
static_cast<OutType*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
static_cast<int>(group_size),
static_cast<int>(num_groups),
static_cast<int>(groups_per_block),
feps,
fmin8,
fmax8,
static_cast<int>(num_groups_per_row),
scale_stride);
}
} else {
LaunchKernel(num_blocks, num_threads, input.device())(
per_token_group_quant_8bit_kernel<DType, OutType, false, false>,
static_cast<const DType*>(input.data_ptr()),
static_cast<OutType*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
static_cast<int>(group_size),
static_cast<int>(num_groups),
static_cast<int>(groups_per_block),
feps,
fmin8,
fmax8,
0,
0);
}
}
} // namespace

67
python/sglang/jit_kernel/per_token_group_quant_8bit.py Normal file
View File

@@ -0,0 +1,67 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
from sglang.jit_kernel.utils import CPP_DTYPE_MAP as OUTPUT_DTYPE_MAP
@cache_once
def _jit_per_token_group_quant_8bit_module(
dtype: torch.dtype, output_type: torch.dtype
) -> Module:
input_args = make_cpp_args(dtype)
out_cpp = OUTPUT_DTYPE_MAP[output_type]
return load_jit(
"per_token_group_quant_8bit",
cuda_files=["gemm/per_token_group_quant_8bit.cuh"],
cuda_wrappers=[
(
"per_token_group_quant_8bit",
f"per_token_group_quant_8bit<{input_args}, {out_cpp}>",
)
],
)
def per_token_group_quant_8bit(
input: torch.Tensor,
output_q: torch.Tensor,
output_s: torch.Tensor,
group_size: int,
eps: float,
fp8_min: float,
fp8_max: float,
scale_ue8m0: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Per-token-group quantization to 8-bit format.
Args:
input: Input tensor to quantize (float, half, or bfloat16).
output_q: Output quantized tensor (e.g., fp8_e4m3 or int8).
output_s: Output scale tensor.
group_size: The size of the group for quantization.
eps: A small value to avoid division by zero.
fp8_min: The minimum value of the 8-bit data type.
fp8_max: The maximum value of the 8-bit data type.
scale_ue8m0: Whether to use UE8M0 format for scales.
"""
module = _jit_per_token_group_quant_8bit_module(input.dtype, output_q.dtype)
module.per_token_group_quant_8bit(
input,
output_q,
output_s,
group_size,
eps,
fp8_min,
fp8_max,
scale_ue8m0,
)
return output_q, output_s

205
python/sglang/jit_kernel/tests/test_per_token_group_quant_8bit.py Normal file
View File

@@ -0,0 +1,205 @@
import itertools
import pytest
import torch
from sglang.srt.utils import is_hip
_is_hip = is_hip()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
from sgl_kernel.test_utils import (
assert_all_close_or_tiny_diff,
create_per_token_group_quant_test_data,
)
from sglang.jit_kernel.per_token_group_quant_8bit import (
per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
)
from sglang.srt.layers.quantization.fp8_kernel import (
create_per_token_group_quant_fp8_output_scale,
)
from sglang.srt.layers.quantization.fp8_kernel import (
per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
)
configs = list(
itertools.product(
[1, 4, 16, 64, 127, 128, 512, 1024, 4096, 8192], # num_tokens
[128, 256, 384, 512, 1024, 1536, 1664, 2048, 4096, 7168, 16384], # hidden_dim
[16, 32, 64, 128], # group_size
[None], # num_ranks
[fp8_type_], # dtype
[
dict(
column_major_scales=False,
scale_tma_aligned=False,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=False,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
],
)
) + list(
itertools.product(
[1, 4, 1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
[2048],
[128],
[8, 16, 32, 48],
[fp8_type_],
[
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="balanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="imbalanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="extreme",
),
],
)
)
@pytest.mark.parametrize(
"num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags", configs
)
def test_per_token_group_quant_with_column_major(
num_tokens,
hidden_dim,
group_size,
num_ranks,
dst_dtype,
flags,
):
arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
if flags["scale_ue8m0"] and (arch_major <= 9):
pytest.skip("Only Blackwell need ue8m0 fusion")
return
if (flags["scale_ue8m0"] and (group_size != 128)) or (
(dst_dtype == torch.int8) and flags["column_major_scales"]
):
pytest.skip()
return
x, masked_m = create_per_token_group_quant_test_data(
num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags
)
execute_kwargs = dict(
x=x,
masked_m=masked_m,
group_size=group_size,
eps=1e-10,
dst_dtype=dst_dtype,
**{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
)
def _postprocess(x_q, x_s):
if masked_m is not None:
print(f"Mask tokens after {masked_m} to be zero")
for i in range(len(masked_m)):
x_q[i, masked_m[i] :, :] = 0
x_s[i, masked_m[i] :, :] = 0
return x_q, x_s
x_q_triton, x_s_triton = _postprocess(
*triton_per_token_group_quant_8bit(**execute_kwargs)
)
fuse_silu_and_mul = False
out_shape = (*x.shape[:-1], x.shape[-1] // (2 if fuse_silu_and_mul else 1))
fp8_dtype = torch.float8_e4m3fn
fp8_max = torch.finfo(fp8_dtype).max
fp8_min = -fp8_max
x_q = torch.empty(out_shape, device=x.device, dtype=fp8_dtype)
x_s = create_per_token_group_quant_fp8_output_scale(
x_shape=out_shape,
device=x.device,
group_size=group_size,
column_major_scales=False,
scale_tma_aligned=False,
scale_ue8m0=False,
)
execute_kwargs = dict(
input=x,
output_q=x_q,
output_s=x_s,
group_size=group_size,
eps=1e-10,
fp8_max=fp8_max,
fp8_min=fp8_min,
)
x_q_sglang, x_s_sglang = _postprocess(
*sglang_per_token_group_quant_8bit(**execute_kwargs)
)
try:
assert_all_close_or_tiny_diff(x_q_triton, x_q_sglang)
torch.testing.assert_close(
x_s_triton.contiguous(),
x_s_sglang.contiguous(),
rtol=1e-3,
atol=1e-5,
msg=lambda message: message + f" {x_s_triton=} {x_s_sglang=}",
)
except AssertionError:
print(
f"{x.shape=} {x_q_triton.shape=} {x_s_triton.shape=} {x_q_sglang.shape=} {x_s_sglang.shape=}"
)
print(f"{x=}")
print(f"{masked_m=}")
print(f"{x_q_triton=}")
print(f"{x_s_triton=}")
print(f"{x_q_sglang=}")
print(f"{x_s_sglang=}")
raise
if __name__ == "__main__":
pytest.main([__file__])

3
python/sglang/jit_kernel/utils.py
View File

@@ -10,7 +10,6 @@ import torch
if TYPE_CHECKING:
from tvm_ffi import Module
F = TypeVar("F", bound=Callable[..., Any])
@@ -73,7 +72,9 @@ class CPPArgList(list[str]):
CPP_DTYPE_MAP = {
torch.float: "fp32_t",
torch.float16: "fp16_t",
torch.float8_e4m3fn: "fp8_e4m3_t",
torch.bfloat16: "bf16_t",
torch.int8: "int8_t",
}

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

@@ -63,6 +63,10 @@ if _is_cuda:
enable_sgl_per_token_group_quant_8bit = False
from sglang.jit_kernel.per_token_group_quant_8bit import (
per_token_group_quant_8bit as sgl_per_token_group_quant_8bit_jit,
)
if _is_hip:
_has_vllm = False
if _use_aiter:
@@ -501,19 +505,31 @@ def sglang_per_token_group_quant_fp8(
if x.shape[0] > 0:
# Temporary
if enable_sgl_per_token_group_quant_8bit:
sgl_per_token_group_quant_8bit(
x,
x_q,
x_s,
group_size,
eps,
fp8_min,
fp8_max,
scale_ue8m0,
fuse_silu_and_mul,
masked_m,
enable_v2=enable_v2,
)
if enable_v2:
sgl_per_token_group_quant_8bit(
x,
x_q,
x_s,
group_size,
eps,
fp8_min,
fp8_max,
scale_ue8m0,
fuse_silu_and_mul,
masked_m,
enable_v2=True,
)
else:
sgl_per_token_group_quant_8bit_jit(
input=x,
output_q=x_q,
output_s=x_s,
group_size=group_size,
eps=eps,
fp8_min=fp8_min,
fp8_max=fp8_max,
scale_ue8m0=scale_ue8m0,
)
else:
assert not enable_v2
sgl_per_token_group_quant_fp8(