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

[Quantization] add humming quantization kernel (#23754)

Browse Source 
Co-authored-by: guzekai01 <zekai01@antgroup.com>
Co-authored-by: Julian Huang <huangzhilin.hzl@gmail.com>
Co-authored-by: 墨楼 <huangzhilin.hzl@antgroup.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Peng Zhang <aniz1905@gmail.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
Jinzhen Lin
2026-07-14 08:42:56 +08:00
committed by GitHub
co-authored by guzekai01 Julian Huang 墨楼 Claude Opus 4.8 Peng Zhang Xiaoyu Zhang
parent 4c997310f5
commit 423b8485fb
33 changed files with 2636 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/common_extension.cc
+1 -1
View File
@@ -161,7 +161,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def(
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! cumsum_buffer, bool "
"pad_sorted_token_ids) -> ()");
"pad_sorted_token_ids, bool ignore_invalid_expert) -> ()");
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
m.def(
sgl-kernel/csrc/common_extension_musa.cc
+1 -1
View File
@@ -109,7 +109,7 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
m.def(
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! cumsum_buffer, bool "
"pad_sorted_token_ids) -> ()");
"pad_sorted_token_ids, bool ignore_invalid_expert) -> ()");
m.impl("moe_align_block_size", torch::kMUSA, &moe_align_block_size);
m.def(
sgl-kernel/csrc/common_extension_rocm.cc
+1 -1
View File
@@ -135,7 +135,7 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
m.def(
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! cumsum_buffer, bool "
"pad_sorted_token_ids) -> ()");
"pad_sorted_token_ids, bool ignore_invalid_expert) -> ()");
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
m.def(
sgl-kernel/csrc/moe/moe_align_kernel.cu
+14 -3
View File
@@ -29,11 +29,13 @@ __global__ void count_and_sort_expert_tokens_kernel(
const scalar_t* __restrict__ topk_ids,
int32_t* __restrict__ sorted_token_ids,
int32_t* __restrict__ cumsum_buffer,
size_t numel) {
size_t numel,
bool ignore_invalid_expert) {
const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
const size_t stride = blockDim.x * gridDim.x;
for (size_t i = tid; i < numel; i += stride) {
if (ignore_invalid_expert && topk_ids[i] < 0) continue;
int32_t expert_id = topk_ids[i] + 1;
int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
sorted_token_ids[rank_post_pad] = i;
@@ -63,6 +65,7 @@ __global__ void moe_align_block_size_kernel(
size_t numel,
int32_t* __restrict__ cumsum,
bool pad_sorted_token_ids,
bool ignore_invalid_expert,
const int32_t scan_size,
int32_t max_num_tokens_padded) {
// Use a separate thread block to populate sorted_token_ids
@@ -95,6 +98,7 @@ __global__ void moe_align_block_size_kernel(
__syncthreads();
for (size_t i = tid; i < numel; i += stride) {
if (ignore_invalid_expert && topk_ids[i] < 0) continue;
int expert_id = topk_ids[i] + 1;
atomicAdd(&shared_counts[expert_id], 1);
}
@@ -242,6 +246,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
int32_t block_size,
size_t numel,
bool pad_sorted_token_ids,
bool ignore_invalid_expert,
int32_t max_num_tokens_padded) {
// Adapted from
// https://github.com/vllm-project/vllm/pull/29642/files#diff-5647b1413f4ae9aacba904eca8f8a8aee9079321eadff4c10101a2c6962dcc53R226
@@ -275,6 +280,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
}
for (size_t i = tid; i < numel; i += stride) {
if (ignore_invalid_expert && topk_ids[i] < 0) continue;
int32_t expert_id = topk_ids[i] + 1;
++tokens_cnts[(tid + 1) * num_experts + expert_id];
}
@@ -307,6 +313,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
}
for (size_t i = tid; i < numel; i += stride) {
if (ignore_invalid_expert && topk_ids[i] < 0) continue;
int32_t expert_id = topk_ids[i] + 1;
int32_t rank_post_pad = tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
sorted_token_ids[rank_post_pad] = i;
@@ -322,7 +329,8 @@ void moe_align_block_size(
torch::Tensor experts_ids,
torch::Tensor num_tokens_post_pad,
torch::Tensor cumsum_buffer,
bool pad_sorted_token_ids) {
bool pad_sorted_token_ids,
bool ignore_invalid_expert) {
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
int threads = 1024;
@@ -348,6 +356,7 @@ void moe_align_block_size(
block_size,
topk_ids.numel(),
pad_sorted_token_ids,
ignore_invalid_expert,
max_num_tokens_padded);
} else {
auto align_kernel = moe_align_block_size_kernel<scalar_t>;
@@ -364,6 +373,7 @@ void moe_align_block_size(
topk_ids.numel(),
cumsum_buffer.data_ptr<int32_t>(),
pad_sorted_token_ids,
ignore_invalid_expert,
scan_size,
max_num_tokens_padded);
@@ -377,7 +387,8 @@ void moe_align_block_size(
topk_ids.data_ptr<scalar_t>(),
sorted_token_ids.data_ptr<int32_t>(),
cumsum_buffer.data_ptr<int32_t>(),
topk_ids.numel());
topk_ids.numel(),
ignore_invalid_expert);
}
});
}
sgl-kernel/include/sgl_kernel_ops.h
+2 -1
View File
@@ -294,7 +294,8 @@ void moe_align_block_size(
torch::Tensor experts_ids,
torch::Tensor num_tokens_post_pad,
torch::Tensor cumsum_buffer,
bool pad_sorted_token_ids);
bool pad_sorted_token_ids,
bool ignore_invalid_expert);
void topk_softmax(
torch::Tensor& topk_weights,
sgl-kernel/python/sgl_kernel/moe.py
+2
View File
@@ -12,6 +12,7 @@ def moe_align_block_size(
num_tokens_post_pad,
cumsum_buffer,
pad_sorted_token_ids=False,
ignore_invalid_expert=False,
):
torch.ops.sgl_kernel.moe_align_block_size.default(
topk_ids,
@@ -22,6 +23,7 @@ def moe_align_block_size(
num_tokens_post_pad,
cumsum_buffer,
pad_sorted_token_ids,
ignore_invalid_expert,
)