From 03b835e7d13e4564e2beb72845486b8dc37606b9 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Mon, 8 Dec 2025 09:24:58 +0800 Subject: [PATCH] Refactor tuning block wise kernel and opt Qwen/Qwen3-VL-32B-Instruct-FP8 (#14141) --- benchmark/kernels/quantization/README.md | 92 +++++++++++ .../quantization/tuning_block_wise_kernel.py | 35 +++- ...dtype=fp8_w8a8,block_shape=[128, 128].json | 26 +++ ...dtype=fp8_w8a8,block_shape=[128, 128].json | 26 +++ ...dtype=fp8_w8a8,block_shape=[128, 128].json | 26 +++ ...dtype=fp8_w8a8,block_shape=[128, 128].json | 26 +++ .../srt/layers/quantization/configs/README.md | 16 ++ .../srt/layers/quantization/fp8_kernel.py | 151 +++++++++++------- 8 files changed, 334 insertions(+), 64 deletions(-) create mode 100644 benchmark/kernels/quantization/README.md create mode 100644 python/sglang/srt/layers/quantization/configs/N=1280,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json create mode 100644 python/sglang/srt/layers/quantization/configs/N=5120,K=1024,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json create mode 100644 python/sglang/srt/layers/quantization/configs/N=5120,K=3200,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json create mode 100644 python/sglang/srt/layers/quantization/configs/N=6400,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json create mode 100644 python/sglang/srt/layers/quantization/configs/README.md diff --git a/benchmark/kernels/quantization/README.md b/benchmark/kernels/quantization/README.md new file mode 100644 index 000000000..acf6f0b0d --- /dev/null +++ b/benchmark/kernels/quantization/README.md @@ -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, ...} +} +``` diff --git a/benchmark/kernels/quantization/tuning_block_wise_kernel.py b/benchmark/kernels/quantization/tuning_block_wise_kernel.py index 1b51e54b7..0a5e7fb53 100644 --- a/benchmark/kernels/quantization/tuning_block_wise_kernel.py +++ b/benchmark/kernels/quantization/tuning_block_wise_kernel.py @@ -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) diff --git a/python/sglang/srt/layers/quantization/configs/N=1280,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1280,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5c0c8d761 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1280,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -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 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=5120,K=1024,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=5120,K=1024,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..15e91cde5 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=5120,K=1024,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -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 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=5120,K=3200,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=5120,K=3200,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c714b7f19 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=5120,K=3200,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -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 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=6400,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=6400,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..f33809b0a --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=6400,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -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 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/README.md b/python/sglang/srt/layers/quantization/configs/README.md new file mode 100644 index 000000000..718c9adb9 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/README.md @@ -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]`) diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 86c0982d2..3f73c0a90 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -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