From 46ae84df159e261fa67d5cce3647f3342fe365a7 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 16 Sep 2026 16:54:19 -0700 Subject: [PATCH] dsv4.1: Hopper FP8 matmul kernels and tuning (#39657) Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> --- ...0,dtype=fp8_w8a8,block_shape=[32, 32].json | 20 +++ ...0,dtype=fp8_w8a8,block_shape=[32, 32].json | 20 +++ ...0,dtype=fp8_w8a8,block_shape=[32, 32].json | 19 +++ ...0,dtype=fp8_w8a8,block_shape=[32, 32].json | 20 +++ ...0,dtype=fp8_w8a8,block_shape=[32, 32].json | 20 +++ ...0,dtype=fp8_w8a8,block_shape=[32, 32].json | 20 +++ ...0,dtype=fp8_w8a8,block_shape=[32, 32].json | 20 +++ .../kernels/ops/quantization/fp8_kernel.py | 156 +++++++++++++++++- 8 files changed, 291 insertions(+), 4 deletions(-) create mode 100644 python/sglang/kernels/ops/quantization/configs/N=1152,K=5120,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json create mode 100644 python/sglang/kernels/ops/quantization/configs/N=1792,K=5120,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json create mode 100644 python/sglang/kernels/ops/quantization/configs/N=25600,K=6144,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json create mode 100644 python/sglang/kernels/ops/quantization/configs/N=4096,K=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json create mode 100644 python/sglang/kernels/ops/quantization/configs/N=5120,K=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json create mode 100644 python/sglang/kernels/ops/quantization/configs/N=5120,K=576,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json create mode 100644 python/sglang/kernels/ops/quantization/configs/N=8192,K=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json diff --git a/python/sglang/kernels/ops/quantization/configs/N=1152,K=5120,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json b/python/sglang/kernels/ops/quantization/configs/N=1152,K=5120,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json new file mode 100644 index 000000000..80ed75f24 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=1152,K=5120,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json @@ -0,0 +1,20 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4, + "SWAP_AB": true, + "SPLIT_K": 16 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=1792,K=5120,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json b/python/sglang/kernels/ops/quantization/configs/N=1792,K=5120,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json new file mode 100644 index 000000000..fb5b60cdf --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=1792,K=5120,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json @@ -0,0 +1,20 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4, + "SWAP_AB": true, + "SPLIT_K": 8 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=25600,K=6144,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json b/python/sglang/kernels/ops/quantization/configs/N=25600,K=6144,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json new file mode 100644 index 000000000..91b3ff0df --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=25600,K=6144,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json @@ -0,0 +1,19 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4, + "SWAP_AB": true + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=4096,K=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json b/python/sglang/kernels/ops/quantization/configs/N=4096,K=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json new file mode 100644 index 000000000..5d34539e4 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=4096,K=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json @@ -0,0 +1,20 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4, + "SWAP_AB": true, + "SPLIT_K": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=5120,K=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json b/python/sglang/kernels/ops/quantization/configs/N=5120,K=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json new file mode 100644 index 000000000..fb5b60cdf --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=5120,K=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json @@ -0,0 +1,20 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4, + "SWAP_AB": true, + "SPLIT_K": 8 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=5120,K=576,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json b/python/sglang/kernels/ops/quantization/configs/N=5120,K=576,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json new file mode 100644 index 000000000..5d34539e4 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=5120,K=576,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json @@ -0,0 +1,20 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4, + "SWAP_AB": true, + "SPLIT_K": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=8192,K=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json b/python/sglang/kernels/ops/quantization/configs/N=8192,K=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json new file mode 100644 index 000000000..b6c100355 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=8192,K=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[32, 32].json @@ -0,0 +1,20 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4, + "SWAP_AB": true, + "SPLIT_K": 2 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/kernels/ops/quantization/fp8_kernel.py b/python/sglang/kernels/ops/quantization/fp8_kernel.py index 5cfe48e09..c6fbd51ac 100644 --- a/python/sglang/kernels/ops/quantization/fp8_kernel.py +++ b/python/sglang/kernels/ops/quantization/fp8_kernel.py @@ -26,6 +26,7 @@ import triton.language as tl from sglang.kernels.jit.utils import is_arch_support_pdl from sglang.kernels.ops.quantization.fp8_utils import fp8_dtype_to_triton from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.runtime_context import get_platform from sglang.srt.utils import ( ceil_align, get_bool_env_var, @@ -1029,6 +1030,133 @@ def _w8a8_block_fp8_matmul( tl.store(c_ptrs, c, mask=c_mask) +@triton.jit +def _w8a8_block_fp8_matmul_hopper( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + needs_masking: tl.constexpr, + SWAP_AB: tl.constexpr = False, + SPLIT_K: tl.constexpr = 1, +): + + pid = tl.program_id(axis=0) + split = tl.program_id(axis=1) + tiles_per_split = tl.cdiv(tl.cdiv(K, BLOCK_SIZE_K), SPLIT_K) + first_tile = split * tiles_per_split + C += split * M * N + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + n_tiles_k_per_group_k = group_k // BLOCK_SIZE_K + + a_ptrs += first_tile * BLOCK_SIZE_K * stride_ak + b_ptrs += first_tile * BLOCK_SIZE_K * stride_bk + As_ptrs += (first_tile // n_tiles_k_per_group_k) * stride_As_k + Bs_ptrs += (first_tile // n_tiles_k_per_group_k) * stride_Bs_k + + # Small-M Hopper configs transpose the MMA so the weight tile occupies M. + if SWAP_AB: + accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32) + else: + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range( + first_tile, tl.minimum(first_tile + tiles_per_split, tl.cdiv(K, BLOCK_SIZE_K)) + ): + 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) + + a_s = tl.load(As_ptrs) + b_s = tl.load(Bs_ptrs) + + scale_step_k = tl.where((k + 1) % n_tiles_k_per_group_k == 0, 1, 0) + if SWAP_AB: + accumulator += ( + tl.dot(tl.trans(b), tl.trans(a)) * b_s[:, None] * a_s[None, :] + ) + else: + 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 SWAP_AB: + accumulator = tl.trans(accumulator) + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +@triton.jit +def _reduce_block_fp8_split_k( + Parts, Out, ELEMENTS: tl.constexpr, SPLITS: tl.constexpr, BLOCK: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + splits = tl.arange(0, SPLITS) + values = tl.load( + Parts + splits[:, None] * ELEMENTS + offsets[None, :], + offsets[None, :] < ELEMENTS, + 0.0, + ) + tl.store(Out + offsets, tl.sum(values, axis=0), offsets < ELEMENTS) + + @triton.jit def _w8a8_block_fp8_matmul_gfx1250( # Pointers to inputs and output @@ -1579,23 +1707,38 @@ def w8a8_block_fp8_matmul_triton( "num_stages": 3, } - if _is_gfx1250: + # Split-K accumulates K in SPLIT_K separate fp32 partials, so its results + # do not match the single-accumulator kernels bit-for-bit. + hopper_tuned = get_platform().is_sm90 and ( + config.get("SWAP_AB", False) or config.get("SPLIT_K", 1) > 1 + ) + if hopper_tuned: + kernel = _w8a8_block_fp8_matmul_hopper + elif _is_gfx1250: config = {**config, "num_stages": 1} kernel = _w8a8_block_fp8_matmul_gfx1250 else: kernel = select_w8a8_block_fp8_matmul_kernel(M, N, config) + split_k = config.get("SPLIT_K", 1) if hopper_tuned else 1 + if split_k > 1: + assert split_k & (split_k - 1) == 0 + partials = torch.empty((split_k, M, N), device=A.device, dtype=torch.float32) + else: + partials = C + 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"]), + blocks = triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv( + N, META["BLOCK_SIZE_N"] ) + return (blocks, split_k) if hopper_tuned else (blocks,) kernel[grid]( A, B, - C, + partials, As, Bs, M, @@ -1617,6 +1760,11 @@ def w8a8_block_fp8_matmul_triton( needs_masking=needs_masking, ) + if split_k > 1: + _reduce_block_fp8_split_k[(triton.cdiv(M * N, 256),)]( + partials, C, M * N, split_k, 256 + ) + return C