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

[sgl-kernel][Feat][B200][2/N] Support MXFP8 Grouped GEMM in Blackwell (#14640)

Browse Source 
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
Qi Yuhang
2026-02-26 11:23:37 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang
parent b3202fe6d0
commit 88ad3b894a
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh
+6 -2
View File
@@ -276,7 +276,9 @@ __global__ void mxfp8_group_quant(
auto scale_factor_shared = make_tensor( auto scale_factor_shared = make_tensor(
make_smem_ptr(shared_memory), make_smem_ptr(shared_memory),
scale_factor_tile_layout); // ((_32,_4), _4):((_16,_4), _1) scale_factor_tile_layout); // ((_32,_4), _4):((_16,_4), _1)
// TODO: Transform Groupwise Schedule into a more efficient Schedule // Transform Groupwise Schedule into Flatten Schedule
uint group_total_tiles = 0;
uint head_cta_id = 0;
for (int g = 0; g < groups; g++) { for (int g = 0; g < groups; g++) {
int m = problem_sizes[g * 3 + 0]; int m = problem_sizes[g * 3 + 0];
int k = problem_sizes[g * 3 + 2]; int k = problem_sizes[g * 3 + 2];
@@ -323,7 +325,9 @@ __global__ void mxfp8_group_quant(
auto tiled_predict_tensor = zipped_divide(predict_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128))) auto tiled_predict_tensor = zipped_divide(predict_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
auto total_tiles = size<1>(tiled_input_tensor); // cdiv(M, 128) * cdiv(K, 128) auto total_tiles = size<1>(tiled_input_tensor); // cdiv(M, 128) * cdiv(K, 128)
decltype(total_tiles) blk_offset = blockIdx.x; group_total_tiles += total_tiles;
auto blk_offset = (blockIdx.x + (gridDim.x - head_cta_id)) % gridDim.x;
head_cta_id = group_total_tiles % gridDim.x;
while (blk_offset < total_tiles) { while (blk_offset < total_tiles) {
auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset)); auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset));
auto current_quant_output_tile = tensor<0>(tiled_quant_output_tensor(_, blk_offset)); auto current_quant_output_tile = tensor<0>(tiled_quant_output_tensor(_, blk_offset));