chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

92
benchmark/kernels/quantization/README.md Normal file
View File

@@ -0,0 +1,92 @@
# W8A8 Block-wise Quantization Kernel Tuning
Auto-tune Triton FP8/INT8 block-wise quantization kernels for optimal performance.
## When to Use Triton FP8 Block-wise Quantization Kernel vs DeepGEMM
**Use Triton FP8 Block-wise Quantization Kernel when:**
- Output dtype is NOT `bfloat16` (e.g., `float16`, `float32`)
- DeepGEMM is disabled (environment variable `SGLANG_ENABLE_JIT_DEEPGEMM=0`)
- Running on GPUs with compute capability < SM90 (DeepGEMM requires SM90+)
- You need cross-platform compatibility (Triton works on both NVIDIA and AMD GPUs)
**Use DeepGEMM when:**
- Output dtype is `bfloat16` AND DeepGEMM is enabled
- Running on NVIDIA GPUs with compute capability >= SM90 (e.g., H100, H200)
- Need maximum performance for production workloads (DeepGEMM is highly optimized for Hopper architecture)
**Note:** DeepGEMM requires CUDA compute capability >= 9.0 (SM90+). It is specifically optimized for NVIDIA Hopper GPUs (H100/H200).
The kernel selection logic in SGLang automatically chooses DeepGEMM when conditions are met (see `w8a8_block_fp8_matmul` function in `fp8_kernel.py`), otherwise falls back to Triton implementation.
## Quick Start
**Default (DeepSeek-V3):**
```bash
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --tp-size 8
```
**Custom Model (specify N and K):**
```bash
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 25600
```
## Parameters
- `--N`, `--K`: Weight matrix dimensions (N=output_dim, K=input_dim). If not specified, uses `--tp-size` for DeepSeek-V3
- `--tp-size`: Tensor parallelism size for DeepSeek-V3 (default: 8)
- `--input-type`: `fp8` or `int8` (default: fp8)
- `--block-n`, `--block-k`: Block quantization granularity (default: 128)
- `--batch-size`: Test single batch size (optional)
## How to Calculate N and K
For a linear layer `y = xW^T` where `x` is (M, K) and `W` is (N, K):
- **N**: Output features (weight matrix output dimension)
- **K**: Input features (weight matrix input dimension)
**Example: Qwen3-VL-32B** (hidden_size=5120, intermediate_size=25600, num_heads=64, num_kv_heads=8, head_dim=128) and TP=1
```bash
# QKV projection: Q(8192) + K(1024) + V(1024) = 10240
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 10240 --K 5120
# MLP gate+up (SwiGLU): 2 * intermediate_size = 51200
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 51200 --K 5120
# MLP down projection
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 25600
# O projection (if separate from QKV)
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 8192
```
If TP=8:
```bash
# QKV projection: Q(8192) + K(1024) + V(1024) = 10240 / TP=8
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 1280 --K 5120
# MLP gate+up (SwiGLU): 2 * intermediate_size = 51200 / TP=8
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 6400 --K 5120
# MLP down projection
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 3200
# O projection (if separate from QKV)
python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 1024
```
## Output
Generates JSON config files saved to `python/sglang/srt/layers/quantization/configs/`:
```
N={N},K={K},device_name={DEVICE},dtype=fp8_w8a8,block_shape=[128,128].json
```
Config maps batch size to optimal kernel parameters:
```json
{
"1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, ...},
"2048": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128, ...}
}
```

35
benchmark/kernels/quantization/tuning_block_wise_kernel.py
View File

@@ -84,6 +84,8 @@ def w8a8_block_matmul(
C_shape = A.shape[:-1] + (N,)
C = A.new_empty(C_shape, dtype=output_dtype)
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"]),
@@ -127,6 +129,7 @@ def w8a8_block_matmul(
Bs.stride(1),
Bs.stride(0),
**config,
needs_masking=needs_masking,
)
return C
@@ -428,7 +431,13 @@ def main(args):
batch_sizes = [args.batch_size]
num_gpus = 1 # If only one batch size, use only one GPU
weight_shapes = get_weight_shapes(args.tp_size)
# Support manual N and K specification
if args.N is not None and args.K is not None:
weight_shapes = [(args.N, args.K)]
print(f"Using manually specified weight shape: N={args.N}, K={args.K}")
else:
weight_shapes = get_weight_shapes(args.tp_size)
print(f"Using predefined weight shapes for TP size {args.tp_size}")
batches_per_gpu = distribute_batch_sizes(batch_sizes, num_gpus)
@@ -453,7 +462,25 @@ def main(args):
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--tp-size", "-tp", type=int, default=8)
parser.add_argument(
"--tp-size",
"-tp",
type=int,
default=8,
help="Tensor parallelism size (ignored if --N and --K are specified)",
)
parser.add_argument(
"--N",
type=int,
default=None,
help="Output dimension of weight matrix (number of columns)",
)
parser.add_argument(
"--K",
type=int,
default=None,
help="Input dimension of weight matrix (number of rows)",
)
parser.add_argument(
"--input-type", type=str, choices=["fp8", "int8"], default="fp8"
)
@@ -471,4 +498,8 @@ if __name__ == "__main__":
)
args = parser.parse_args()
# Validate arguments
if (args.N is None) != (args.K is None):
parser.error("--N and --K must be specified together or not at all")
main(args)

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