Skip to content

Expand deep_gemm entrypoint to support more FP8 recipes. #17294

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zianglih wants to merge 1 commit into from
Closed

Expand deep_gemm entrypoint to support more FP8 recipes. #17294

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 105 additions & 22 deletions python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand All @@ -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.
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -212,69 +224,131 @@ 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}")
Comment on lines +237 to +242

@gemini-code-assist gemini-code-assist Bot Jan 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This logic for handling sf_dtype is duplicated in _empty_block_fp8 (lines 263-268). To improve maintainability and reduce code duplication, consider extracting this logic into a shared helper function. This would also be a good place to make the handling of sf_dtype=None more explicit by defaulting it to torch.float32 at the beginning of the helper.

Comment on lines +237 to +242

@gemini-code-assist gemini-code-assist Bot Jan 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This logic for determining sf_storage_elements_per_scale is duplicated in _empty_block_fp8 (lines 263-268). To improve maintainability and reduce code duplication, consider extracting this logic into a separate helper function.

For example:

def _get_sf_storage_elements_per_scale(sf_dtype: Optional[torch.dtype]) -> int:
    if sf_dtype is None or sf_dtype == torch.float32:
        return 1
    if sf_dtype == torch.int:
        return 4
    raise ValueError(f"Unimplemented sf_dtype: {sf_dtype}")

You can then call this helper function in both _empty_token_fp8 and _empty_block_fp8.

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,

@gemini-code-assist gemini-code-assist Bot Jan 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

While passing dtype=None to torch.empty defaults to torch.float32, making this explicit improves code clarity and robustness. It's better to explicitly handle the None case for sf_dtype. This also applies to _empty_block_fp8 on line 278.

Suggested change
dtype=sf_dtype,
dtype=sf_dtype or torch.float32,

),
)


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,
Comment thread
zianglih marked this conversation as resolved.
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
Comment thread
zianglih marked this conversation as resolved.

def execute(self, m):
deep_gemm.m_grouped_fp8_gemm_nt_contiguous(
(self.lhs_q[:m], self.lhs_s[: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
Comment thread
zianglih marked this conversation as resolved.

def execute(self, m):
deep_gemm.fp8_m_grouped_gemm_nt_masked(
Expand All @@ -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
22 changes: 18 additions & 4 deletions python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
):
Expand All @@ -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
Expand All @@ -52,6 +54,7 @@ def grouped_gemm_nt_f8f8bf16_masked(
out,
masked_m,
expected_m,
recipe=recipe,
**(
dict(
enable_overlap=True,
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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,
)


Expand Down
2 changes: 1 addition & 1 deletion python/sglang/srt/layers/quantization/fp8_kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading