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

Refactor tuning block wise kernel and opt Qwen/Qwen3-VL-32B-Instruct-FP8 (#14141)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-12-08 09:24:58 +08:00
committed by GitHub
parent aff1238ef2
commit 03b835e7d1
8 changed files with 334 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/sglang/srt/layers/quantization/configs/N=1280,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json Normal file
View File

@@ -0,0 +1,26 @@
{
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 4
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3
}
}

26
python/sglang/srt/layers/quantization/configs/N=5120,K=1024,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json Normal file
View File

@@ -0,0 +1,26 @@
{
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2
}
}

26
python/sglang/srt/layers/quantization/configs/N=5120,K=3200,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json Normal file
View File

@@ -0,0 +1,26 @@
{
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 4
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 4
}
}

26
python/sglang/srt/layers/quantization/configs/N=6400,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json Normal file
View File

@@ -0,0 +1,26 @@
{
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3
}
}

16
python/sglang/srt/layers/quantization/configs/README.md Normal file
View File

@@ -0,0 +1,16 @@
# W8A8 Block FP8 Kernel Configurations
This directory contains optimized kernel configurations for the W8A8 block FP8 matrix multiplication kernel.
## Configuration File Format
Configuration files are named using the following pattern:
```
N={N},K={K},device_name={DEVICE_NAME},dtype=fp8_w8a8,block_shape=[{BLOCK_N},{BLOCK_K}].json
```
Where:
- `N`: Output dimension (number of columns in weight matrix)
- `K`: Input dimension (number of columns in activation matrix)
- `DEVICE_NAME`: GPU device name with spaces replaced by underscores (e.g., `NVIDIA_H100_80GB_HBM3`)
- `BLOCK_N`, `BLOCK_K`: Block quantization granularity (typically `[128,128]`)

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

@@ -719,6 +719,7 @@ def _w8a8_block_fp8_matmul(
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
needs_masking: tl.constexpr,
):
"""Triton-accelerated function used to perform linear operations (dot
product) on input tensors `A` and `B` with block-wise quantization, and store the result in output
@@ -744,20 +745,25 @@ def _w8a8_block_fp8_matmul(
As_ptrs = As + offs_am * stride_As_m
offs_bsn = offs_bn // group_n
Bs_ptrs = Bs + offs_bsn * stride_Bs_n
scale_step_k = BLOCK_SIZE_K // group_k
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
if needs_masking:
a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
else:
a = tl.load(a_ptrs)
b = tl.load(b_ptrs)
k_start = k * BLOCK_SIZE_K
offs_ks = k_start // group_k
a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
a_s = tl.load(As_ptrs)
b_s = tl.load(Bs_ptrs)
accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
As_ptrs += scale_step_k * stride_As_k
Bs_ptrs += scale_step_k * stride_Bs_k
if C.dtype.element_ty == tl.bfloat16:
c = accumulator.to(tl.bfloat16)
@@ -804,6 +810,7 @@ def _w8a8_block_fp8_matmul_unrolledx4(
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
needs_masking: tl.constexpr,
):
"""Triton-accelerated function used to perform linear operations (dot
product) on input tensors `A` and `B` with block-wise quantization, and store the result in output
@@ -829,94 +836,111 @@ def _w8a8_block_fp8_matmul_unrolledx4(
As_ptrs = As + offs_am * stride_As_m
offs_bsn = offs_bn // group_n
Bs_ptrs = Bs + offs_bsn * stride_Bs_n
scale_step_k = BLOCK_SIZE_K // group_k
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
# manually unroll to 4 iterations
UNROLL_FACTOR = 4
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * UNROLL_FACTOR)):
# 1st iteration
a = tl.load(
a_ptrs,
mask=offs_k[None, :] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K,
other=0.0,
)
b = tl.load(
b_ptrs,
mask=offs_k[:, None] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K,
other=0.0,
)
if needs_masking:
a = tl.load(
a_ptrs,
mask=offs_k[None, :] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K,
other=0.0,
)
b = tl.load(
b_ptrs,
mask=offs_k[:, None] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K,
other=0.0,
)
else:
a = tl.load(a_ptrs)
b = tl.load(b_ptrs)
k_start = (k * UNROLL_FACTOR) * BLOCK_SIZE_K
offs_ks = k_start // group_k
a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
a_s = tl.load(As_ptrs)
b_s = tl.load(Bs_ptrs)
accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
As_ptrs += scale_step_k * stride_As_k
Bs_ptrs += scale_step_k * stride_Bs_k
# 2nd iteration
a = tl.load(
a_ptrs,
mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K,
other=0.0,
)
b = tl.load(
b_ptrs,
mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K,
other=0.0,
)
if needs_masking:
a = tl.load(
a_ptrs,
mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K,
other=0.0,
)
b = tl.load(
b_ptrs,
mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K,
other=0.0,
)
else:
a = tl.load(a_ptrs)
b = tl.load(b_ptrs)
k_start = k_start + BLOCK_SIZE_K
offs_ks = k_start // group_k
a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
a_s = tl.load(As_ptrs)
b_s = tl.load(Bs_ptrs)
accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
As_ptrs += scale_step_k * stride_As_k
Bs_ptrs += scale_step_k * stride_Bs_k
# 3rd iteration
a = tl.load(
a_ptrs,
mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K,
other=0.0,
)
b = tl.load(
b_ptrs,
mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K,
other=0.0,
)
if needs_masking:
a = tl.load(
a_ptrs,
mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K,
other=0.0,
)
b = tl.load(
b_ptrs,
mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K,
other=0.0,
)
else:
a = tl.load(a_ptrs)
b = tl.load(b_ptrs)
k_start = k_start + BLOCK_SIZE_K
offs_ks = k_start // group_k
a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
a_s = tl.load(As_ptrs)
b_s = tl.load(Bs_ptrs)
accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
As_ptrs += scale_step_k * stride_As_k
Bs_ptrs += scale_step_k * stride_Bs_k
# 4th iteration
a = tl.load(
a_ptrs,
mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K,
other=0.0,
)
b = tl.load(
b_ptrs,
mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K,
other=0.0,
)
if needs_masking:
a = tl.load(
a_ptrs,
mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K,
other=0.0,
)
b = tl.load(
b_ptrs,
mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K,
other=0.0,
)
else:
a = tl.load(a_ptrs)
b = tl.load(b_ptrs)
k_start = k_start + BLOCK_SIZE_K
offs_ks = k_start // group_k
a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
a_s = tl.load(As_ptrs)
b_s = tl.load(Bs_ptrs)
accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
As_ptrs += scale_step_k * stride_As_k
Bs_ptrs += scale_step_k * stride_Bs_k
if C.dtype.element_ty == tl.bfloat16:
c = accumulator.to(tl.bfloat16)
@@ -1111,6 +1135,8 @@ def w8a8_block_fp8_matmul_triton(
"num_stages": 3,
}
needs_masking = bool(K % config["BLOCK_SIZE_K"] != 0)
def grid(META):
return (
triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
@@ -1140,6 +1166,7 @@ def w8a8_block_fp8_matmul_triton(
Bs.stride(1),
Bs.stride(0),
**config,
needs_masking=needs_masking,
)
return C