515 lines
19 KiB
Python
515 lines
19 KiB
Python
import logging
|
|
import os
|
|
import time
|
|
from contextlib import contextmanager, nullcontext
|
|
from enum import IntEnum, auto
|
|
from typing import Dict, List, Tuple
|
|
|
|
import torch
|
|
from tqdm import tqdm
|
|
|
|
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
|
|
disable_symmetric_memory_context,
|
|
restore_symmetric_memory_context,
|
|
)
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.srt.utils import ceil_align, ceil_div, get_available_gpu_memory, is_musa
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_is_musa = is_musa()
|
|
|
|
if ENABLE_JIT_DEEPGEMM:
|
|
import deep_gemm
|
|
|
|
|
|
_BUILTIN_M_LIST = list(range(1, 1024 * 16 + 1))
|
|
_ENABLE_JIT_DEEPGEMM_PRECOMPILE = envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.get()
|
|
_DO_COMPILE_ALL = True
|
|
_IS_FIRST_RANK_ON_NODE = envs.SGLANG_IS_FIRST_RANK_ON_NODE.get()
|
|
_IN_PRECOMPILE_STAGE = envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get()
|
|
_FAST_WARMUP = envs.SGLANG_JIT_DEEPGEMM_FAST_WARMUP.get()
|
|
|
|
# Force redirect deep_gemm cache_dir
|
|
os.environ["DG_JIT_CACHE_DIR"] = os.getenv(
|
|
"SGLANG_DG_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm")
|
|
)
|
|
|
|
# Refer to https://github.com/deepseek-ai/DeepGEMM/commit/d75b218b7b8f4a5dd5406ac87905039ead3ae42f
|
|
# NVRTC may have performance loss with some cases.
|
|
# And NVCC JIT speed is also 9x faster in the ref commit
|
|
os.environ["DG_JIT_USE_NVRTC"] = os.getenv("SGL_DG_USE_NVRTC", "0")
|
|
|
|
|
|
def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs):
|
|
global _BUILTIN_M_LIST
|
|
global _DO_COMPILE_ALL
|
|
global _IS_FIRST_RANK_ON_NODE
|
|
|
|
_BUILTIN_M_LIST = []
|
|
|
|
if _FAST_WARMUP:
|
|
# In fast warmup mode, only compile a small set of typical Ms
|
|
|
|
# First cover all the small bs to ensure decode performance
|
|
_BUILTIN_M_LIST += list(range(1, 1025))
|
|
|
|
# Then cover larger batch sizes with gradually increasing steps
|
|
# For example, when chunekd prefill size is 16384
|
|
# The sampled Ms would be:
|
|
# 1024, 1026, ... 2046 (step 2)
|
|
# 2048, 2052, ... 4092 (step 4)
|
|
# 4096, 5004, ... 8184 (step 8)
|
|
# 8192, 9008, ... 16384 (step 16)
|
|
# Totally 1024 + 1024 / 2 + 2048 / 4 + 4096 / 8 + 8192 / 16 = 3072 kernels
|
|
next_m, sample_step = 1024, 2
|
|
max_prefill_bs = (
|
|
min(server_args.chunked_prefill_size, 32 * 1024)
|
|
if server_args.chunked_prefill_size >= 1
|
|
else 16 * 1024
|
|
)
|
|
while next_m < max_prefill_bs:
|
|
_BUILTIN_M_LIST += list(range(next_m, 2 * next_m, sample_step))
|
|
next_m = next_m * 2
|
|
sample_step = sample_step * 2
|
|
_BUILTIN_M_LIST.append(max_prefill_bs)
|
|
_BUILTIN_M_LIST = sorted(list(set(_BUILTIN_M_LIST)))
|
|
else:
|
|
# When fast warmup isn't enabled, generate m_max and compile all the covered Ms.
|
|
m_max = 1024 * 16
|
|
if server_args.chunked_prefill_size < 1:
|
|
m_max = 1024 * 64
|
|
elif server_args.chunked_prefill_size > 8192:
|
|
m_max = server_args.chunked_prefill_size * 2
|
|
m_max = min(1024 * 128, m_max)
|
|
_BUILTIN_M_LIST += list(range(1, m_max + 1))
|
|
|
|
_IS_FIRST_RANK_ON_NODE = server_args.base_gpu_id == gpu_id
|
|
|
|
# Check if is the first rank on node.
|
|
# Default each rank will try compile all Ms to
|
|
# load all symbols at the launch stages.
|
|
# Avoid loading symbols at the serving stages.
|
|
_DO_COMPILE_ALL = _IS_FIRST_RANK_ON_NODE
|
|
|
|
|
|
class DeepGemmKernelType(IntEnum):
|
|
GROUPED_GEMM_NT_F8F8BF16_MASKED = auto()
|
|
GROUPED_GEMM_NT_F8F8BF16_CONTIG = auto()
|
|
GROUPED_GEMM_NT_BF16_MASKED = auto()
|
|
GROUPED_GEMM_NT_BF16_CONTIG = auto()
|
|
GEMM_NT_F8F8BF16 = auto()
|
|
GEMM_NT_BF16BF16F32 = auto()
|
|
TF32_HC_PRENORM_GEMM = auto()
|
|
|
|
|
|
_INITIALIZATION_DICT: Dict[Tuple[DeepGemmKernelType, int, int, int], bool] = dict()
|
|
|
|
|
|
# TODO improve code
|
|
def _maybe_compile_deep_gemm_one_type_all(
|
|
kernel_type: DeepGemmKernelType,
|
|
n: int,
|
|
k: int,
|
|
num_groups: int,
|
|
) -> None:
|
|
global _INITIALIZATION_DICT
|
|
global _BUILTIN_M_LIST
|
|
|
|
query_key = (kernel_type, n, k, num_groups)
|
|
if (
|
|
_ENABLE_JIT_DEEPGEMM_PRECOMPILE
|
|
and _DO_COMPILE_ALL
|
|
and _INITIALIZATION_DICT.get(query_key) is None
|
|
):
|
|
_INITIALIZATION_DICT[query_key] = True
|
|
|
|
# TODO maybe improve logs
|
|
if not _IN_PRECOMPILE_STAGE and _IS_FIRST_RANK_ON_NODE:
|
|
logger.warning(
|
|
"Entering DeepGEMM JIT Pre-Compile session. "
|
|
"It may take a long time (typically 10-20 mins) "
|
|
"if you have not run `sglang.compile_deep_gemm`. "
|
|
"It is recommended to run `sglang.compile_deep_gemm` with same args as `sglang.launch_server`"
|
|
" for pre-compilation to reduce the overhead if you have not run it before. "
|
|
"For example: "
|
|
"`python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code`"
|
|
)
|
|
|
|
logger.info(
|
|
f"Try DeepGEMM JIT Compiling for "
|
|
f"<{kernel_type.name}> N={n}, K={k}, num_groups={num_groups} with all Ms."
|
|
f"{' It only takes a little time (typically 1 sec) if you have run `python3 -m sglang.compile_deep_gemm`. ' if not _IN_PRECOMPILE_STAGE else ''}"
|
|
)
|
|
|
|
_compile_deep_gemm_one_type_all(
|
|
kernel_type=kernel_type,
|
|
n=n,
|
|
k=k,
|
|
num_groups=num_groups,
|
|
m_list=_BUILTIN_M_LIST,
|
|
)
|
|
|
|
|
|
# NOTE(alcanderian): get_num_sms should be change when 2-batch-overlap is introduced
|
|
def _compile_deep_gemm_one_type_all(
|
|
kernel_type: DeepGemmKernelType,
|
|
n: int,
|
|
k: int,
|
|
num_groups: int,
|
|
m_list: List[int],
|
|
) -> None:
|
|
# Symmetric memory allocation performs a collective operation across all the GPUs.
|
|
# Temporary disable symmetric memory during compilation since it only runs on the first rank.
|
|
saved_context = disable_symmetric_memory_context()
|
|
try:
|
|
if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG:
|
|
m_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout()
|
|
m_list = sorted(list(set(m for m in m_list if m % m_alignment == 0)))
|
|
elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_BF16_CONTIG:
|
|
m_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout()
|
|
m_list = sorted(list(set(m for m in m_list if m % m_alignment == 0)))
|
|
|
|
# Here the precompilation is only run on the first rank, so gpu_id should be 0
|
|
memory_budget = get_available_gpu_memory(device="cuda", gpu_id=0)
|
|
|
|
# If the memory budget is less memory requirement, we need to reduce max_m to avoid out of memory, which might further cause hanging during warmup
|
|
max_m = max(m_list)
|
|
required_memory = _BaseWarmupExecutor.get_memory_requirement(
|
|
kernel_type, max_m=max_m, n=n, k=k, num_groups=num_groups
|
|
)
|
|
logger.info(
|
|
f"Required memory for warmup: {required_memory}GB, Available memory: {memory_budget}GB"
|
|
)
|
|
if memory_budget < required_memory:
|
|
# TODO: Maybe compute the max_m based on the memory budget
|
|
while (
|
|
_BaseWarmupExecutor.get_memory_requirement(
|
|
kernel_type, max_m=max_m, n=n, k=k, num_groups=num_groups
|
|
)
|
|
> memory_budget
|
|
and max_m > 4096
|
|
):
|
|
max_m = max_m // 2
|
|
logger.warning(
|
|
f"Available memory {memory_budget}GB is less than required memory {required_memory}GB for warmup, reducing max_m to {max_m} to avoid out of memory"
|
|
)
|
|
m_list = [m for m in m_list if m <= max_m]
|
|
|
|
# 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
|
|
)
|
|
|
|
has_compile_mode_api = hasattr(deep_gemm, "get_compile_mode") and hasattr(
|
|
deep_gemm, "set_compile_mode"
|
|
)
|
|
if has_compile_mode_api:
|
|
old_compile_mode = deep_gemm.get_compile_mode()
|
|
deep_gemm.set_compile_mode(1)
|
|
|
|
# TODO can use multi thread
|
|
for m in tqdm(m_list, desc="DeepGEMM warmup"):
|
|
executor.execute(m=m)
|
|
if has_compile_mode_api:
|
|
deep_gemm.set_compile_mode(old_compile_mode)
|
|
|
|
# clean up input buffers
|
|
torch.cuda.current_stream().synchronize()
|
|
del executor
|
|
torch.cuda.empty_cache()
|
|
finally:
|
|
# Restore symmetric memory context
|
|
restore_symmetric_memory_context(saved_context)
|
|
|
|
|
|
class _BaseWarmupExecutor:
|
|
@staticmethod
|
|
def create(kernel_type: DeepGemmKernelType, **kwargs):
|
|
return {
|
|
DeepGemmKernelType.GEMM_NT_F8F8BF16: _NormalWarmupExecutor,
|
|
DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: _GroupedContWarmupExecutor,
|
|
DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED: _GroupedMaskedWarmupExecutor,
|
|
DeepGemmKernelType.GEMM_NT_BF16BF16F32: _BF16F32WarmupExecutor,
|
|
DeepGemmKernelType.GROUPED_GEMM_NT_BF16_CONTIG: _BF16GroupedContWarmupExecutor,
|
|
DeepGemmKernelType.GROUPED_GEMM_NT_BF16_MASKED: _BF16GroupedMaskedWarmupExecutor,
|
|
DeepGemmKernelType.TF32_HC_PRENORM_GEMM: _TF32HcPrenormWarmupExecutor,
|
|
}[kernel_type](**kwargs)
|
|
|
|
@staticmethod
|
|
def get_memory_requirement(
|
|
kernel_type: DeepGemmKernelType, max_m: int, n: int, k: int, num_groups: int
|
|
) -> int:
|
|
# Return the required memory space in GB for warmup executor
|
|
_GB = 1 << 30
|
|
if kernel_type == DeepGemmKernelType.GEMM_NT_F8F8BF16:
|
|
return (max_m * k + n * k + max_m * n * 2) / _GB
|
|
elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG:
|
|
return (max_m * k + num_groups * n * k + max_m * 4 + max_m * n * 2) / _GB
|
|
elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_BF16_CONTIG:
|
|
return (
|
|
max_m * k * 2 + num_groups * n * k * 2 + max_m * 4 + max_m * n * 2
|
|
) / _GB
|
|
elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED:
|
|
return (
|
|
num_groups * max_m * k
|
|
+ num_groups * n * k
|
|
+ num_groups * 4
|
|
+ num_groups * max_m * n * 2
|
|
) / _GB
|
|
elif kernel_type == DeepGemmKernelType.GEMM_NT_BF16BF16F32:
|
|
# bf16 lhs + bf16 rhs + fp32 out
|
|
return (max_m * k * 2 + n * k * 2 + max_m * n * 4) / _GB
|
|
elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_BF16_MASKED:
|
|
return (
|
|
num_groups * max_m * k * 2
|
|
+ num_groups * n * k * 2
|
|
+ num_groups * 4
|
|
+ num_groups * max_m * n * 2
|
|
) / _GB
|
|
elif kernel_type == DeepGemmKernelType.TF32_HC_PRENORM_GEMM:
|
|
# The generic hook's fourth dimension is num_splits for MHC.
|
|
# A value of 0 represents DeepGEMM's unsplit num_splits=None path.
|
|
num_splits = num_groups if num_groups > 0 else 1
|
|
return (max_m * k * 2 + n * k * 4 + num_splits * max_m * (n + 1) * 4) / _GB
|
|
else:
|
|
raise ValueError(f"Invalid kernel type: {kernel_type}")
|
|
|
|
def execute(self, m):
|
|
raise NotImplementedError
|
|
|
|
|
|
def _empty_token_fp8(size):
|
|
*dims, k = size
|
|
return (
|
|
torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn),
|
|
torch.empty(
|
|
(*dims, ceil_div(k, _BLOCK_SIZE)), device="cuda", dtype=torch.float32
|
|
),
|
|
)
|
|
|
|
|
|
def _empty_block_fp8(size):
|
|
*dims, n, k = size
|
|
return (
|
|
torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn),
|
|
torch.empty(
|
|
(*dims, ceil_div(n, _BLOCK_SIZE), ceil_div(k, _BLOCK_SIZE)),
|
|
device="cuda",
|
|
dtype=torch.float32,
|
|
),
|
|
)
|
|
|
|
|
|
_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))
|
|
self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16)
|
|
|
|
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],
|
|
)
|
|
|
|
|
|
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))
|
|
self.m_indices = torch.zeros((max_m,), device="cuda", dtype=torch.int32)
|
|
self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16)
|
|
|
|
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],
|
|
self.m_indices[:m],
|
|
)
|
|
|
|
|
|
class _BF16GroupedContWarmupExecutor(_BaseWarmupExecutor):
|
|
def __init__(self, max_m: int, n: int, k: int, num_groups: int):
|
|
self.a = torch.empty((max_m, k), device="cuda", dtype=torch.bfloat16)
|
|
self.b = torch.empty((num_groups, n, k), device="cuda", dtype=torch.bfloat16)
|
|
self.m_indices = torch.zeros((max_m,), device="cuda", dtype=torch.int32)
|
|
self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16)
|
|
|
|
def execute(self, m):
|
|
deep_gemm.m_grouped_bf16_gemm_nt_contiguous(
|
|
self.a[:m],
|
|
self.b,
|
|
self.out[:m],
|
|
self.m_indices[:m],
|
|
)
|
|
|
|
|
|
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))
|
|
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
|
|
)
|
|
|
|
def execute(self, m):
|
|
deep_gemm.fp8_m_grouped_gemm_nt_masked(
|
|
(self.lhs_q, self.lhs_s),
|
|
(self.rhs_q, self.rhs_s),
|
|
self.out,
|
|
masked_m=self.masked_m,
|
|
# DeepGEMM uses `expect_m` instead of input shape for `get_best_config`
|
|
expected_m=m,
|
|
)
|
|
|
|
|
|
class _BF16F32WarmupExecutor(_BaseWarmupExecutor):
|
|
def __init__(self, max_m: int, n: int, k: int, num_groups: int):
|
|
self.lhs = torch.empty((max_m, k), device="cuda", dtype=torch.bfloat16)
|
|
self.rhs = torch.empty((n, k), device="cuda", dtype=torch.bfloat16)
|
|
self.out = torch.empty((max_m, n), device="cuda", dtype=torch.float32)
|
|
|
|
def execute(self, m):
|
|
deep_gemm.bf16_gemm_nt(self.lhs[:m], self.rhs, self.out[:m])
|
|
|
|
|
|
class _BF16GroupedMaskedWarmupExecutor(_BaseWarmupExecutor):
|
|
def __init__(self, max_m: int, n: int, k: int, num_groups: int):
|
|
self.a = torch.empty(
|
|
(num_groups, max_m, k), device="cuda", dtype=torch.bfloat16
|
|
)
|
|
self.b = torch.empty((num_groups, n, k), device="cuda", dtype=torch.bfloat16)
|
|
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
|
|
)
|
|
|
|
def execute(self, m):
|
|
deep_gemm.m_grouped_bf16_gemm_nt_masked(
|
|
self.a,
|
|
self.b,
|
|
self.out,
|
|
masked_m=self.masked_m,
|
|
# DeepGEMM uses `expect_m` instead of input shape for `get_best_config`
|
|
expected_m=m,
|
|
)
|
|
|
|
|
|
class _TF32HcPrenormWarmupExecutor(_BaseWarmupExecutor):
|
|
def __init__(self, max_m: int, n: int, k: int, num_groups: int):
|
|
self.x = torch.empty((max_m, k), device="cuda", dtype=torch.bfloat16)
|
|
self.fn = torch.empty((n, k), device="cuda", dtype=torch.float32)
|
|
self.n = n
|
|
# The generic warmup executor's num_groups argument is num_splits here.
|
|
# A value of 0 represents DeepGEMM's unsplit num_splits=None path.
|
|
self.num_splits = num_groups if num_groups > 0 else None
|
|
|
|
def execute(self, m):
|
|
if self.num_splits is None:
|
|
out = torch.empty((m, self.n), device="cuda", dtype=torch.float32)
|
|
sqrsum = torch.empty((m,), device="cuda", dtype=torch.float32)
|
|
else:
|
|
# Slicing the middle dimension of a preallocated
|
|
# (num_splits, max_m, n) output would create a strided view.
|
|
out = torch.empty(
|
|
(self.num_splits, m, self.n), device="cuda", dtype=torch.float32
|
|
)
|
|
sqrsum = torch.empty(
|
|
(self.num_splits, m), device="cuda", dtype=torch.float32
|
|
)
|
|
deep_gemm.tf32_hc_prenorm_gemm(
|
|
self.x[:m],
|
|
self.fn,
|
|
out,
|
|
sqrsum,
|
|
num_splits=self.num_splits,
|
|
)
|
|
|
|
|
|
def deep_gemm_execution_hook(
|
|
m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType
|
|
):
|
|
if _is_musa:
|
|
return nullcontext()
|
|
|
|
return _deep_gemm_execution_hook(m, n, k, num_groups, kernel_type)
|
|
|
|
|
|
@contextmanager
|
|
def _deep_gemm_execution_hook(
|
|
m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType
|
|
):
|
|
if m > 0:
|
|
_maybe_compile_deep_gemm_one_type_all(kernel_type, n, k, num_groups)
|
|
yield
|
|
|
|
|
|
def pp_parallel_deep_gemm_warmup(model_runner) -> None:
|
|
"""Run per-PP-rank dummy DECODE+EXTEND forwards so each rank's
|
|
DeepGEMM JIT compiles in parallel instead of serially via the warmup
|
|
/generate flowing through the pipeline. Opt-in via
|
|
SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.
|
|
"""
|
|
# n_splits ~= n_sms / ceil(bs/block_m) with block_m=64; sweep 5 bs to
|
|
# cover the brackets real /generate hits (smallest decode shape,
|
|
# mid-low, two mid, and n_splits=1 for ~5K+ token prefill). Ceil-align
|
|
# to attn_cp_size for DSA prefill CP's seq_len % cp_size == 0 assert.
|
|
n_sms = torch.cuda.get_device_properties(model_runner.device).multi_processor_count
|
|
block_m = 64
|
|
cp = max(model_runner.attn_cp_size, 1)
|
|
batch_sizes = sorted(
|
|
{
|
|
ceil_align(bs, cp)
|
|
for bs in (
|
|
1,
|
|
2 * block_m,
|
|
max(n_sms // 8, 2) * block_m,
|
|
max(n_sms // 4, 4) * block_m,
|
|
n_sms * block_m,
|
|
)
|
|
}
|
|
)
|
|
|
|
# In PD, prefill-only nodes never decode (indexer would OOM at large
|
|
# bs) and decode-only nodes never extend.
|
|
disagg_mode = model_runner.server_args.disaggregation_mode
|
|
run_decode = model_runner.is_generation and disagg_mode != "prefill"
|
|
run_extend = disagg_mode != "decode"
|
|
|
|
logger.info(
|
|
"PP-parallel DeepGEMM warmup start "
|
|
"(pp_rank=%d, tp_rank=%d, batch_sizes=%s, disagg=%s).",
|
|
model_runner.pp_rank,
|
|
model_runner.tp_rank,
|
|
batch_sizes,
|
|
disagg_mode,
|
|
)
|
|
|
|
t0 = time.perf_counter()
|
|
with torch.inference_mode():
|
|
for bs in batch_sizes:
|
|
if run_decode:
|
|
model_runner._dummy_run(
|
|
batch_size=bs, forward_mode_override=ForwardMode.DECODE
|
|
)
|
|
if run_extend:
|
|
model_runner._dummy_run(
|
|
batch_size=bs, forward_mode_override=ForwardMode.EXTEND
|
|
)
|
|
|
|
logger.info(
|
|
"PP-parallel DeepGEMM warmup done in %.2fs (pp_rank=%d).",
|
|
time.perf_counter() - t0,
|
|
model_runner.pp_rank,
|
|
)
|