diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 5e25e56a239c..92d6dab80bfa 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -2,7 +2,7 @@ import os from contextlib import contextmanager from enum import IntEnum, auto -from typing import Dict, List, Tuple +from typing import Dict, List, Optional, Tuple import torch from tqdm import tqdm @@ -77,11 +77,13 @@ def _maybe_compile_deep_gemm_one_type_all( n: int, k: int, num_groups: int, + recipe: Optional[tuple[int, int, int]], + sf_dtype: Optional[torch.dtype], ) -> None: global _INITIALIZATION_DICT global _BUILTIN_M_LIST - query_key = (kernel_type, n, k, num_groups) + query_key = (kernel_type, n, k, num_groups, recipe, sf_dtype) if ( _ENABLE_JIT_DEEPGEMM_PRECOMPILE and _DO_COMPILE_ALL @@ -112,6 +114,8 @@ def _maybe_compile_deep_gemm_one_type_all( n=n, k=k, num_groups=num_groups, + recipe=recipe, + sf_dtype=sf_dtype, m_list=_BUILTIN_M_LIST, ) @@ -122,6 +126,8 @@ def _compile_deep_gemm_one_type_all( n: int, k: int, num_groups: int, + recipe: Optional[tuple[int, int, int]], + sf_dtype: Optional[torch.dtype], m_list: List[int], ) -> None: # Symmetric memory allocation performs a collective operation across all the GPUs. @@ -160,7 +166,13 @@ def _compile_deep_gemm_one_type_all( # Need some methods to estimate needed memory for warmup executor = _BaseWarmupExecutor.create( - kernel_type, max_m=max_m, n=n, k=k, num_groups=num_groups + kernel_type, + max_m=max_m, + n=n, + k=k, + num_groups=num_groups, + recipe=recipe, + sf_dtype=sf_dtype, ) old_compile_mode = deep_gemm.get_compile_mode() @@ -212,51 +224,101 @@ def execute(self, m): raise NotImplementedError -def _empty_token_fp8(size): +def _empty_token_fp8( + size, + recipe: Optional[tuple[int, int, int]] = None, + sf_dtype: Optional[torch.dtype] = None, +): *dims, k = size + if recipe is None: + block_k = 128 + else: + _, _, block_k = recipe + if sf_dtype is None or sf_dtype == torch.float32: + sf_storage_elements_per_scale = 1 + elif sf_dtype == torch.int: + sf_storage_elements_per_scale = 4 + else: + raise ValueError(f"Unimplemented sf_dtype: {sf_dtype}") return ( torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn), torch.empty( - (*dims, ceil_div(k, _BLOCK_SIZE)), device="cuda", dtype=torch.float32 + (*dims, ceil_div(k, block_k * sf_storage_elements_per_scale)), + device="cuda", + dtype=sf_dtype, ), ) -def _empty_block_fp8(size): +def _empty_block_fp8( + size, + recipe: Optional[tuple[int, int, int]] = None, + sf_dtype: Optional[torch.dtype] = None, +): *dims, n, k = size + if recipe is None: + block_n = block_k = 128 + else: + _, block_n, block_k = recipe + if sf_dtype is None or sf_dtype == torch.float32: + sf_storage_elements_per_scale = 1 + elif sf_dtype == torch.int: + sf_storage_elements_per_scale = 4 + else: + raise ValueError(f"Unimplemented sf_dtype: {sf_dtype}") return ( torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn), torch.empty( - (*dims, ceil_div(n, _BLOCK_SIZE), ceil_div(k, _BLOCK_SIZE)), + ( + *dims, + ceil_div(n, block_n * sf_storage_elements_per_scale), + ceil_div(k, block_k * sf_storage_elements_per_scale), + ), device="cuda", - dtype=torch.float32, + dtype=sf_dtype, ), ) -_BLOCK_SIZE = 128 - - class _NormalWarmupExecutor(_BaseWarmupExecutor): - def __init__(self, max_m: int, n: int, k: int, num_groups: int): - self.lhs_q, self.lhs_s = _empty_token_fp8((max_m, k)) - self.rhs_q, self.rhs_s = _empty_block_fp8((n, k)) + def __init__( + self, + max_m: int, + n: int, + k: int, + num_groups: int, + recipe: Optional[tuple[int, int, int]], + sf_dtype: Optional[torch.dtype], + ): + self.lhs_q, self.lhs_s = _empty_token_fp8((max_m, k), recipe, sf_dtype) + self.rhs_q, self.rhs_s = _empty_block_fp8((n, k), recipe, sf_dtype) self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16) + self.recipe = recipe def execute(self, m): deep_gemm.fp8_gemm_nt( (self.lhs_q[:m], self.lhs_s[:m]), (self.rhs_q, self.rhs_s), self.out[:m], + recipe=self.recipe, ) class _GroupedContWarmupExecutor(_BaseWarmupExecutor): - def __init__(self, max_m: int, n: int, k: int, num_groups: int): - self.lhs_q, self.lhs_s = _empty_token_fp8((max_m, k)) - self.rhs_q, self.rhs_s = _empty_block_fp8((num_groups, n, k)) + def __init__( + self, + max_m: int, + n: int, + k: int, + num_groups: int, + recipe: Optional[tuple[int, int, int]], + sf_dtype: Optional[torch.dtype], + ): + self.lhs_q, self.lhs_s = _empty_token_fp8((max_m, k), recipe, sf_dtype) + self.rhs_q, self.rhs_s = _empty_block_fp8((num_groups, n, k), recipe, sf_dtype) self.m_indices = torch.zeros((max_m,), device="cuda", dtype=torch.int32) self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16) + self.recipe = recipe def execute(self, m): deep_gemm.m_grouped_fp8_gemm_nt_contiguous( @@ -264,17 +326,29 @@ def execute(self, m): (self.rhs_q, self.rhs_s), self.out[:m], m_indices=self.m_indices[:m], + recipe=self.recipe, ) class _GroupedMaskedWarmupExecutor(_BaseWarmupExecutor): - def __init__(self, max_m: int, n: int, k: int, num_groups: int): - self.lhs_q, self.lhs_s = _empty_token_fp8((num_groups, max_m, k)) - self.rhs_q, self.rhs_s = _empty_block_fp8((num_groups, n, k)) + def __init__( + self, + max_m: int, + n: int, + k: int, + num_groups: int, + recipe: Optional[tuple[int, int, int]], + sf_dtype: Optional[torch.dtype], + ): + self.lhs_q, self.lhs_s = _empty_token_fp8( + (num_groups, max_m, k), recipe, sf_dtype + ) + self.rhs_q, self.rhs_s = _empty_block_fp8((num_groups, n, k), recipe, sf_dtype) self.masked_m = torch.zeros((num_groups,), device="cuda", dtype=torch.int32) self.out = torch.empty( (num_groups, max_m, n), device="cuda", dtype=torch.bfloat16 ) + self.recipe = recipe def execute(self, m): deep_gemm.fp8_m_grouped_gemm_nt_masked( @@ -284,13 +358,22 @@ def execute(self, m): masked_m=self.masked_m, # DeepGEMM uses `expect_m` instead of input shape for `get_best_config` expected_m=m, + recipe=self.recipe, ) @contextmanager def deep_gemm_execution_hook( - m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType + m: int, + n: int, + k: int, + num_groups: int, + recipe: Optional[tuple[int, int, int]], + sf_dtype: Optional[torch.dtype], + kernel_type: DeepGemmKernelType, ): if m > 0: - _maybe_compile_deep_gemm_one_type_all(kernel_type, n, k, num_groups) + _maybe_compile_deep_gemm_one_type_all( + kernel_type, n, k, num_groups, recipe, sf_dtype + ) yield diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 88d0a959b156..8c931f5ac2af 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -29,6 +29,8 @@ def grouped_gemm_nt_f8f8bf16_masked( out: torch.Tensor, masked_m: torch.Tensor, expected_m: int, + recipe: Optional[tuple[int, int, int]] = None, + sf_dtype: Optional[torch.dtype] = None, overlap_args: Optional[Any] = None, max_block_n: int = 256, ): @@ -40,7 +42,7 @@ def grouped_gemm_nt_f8f8bf16_masked( _sanity_check_input(rhs) with compile_utils.deep_gemm_execution_hook( - expected_m, n, k, num_groups, kernel_type + expected_m, n, k, num_groups, recipe, sf_dtype, kernel_type ): with configure_deep_gemm_num_sms( overlap_args.num_sms if overlap_args is not None else None @@ -52,6 +54,7 @@ def grouped_gemm_nt_f8f8bf16_masked( out, masked_m, expected_m, + recipe=recipe, **( dict( enable_overlap=True, @@ -69,6 +72,8 @@ def grouped_gemm_nt_f8f8bf16_contig( rhs: Tuple[torch.Tensor, torch.Tensor], out: torch.Tensor, m_indices: torch.Tensor, + recipe: Optional[tuple[int, int, int]] = None, + sf_dtype: Optional[torch.dtype] = None, ): m, k = lhs[0].shape num_groups, n, _ = rhs[0].shape @@ -77,14 +82,20 @@ def grouped_gemm_nt_f8f8bf16_contig( _sanity_check_input(lhs) _sanity_check_input(rhs) - with compile_utils.deep_gemm_execution_hook(m, n, k, num_groups, kernel_type): - deep_gemm.m_grouped_fp8_gemm_nt_contiguous(lhs, rhs, out, m_indices) + with compile_utils.deep_gemm_execution_hook( + m, n, k, num_groups, recipe, sf_dtype, kernel_type + ): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + lhs, rhs, out, m_indices, recipe=recipe + ) def gemm_nt_f8f8bf16( lhs: Tuple[torch.Tensor, torch.Tensor], rhs: Tuple[torch.Tensor, torch.Tensor], out: torch.Tensor, + recipe: Optional[tuple[int, int, int]] = None, + sf_dtype: Optional[torch.dtype] = None, ): m, k = lhs[0].shape n, _ = rhs[0].shape @@ -94,11 +105,14 @@ def gemm_nt_f8f8bf16( _sanity_check_input(lhs) _sanity_check_input(rhs) - with compile_utils.deep_gemm_execution_hook(m, n, k, num_groups, kernel_type): + with compile_utils.deep_gemm_execution_hook( + m, n, k, num_groups, recipe, sf_dtype, kernel_type + ): deep_gemm.fp8_gemm_nt( lhs, rhs, out, + recipe=recipe, ) diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 7701f9757f52..98115c479eb4 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -428,7 +428,7 @@ def create_per_token_group_quant_fp8_output_scale( if scale_ue8m0: assert column_major_scales and scale_tma_aligned *x_batch, x_q_mn, x_q_k = x_shape - x_s_mn, x_s_k = x_q_mn, x_q_k // 128 + x_s_mn, x_s_k = x_q_mn, x_q_k // group_size aligned_mn = ceil_align(x_s_mn, 4) aligned_k = ceil_align(x_s_k, 4) # TODO(FIXME): Fix cuda kernel and recover here to empty.