diff --git a/python/sglang/kernels/ops/quantization/configs/N=24576,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json b/python/sglang/kernels/ops/quantization/configs/N=24576,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json new file mode 100644 index 000000000..5c6d3e323 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=24576,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json @@ -0,0 +1,103 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 1024, + "num_warps": 2, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 512, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 8, + "num_stages": 1 + }, + "256": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 8, + "num_stages": 1 + }, + "512": null, + "1024": null, + "2048": null, + "4096": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 2 + }, + "8192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 2 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=4096,K=12288,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json b/python/sglang/kernels/ops/quantization/configs/N=4096,K=12288,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json new file mode 100644 index 000000000..3f7848c44 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=4096,K=12288,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json @@ -0,0 +1,121 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 512, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 8, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 8, + "num_stages": 4 + }, + "192": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 3 + }, + "8192": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 2 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=4096,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json b/python/sglang/kernels/ops/quantization/configs/N=4096,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json new file mode 100644 index 000000000..e9e1aec53 --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=4096,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json @@ -0,0 +1,121 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 1 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "num_warps": 2, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 512, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "192": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 4, + "num_stages": 2 + }, + "8192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 2 + } +} diff --git a/python/sglang/kernels/ops/quantization/configs/N=6144,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json b/python/sglang/kernels/ops/quantization/configs/N=6144,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json new file mode 100644 index 000000000..4181a9c7e --- /dev/null +++ b/python/sglang/kernels/ops/quantization/configs/N=6144,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json @@ -0,0 +1,121 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 512, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 1024, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 512, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 2, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "num_warps": 2, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 2 + }, + "192": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "num_warps": 8, + "num_stages": 3 + }, + "8192": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/kernels/ops/quantization/fp8_kernel.py b/python/sglang/kernels/ops/quantization/fp8_kernel.py index 6e3253fc4..f232e548c 100644 --- a/python/sglang/kernels/ops/quantization/fp8_kernel.py +++ b/python/sglang/kernels/ops/quantization/fp8_kernel.py @@ -1263,6 +1263,57 @@ def get_w8a8_block_fp8_configs( return None +@functools.lru_cache +def get_w8a8_channelwise_fp8_configs(N: int, K: int) -> Optional[Dict[int, Any]]: + """ + Return tuned Triton configurations for the per-token/per-channel w8a8 fp8 + GEMM (`triton_scaled_mm`), for one weight shape on the current device. + + The return value maps an irregular grid of token counts M to + `triton_scaled_mm` tile configs; the closest M in the grid should be picked + for a given batch. A null entry means CUTLASS was faster at that M, so the + grid keeps the point (to stop a neighbouring M snapping onto it) but the + caller must fall back. A missing file / shape likewise means "keep the + default". + """ + # Intentional: the tuned lookup (host-side device name + file I/O) isn't + # traceable, so under torch.compile return None and fall back to CUTLASS -- + # same as mainline's block-FP8 get_w8a8_block_fp8_configs. + if torch._dynamo.is_compiling(): + return None + + device_name = get_device_name().replace(" ", "_") + json_file_name = ( + f"N={N},K={K},device_name={device_name},dtype=fp8_w8a8_channelwise.json" + ) + config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name + ) + if not os.path.exists(config_file_path): + return None + + with open(config_file_path) as f: + log_info_on_rank0( + logger, + f"Using configuration from {config_file_path} for W8A8 channelwise FP8 GEMM.", + ) + return {int(key): val for key, val in json.load(f).items()} + + +def get_w8a8_channelwise_fp8_config(N: int, K: int, M: int) -> Optional[Dict[str, int]]: + """Tuned `triton_scaled_mm` config for this GEMM, or None to keep the default. + + N / K are the weight's output / contraction dims and M the token count. None + means no tuned config applies -- the caller keeps its existing kernel choice + -- either because this device / shape was never tuned, or the nearest tuned + M is a point where CUTLASS won. The caller gates on the feature flag. + """ + configs = get_w8a8_channelwise_fp8_configs(N, K) + if not configs: + return None + return configs[min(configs.keys(), key=lambda x: abs(x - M))] + + def select_w8a8_block_fp8_matmul_kernel(M, N, META): return _w8a8_block_fp8_matmul @@ -2124,6 +2175,8 @@ def triton_scaled_mm( block_size_n: int = 32, block_size_k: int = 32, use_heuristic=True, + num_warps: Optional[int] = None, + num_stages: Optional[int] = None, ) -> torch.Tensor: M, K = input.shape N = weight.shape[1] @@ -2163,14 +2216,23 @@ def triton_scaled_mm( tile_shape = (64, 128, 128) else: tile_shape = (128, 128, 128) - - block_size_m, block_size_n, block_size_k = tile_shape + block_size_m, block_size_n, block_size_k = tile_shape + # else: use the block_size_{m,n,k} the caller passed (e.g. a tuned config). block_size_sa = 1 if has_scalar(scale_a) else block_size_m block_size_sb = 1 if has_scalar(scale_b) else block_size_n accumulator_dtype = tl.float32 if input.is_floating_point() else tl.int32 + # num_warps / num_stages are triton.jit launch kwargs, and triton rejects an + # explicit None -- pass them only when the caller set them, so the default + # path keeps triton's own defaults. + launch_kwargs = {} + if num_warps is not None: + launch_kwargs["num_warps"] = num_warps + if num_stages is not None: + launch_kwargs["num_stages"] = num_stages + # A = input, B = weight, C = result # A = M x K, B = K x N, C = M x N scaled_mm_kernel[grid]( @@ -2195,6 +2257,7 @@ def triton_scaled_mm( BLOCK_SIZE_K=block_size_k, BLOCK_SIZE_SCALE_A=block_size_sa, BLOCK_SIZE_SCALE_B=block_size_sb, + **launch_kwargs, ) return result.to(out_dtype) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 09c406584..64e7826c4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -847,6 +847,15 @@ class Envs: SGLANG_QUANT_ALLOW_DOWNCASTING = EnvBool(False) SGLANG_FP8_IGNORED_LAYERS = EnvStr("") SGLANG_FP4_IGNORED_LAYERS = EnvStr("") + # On by default; set SGLANG_ENABLE_FP8_GEMM_CONFIG_TUNE=0 as a kill switch. + # Consults the tuned per-(N, K, M) Triton tile config table in + # apply_fp8_linear. When a tuned config exists for this GPU / weight shape / + # token count, run the Triton w8a8 FP8 GEMM with it; otherwise keep the + # default CUTLASS path. Only takes effect on a GPU with a matching + # dtype=fp8_w8a8_channelwise config JSON under + # kernels/ops/quantization/configs/ (currently L40S), so it is a no-op on + # any other GPU / untuned shape even when enabled. + SGLANG_ENABLE_FP8_GEMM_CONFIG_TUNE = EnvBool(True) # =================================================================== # Humming quantization diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index ec9987b10..4e1060fd9 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -11,6 +11,7 @@ from sglang.kernels.ops.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8, sglang_per_token_group_quant_fp8_row_padded, ) +from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil from sglang.srt.runtime_context import get_exec, get_parallel @@ -23,6 +24,7 @@ from sglang.kernels.ops.quantization.fp8_kernel import ( fp8_dtype, fp8_max, fp8_min, + get_w8a8_channelwise_fp8_config, is_fp8_fnuz, per_token_group_quant_fp8, scaled_fp8_quant, @@ -1788,6 +1790,12 @@ def apply_fp8_linear( use_cutlass_channelwise_gemm = ( channelwise_cutlass and cutlass_compatible_b and not use_triton_w8a8_fp8_kernel ) + # Consider a tuned Triton tile only where the shape would otherwise go to + # CUTLASS (that is the path the offline sweep tuned against). On by default; + # SGLANG_ENABLE_FP8_GEMM_CONFIG_TUNE=0 is the kill switch. + use_tuned_triton_channelwise = ( + use_cutlass_channelwise_gemm and envs.SGLANG_ENABLE_FP8_GEMM_CONFIG_TUNE.get() + ) native_scalar_a_scale = use_cutlass_channelwise_gemm and ( _is_sm90_supported or _is_sm100_supported or _is_sm120_supported ) @@ -1866,12 +1874,38 @@ def apply_fp8_linear( ) if channelwise_cutlass: + # A tuned config exists only for shapes where tuned Triton beat the + # CUTLASS dispatch on this GPU; otherwise this is None and the backend + # choice below is unchanged. weight is [K, N] here. + tuned_config = ( + get_w8a8_channelwise_fp8_config( + N=weight.shape[1], K=weight.shape[0], M=qinput.shape[0] + ) + if use_tuned_triton_channelwise + else None + ) if not use_cutlass_channelwise_gemm: # Massage the input to be 2D qinput = qinput.view(-1, qinput.shape[-1]) output = triton_scaled_mm( qinput, weight, x_scale, weight_scale, output_dtype, bias ) + elif tuned_config is not None: + qinput = qinput.view(-1, qinput.shape[-1]) + output = triton_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + output_dtype, + bias, + block_size_m=tuned_config["BLOCK_SIZE_M"], + block_size_n=tuned_config["BLOCK_SIZE_N"], + block_size_k=tuned_config["BLOCK_SIZE_K"], + use_heuristic=False, + num_warps=tuned_config["num_warps"], + num_stages=tuned_config["num_stages"], + ) else: output = fp8_scaled_mm( qinput,