[DSV4] perf: Make FP8 quant output tensor contiguous (#27926)
Co-authored-by: liqichao <liqichao@baidu.com> Co-authored-by: chenbong <bhchen@stu.xmu.edu.cn>
This commit is contained in:
Kaixi
co-authored by
liqichao
chenbong
liqichao
chenbong
parent
b363249423
commit
d7dcdf3efd
@@ -0,0 +1,169 @@
|
||||
// DeepSeek-V4 wo_a activation quantization for DeepGEMM fp8_einsum.
|
||||
//
|
||||
// This is intentionally narrower than the generic per_token_group_quant_8bit_v2
|
||||
// kernel: input is a [T, G, D] view with contiguous hidden groups, output_q is
|
||||
// contiguous [T, G, D], group_size is fixed to 128, scales are fp32 UE8M0
|
||||
// power-of-two values, and output_s is a logical [T, G, D/128] view backed by
|
||||
// group-major [G, T, D/128] storage.
|
||||
//
|
||||
// The generic kernel cannot read the strided DSV4 view while producing
|
||||
// contiguous [T, G, D] codes and group-major scales without an extra full-tensor
|
||||
// copy.
|
||||
#include <sgl_kernel/tensor.h> // TensorMatcher, SymbolicSize/Device
|
||||
#include <sgl_kernel/utils.h> // RuntimeCheck
|
||||
|
||||
#include <sgl_kernel/utils.cuh> // fp8 aliases, PDL helpers
|
||||
#include <sgl_kernel/warp.cuh> // warp::reduce_max
|
||||
|
||||
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh> // UE8M0 and FP8 helpers
|
||||
|
||||
#include <tvm/ffi/container/tensor.h> // tvm::ffi::TensorView
|
||||
|
||||
#include <cstdint>
|
||||
#include <cuda_fp8.h>
|
||||
|
||||
namespace {
|
||||
|
||||
using deepseek_v4::fp8::cast_to_ue8m0;
|
||||
using deepseek_v4::fp8::inv_scale_ue8m0;
|
||||
using deepseek_v4::fp8::pack_fp8;
|
||||
|
||||
constexpr float LOCAL_ABSMAX_ABS = 1e-10f;
|
||||
constexpr uint32_t GROUP_SIZE = 128;
|
||||
constexpr uint32_t THREADS_PER_GROUP = 8;
|
||||
constexpr uint32_t SUBWARPS_PER_BLOCK = 16;
|
||||
constexpr uint32_t INPUT_VEC_NUM_BYTES = 32;
|
||||
constexpr uint32_t INPUT_INT4_SIZE = INPUT_VEC_NUM_BYTES / sizeof(int4);
|
||||
|
||||
template <int THREADS_PER_SUBWARP>
|
||||
SGL_DEVICE float GroupReduceMax(float val) {
|
||||
static_assert(
|
||||
(THREADS_PER_SUBWARP & (THREADS_PER_SUBWARP - 1)) == 0 && THREADS_PER_SUBWARP <= 16 && THREADS_PER_SUBWARP >= 1,
|
||||
"THREADS_PER_SUBWARP must be 1, 2, 4, 8, or 16");
|
||||
// Tail subwarps can be inactive at the bounds check, so reduce with only the
|
||||
// current subgroup's lanes rather than a full-warp mask.
|
||||
constexpr device::warp::mask_t kSub = (device::warp::mask_t{1} << THREADS_PER_SUBWARP) - 1;
|
||||
const device::warp::mask_t mask = kSub << (THREADS_PER_SUBWARP * ((threadIdx.x % 32) / THREADS_PER_SUBWARP));
|
||||
return device::warp::reduce_max<THREADS_PER_SUBWARP>(val, mask);
|
||||
}
|
||||
|
||||
template <typename T, bool kUsePDL>
|
||||
__global__ void fp8_wo_a_group_major_quant_ue8m0_kernel(
|
||||
const T* __restrict__ input,
|
||||
fp8_e4m3_t* __restrict__ output_q,
|
||||
float* __restrict__ output_s,
|
||||
int64_t total_scale_groups,
|
||||
int64_t num_tokens,
|
||||
int hidden_dim_groups,
|
||||
int num_outer_groups,
|
||||
int64_t input_stride_t) {
|
||||
device::PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
const int64_t subwarp_id = threadIdx.x / THREADS_PER_GROUP;
|
||||
const int lane_id = threadIdx.x % THREADS_PER_GROUP;
|
||||
const int64_t group_id = static_cast<int64_t>(blockIdx.x) * SUBWARPS_PER_BLOCK + subwarp_id;
|
||||
if (group_id < total_scale_groups) {
|
||||
const int hidden_group = group_id % hidden_dim_groups;
|
||||
const int64_t token_outer = group_id / hidden_dim_groups;
|
||||
const int outer_idx = token_outer % num_outer_groups;
|
||||
const int64_t token_idx = token_outer / num_outer_groups;
|
||||
|
||||
constexpr uint32_t INPUT_VEC_SIZE = INPUT_VEC_NUM_BYTES / sizeof(T);
|
||||
static_assert(INPUT_VEC_SIZE * THREADS_PER_GROUP == GROUP_SIZE);
|
||||
|
||||
const int64_t input_group_start_offset =
|
||||
token_idx * input_stride_t + outer_idx * GROUP_SIZE * hidden_dim_groups + hidden_group * GROUP_SIZE;
|
||||
const int64_t output_group_start_offset = group_id * GROUP_SIZE;
|
||||
|
||||
int4 input_int4[INPUT_INT4_SIZE];
|
||||
T* input_vec = reinterpret_cast<T*>(input_int4);
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < INPUT_INT4_SIZE; ++j) {
|
||||
input_int4[j] = reinterpret_cast<const int4*>(input + input_group_start_offset + lane_id * INPUT_VEC_SIZE)[j];
|
||||
}
|
||||
|
||||
float local_absmax = LOCAL_ABSMAX_ABS;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < INPUT_VEC_SIZE; ++j) {
|
||||
const float val = static_cast<float>(input_vec[j]);
|
||||
local_absmax = fmaxf(local_absmax, fabsf(val));
|
||||
}
|
||||
|
||||
local_absmax = GroupReduceMax<THREADS_PER_GROUP>(local_absmax);
|
||||
|
||||
constexpr float kFp8MaxInv = 1.0f / kFP8E4M3Max;
|
||||
const int32_t scale_ue8m0 = cast_to_ue8m0(local_absmax * kFp8MaxInv);
|
||||
const float y_scale = inv_scale_ue8m0(scale_ue8m0);
|
||||
const float y_scale_inv = __uint_as_float(static_cast<uint32_t>(scale_ue8m0) << 23);
|
||||
|
||||
int4 output_buf;
|
||||
auto* output_buf_ptr = reinterpret_cast<fp8x2_e4m3_t*>(&output_buf);
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < INPUT_VEC_SIZE; j += 2) {
|
||||
output_buf_ptr[j / 2] =
|
||||
pack_fp8(static_cast<float>(input_vec[j]) * y_scale, static_cast<float>(input_vec[j + 1]) * y_scale);
|
||||
}
|
||||
|
||||
*reinterpret_cast<int4*>(output_q + output_group_start_offset + lane_id * INPUT_VEC_SIZE) = output_buf;
|
||||
|
||||
if (lane_id == 0) {
|
||||
output_s[(outer_idx * num_tokens + token_idx) * hidden_dim_groups + hidden_group] = y_scale_inv;
|
||||
}
|
||||
}
|
||||
|
||||
device::PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
template <typename T, bool kUsePDL>
|
||||
struct FP8WoAGroupMajorQuantUE8M0Kernel {
|
||||
static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
|
||||
using namespace host;
|
||||
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
auto TSize = SymbolicSize{"num_tokens"};
|
||||
auto GSize = SymbolicSize{"num_outer_groups"};
|
||||
auto DSize = SymbolicSize{"hidden_dim"};
|
||||
auto SSize = SymbolicSize{"hidden_dim_groups"};
|
||||
|
||||
TensorMatcher({TSize, GSize, DSize}).with_strides({-1, DSize, 1}).with_dtype<T>().with_device(device).verify(input);
|
||||
TensorMatcher({TSize, GSize, DSize}).with_dtype<fp8_e4m3_t>().with_device(device).verify(output_q);
|
||||
TensorMatcher({GSize, TSize, SSize}).with_dtype<float>().with_device(device).verify(output_s);
|
||||
|
||||
const auto num_tokens = TSize.unwrap();
|
||||
const auto num_outer_groups = GSize.unwrap();
|
||||
const auto hidden_dim = DSize.unwrap();
|
||||
const auto hidden_dim_groups = SSize.unwrap();
|
||||
const auto input_stride_t = input.stride(0);
|
||||
constexpr int64_t kInputAlignElements = sizeof(int4) / sizeof(T);
|
||||
|
||||
RuntimeCheck(hidden_dim % GROUP_SIZE == 0, "hidden_dim must be divisible by 128");
|
||||
RuntimeCheck(hidden_dim_groups == hidden_dim / GROUP_SIZE, "output_s hidden dim mismatch");
|
||||
RuntimeCheck(
|
||||
reinterpret_cast<uintptr_t>(input.data_ptr()) % sizeof(int4) == 0,
|
||||
"input base pointer must be 16-byte aligned");
|
||||
RuntimeCheck(
|
||||
num_tokens <= 1 || input_stride_t % kInputAlignElements == 0,
|
||||
"input token stride must preserve 16-byte vector-load alignment");
|
||||
|
||||
const int64_t total_scale_groups = num_tokens * num_outer_groups * hidden_dim_groups;
|
||||
if (total_scale_groups == 0) return;
|
||||
|
||||
const auto grid = dim3((total_scale_groups + SUBWARPS_PER_BLOCK - 1) / SUBWARPS_PER_BLOCK);
|
||||
const auto block = dim3(SUBWARPS_PER_BLOCK * THREADS_PER_GROUP);
|
||||
host::LaunchKernel(grid, block, device.unwrap())
|
||||
.enable_pdl(kUsePDL)(
|
||||
fp8_wo_a_group_major_quant_ue8m0_kernel<T, kUsePDL>,
|
||||
static_cast<const T*>(input.data_ptr()),
|
||||
static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
|
||||
static_cast<float*>(output_s.data_ptr()),
|
||||
total_scale_groups,
|
||||
static_cast<int64_t>(num_tokens),
|
||||
static_cast<int>(hidden_dim_groups),
|
||||
static_cast<int>(num_outer_groups),
|
||||
static_cast<int64_t>(input_stride_t));
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
@@ -19,6 +19,7 @@ from .elementwise import (
|
||||
fused_q_norm_rope,
|
||||
fused_rope_inplace,
|
||||
)
|
||||
from .fp8_wo_a import sglang_per_token_group_quant_fp8_dsv4_wo_a
|
||||
from .gemm import linear_bf16_fp32
|
||||
from .moe import (
|
||||
hash_topk,
|
||||
@@ -45,6 +46,7 @@ __all__ = [
|
||||
"fused_q_indexer_rope_hadamard_fp4_quant",
|
||||
"fused_q_indexer_rope_hadamard_quant",
|
||||
"fused_k_norm_rope_flashmla",
|
||||
"sglang_per_token_group_quant_fp8_dsv4_wo_a",
|
||||
"make_name",
|
||||
"linear_bf16_fp32",
|
||||
"get_paged_mqa_logits_metadata",
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
from .utils import make_name
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
_GROUP_SIZE = 128
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_module(in_dtype: torch.dtype, use_pdl: bool) -> Module:
|
||||
args = make_cpp_args(in_dtype, use_pdl)
|
||||
return load_jit(
|
||||
make_name("fp8_wo_a_group_major_quant_ue8m0"),
|
||||
*args,
|
||||
cuda_files=["deepseek_v4/fp8_wo_a_group_major_quant.cuh"],
|
||||
cuda_wrappers=[
|
||||
(
|
||||
"fp8_wo_a_group_major_quant_ue8m0",
|
||||
f"FP8WoAGroupMajorQuantUE8M0Kernel<{args}>::run",
|
||||
)
|
||||
],
|
||||
# Match the AOT/JIT v2 quant path's fast-math build so FP8 rounding stays
|
||||
# bit-identical for the DSV4 wo_a replacement.
|
||||
extra_cuda_cflags=["--use_fast_math"],
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(
|
||||
op_name="fp8_wo_a_group_major_quant_ue8m0",
|
||||
mutates_args=["output_q", "output_s"],
|
||||
)
|
||||
def _fp8_wo_a_group_major_quant_ue8m0_custom_op(
|
||||
input: torch.Tensor,
|
||||
output_q: torch.Tensor,
|
||||
output_s: torch.Tensor,
|
||||
) -> None:
|
||||
"""Opaque custom-op boundary for the DeepSeek-V4 wo_a quant JIT kernel."""
|
||||
assert input.dtype in (torch.bfloat16, torch.float16)
|
||||
|
||||
module = _jit_module(input.dtype, is_arch_support_pdl())
|
||||
module.fp8_wo_a_group_major_quant_ue8m0(input, output_q, output_s)
|
||||
|
||||
|
||||
@debug_kernel_api
|
||||
def fp8_wo_a_group_major_quant_ue8m0(
|
||||
input: torch.Tensor,
|
||||
output_q: torch.Tensor,
|
||||
output_s: torch.Tensor,
|
||||
) -> None:
|
||||
_fp8_wo_a_group_major_quant_ue8m0_custom_op(input, output_q, output_s)
|
||||
|
||||
|
||||
def sglang_per_token_group_quant_fp8_dsv4_wo_a(
|
||||
x: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Quantize DSV4 wo_a activations for DeepGEMM fp8_einsum.
|
||||
|
||||
The input is a [T, G, D] bf16/fp16 tensor whose hidden dimension is
|
||||
contiguous. The output codes are contiguous [T, G, D] fp8 values. The scale
|
||||
tensor is returned as logical [T, G, D/128] fp32 UE8M0 values backed by
|
||||
contiguous [G, T, D/128] storage, so each group/head [T, S] panel is
|
||||
contiguous for the DeepGEMM recipe=(1, 1, 128) consumer. Group size is fixed
|
||||
to 128 and the absmax floor is fixed to 1e-10.
|
||||
"""
|
||||
num_tokens, num_groups, hidden = x.shape
|
||||
hidden_groups = hidden // _GROUP_SIZE
|
||||
x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn)
|
||||
x_s_storage = torch.empty(
|
||||
(num_groups, num_tokens, hidden_groups),
|
||||
device=x.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
if x.numel() > 0:
|
||||
fp8_wo_a_group_major_quant_ue8m0(x, x_q, x_s_storage)
|
||||
|
||||
# DeepGEMM fp8_einsum consumes each group/head [T, S] scale panel contiguously.
|
||||
return x_q, x_s_storage.transpose(0, 1)
|
||||
@@ -24,6 +24,7 @@ from sglang.jit_kernel.dsv4 import (
|
||||
fused_norm_rope_inplace,
|
||||
fused_q_norm_rope,
|
||||
fused_rope_inplace,
|
||||
sglang_per_token_group_quant_fp8_dsv4_wo_a,
|
||||
)
|
||||
from sglang.srt.compilation.compilation_config import register_split_op
|
||||
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
|
||||
@@ -76,7 +77,6 @@ from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||
from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||
from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8
|
||||
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
|
||||
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
|
||||
from sglang.srt.layers.utils.cp_utils import (
|
||||
@@ -1090,15 +1090,11 @@ class MQALayer(nn.Module):
|
||||
|
||||
T, G, D = o.shape
|
||||
R = self.o_lora_rank
|
||||
o_fp8, o_s = sglang_per_token_group_quant_fp8(
|
||||
o.reshape(T * G, D).contiguous(),
|
||||
group_size=128,
|
||||
scale_ue8m0=True,
|
||||
)
|
||||
o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o)
|
||||
output = torch.empty(T, G, R, device=o.device, dtype=torch.bfloat16)
|
||||
deep_gemm.fp8_einsum(
|
||||
"bhr,hdr->bhd",
|
||||
(o_fp8.view(T, G, D), o_s.view(T, G, -1)),
|
||||
(o_fp8, o_s),
|
||||
(self.wo_a.weight.view(G, R, D), self.wo_a.weight_scale_inv.data),
|
||||
output,
|
||||
recipe=(1, 1, 128),
|
||||
|
||||
Reference in New Issue
Block a user