477 lines
20 KiB
Python
477 lines
20 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.
|
|
# ==============================================================================
|
|
"""No-cuda-graph phase runner; the eager dual of BaseCudaGraphRunner."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
from dataclasses import replace
|
|
from typing import TYPE_CHECKING, Any, Tuple, Union
|
|
|
|
import torch
|
|
|
|
from sglang.srt.dllm.config import DllmConfig
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.cp.utils import (
|
|
cp_gather_after_forward,
|
|
cp_shard_model_inputs,
|
|
get_cp_strategy,
|
|
is_cp_active,
|
|
prepare_cp_forward,
|
|
)
|
|
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
|
|
from sglang.srt.model_executor.cuda_graph_buffer_registry import (
|
|
build_eager_registry,
|
|
)
|
|
from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
|
|
create_chunked_prefix_cache_kv_indices,
|
|
)
|
|
from sglang.srt.model_executor.forward_batch_info import (
|
|
ForwardBatch,
|
|
ForwardMode,
|
|
PPProxyTensors,
|
|
)
|
|
from sglang.srt.model_executor.forward_context import (
|
|
ForwardContext,
|
|
forward_context,
|
|
get_req_to_token_pool,
|
|
get_token_to_kv_pool,
|
|
)
|
|
from sglang.srt.model_executor.runner.base_runner import BaseRunner
|
|
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
|
|
enable_tc_piecewise_cuda_graph,
|
|
set_tc_piecewise_forward_context,
|
|
)
|
|
from sglang.srt.model_executor.runner_utils import (
|
|
maybe_publish_prefill_shared_read_done,
|
|
)
|
|
from sglang.srt.runtime_context import (
|
|
get_exec,
|
|
get_parallel,
|
|
get_spec,
|
|
max_prefill_buffer_tokens,
|
|
max_speculative_num_draft_tokens,
|
|
)
|
|
from sglang.srt.utils import is_hip, is_npu
|
|
from sglang.srt.utils.common import (
|
|
ceil_align,
|
|
get_eager_max_batch_size,
|
|
require_mlp_sync,
|
|
)
|
|
from sglang.srt.utils.device_timer import device_timer_ctx
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_is_hip = is_hip()
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
|
from sglang.srt.model_executor.model_runner import ModelRunner
|
|
|
|
|
|
class EagerRunner(BaseRunner):
|
|
def __init__(self, model_runner: ModelRunner) -> None:
|
|
super().__init__(model_runner)
|
|
mr = model_runner
|
|
sa = mr.server_args
|
|
# Built first so the cg runners coalesce onto its buffers via the shared
|
|
# input pool; size to the largest tokens/req across modes the worker hits.
|
|
num_tokens_per_req = 1
|
|
if mr.spec_algorithm.is_speculative():
|
|
# speculative_adaptive can grow draft tokens at runtime; size to the max.
|
|
num_draft_tokens = max_speculative_num_draft_tokens() or 1
|
|
if mr.is_draft_worker:
|
|
num_tokens_per_req = max(
|
|
get_spec().speculative_eagle_topk or 1,
|
|
num_draft_tokens,
|
|
(
|
|
2 * (get_spec().speculative_num_steps or 0)
|
|
if get_spec().enable_multi_layer_eagle
|
|
else 0
|
|
),
|
|
)
|
|
else:
|
|
num_tokens_per_req = mr.decode_num_tokens_per_req(
|
|
num_draft_tokens=num_draft_tokens
|
|
)
|
|
else:
|
|
dllm_config = DllmConfig.from_server_args(sa)
|
|
if dllm_config is not None:
|
|
# dLLM runs block_size tokens/request (DLLM_EXTEND).
|
|
num_tokens_per_req = dllm_config.block_size
|
|
max_bs = mr.max_running_requests
|
|
if (
|
|
mr.is_draft_worker
|
|
and mr.spec_algorithm.is_frozen_kv_mtp()
|
|
and get_spec().speculative_eagle_topk > 1
|
|
):
|
|
# Frozen-KV MTP expands the draft batch by topk on the bs axis
|
|
# (expand_for_topk_draft) before the eager fallback.
|
|
max_bs *= get_spec().speculative_eagle_topk
|
|
# Mirror prepare_mlp_sync_batch padding so the registry holds what load_batch copies.
|
|
max_bs = get_eager_max_batch_size(max_bs)
|
|
prefill_ceiling = max(mr.max_total_num_tokens, max_prefill_buffer_tokens())
|
|
max_num_token = max(prefill_ceiling, max_bs * num_tokens_per_req)
|
|
if require_mlp_sync():
|
|
from sglang.srt.layers.cp.padding import get_cp_padding_align_size
|
|
|
|
max_num_token = ceil_align(max_num_token, self.attn_tp_size)
|
|
max_num_token = ceil_align(max_num_token, get_cp_padding_align_size())
|
|
self._eager_max_bs = max_bs
|
|
self._eager_num_tokens_per_req = num_tokens_per_req
|
|
is_encoder_decoder = mr.model_config.is_encoder_decoder
|
|
self._eager_registry = build_eager_registry(
|
|
device=mr.device,
|
|
max_bs=max_bs,
|
|
max_num_token=max_num_token,
|
|
cache_loc_dtype=torch.int64,
|
|
enable_mamba_track=(
|
|
get_exec().mamba.enable_mamba_extra_buffer
|
|
and mr.spec_algorithm.is_none()
|
|
),
|
|
is_encoder_decoder=is_encoder_decoder,
|
|
encoder_len_fill_value=(
|
|
getattr(mr.model_config.hf_config, "max_source_positions", 0)
|
|
if is_encoder_decoder
|
|
else 0
|
|
),
|
|
encoder_lens_dtype=(
|
|
torch.int64 if torch.device(mr.device).type == "cpu" else torch.int32
|
|
),
|
|
dp_size=get_parallel().dp_size,
|
|
)
|
|
# Eager has no capture step, so warm up here (run-once via mr._kernel_warmed_up).
|
|
self.warmup()
|
|
|
|
def _autotune_buffers(self) -> Tuple[Any, int]:
|
|
"""Decode-shaped dummy buffers (bs * num_tokens_per_req) for the warmup
|
|
flashinfer-autotune forward.
|
|
|
|
flashinfer's MoE autotuner times candidate tactics against the buffer it
|
|
is given, so it must match the live decode shape for the cached tactic to
|
|
be optimal at decode. The eager input registry spans the prefill token
|
|
ceiling; the dummy run only needs the decode-sized slice.
|
|
"""
|
|
mr = self.model_runner
|
|
num_tokens_per_req = 1
|
|
if mr.spec_algorithm.is_speculative():
|
|
num_tokens_per_req = mr.decode_num_tokens_per_req()
|
|
return (
|
|
self._alloc_dummy_decode_buffers(
|
|
self._eager_max_bs, num_tokens_per_req=num_tokens_per_req
|
|
),
|
|
self._eager_max_bs,
|
|
)
|
|
|
|
def can_run_graph(self, forward_batch: ForwardBatch) -> bool:
|
|
# Eager never runs a cuda graph; callers dispatch on isinstance(...,
|
|
# EagerRunner) and must not route an eager batch into a replay branch.
|
|
return False
|
|
|
|
def load_batch(
|
|
self, forward_batch: ForwardBatch, pp_proxy_tensors=None, **kwargs
|
|
) -> ForwardBatch:
|
|
"""Copy the live batch into the fixed-max eager static buffers (sliced to
|
|
this batch's shape) — the eager counterpart of the cuda-graph runners'
|
|
load_batch."""
|
|
if envs.SGLANG_EAGER_INPUT_NO_COPY.get():
|
|
return replace(forward_batch)
|
|
raw_bs = forward_batch.batch_size
|
|
if forward_batch.input_ids is not None:
|
|
raw_num_tokens = forward_batch.input_ids.shape[0]
|
|
elif forward_batch.input_embeds is not None:
|
|
raw_num_tokens = forward_batch.input_embeds.shape[0]
|
|
else:
|
|
raw_num_tokens = 0
|
|
registry = self._eager_registry
|
|
registry.fill_from(
|
|
forward_batch,
|
|
raw_bs=raw_bs,
|
|
padded_bs=raw_bs,
|
|
raw_num_tokens=raw_num_tokens,
|
|
padded_num_tokens=raw_num_tokens,
|
|
pp_proxy_tensors=pp_proxy_tensors,
|
|
)
|
|
return registry.extract_buffer(
|
|
padded_bs=raw_bs,
|
|
padded_num_tokens=raw_num_tokens,
|
|
forward_batch_template=forward_batch,
|
|
)
|
|
|
|
def execute(
|
|
self, forward_batch: ForwardBatch, pp_proxy_tensors=None, **kwargs
|
|
) -> Any:
|
|
mode = forward_batch.forward_mode
|
|
if mode.is_mixed() and not is_npu() and get_cp_strategy() is None:
|
|
# A mixed batch is extend-shaped (decode tails are 1-token
|
|
# extends); run it as EXTEND. NPU keeps MIXED for its dedicated
|
|
# kernel; CP keeps it to skip the zigzag split.
|
|
forward_batch.forward_mode = ForwardMode.EXTEND
|
|
mode = ForwardMode.EXTEND
|
|
if mode.is_decode():
|
|
return self._execute_decode(forward_batch, pp_proxy_tensors)
|
|
if mode.is_idle():
|
|
return self._execute_idle(forward_batch, pp_proxy_tensors)
|
|
if mode.is_extend(include_draft_extend_v2=True):
|
|
return self._execute_extend(forward_batch, pp_proxy_tensors)
|
|
raise ValueError(f"Invalid forward mode for eager runner: {mode}")
|
|
|
|
def _resolve_decode_pdmux(
|
|
self,
|
|
) -> Tuple[Any, contextlib.AbstractContextManager]:
|
|
"""Resolve the (attn_backend, forward_context) the eager decode forward
|
|
runs under. PDmux selects a per-stream backend and publishes it via an
|
|
active ForwardContext; non-pdmux uses attn_backend + the ambient ctx."""
|
|
model_runner = self.model_runner
|
|
if self.enable_pdmux:
|
|
return model_runner.decode_attn_backend, forward_context(
|
|
ForwardContext(attn_backend=model_runner.decode_attn_backend)
|
|
)
|
|
return model_runner.attn_backend, contextlib.nullcontext()
|
|
|
|
def _execute_decode(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
pp_proxy_tensors=None,
|
|
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
|
|
model_runner = self.model_runner
|
|
enable_pdmux = self.enable_pdmux
|
|
attn_backend, pdmux_ctx = self._resolve_decode_pdmux()
|
|
if not enable_pdmux:
|
|
forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
|
|
if forward_batch.needs_forward_metadata_init():
|
|
if hasattr(model_runner.model, "prepare_forward_batch"):
|
|
# Prepare model-specific attention metadata before planning,
|
|
# e.g. Moss-VL's prefill cross-attention custom mask.
|
|
model_runner.model.prepare_forward_batch(forward_batch)
|
|
attn_backend.init_forward_metadata(forward_batch)
|
|
# FIXME: add pp_proxy_tensors arg to all models
|
|
kwargs = model_runner._pp_kwargs(pp_proxy_tensors)
|
|
|
|
ctx = device_timer_ctx(model_runner.device_timer, "decode")
|
|
|
|
with ctx, pdmux_ctx:
|
|
return model_runner.model.forward(
|
|
forward_batch.input_ids,
|
|
forward_batch.positions,
|
|
forward_batch,
|
|
**kwargs,
|
|
)
|
|
|
|
def _execute_extend(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
pp_proxy_tensors=None,
|
|
) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
|
|
model_runner = self.model_runner
|
|
kwargs = model_runner._extend_forward_kwargs(forward_batch, pp_proxy_tensors)
|
|
|
|
if not self.enable_pdmux:
|
|
forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
|
|
|
|
cp_active = is_cp_active(forward_batch)
|
|
if cp_active:
|
|
prepare_cp_forward(forward_batch)
|
|
|
|
# Target verify can arrive with ``forward_metadata_ready`` set by an
|
|
# upstream/speculative planning step. That mark does not initialize
|
|
# the final target hybrid backend, and unlike a graph replay eager has
|
|
# no static metadata load to fill the gap. Re-plan target verify from
|
|
# the final batch every time; eager metadata is intentionally derived
|
|
# directly from the live ``spec_info`` tensors.
|
|
if (
|
|
forward_batch.needs_forward_metadata_init()
|
|
or cp_active
|
|
or forward_batch.forward_mode.is_target_verify()
|
|
):
|
|
if model_runner.attn_dcp_size > 1 and hasattr(
|
|
model_runner.model, "prepare_context_parallel_metadata_for_dcp"
|
|
):
|
|
# prepare kv cache buffer for dcp to gather kv cache
|
|
forward_batch.attn_dcp_metadata = (
|
|
model_runner.model.prepare_context_parallel_metadata_for_dcp(
|
|
forward_batch.seq_lens,
|
|
forward_batch.extend_prefix_lens,
|
|
forward_batch.extend_prefix_lens_cpu,
|
|
forward_batch.extend_seq_lens,
|
|
forward_batch.req_pool_indices,
|
|
get_req_to_token_pool().req_to_token,
|
|
forward_batch.seq_lens_sum,
|
|
get_token_to_kv_pool().get_kv_buffer_shape()[0],
|
|
model_runner.kv_cache_dtype,
|
|
model_runner.device,
|
|
create_chunked_prefix_cache_kv_indices,
|
|
)
|
|
)
|
|
if hasattr(model_runner.model, "prepare_forward_batch"):
|
|
# Prepare model-specific attention metadata before planning,
|
|
# e.g. Moss-VL's prefill cross-attention custom mask.
|
|
model_runner.model.prepare_forward_batch(forward_batch)
|
|
model_runner.attn_backend.init_forward_metadata(forward_batch)
|
|
model_runner.attn_backend.prepare_prefill_shared_read_snapshot(
|
|
forward_batch,
|
|
num_qo_tokens=len(forward_batch.input_ids),
|
|
)
|
|
maybe_publish_prefill_shared_read_done(
|
|
model_runner,
|
|
forward_batch,
|
|
torch.get_device_module(model_runner.device),
|
|
)
|
|
|
|
if not cp_active:
|
|
forward_batch.attn_cp_metadata = None
|
|
|
|
category = (
|
|
"target_verify"
|
|
if forward_batch.forward_mode.is_target_verify()
|
|
else "extend"
|
|
)
|
|
with device_timer_ctx(model_runner.device_timer, category):
|
|
pcg_runner = model_runner.prefill_cuda_graph_runner
|
|
if (
|
|
_is_hip
|
|
and pcg_runner is not None
|
|
and not isinstance(pcg_runner, EagerRunner)
|
|
and not cp_active
|
|
):
|
|
# HIP PCG eager fallback: enter the PCG context so Dynamo guards
|
|
# and PCG-specific MoE/attention paths stay consistent.
|
|
with (
|
|
enable_tc_piecewise_cuda_graph(),
|
|
set_tc_piecewise_forward_context(
|
|
forward_batch,
|
|
model_runner.attention_layers,
|
|
getattr(model_runner.model, "quant_config", None),
|
|
model_runner.moe_layers,
|
|
model_runner.moe_fusions,
|
|
dsa_indexers=model_runner.dsa_indexers,
|
|
mha_companion_layers=model_runner.mha_companion_layers,
|
|
),
|
|
):
|
|
ret = model_runner.model.forward(
|
|
forward_batch.input_ids,
|
|
forward_batch.positions,
|
|
forward_batch,
|
|
**kwargs,
|
|
)
|
|
elif cp_active:
|
|
ret = self._execute_extend_cp(forward_batch, kwargs)
|
|
else:
|
|
ret = model_runner.model.forward(
|
|
forward_batch.input_ids,
|
|
forward_batch.positions,
|
|
forward_batch,
|
|
**kwargs,
|
|
)
|
|
return ret
|
|
|
|
def _execute_extend_cp(
|
|
self, forward_batch: ForwardBatch, kwargs: dict
|
|
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
|
|
"""CP extend: shard inputs at the model boundary, run the body on the
|
|
rank-local slice, then gather hidden states before the logits step.
|
|
"""
|
|
model = self.model_runner.model
|
|
|
|
input_ids = forward_batch.input_ids
|
|
input_embeds = kwargs.get("input_embeds")
|
|
# Multimodal spans must be embedded in global token order, before CP
|
|
# slicing. The model may also normalize image hash IDs for its router.
|
|
prepare_inputs = getattr(model, "prepare_language_model_inputs", None)
|
|
if prepare_inputs is not None:
|
|
input_ids, input_embeds = prepare_inputs(
|
|
input_ids, forward_batch, input_embeds
|
|
)
|
|
if input_embeds is None:
|
|
input_embeds = model.get_input_embeddings()(input_ids)
|
|
with cp_shard_model_inputs(
|
|
input_embeds,
|
|
forward_batch.positions,
|
|
forward_batch,
|
|
input_ids,
|
|
) as (sharded_input_embeds, sharded_positions, model_input_ids):
|
|
model_kwargs = {"input_embeds": sharded_input_embeds}
|
|
if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None:
|
|
model_kwargs["pp_proxy_tensors"] = pp_proxy_tensors
|
|
hidden_states = model.model(
|
|
model_input_ids,
|
|
sharded_positions,
|
|
forward_batch,
|
|
**model_kwargs,
|
|
)
|
|
capture_aux_hidden_states = getattr(model, "capture_aux_hidden_states", False)
|
|
aux_hidden_states = None
|
|
if capture_aux_hidden_states:
|
|
hidden_states, aux_hidden_states = hidden_states
|
|
|
|
if not model.pp_group.is_last_rank:
|
|
return (
|
|
(hidden_states, aux_hidden_states)
|
|
if capture_aux_hidden_states
|
|
else hidden_states
|
|
)
|
|
|
|
stream = torch.cuda.current_stream()
|
|
hidden_states = cp_gather_after_forward(hidden_states, forward_batch, stream)
|
|
# DSpark aux tensors ride the same CP token split; gather them the same way.
|
|
if aux_hidden_states is not None:
|
|
if isinstance(aux_hidden_states, torch.Tensor):
|
|
aux_hidden_states = cp_gather_after_forward(
|
|
aux_hidden_states, forward_batch, stream
|
|
)
|
|
else:
|
|
aux_hidden_states = [
|
|
cp_gather_after_forward(aux, forward_batch, stream)
|
|
for aux in aux_hidden_states
|
|
]
|
|
logits_kwargs = {}
|
|
# DSV4 returns (hidden_states, hidden_states_before_norm) from its model body.
|
|
if isinstance(hidden_states, tuple):
|
|
hidden_states, hidden_states_before_norm = hidden_states
|
|
# Mirror DeepseekV4ForCausalLM.forward: drop pre_hc_head when
|
|
# DSpark aux capture is on, else it overrides the packed aux.
|
|
if aux_hidden_states is None:
|
|
logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
|
|
return model.logits_processor(
|
|
input_ids,
|
|
hidden_states,
|
|
model.lm_head,
|
|
forward_batch,
|
|
aux_hidden_states,
|
|
**logits_kwargs,
|
|
)
|
|
|
|
def _execute_idle(
|
|
self, forward_batch: ForwardBatch, pp_proxy_tensors=None
|
|
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
|
|
model_runner = self.model_runner
|
|
# Padded idle (DP-attn MLP sync) needs metadata reinit; unpadded must
|
|
# drop stale forward_metadata to avoid an SWA use-after-free on req_pool.
|
|
if forward_batch.batch_size > 0:
|
|
if not self.enable_pdmux:
|
|
forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
|
|
model_runner.attn_backend.init_forward_metadata(forward_batch)
|
|
else:
|
|
model_runner.attn_backend.forward_metadata = None
|
|
|
|
kwargs = model_runner._pp_kwargs(pp_proxy_tensors)
|
|
with device_timer_ctx(model_runner.device_timer, "idle"):
|
|
return model_runner.model.forward(
|
|
forward_batch.input_ids,
|
|
forward_batch.positions,
|
|
forward_batch,
|
|
**kwargs,
|
|
)
|