731 lines
28 KiB
Python
731 lines
28 KiB
Python
# Copyright 2023-2026 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Base class shared by EagerRunner and BaseCudaGraphRunner."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import datetime
|
|
import hashlib
|
|
import inspect
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import TYPE_CHECKING, Any, Optional, Tuple
|
|
|
|
import torch
|
|
|
|
from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
|
|
from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers import deep_gemm_wrapper
|
|
from sglang.srt.layers.dp_attention import (
|
|
DpPaddingMode,
|
|
set_dp_buffer_len,
|
|
set_is_extend_in_batch,
|
|
)
|
|
from sglang.srt.model_executor.forward_batch_info import (
|
|
CaptureHiddenMode,
|
|
ForwardBatch,
|
|
ForwardMode,
|
|
NgramEmbeddingInfo,
|
|
PPProxyTensors,
|
|
)
|
|
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
|
|
from sglang.srt.runtime_context import get_parallel
|
|
from sglang.srt.speculative.spec_info import create_dummy_verify_input
|
|
from sglang.srt.utils import (
|
|
empty_context,
|
|
log_info_on_rank0,
|
|
require_attn_tp_gather,
|
|
require_gathered_buffer,
|
|
require_mlp_tp_gather,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.model_executor.model_runner import ModelRunner
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _allocate_decode_buffers(
|
|
*,
|
|
device: torch.device,
|
|
max_bs: int,
|
|
max_num_token: int,
|
|
hidden_size: int,
|
|
vocab_size: int,
|
|
dtype: torch.dtype,
|
|
dp_size: int,
|
|
pp_size: int,
|
|
is_encoder_decoder: bool,
|
|
require_mlp_tp_gather: bool,
|
|
seq_len_fill_value: int,
|
|
encoder_len_fill_value: int,
|
|
num_tokens_per_bs: int,
|
|
cache_loc_dtype: torch.dtype,
|
|
enable_mamba_track: bool,
|
|
ne_token_table: Optional[torch.Tensor] = None,
|
|
hc_hidden_size: Optional[int] = None,
|
|
pp_proxy_topk_size: Optional[int] = None,
|
|
) -> SimpleNamespace:
|
|
"""Allocate the FB-shared decode buffers."""
|
|
with torch.device(device):
|
|
input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
|
|
input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
|
|
req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
|
|
seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int64)
|
|
out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
|
|
positions = torch.zeros((max_num_token,), dtype=torch.int64)
|
|
mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
|
|
num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
|
|
custom_mask = torch.ones(
|
|
(max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
|
|
dtype=torch.bool,
|
|
)
|
|
next_token_logits_buffer = torch.zeros(
|
|
(max_num_token, vocab_size),
|
|
dtype=torch.float,
|
|
)
|
|
mamba_track_indices = (
|
|
torch.zeros((max_bs,), dtype=torch.int64) if enable_mamba_track else None
|
|
)
|
|
mamba_track_mask = (
|
|
torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
|
|
)
|
|
|
|
if pp_size > 1:
|
|
# mHC (e.g. DSV4) flattens residual into hidden_states (size = hc_hidden_size).
|
|
is_mhc = hc_hidden_size is not None
|
|
hs = hc_hidden_size if is_mhc else hidden_size
|
|
pp_proxy_tensors = {
|
|
"hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
|
|
}
|
|
if not is_mhc:
|
|
pp_proxy_tensors["residual"] = torch.zeros(
|
|
(max_bs, hidden_size), dtype=dtype
|
|
)
|
|
if pp_proxy_topk_size is not None:
|
|
pp_proxy_tensors["topk_indices"] = torch.zeros(
|
|
(max_num_token, pp_proxy_topk_size), dtype=torch.int32
|
|
)
|
|
else:
|
|
pp_proxy_tensors = None
|
|
|
|
if is_encoder_decoder:
|
|
encoder_lens = torch.full(
|
|
(max_bs,), encoder_len_fill_value, dtype=torch.int32
|
|
)
|
|
else:
|
|
encoder_lens = None
|
|
|
|
if require_mlp_tp_gather:
|
|
global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
|
|
global_num_tokens_for_logprob_gpu = torch.zeros(
|
|
(dp_size,), dtype=torch.int32
|
|
)
|
|
else:
|
|
global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
|
|
global_num_tokens_for_logprob_gpu = torch.zeros((1,), dtype=torch.int32)
|
|
|
|
ngram_embedding_info = (
|
|
NgramEmbeddingInfo(
|
|
token_table=ne_token_table,
|
|
column_starts=torch.zeros([max_bs], dtype=torch.int32),
|
|
req_lens=torch.ones([max_bs], dtype=torch.int32),
|
|
out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
|
|
out_req_lens=torch.ones([max_bs], dtype=torch.int32),
|
|
)
|
|
if ne_token_table is not None
|
|
else None
|
|
)
|
|
|
|
if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get():
|
|
rids_int = torch.zeros((max_bs,), dtype=torch.int64)
|
|
bootstrap_room_ids_int = torch.full((max_bs,), -1, dtype=torch.int64)
|
|
else:
|
|
rids_int = None
|
|
bootstrap_room_ids_int = None
|
|
|
|
seq_lens_cpu = torch.full(
|
|
(max_bs,),
|
|
seq_len_fill_value,
|
|
dtype=torch.int64,
|
|
device="cpu",
|
|
)
|
|
|
|
return SimpleNamespace(
|
|
input_ids=input_ids,
|
|
input_embeds=input_embeds,
|
|
req_pool_indices=req_pool_indices,
|
|
seq_lens=seq_lens,
|
|
seq_lens_cpu=seq_lens_cpu,
|
|
out_cache_loc=out_cache_loc,
|
|
positions=positions,
|
|
mrope_positions=mrope_positions,
|
|
num_token_non_padded=num_token_non_padded,
|
|
custom_mask=custom_mask,
|
|
next_token_logits_buffer=next_token_logits_buffer,
|
|
mamba_track_indices=mamba_track_indices,
|
|
mamba_track_mask=mamba_track_mask,
|
|
encoder_lens=encoder_lens,
|
|
global_num_tokens_gpu=global_num_tokens_gpu,
|
|
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
|
|
pp_proxy_tensors=pp_proxy_tensors,
|
|
ngram_embedding_info=ngram_embedding_info,
|
|
rids_int=rids_int,
|
|
bootstrap_room_ids_int=bootstrap_room_ids_int,
|
|
)
|
|
|
|
|
|
class BaseRunner(ABC):
|
|
def __init__(self, model_runner: ModelRunner) -> None:
|
|
self.model_runner = model_runner
|
|
self.device = model_runner.device
|
|
self.device_module = torch.get_device_module(self.device)
|
|
self.tp_size = model_runner.server_args.tp_size
|
|
self.dp_size = model_runner.server_args.dp_size
|
|
self.pp_size = model_runner.server_args.pp_size
|
|
self.attn_tp_size = get_parallel().attn_tp_size
|
|
self.attn_tp_rank = get_parallel().attn_tp_rank
|
|
self.tbo_plugin = TboCudaGraphRunnerPlugin()
|
|
|
|
def warmup(self) -> None:
|
|
"""Run kernel warmup + autotune once, gated by mr._kernel_warmed_up."""
|
|
mr = self.model_runner
|
|
if getattr(mr, "_kernel_warmed_up", False):
|
|
return
|
|
mr._kernel_warmed_up = True
|
|
|
|
if mr.device != "cuda":
|
|
return
|
|
|
|
self._pre_initialize_flashinfer_allreduce_workspace()
|
|
|
|
if self._should_run_flashinfer_autotune():
|
|
buffers, batch_size = self._autotune_buffers()
|
|
assert (
|
|
buffers is not None
|
|
), "_autotune_buffers() must return a reusable buffer set for autotune"
|
|
self._flashinfer_autotune(buffers=buffers, batch_size=batch_size)
|
|
|
|
if (
|
|
envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get()
|
|
and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
|
|
and mr.pp_size > 1
|
|
and not mr.spec_algorithm.is_speculative()
|
|
):
|
|
from sglang.srt.layers.deep_gemm_wrapper.compile_utils import (
|
|
pp_parallel_deep_gemm_warmup,
|
|
)
|
|
|
|
pp_parallel_deep_gemm_warmup(self)
|
|
|
|
def _pre_initialize_flashinfer_allreduce_workspace(self):
|
|
"""Allocate flashinfer allreduce workspaces; must run before CG capture
|
|
to keep broadcasts/barriers outside the capture context (else deadlock
|
|
with custom_all_reduce.register_graph_buffers).
|
|
"""
|
|
mr = self.model_runner
|
|
if mr.server_args.flashinfer_allreduce_fusion_backend is None:
|
|
return
|
|
|
|
from sglang.srt.layers.communicator import FUSE_ALLREDUCE_MAX_BATCH_SIZE
|
|
from sglang.srt.layers.flashinfer_comm_fusion import pre_initialize_workspaces
|
|
|
|
pre_initialize_workspaces(
|
|
max_token_num=FUSE_ALLREDUCE_MAX_BATCH_SIZE,
|
|
hidden_dim=mr.model_config.hidden_size,
|
|
dtype=mr.dtype,
|
|
)
|
|
|
|
def _should_run_flashinfer_autotune(self) -> bool:
|
|
"""Check if flashinfer autotune should be run."""
|
|
mr = self.model_runner
|
|
if mr.server_args.disable_flashinfer_autotune:
|
|
return False
|
|
|
|
# CuteDSL v1 (cutedsl runner + deepep a2a) bypasses MoeRunner and must not
|
|
# be autotuned -- its _dummy_run would dispatch more tokens per rank than
|
|
# SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK, tripping a DeepEP assert.
|
|
# Read server_args directly to avoid depending on initialize_moe_config()
|
|
# having already populated the MoE backend globals.
|
|
if (
|
|
mr.server_args.moe_runner_backend == "flashinfer_cutedsl"
|
|
and mr.server_args.moe_a2a_backend == "deepep"
|
|
):
|
|
return False
|
|
|
|
backend_str = mr.server_args.moe_runner_backend
|
|
|
|
# TODO smor- support other cases for flashinfer autotune, such as, mamba backend
|
|
|
|
moe_needs_autotune = backend_str in [
|
|
"flashinfer_trtllm",
|
|
"flashinfer_trtllm_routed",
|
|
"flashinfer_mxfp4",
|
|
"flashinfer_cutedsl",
|
|
"flashinfer_cutlass",
|
|
]
|
|
|
|
from sglang.srt.layers.quantization.fp4_utils import (
|
|
get_fp4_gemm_runner_backend,
|
|
)
|
|
|
|
model_uses_fp4 = mr.model_config.quantization in (
|
|
"modelopt_fp4",
|
|
"modelopt_mixed",
|
|
)
|
|
fp4_gemm_needs_autotune = model_uses_fp4 and (
|
|
get_fp4_gemm_runner_backend().is_flashinfer_cutlass()
|
|
or get_fp4_gemm_runner_backend().is_flashinfer_cutedsl()
|
|
)
|
|
|
|
from sglang.srt.layers.quantization.fp8_utils import (
|
|
get_fp8_gemm_runner_backend,
|
|
)
|
|
from sglang.srt.utils import is_sm100_supported
|
|
|
|
model_uses_modelopt_fp8 = mr.model_config.quantization in (
|
|
"modelopt",
|
|
"modelopt_fp8",
|
|
"modelopt_mixed",
|
|
)
|
|
# Online MXFP8 (microscaling) linears dispatch to flashinfer's
|
|
# ``mm_mxfp8``, which the flashinfer fp8 autotune dummy run does not
|
|
# exercise correctly -- it triggers an illegal memory access inside the
|
|
# mxfp8 cutlass cubin. The mxfp8 gemm is fixed-config and needs no
|
|
# tuning, so skip autotune for these models.
|
|
model_uses_mxfp8 = "mxfp8" in (mr.model_config.quantization or "")
|
|
fp8_gemm_needs_autotune = not model_uses_mxfp8 and (
|
|
get_fp8_gemm_runner_backend().is_flashinfer_cutlass()
|
|
or (model_uses_modelopt_fp8 and is_sm100_supported())
|
|
)
|
|
|
|
if not (
|
|
moe_needs_autotune or fp4_gemm_needs_autotune or fp8_gemm_needs_autotune
|
|
):
|
|
return False
|
|
|
|
major, _ = torch.cuda.get_device_capability()
|
|
if major < 9:
|
|
return False
|
|
|
|
if mr.spec_algorithm.is_speculative():
|
|
return not mr.is_draft_worker
|
|
|
|
return True
|
|
|
|
def _flashinfer_autotune(self, *, buffers, batch_size):
|
|
"""Run flashinfer autotune.
|
|
|
|
buffers / batch_size: a prepared static decode-buffer set and its bs,
|
|
reused for the dummy forward instead of allocating a throwaway set.
|
|
Supplied by warmup() (the decode runner's captured buffers when a graph
|
|
runner exists; a freshly-allocated dummy set in the eager path).
|
|
"""
|
|
from flashinfer.autotuner import autotune
|
|
|
|
from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
|
|
|
|
mr = self.model_runner
|
|
cache_path = self._flashinfer_autotune_cache_path()
|
|
if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get():
|
|
autotune_cache = cache_path
|
|
logger.info("Running FlashInfer autotune with cache: %s", autotune_cache)
|
|
else:
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
runs_dir = cache_path.parent / "runs"
|
|
runs_dir.mkdir(parents=True, exist_ok=True)
|
|
autotune_cache = (
|
|
runs_dir / f"{cache_path.stem}.{timestamp}{cache_path.suffix}"
|
|
)
|
|
logger.info(
|
|
"Running FlashInfer autotune (cache reuse DISABLED via "
|
|
"SGLANG_FLASHINFER_AUTOTUNE_CACHE=0); writing fresh result to: %s",
|
|
autotune_cache,
|
|
)
|
|
|
|
# Run warmup on the non-default stream to avoid NCCL 2.29+ cudaMemcpyBatchAsync
|
|
# calls on default stream (unsupported by CUDA) when --enable-symm-mem is used.
|
|
mr.forward_stream.wait_stream(torch.cuda.current_stream())
|
|
with torch.get_device_module(mr.device).stream(mr.forward_stream):
|
|
with (
|
|
torch.inference_mode(),
|
|
autotune(True, cache=str(autotune_cache)),
|
|
autotune_dummy_run_mode(),
|
|
):
|
|
self._dummy_run(batch_size=batch_size, buffers=buffers)
|
|
torch.cuda.current_stream().wait_stream(mr.forward_stream)
|
|
logger.info("FlashInfer autotune completed.")
|
|
|
|
def _flashinfer_autotune_cache_path(self) -> Path:
|
|
import flashinfer
|
|
|
|
mr = self.model_runner
|
|
major, minor = torch.cuda.get_device_capability(mr.device)
|
|
arch = f"sm{major}{minor}"
|
|
flashinfer_version = getattr(flashinfer, "__version__", "unknown")
|
|
|
|
server_args = mr.server_args
|
|
model_key = "|".join(
|
|
[
|
|
str(server_args.model_path),
|
|
str(mr.dtype),
|
|
str(server_args.quantization),
|
|
str(server_args.moe_runner_backend),
|
|
str(mr.tp_size),
|
|
str(mr.pp_size),
|
|
str(mr.dp_size),
|
|
str(mr.moe_ep_size),
|
|
str(mr.model_config.hf_config.__class__.__name__),
|
|
]
|
|
)
|
|
cache_key = hashlib.sha256(model_key.encode()).hexdigest()[:16]
|
|
cache_dir = (
|
|
Path(envs.SGLANG_CACHE_DIR.get())
|
|
/ "flashinfer"
|
|
/ "autotune"
|
|
/ flashinfer_version
|
|
/ arch
|
|
/ cache_key
|
|
)
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
return (
|
|
cache_dir / f"rank_tp{mr.tp_rank}_pp{mr.pp_rank}_dp{mr.dp_rank or 0}.json"
|
|
)
|
|
|
|
def _alloc_dummy_decode_buffers(self, max_bs: int, *, num_tokens_per_bs: int = 1):
|
|
"""Allocate one static decode-buffer set for a dummy forward, sized to
|
|
(max_bs, max_bs * num_tokens_per_bs).
|
|
|
|
The PP-parallel DeepGEMM warmup sweeps batch sizes far larger than any
|
|
runner's max_bs (up to ~n_sms*block_m), so no pre-allocated runner buffer
|
|
set fits; it builds one here and hands it to _dummy_run (reused across the
|
|
sweep; _dummy_run slices it per shape). The flashinfer autotune does NOT
|
|
use this -- it reuses an existing runner's buffers via _autotune_buffers
|
|
(the eager input registry, or the decode cuda-graph runner's captured
|
|
buffers).
|
|
"""
|
|
mr = self.model_runner
|
|
return _allocate_decode_buffers(
|
|
device=mr.device,
|
|
max_bs=max_bs,
|
|
max_num_token=max_bs * num_tokens_per_bs,
|
|
hidden_size=mr.model_config.hidden_size,
|
|
vocab_size=mr.model_config.vocab_size,
|
|
dtype=mr.model_config.dtype,
|
|
dp_size=mr.server_args.dp_size,
|
|
pp_size=mr.server_args.pp_size,
|
|
is_encoder_decoder=mr.model_config.is_encoder_decoder,
|
|
require_mlp_tp_gather=require_mlp_tp_gather(mr.server_args),
|
|
seq_len_fill_value=mr.attn_backend.get_cuda_graph_seq_len_fill_value(),
|
|
encoder_len_fill_value=(
|
|
getattr(mr.model_config.hf_config, "max_source_positions", 0)
|
|
if mr.model_config.is_encoder_decoder
|
|
else 0
|
|
),
|
|
num_tokens_per_bs=num_tokens_per_bs,
|
|
cache_loc_dtype=torch.int64,
|
|
enable_mamba_track=False,
|
|
hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None),
|
|
pp_proxy_topk_size=mr.get_pp_proxy_topk_size(),
|
|
)
|
|
|
|
def _dummy_run(
|
|
self,
|
|
batch_size: int,
|
|
run_ctx=None,
|
|
forward_mode_override: Optional[ForwardMode] = None,
|
|
*,
|
|
buffers,
|
|
):
|
|
"""Run a dummy forward pass for warmup/profiling.
|
|
|
|
forward_mode_override forces EXTEND/DECODE regardless of
|
|
is_generation (used by the PP-parallel DeepGEMM warmup).
|
|
|
|
buffers: a prepared static buffer set (or lightweight adapter exposing
|
|
the same fields), sized >= this dummy shape, which _dummy_run slices to
|
|
(batch_size, num_tokens). The caller owns the shape and the allocation --
|
|
the flashinfer autotune reuses an existing runner's buffers via
|
|
_autotune_buffers (the eager input registry, or the decode cuda-graph
|
|
runner's captured buffers); the PP-DeepGEMM warmup builds one via
|
|
_alloc_dummy_decode_buffers. _dummy_run never allocates and never re-pads
|
|
(autotune must run at the reused shape; the PP warmup pre-pads and sizes
|
|
its buffer to match). next_token_logits_buffer is optional -- a live
|
|
autotune forward returns logits fresh, so the eager-reuse path passes
|
|
None (only the PP warmup set still carries one).
|
|
"""
|
|
mr = self.model_runner
|
|
if forward_mode_override is not None:
|
|
capture_forward_mode = forward_mode_override
|
|
elif mr.is_generation:
|
|
capture_forward_mode = ForwardMode.DECODE
|
|
else:
|
|
capture_forward_mode = ForwardMode.EXTEND
|
|
capture_hidden_mode = CaptureHiddenMode.NULL
|
|
num_tokens_per_bs = 1
|
|
if mr.spec_algorithm.is_speculative():
|
|
if mr.is_draft_worker:
|
|
if not mr.spec_algorithm.supports_target_verify_for_draft():
|
|
raise RuntimeError("This should not happen")
|
|
capture_forward_mode = ForwardMode.TARGET_VERIFY
|
|
num_tokens_per_bs = (
|
|
mr.spec_algorithm.get_num_tokens_per_bs_for_target_verify(
|
|
mr.server_args.speculative_num_draft_tokens, mr.is_draft_worker
|
|
)
|
|
)
|
|
|
|
if mr.server_args.enable_return_hidden_states:
|
|
capture_hidden_mode = CaptureHiddenMode.FULL
|
|
|
|
num_tokens = batch_size * num_tokens_per_bs
|
|
|
|
# Caller owns the shape: passes a static buffer >= the dummy shape; no
|
|
# allocation, no re-padding (would overflow the reused buffers).
|
|
assert (
|
|
buffers is not None
|
|
and num_tokens <= buffers.input_ids.shape[0]
|
|
and batch_size <= buffers.seq_lens.shape[0]
|
|
), (
|
|
f"_dummy_run needs a static buffer >= (num_tokens={num_tokens}, "
|
|
f"batch_size={batch_size}); got "
|
|
+ (
|
|
"None"
|
|
if buffers is None
|
|
else f"(input_ids={buffers.input_ids.shape[0]}, "
|
|
f"seq_lens={buffers.seq_lens.shape[0]})"
|
|
)
|
|
)
|
|
|
|
seq_len_fill_value = mr.attn_backend.get_cuda_graph_seq_len_fill_value()
|
|
|
|
if mr.server_args.enable_torch_compile:
|
|
set_torch_compile_config()
|
|
should_disable_torch_compile = not getattr(
|
|
mr.model, "_can_torch_compile", True
|
|
)
|
|
if should_disable_torch_compile:
|
|
log_info_on_rank0(
|
|
logger,
|
|
"Transformers backend model reports it is not torch.compile "
|
|
"compatible (e.g. dynamic rope scaling). Disabling torch.compile.",
|
|
)
|
|
mr.server_args.enable_torch_compile = False
|
|
|
|
# NOTE: aux hidden state capture (eagle3/dflash) is already
|
|
# configured by init_aux_hidden_state_capture() in initialize().
|
|
|
|
require_mlp_tp_gather_ = require_mlp_tp_gather(mr.server_args)
|
|
if require_gathered_buffer(mr.server_args):
|
|
assert require_mlp_tp_gather_ or require_attn_tp_gather(mr.server_args)
|
|
|
|
input_ids = buffers.input_ids[:num_tokens]
|
|
positions = buffers.positions[:num_tokens]
|
|
out_cache_loc = buffers.out_cache_loc[:num_tokens]
|
|
# Eager-reuse drops the logits buffer; only buffer sets that carry one slice it.
|
|
next_token_logits_buffer = (
|
|
buffers.next_token_logits_buffer[:num_tokens]
|
|
if buffers.next_token_logits_buffer is not None
|
|
else None
|
|
)
|
|
mrope_positions = buffers.mrope_positions[:, :num_tokens]
|
|
req_pool_indices = buffers.req_pool_indices[:batch_size]
|
|
seq_lens = buffers.seq_lens[:batch_size]
|
|
seq_lens_cpu = buffers.seq_lens_cpu[:batch_size]
|
|
encoder_lens = (
|
|
buffers.encoder_lens[:batch_size]
|
|
if buffers.encoder_lens is not None
|
|
else None
|
|
)
|
|
|
|
buffers.num_token_non_padded[...] = num_tokens
|
|
|
|
# For extend mode
|
|
if capture_forward_mode == ForwardMode.EXTEND:
|
|
extend_prefix_lens_cpu = [0] * batch_size
|
|
extend_seq_lens_cpu = [seq_len_fill_value] * batch_size
|
|
extend_num_tokens = num_tokens
|
|
extend_seq_lens = torch.full(
|
|
(batch_size,), seq_len_fill_value, dtype=torch.int32, device=mr.device
|
|
)
|
|
extend_prefix_lens = torch.zeros(
|
|
(batch_size,), dtype=torch.int32, device=mr.device
|
|
)
|
|
extend_start_loc = torch.arange(
|
|
0, num_tokens, num_tokens_per_bs, dtype=torch.int32, device=mr.device
|
|
)
|
|
else:
|
|
extend_prefix_lens_cpu = None
|
|
extend_seq_lens_cpu = None
|
|
extend_num_tokens = None
|
|
extend_seq_lens = None
|
|
extend_prefix_lens = None
|
|
extend_start_loc = None
|
|
|
|
if mr.server_args.pp_size > 1:
|
|
# PP0 already cp-split hidden_states before send.
|
|
pp_hidden_tokens = num_tokens
|
|
if (
|
|
capture_forward_mode == ForwardMode.EXTEND
|
|
and mr.pp_rank != 0
|
|
and mr.attn_cp_size > 1
|
|
):
|
|
pp_hidden_tokens = num_tokens // mr.attn_cp_size
|
|
pp_proxy_tensors = PPProxyTensors(
|
|
{k: v[:pp_hidden_tokens] for k, v in buffers.pp_proxy_tensors.items()}
|
|
)
|
|
|
|
if require_mlp_tp_gather_:
|
|
global_num_tokens_cpu = [num_tokens] * mr.server_args.dp_size
|
|
elif require_attn_tp_gather(mr.server_args):
|
|
global_num_tokens_cpu = [num_tokens]
|
|
else:
|
|
global_num_tokens_cpu = None
|
|
|
|
if global_num_tokens_cpu is not None:
|
|
global_dp_buffer_len = sum(global_num_tokens_cpu)
|
|
num_tokens_tensor = torch.tensor(
|
|
global_num_tokens_cpu, dtype=torch.int32, device=mr.device
|
|
)
|
|
buffers.global_num_tokens_gpu.copy_(num_tokens_tensor)
|
|
buffers.global_num_tokens_for_logprob_gpu.copy_(num_tokens_tensor)
|
|
else:
|
|
global_dp_buffer_len = None
|
|
global_num_tokens_cpu = None
|
|
|
|
spec_info = create_dummy_verify_input(
|
|
mr.spec_algorithm,
|
|
mr.server_args,
|
|
buffers.custom_mask,
|
|
num_tokens_per_bs,
|
|
mr.is_draft_worker,
|
|
)
|
|
if spec_info is not None and (
|
|
mr.spec_algorithm.is_eagle() or mr.spec_algorithm.is_standalone()
|
|
):
|
|
# MTP models (e.g. deepseek_nextn) read spec_info.hidden_states
|
|
# during forward; provide a dummy so warmup doesn't crash.
|
|
spec_info.hidden_states = torch.zeros(
|
|
(num_tokens, mr.model_config.hidden_size),
|
|
dtype=mr.dtype,
|
|
device=mr.device,
|
|
)
|
|
if capture_hidden_mode != CaptureHiddenMode.FULL:
|
|
capture_hidden_mode = (
|
|
spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
|
|
)
|
|
|
|
if mr.server_args.enable_lora:
|
|
lora_ids = [None] * batch_size
|
|
else:
|
|
lora_ids = None
|
|
|
|
forward_batch = ForwardBatch(
|
|
forward_mode=capture_forward_mode,
|
|
batch_size=batch_size,
|
|
input_ids=input_ids,
|
|
req_pool_indices=req_pool_indices,
|
|
seq_lens=seq_lens,
|
|
seq_lens_cpu=seq_lens_cpu,
|
|
next_token_logits_buffer=next_token_logits_buffer,
|
|
orig_seq_lens=seq_lens,
|
|
out_cache_loc=out_cache_loc,
|
|
seq_lens_sum=seq_lens.sum().item(),
|
|
encoder_lens=encoder_lens,
|
|
return_logprob=False,
|
|
positions=positions,
|
|
extend_num_tokens=extend_num_tokens,
|
|
extend_seq_lens=extend_seq_lens,
|
|
extend_prefix_lens=extend_prefix_lens,
|
|
extend_start_loc=extend_start_loc,
|
|
extend_prefix_lens_cpu=extend_prefix_lens_cpu,
|
|
extend_seq_lens_cpu=extend_seq_lens_cpu,
|
|
global_num_tokens_gpu=buffers.global_num_tokens_gpu,
|
|
global_num_tokens_cpu=global_num_tokens_cpu,
|
|
global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu,
|
|
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
|
|
global_dp_buffer_len=global_dp_buffer_len,
|
|
mrope_positions=mrope_positions,
|
|
spec_algorithm=mr.spec_algorithm,
|
|
spec_info=spec_info,
|
|
capture_hidden_mode=capture_hidden_mode,
|
|
num_token_non_padded=buffers.num_token_non_padded,
|
|
global_forward_mode=capture_forward_mode,
|
|
lora_ids=lora_ids,
|
|
)
|
|
|
|
if lora_ids is not None:
|
|
mr.lora_manager.prepare_lora_batch(forward_batch)
|
|
|
|
mr.attn_backend.init_forward_metadata(forward_batch)
|
|
|
|
def run_once():
|
|
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
|
|
set_dp_buffer_len(
|
|
global_dp_buffer_len,
|
|
num_tokens,
|
|
forward_batch.dp_padding_mode.is_max_len(),
|
|
global_num_tokens_cpu,
|
|
)
|
|
set_is_extend_in_batch(False)
|
|
|
|
kwargs = {}
|
|
if (
|
|
mr.server_args.pp_size > 1
|
|
and "pp_proxy_tensors" in inspect.signature(mr.model.forward).parameters
|
|
):
|
|
kwargs["pp_proxy_tensors"] = PPProxyTensors(
|
|
{k: v.clone() for k, v in pp_proxy_tensors.tensors.items()}
|
|
)
|
|
if not mr.is_generation:
|
|
kwargs["get_embedding"] = True
|
|
|
|
logits_output_or_pp_proxy_tensors = mr.model.forward(
|
|
input_ids,
|
|
forward_batch.positions,
|
|
forward_batch,
|
|
**kwargs,
|
|
)
|
|
return logits_output_or_pp_proxy_tensors
|
|
|
|
torch.get_device_module(mr.device).synchronize()
|
|
mr.tp_group.barrier()
|
|
with forward_context(ForwardContext(attn_backend=mr.attn_backend)):
|
|
with torch.inference_mode(), run_ctx or empty_context():
|
|
run_once()
|
|
|
|
def _autotune_buffers(self) -> Tuple[Optional[Any], Optional[int]]:
|
|
"""Return (buffers, bs) for the autotune dummy forward to reuse; the
|
|
EagerRunner and DecodeCudaGraphRunner override this."""
|
|
return None, None
|
|
|
|
@abstractmethod
|
|
def can_run_graph(self, forward_batch: ForwardBatch) -> bool: ...
|
|
|
|
@abstractmethod
|
|
def load_batch(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
**kwargs,
|
|
) -> Any: ...
|
|
|
|
@abstractmethod
|
|
def execute(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
**kwargs,
|
|
) -> Any: ...
|