b328baec7c
The failing runs showed EAGLE accept length collapsing when draft cache-hit suffixes used the new partial-current splice path. This keeps target partial-current reuse enabled, but returns EAGLE/NextN draft cache-hit suffixes to the previous full-materialize path with an explicit fallback warning until the draft splice path has value-level ETE proof.\n\nThe same change set also tightens the page-granular CP HiCache contract for scheduler-visible hits and makes the prefill-to-decode EAGLE handoff observable without cloning hot-path metadata. Exact non-page CP hits are floored to a page boundary for new scheduling decisions, while internal unfinished-request refresh keeps its exact accounting.\n\nConstraint: CP shared KV and HiCache operate at page granularity; exposing token-precise CP tails to scheduler-visible cache hits can force non-page partial materialization.\nConstraint: EAGLE/NextN draft has only one executable layer, so draft prefetch and draft partial-current splice need a separate correctness contract from target layers.\nRejected: Keep draft partial-current splice enabled | remote logs correlate it with avg accept length around 0.068 and median 0.\nRejected: Clone decode metadata tensors on transfer | slot ownership until process_prebuilt consumes them avoids extra hot-path copies.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable draft partial-current reuse without metadata/draft-KV value checks and ETE accept-length evidence.\nTested: g0034 container py_compile for touched modules.\nTested: g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 183 passed, 5 warnings, 2 subtests passed.\nNot-tested: Fresh ETE accept-length run after this exact commit; requires user-driven traffic restart.
243 lines
10 KiB
Python
243 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from http import HTTPStatus
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
|
|
from sglang.srt.disaggregation.utils import (
|
|
eagle_accept_debug_should_log,
|
|
eagle_accept_debug_tensor_digest,
|
|
)
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode
|
|
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_EAGLE_ACCEPT_PREBUILT_DEBUG_COUNTER = 0
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.overlap_utils import FutureMap
|
|
from sglang.srt.managers.schedule_batch import ScheduleBatch
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
|
|
class ScheduleBatchDisaggregationDecodeMixin:
|
|
|
|
def prepare_for_prebuilt(self: ScheduleBatch):
|
|
"""
|
|
Prepare a prebuilt extend by populate metadata
|
|
Adapted from .prepare_for_extend().
|
|
"""
|
|
|
|
self.forward_mode = ForwardMode.PREBUILT
|
|
reqs = self.reqs
|
|
input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
|
|
extend_num_tokens = sum(len(ids) for ids in input_ids)
|
|
seq_lens = []
|
|
pre_lens = []
|
|
req_pool_indices = []
|
|
|
|
# Pre-calculate total size
|
|
total_size = sum(req.extend_input_len for req in reqs)
|
|
out_cache_loc = torch.empty(total_size, dtype=torch.int64, device=self.device)
|
|
|
|
# Fill the tensor in one pass
|
|
offset = 0
|
|
for i, req in enumerate(reqs):
|
|
req_pool_indices.append(req.req_pool_idx)
|
|
|
|
pre_len = len(req.prefix_indices)
|
|
chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][
|
|
: req.extend_input_len
|
|
]
|
|
assert (
|
|
offset + req.extend_input_len <= total_size
|
|
), f"Exceeds total size: offset={offset}, req.extend_input_len={req.extend_input_len}, total_size={total_size}"
|
|
out_cache_loc[offset : offset + req.extend_input_len] = chunk
|
|
offset += req.extend_input_len
|
|
|
|
seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1)
|
|
seq_lens.append(seq_len)
|
|
if len(req.output_ids) == 0:
|
|
assert (
|
|
seq_len - pre_len == req.extend_input_len
|
|
), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_input_len={req.extend_input_len}"
|
|
|
|
if not req.retracted_stain:
|
|
req.cached_tokens += pre_len - req.already_computed
|
|
req.already_computed = seq_len
|
|
req.is_retracted = False
|
|
pre_lens.append(pre_len)
|
|
req.extend_logprob_start_len = 0
|
|
|
|
if envs.SGLANG_EAGLE_ACCEPT_DEBUG.get():
|
|
global _EAGLE_ACCEPT_PREBUILT_DEBUG_COUNTER
|
|
_EAGLE_ACCEPT_PREBUILT_DEBUG_COUNTER += 1
|
|
counter = _EAGLE_ACCEPT_PREBUILT_DEBUG_COUNTER
|
|
if (
|
|
counter <= 16
|
|
or counter % 256 == 0
|
|
or (pre_len > 0 and (counter <= 128 or counter % 128 == 0))
|
|
):
|
|
logger.warning(
|
|
"[EAGLE_ACCEPT_DEBUG][prebuilt_prepare] rid=%s "
|
|
"pre_len=%s extend_input_len=%s fill_len=%s origin_len=%s "
|
|
"output_len=%s seq_len=%s cached_tokens=%s "
|
|
"req_pool_idx=%s out_chunk_start=%s expected_suffix_start=%s "
|
|
"out_chunk_len=%s",
|
|
str(getattr(req, "rid", ""))[:8],
|
|
pre_len,
|
|
req.extend_input_len,
|
|
len(req.fill_ids),
|
|
len(req.origin_input_ids),
|
|
len(req.output_ids),
|
|
seq_len,
|
|
int(getattr(req, "cached_tokens", 0) or 0),
|
|
req.req_pool_idx,
|
|
0,
|
|
pre_len,
|
|
int(chunk.numel()),
|
|
)
|
|
|
|
extend_input_logprob_token_ids = None
|
|
|
|
# Set fields
|
|
self.input_ids = torch.tensor(
|
|
sum(input_ids, []), dtype=torch.int32, device=self.device
|
|
)
|
|
self.req_pool_indices = torch.tensor(
|
|
req_pool_indices, dtype=torch.int64, device=self.device
|
|
)
|
|
self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, device=self.device)
|
|
self.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64)
|
|
self.orig_seq_lens = torch.tensor(
|
|
seq_lens, dtype=torch.int32, device=self.device
|
|
)
|
|
self.out_cache_loc = out_cache_loc
|
|
self.seq_lens_sum = sum(seq_lens)
|
|
|
|
if self.return_logprob:
|
|
self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
|
|
self.token_ids_logprobs = [r.token_ids_logprob for r in reqs]
|
|
|
|
self.extend_num_tokens = extend_num_tokens
|
|
self.prefix_lens = [len(r.prefix_indices) for r in reqs]
|
|
self.extend_lens = [r.extend_input_len for r in reqs]
|
|
self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs]
|
|
self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
|
|
self.multimodal_inputs = [r.multimodal_inputs for r in reqs]
|
|
|
|
# Build sampling info
|
|
self.sampling_info = SamplingBatchInfo.from_schedule_batch(
|
|
self,
|
|
self.model_config.vocab_size,
|
|
)
|
|
|
|
def process_prebuilt(
|
|
self: ScheduleBatch,
|
|
server_args: ServerArgs,
|
|
future_map: FutureMap,
|
|
):
|
|
"""Assign the buffered last input id to schedule batch"""
|
|
self.output_ids = []
|
|
for req in self.reqs:
|
|
self.output_ids.append(req.output_ids[-1])
|
|
self.tree_cache.cache_unfinished_req(req)
|
|
if req.grammar is not None:
|
|
# FIXME: this try-except block is for handling unexpected xgrammar issue.
|
|
try:
|
|
# if it is not None, then the grammar is from a retracted request, and we should not
|
|
# accept the token as it's already accepted
|
|
if req.grammar.current_token is None:
|
|
req.grammar.accept_token(req.output_ids[-1])
|
|
except ValueError as e:
|
|
from sglang.srt.managers.schedule_batch import FINISH_ABORT
|
|
|
|
# Grammar accept_token can raise ValueError if the token is not in the grammar.
|
|
# This can happen if the grammar is not set correctly or the token is invalid.
|
|
# Use to_finish (not finished_reason) so that process_batch_result_prebuilt
|
|
# handles the release via check_finished -> release_kv_cache in one place.
|
|
error_message = f"Grammar accept_token failed for req {req.rid} with token {req.output_ids[-1]}: {e}"
|
|
req.to_finish = FINISH_ABORT(
|
|
error_message, HTTPStatus.INTERNAL_SERVER_ERROR
|
|
)
|
|
req.grammar.finished = req.finished()
|
|
self.output_ids = torch.tensor(self.output_ids, device=self.device)
|
|
|
|
# Simulate the eagle run.
|
|
if self.spec_algorithm.is_eagle():
|
|
num_states = server_args.speculative_eagle_topk
|
|
if server_args.enable_multi_layer_eagle:
|
|
num_states *= server_args.speculative_num_steps
|
|
topk_p = torch.stack(
|
|
[
|
|
torch.as_tensor(
|
|
req.output_topk_p[:num_states],
|
|
device=self.device,
|
|
dtype=torch.float32,
|
|
)
|
|
for req in self.reqs
|
|
],
|
|
dim=0,
|
|
)
|
|
topk_index = torch.stack(
|
|
[
|
|
torch.as_tensor(
|
|
req.output_topk_index[:num_states],
|
|
device=self.device,
|
|
dtype=torch.int64,
|
|
)
|
|
for req in self.reqs
|
|
],
|
|
dim=0,
|
|
)
|
|
|
|
hidden_states_list = [req.hidden_states_tensor for req in self.reqs]
|
|
hidden_states = torch.stack(hidden_states_list, dim=0).to(self.device)
|
|
if eagle_accept_debug_should_log("prebuilt_state"):
|
|
req0 = self.reqs[0] if self.reqs else None
|
|
logger.warning(
|
|
"[EAGLE_ACCEPT_DEBUG][prebuilt_state] rid=%s bs=%s "
|
|
"num_states=%s output_ids=%s seq_lens=%s topk_p=%s "
|
|
"topk_index=%s hidden=%s",
|
|
str(getattr(req0, "rid", ""))[:8] if req0 is not None else None,
|
|
len(self.reqs),
|
|
num_states,
|
|
self.output_ids[: min(4, self.output_ids.numel())].detach()
|
|
.cpu()
|
|
.tolist(),
|
|
self.seq_lens[: min(4, self.seq_lens.numel())].detach()
|
|
.cpu()
|
|
.tolist(),
|
|
eagle_accept_debug_tensor_digest(topk_p[: min(1, topk_p.shape[0])]),
|
|
eagle_accept_debug_tensor_digest(
|
|
topk_index[: min(1, topk_index.shape[0])]
|
|
),
|
|
eagle_accept_debug_tensor_digest(
|
|
hidden_states[: min(1, hidden_states.shape[0])]
|
|
),
|
|
)
|
|
|
|
# local import to avoid circular import
|
|
from sglang.srt.speculative.eagle_info import EagleDraftInput
|
|
|
|
spec_info = EagleDraftInput(
|
|
topk_p=topk_p,
|
|
topk_index=topk_index,
|
|
hidden_states=hidden_states,
|
|
verified_id=self.output_ids,
|
|
new_seq_lens=self.seq_lens,
|
|
)
|
|
spec_info.prepare_for_extend(self)
|
|
spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
|
|
if self.enable_overlap:
|
|
spec_info.future_indices = future_map.alloc_future_indices(
|
|
len(self.seq_lens)
|
|
)
|
|
future_map.store_to_map_for_new_batch(
|
|
spec_info.future_indices, spec_info
|
|
)
|
|
self.spec_info = spec_info
|