minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[DSV4] Emit TMA-aligned UE8M0 scales for FP8 einsum (#34277)

Browse Source
This commit is contained in:
Kaixi
2026-08-16 22:06:41 -07:00
committed by GitHub
parent 0fb040cbeb
commit f3225bceb3
3 changed files with 113 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh
+48 -13
View File
@@ -2,9 +2,10 @@
//
// 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.
// contiguous [T, G, D], group_size is fixed to 128, and output_s stores four
// packed UE8M0 exponent bytes per int32. Its physical layout is
// [G, ceil((D/128)/4), align_up(T, 4)]; the Python wrapper returns the logical
// [T, G, ceil((D/128)/4)] view consumed natively by DeepGEMM.
//
// 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
@@ -34,6 +35,8 @@ 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);
constexpr int UE8M0_SCALES_PER_PACK = 4;
static_assert(UE8M0_SCALES_PER_PACK == sizeof(int32_t));
template <int THREADS_PER_SUBWARP>
SGL_DEVICE float GroupReduceMax(float val) {
@@ -51,11 +54,12 @@ 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,
int32_t* __restrict__ output_s,
int64_t total_scale_groups,
int64_t num_tokens,
int hidden_dim_groups,
int packed_hidden_dim_groups,
int num_outer_groups,
int64_t aligned_num_tokens,
int64_t input_stride_t) {
device::PDLWaitPrimary<kUsePDL>();
@@ -95,7 +99,6 @@ __global__ void fp8_wo_a_group_major_quant_ue8m0_kernel(
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);
@@ -108,7 +111,23 @@ __global__ void fp8_wo_a_group_major_quant_ue8m0_kernel(
*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;
const int hidden_pack_idx = hidden_group / UE8M0_SCALES_PER_PACK;
const int pack_byte_idx = hidden_group % UE8M0_SCALES_PER_PACK;
const int64_t scale_word_offset =
(static_cast<int64_t>(outer_idx) * packed_hidden_dim_groups + hidden_pack_idx) * aligned_num_tokens +
token_idx;
auto* scale_output =
reinterpret_cast<uint8_t*>(output_s) + scale_word_offset * UE8M0_SCALES_PER_PACK + pack_byte_idx;
*scale_output = static_cast<uint8_t>(scale_ue8m0);
// DeepGEMM consumes complete int32 packs. Zero bytes without a matching
// hidden group so allocator garbage cannot become an activation scale.
if (hidden_group == hidden_dim_groups - 1) {
#pragma unroll
for (int byte_idx = pack_byte_idx + 1; byte_idx < UE8M0_SCALES_PER_PACK; ++byte_idx) {
scale_output[byte_idx - pack_byte_idx] = 0;
}
}
}
}
@@ -125,24 +144,39 @@ struct FP8WoAGroupMajorQuantUE8M0Kernel {
auto TSize = SymbolicSize{"num_tokens"};
auto GSize = SymbolicSize{"num_outer_groups"};
auto DSize = SymbolicSize{"hidden_dim"};
auto SSize = SymbolicSize{"hidden_dim_groups"};
auto PSize = SymbolicSize{"packed_hidden_dim_groups"};
auto ASize = SymbolicSize{"aligned_num_tokens"};
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);
TensorMatcher({GSize, PSize, ASize}).with_dtype<int32_t>().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 packed_hidden_dim_groups = PSize.unwrap();
const auto aligned_num_tokens = ASize.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");
const auto hidden_dim_groups = hidden_dim / GROUP_SIZE;
RuntimeCheck(
packed_hidden_dim_groups == (hidden_dim_groups + UE8M0_SCALES_PER_PACK - 1) / UE8M0_SCALES_PER_PACK,
"packed output_s hidden dim mismatch");
RuntimeCheck(
aligned_num_tokens == (num_tokens + UE8M0_SCALES_PER_PACK - 1) / UE8M0_SCALES_PER_PACK * UE8M0_SCALES_PER_PACK,
"output_s token dim must be TMA-aligned");
RuntimeCheck(
output_s.stride(0) == packed_hidden_dim_groups * aligned_num_tokens &&
output_s.stride(1) == aligned_num_tokens && output_s.stride(2) == 1,
"output_s must use contiguous [group, packed_hidden, aligned_token] storage");
RuntimeCheck(
reinterpret_cast<uintptr_t>(input.data_ptr()) % sizeof(int4) == 0,
"input base pointer must be 16-byte aligned");
RuntimeCheck(
reinterpret_cast<uintptr_t>(output_s.data_ptr()) % sizeof(int4) == 0,
"output_s 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");
@@ -157,11 +191,12 @@ struct FP8WoAGroupMajorQuantUE8M0Kernel {
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()),
static_cast<int32_t*>(output_s.data_ptr()),
total_scale_groups,
static_cast<int64_t>(num_tokens),
static_cast<int>(hidden_dim_groups),
static_cast<int>(packed_hidden_dim_groups),
static_cast<int>(num_outer_groups),
static_cast<int64_t>(aligned_num_tokens),
static_cast<int64_t>(input_stride_t));
}
};
python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py
+21 -9
View File
@@ -19,6 +19,7 @@ if TYPE_CHECKING:
from tvm_ffi.module import Module
_GROUP_SIZE = 128
_UE8M0_SCALES_PER_PACK = 4
@cache_once
@@ -71,23 +72,34 @@ def sglang_per_token_group_quant_fp8_dsv4_wo_a(
"""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.
contiguous. The output codes are contiguous [T, G, D] fp8 values. Four
UE8M0 scale exponent bytes are packed into each int32. The returned logical
scale tensor has shape [T, G, ceil((D/128)/4)] and DeepGEMM's native
TMA-aligned strides, so fp8_einsum can consume it without a scale-layout
conversion kernel. 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
packed_hidden_groups = (
hidden_groups + _UE8M0_SCALES_PER_PACK - 1
) // _UE8M0_SCALES_PER_PACK
aligned_num_tokens = (
(num_tokens + _UE8M0_SCALES_PER_PACK - 1)
// _UE8M0_SCALES_PER_PACK
* _UE8M0_SCALES_PER_PACK
)
x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn)
x_s_storage = torch.empty(
(num_groups, num_tokens, hidden_groups),
(num_groups, packed_hidden_groups, aligned_num_tokens),
device=x.device,
dtype=torch.float32,
dtype=torch.int32,
)
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)
# DeepGEMM permutes this to [G, T, packed_hidden_groups], where tokens are
# contiguous and the packed-hidden stride is aligned_num_tokens.
x_s = x_s_storage.transpose(-1, -2)[:, :num_tokens, :].transpose(0, 1)
return x_q, x_s