chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Stabilize EAGLE draft cache hits under CP HiCache

Browse Source 
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.
This commit is contained in:
laoyao0822
2026-05-30 22:31:43 +08:00
parent 10296a5fef
commit b328baec7c
14 changed files with 920 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+21
View File
@@ -42,6 +42,8 @@ from sglang.srt.disaggregation.utils import (
ReqToMetadataIdxAllocator,
TransferBackend,
append_cp_draft_state_buffers,
eagle_accept_debug_should_log,
eagle_accept_debug_tensor_digest,
get_kv_class,
is_mla_backend,
poll_and_all_reduce,
@@ -1172,6 +1174,25 @@ class DecodeTransferQueue:
return True
# Case 3: Success - commit the transfer
if (
not self.spec_algorithm.is_none()
and eagle_accept_debug_should_log("metadata_get")
):
logger.warning(
"[EAGLE_ACCEPT_DEBUG][metadata_get] rid=%s room=%s idx=%s "
"output_id=%s cached_tokens=%s actual_room=%s topk_p=%s "
"topk_index=%s hidden=%s",
str(getattr(decode_req.req, "rid", ""))[:8],
decode_req.req.bootstrap_room,
idx,
int(output_id[0].item()),
int(cached_tokens[0].item()),
int(actual_room),
eagle_accept_debug_tensor_digest(output_topk_p[:1]),
eagle_accept_debug_tensor_digest(output_topk_index[:1]),
eagle_accept_debug_tensor_digest(output_hidden_states),
)
decode_req.req.output_ids.append(output_id[0].item())
decode_req.req.cached_tokens = cached_tokens[0].item()
if not self.spec_algorithm.is_none():
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+59 -1
View File
@@ -6,10 +6,16 @@ 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
@@ -42,6 +48,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
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
]
@@ -51,7 +58,6 @@ class ScheduleBatchDisaggregationDecodeMixin:
out_cache_loc[offset : offset + req.extend_input_len] = chunk
offset += req.extend_input_len
pre_len = len(req.prefix_indices)
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:
@@ -66,6 +72,35 @@ class ScheduleBatchDisaggregationDecodeMixin:
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
@@ -161,6 +196,29 @@ class ScheduleBatchDisaggregationDecodeMixin:
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
python/sglang/srt/disaggregation/utils.py
+81
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import os
import random
import logging
from collections import deque
from contextlib import nullcontext
from enum import Enum
@@ -14,6 +15,8 @@ import torch.distributed as dist
from sglang.srt.environ import envs
from sglang.srt.utils import is_npu
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.disaggregation.base.conn import KVArgs
from sglang.srt.disaggregation.common.conn import (
@@ -84,6 +87,63 @@ def poll_and_all_reduce_attn_cp_tp_group(
# Metadata Buffers
#########################
_EAGLE_ACCEPT_DEBUG_COUNTERS = {}
def eagle_accept_debug_should_log(
key: str,
*,
first: int = 16,
every: int = 256,
) -> bool:
if not envs.SGLANG_EAGLE_ACCEPT_DEBUG.get():
return False
count = _EAGLE_ACCEPT_DEBUG_COUNTERS.get(key, 0) + 1
_EAGLE_ACCEPT_DEBUG_COUNTERS[key] = count
return count <= first or (every > 0 and count % every == 0)
def eagle_accept_debug_tensor_digest(tensor: Any, *, sample: int = 64) -> str:
"""Small deterministic tensor summary for EAGLE handoff debugging.
This intentionally samples only a prefix. It is enabled only under
SGLANG_EAGLE_ACCEPT_DEBUG and is for handoff equality checks, not for full
numerical validation.
"""
if tensor is None:
return "None"
try:
t = torch.as_tensor(tensor)
except Exception as exc: # pragma: no cover - defensive debug helper
return f"unavailable({type(exc).__name__})"
shape = tuple(t.shape)
dtype = str(t.dtype).replace("torch.", "")
numel = int(t.numel())
if numel == 0:
return f"shape={shape} dtype={dtype} numel=0"
flat = t.detach().reshape(-1)
sample_count = min(sample, numel)
try:
sample_cpu = flat[:sample_count].cpu()
head_cpu = sample_cpu[: min(8, sample_count)]
head = head_cpu.tolist()
if sample_cpu.is_floating_point():
checksum = float(sample_cpu.float().sum().item())
abs_checksum = float(sample_cpu.float().abs().sum().item())
return (
f"shape={shape} dtype={dtype} sample={sample_count} "
f"sum={checksum:.6g} abs={abs_checksum:.6g} head={head}"
)
checksum = int(sample_cpu.to(torch.int64).sum().item())
return (
f"shape={shape} dtype={dtype} sample={sample_count} "
f"sum={checksum} head={head}"
)
except Exception as exc: # pragma: no cover - defensive debug helper
return f"shape={shape} dtype={dtype} numel={numel} digest_error={type(exc).__name__}"
def append_cp_draft_state_buffers(
kv_args: Any,
@@ -307,6 +367,27 @@ class MetadataBuffers:
self.output_hidden_states[req.metadata_buffer_index].copy_(
req.hidden_states_tensor
)
if eagle_accept_debug_should_log("metadata_set"):
logger.warning(
"[EAGLE_ACCEPT_DEBUG][metadata_set] rid=%s room=%s idx=%s "
"output_id=%s cached_tokens=%s topk=%s topk_p=%s "
"topk_index=%s hidden=%s",
str(getattr(req, "rid", ""))[:8],
getattr(req, "bootstrap_room", None),
req.metadata_buffer_index,
req.output_ids[0] if req.output_ids else None,
req.cached_tokens,
topk,
eagle_accept_debug_tensor_digest(
self.output_topk_p[req.metadata_buffer_index, :topk]
),
eagle_accept_debug_tensor_digest(
self.output_topk_index[req.metadata_buffer_index, :topk]
),
eagle_accept_debug_tensor_digest(
self.output_hidden_states[req.metadata_buffer_index]
),
)
# Store bootstrap_room for validation on decode side
self.bootstrap_room[req.metadata_buffer_index, 0] = (
req.bootstrap_room if req.bootstrap_room is not None else 0
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+25 -7
View File
@@ -978,20 +978,38 @@ def can_reuse_current_extend_kv(forward_batch) -> bool:
def should_reuse_current_extend_kv(forward_batch) -> bool:
"""Return whether MLA should splice current extend KV into materialized KV.
The contract is model-role agnostic: target and prefill-time EAGLE/NextN
draft both write only the CP-owned suffix rows into the persistent pool, so
cache-hit attention should materialize the page-aligned prefix and splice the
freshly computed current suffix explicitly. Draft async prefetch remains
disabled elsewhere because NextN has no next layer to prefetch, but disabling
that optimization must not force draft cache-hit suffixes back to full
materialization from the pool.
Current-only reuse is safe for both target and draft because there is no
cached prefix to compose. Partial current reuse is currently a target-model
contract only. EAGLE/NextN draft cache-hit suffixes keep using the older
full-materialize path until the draft splice path has value-level ETE proof;
the 2026-05-30 accept-length regression correlated with enabling that draft
partial splice path.
"""
if not cp_shared_kv_current_reuse_enabled():
return False
current_only = is_current_only_extend_batch(forward_batch)
if current_only:
return True
partial_current = can_reuse_current_extend_kv(forward_batch)
if partial_current and cp_shared_kv_is_draft_input(forward_batch):
prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
extend_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
_log_current_reuse_fallback(
"draft_partial_current_reuse_disabled",
"cache-hit EAGLE/NextN draft uses full materialize instead of "
"partial current reuse. prefix_lens=%s extend_lens=%s",
[int(x) for x in prefix_lens_cpu]
if prefix_lens_cpu is not None
else None,
[int(x) for x in extend_lens_cpu]
if extend_lens_cpu is not None
else None,
)
return False
return current_only or partial_current
python/sglang/srt/layers/attention/nsa_backend.py
+66
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import logging
from dataclasses import dataclass
from enum import IntEnum, auto
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeAlias
@@ -68,6 +69,8 @@ if TYPE_CHECKING:
_is_hip = is_hip()
logger = logging.getLogger(__name__)
_EAGLE_ACCEPT_DRAFT_MLA_PATH_DEBUG_COUNTS: Dict[Tuple[int, int], int] = {}
if _is_hip:
from sglang.srt.layers.attention.nsa.triton_kernel import get_valid_kv_indices
@@ -1737,6 +1740,10 @@ class NativeSparseAttnBackend(
and topk_transform_method == TopkTransformMethod.PAGED
):
assert forward_batch.cp_shared_kv_layout is not None
is_draft_mla_input = cp_shared_kv_is_draft_input(forward_batch)
eagle_draft_mla_branch = "not_selected"
eagle_draft_mla_used_prefetch = False
eagle_draft_mla_current_rows = None
mla_prefetcher = getattr(
forward_batch, "cp_shared_kv_mla_prefetcher", None
)
@@ -1789,6 +1796,7 @@ class NativeSparseAttnBackend(
)
if is_current_only_extend_batch(forward_batch):
eagle_draft_mla_branch = "current_only"
current_mask, page_table_1 = build_current_loc_remap(
logical_page_table_1,
forward_batch.out_cache_loc,
@@ -1846,8 +1854,11 @@ class NativeSparseAttnBackend(
current_remap_logical_page_capacity=current_remap_logical_page_capacity,
)
if prefetched_kv is not None:
eagle_draft_mla_branch = "partial_current_prefetch"
eagle_draft_mla_used_prefetch = True
kv_cache, page_table_1 = prefetched_kv
else:
eagle_draft_mla_branch = "partial_current_sync"
prefix_lens_cpu = getattr(
forward_batch, "extend_prefix_lens_cpu", None
)
@@ -1969,7 +1980,9 @@ class NativeSparseAttnBackend(
tensor_debug_checksum(k),
tensor_debug_checksum(k_rope),
)
eagle_draft_mla_current_rows = int(current_kv_cache.shape[0])
else:
eagle_draft_mla_branch = "full_materialize"
prefetched_kv = None
if mla_prefetcher is not None:
prefetched_kv = mla_prefetcher.consume(
@@ -1978,6 +1991,8 @@ class NativeSparseAttnBackend(
logical_locs=page_table_1,
)
if prefetched_kv is not None:
eagle_draft_mla_branch = "full_materialize_prefetch"
eagle_draft_mla_used_prefetch = True
kv_cache, page_table_1 = prefetched_kv
else:
slot_remap = get_or_build_shared_token_kv_slot_remap(
@@ -1998,6 +2013,57 @@ class NativeSparseAttnBackend(
nvtx_source="mla.full_materialize",
nvtx_layer_id=layer.layer_id,
)
if (
envs.SGLANG_EAGLE_ACCEPT_DEBUG.get()
and is_draft_mla_input
and int(layer.layer_id) == 0
):
debug_key = (
int(forward_batch.cp_shared_kv_layout.cp_rank),
int(layer.layer_id),
)
debug_count = (
_EAGLE_ACCEPT_DRAFT_MLA_PATH_DEBUG_COUNTS.get(debug_key, 0) + 1
)
_EAGLE_ACCEPT_DRAFT_MLA_PATH_DEBUG_COUNTS[debug_key] = debug_count
if (
debug_count <= 16
or eagle_draft_mla_branch in ("full_materialize", "not_selected")
or debug_count % 256 == 0
):
prefix_lens_cpu = getattr(
forward_batch, "extend_prefix_lens_cpu", None
)
extend_lens_cpu = getattr(
forward_batch, "extend_seq_lens_cpu", None
)
out_cache_loc = getattr(forward_batch, "out_cache_loc", None)
logger.warning(
"[EAGLE_ACCEPT_DEBUG][draft_mla_path] cp_rank=%s "
"layer=%s count=%s branch=%s used_prefetch=%s "
"has_prefetcher=%s can_current_reuse=%s prefix_lens=%s "
"extend_lens=%s current_rows=%s kv_rows=%s "
"page_table_shape=%s out_cache_loc_shape=%s",
forward_batch.cp_shared_kv_layout.cp_rank,
layer.layer_id,
debug_count,
eagle_draft_mla_branch,
eagle_draft_mla_used_prefetch,
mla_prefetcher is not None,
can_reuse_current_kv,
[int(x) for x in prefix_lens_cpu]
if prefix_lens_cpu is not None
else None,
[int(x) for x in extend_lens_cpu]
if extend_lens_cpu is not None
else None,
eagle_draft_mla_current_rows,
int(kv_cache.shape[0]) if kv_cache is not None else None,
tuple(page_table_1.shape)
if page_table_1 is not None
else None,
tuple(out_cache_loc.shape) if out_cache_loc is not None else None,
)
if mla_prefetcher is not None and cp_shared_kv_should_prefetch_next_layer(
forward_batch, layer.layer_id
):
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+7
View File
@@ -63,6 +63,13 @@ class SchedulerOutputProcessorMixin:
)
allocator.free(idx)
req.metadata_buffer_index = -1
# The EAGLE handoff tensors are views into the reusable metadata slot.
# process_prebuilt consumes them into the batch before this helper is
# called; keeping the views afterwards makes the request observe a
# future transfer that reuses the same slot.
req.output_topk_p = None
req.output_topk_index = None
req.hidden_states_tensor = None
def _maybe_log_eagle_accept_debug(
self: Scheduler,
python/sglang/srt/mem_cache/base_prefix_cache.py
+1
View File
@@ -37,6 +37,7 @@ class MatchPrefixParams:
"""Unified parameters for match_prefix across different cache types"""
key: RadixKey
cp_floor_exact: bool = True
# Mamba specific
cow_mamba: bool = False
python/sglang/srt/mem_cache/hiradix_cache.py
+40 -7
View File
@@ -2109,7 +2109,7 @@ class HiRadixCache(RadixCache):
child = node.children[child_key]
prefix_len = self.key_match_fn(child.key, key)
prefix_len = self._cp_floor_exact_valid_tail_extension_len(
child, prefix_len, len(key)
child, prefix_len, len(key), floor_exact_key=True
)
if prefix_len <= 0:
break
@@ -2144,14 +2144,35 @@ class HiRadixCache(RadixCache):
return prefix_len // self.page_size * self.page_size
def _cp_floor_exact_valid_tail_extension_len(
self, child: TreeNode, prefix_len: int, request_len: int
self,
child: TreeNode,
prefix_len: int,
request_len: int,
*,
floor_exact_key: bool = False,
) -> int:
"""Floor exact CP valid-tail hits to the previous physical page.
CP HiCache may keep radix keys at scheduler-visible valid lengths while
the underlying target/draft pools and host reservations are page-owned.
Even when the incoming radix key exactly equals a non-page-aligned child
key, the scheduler will still compute the current suffix token
(`max_prefix_len = input_len - 1`). Exposing the sub-page tail as a
protected prefix lets later backup/current-reuse paths start inside a
page. Sacrifice that tail and let the new request re-own it.
`floor_exact_key` is enabled for scheduler-visible prefix matching and
prepared-backup probing. Internal cache insertion refreshes can keep
exact sub-page tails for the current request so they do not immediately
invalidate their own just-inserted prefix.
"""
if (
not self._uses_cp_hicache
or self.page_size <= 1
or prefix_len <= 0
or prefix_len != len(child.key)
or prefix_len >= request_len
or (prefix_len >= request_len and not floor_exact_key)
or prefix_len % self.page_size == 0
):
return prefix_len
@@ -3505,7 +3526,11 @@ class HiRadixCache(RadixCache):
deferred_node = None
try:
value, last_node = self._match_prefix_helper(self.root_node, key)
value, last_node = self._match_prefix_helper(
self.root_node,
key,
floor_exact_key=getattr(params, "cp_floor_exact", True),
)
except HiCachePendingBackupSplit as exc:
value = []
last_node = exc.node.parent if exc.node.parent is not None else self.root_node
@@ -3633,7 +3658,9 @@ class HiRadixCache(RadixCache):
return matched_length
def _match_prefix_helper(self, node: TreeNode, key: RadixKey):
def _match_prefix_helper(
self, node: TreeNode, key: RadixKey, *, floor_exact_key: bool = True
):
node.last_access_time = time.monotonic()
child_key = self.get_child_key_fn(key)
value = []
@@ -3646,7 +3673,10 @@ class HiRadixCache(RadixCache):
child.pin_expiry = time.monotonic() + child.pin_ttl
raw_prefix_len = self.key_match_fn(child.key, key)
prefix_len = self._cp_floor_exact_valid_tail_extension_len(
child, raw_prefix_len, len(key)
child,
raw_prefix_len,
len(key),
floor_exact_key=floor_exact_key,
)
stop_after_page_floor = prefix_len != raw_prefix_len
prune_stale_tail_after_split = stop_after_page_floor
@@ -3765,7 +3795,10 @@ class HiRadixCache(RadixCache):
node.priority = max(node.priority, priority)
raw_prefix_len = self.key_match_fn(node.key, key)
prefix_len = self._cp_floor_exact_valid_tail_extension_len(
node, raw_prefix_len, len(key)
node,
raw_prefix_len,
len(key),
floor_exact_key=prepared_cp_backup is not None,
)
stop_after_page_floor = prefix_len != raw_prefix_len
if prefix_len <= 0:
python/sglang/srt/mem_cache/radix_cache.py
+3 -1
View File
@@ -653,7 +653,9 @@ class RadixCache(BasePrefixCache):
)
# The prefix indices could be updated, reuse it
match_result = self.match_prefix(MatchPrefixParams(key=radix_key))
match_result = self.match_prefix(
MatchPrefixParams(key=radix_key, cp_floor_exact=False)
)
new_indices, new_last_node = (
match_result.device_indices,
match_result.last_device_node,
python/sglang/srt/models/deepseek_nextn.py
+78
View File
@@ -63,6 +63,19 @@ logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
_is_npu = is_npu()
_EAGLE_ACCEPT_CP_DRAFT_HIDDEN_DEBUG_COUNTS = {}
def _log_eagle_accept_cp_draft_hidden_debug(key: str, message: str, *args):
if not envs.SGLANG_EAGLE_ACCEPT_DEBUG.get():
return
count = _EAGLE_ACCEPT_CP_DRAFT_HIDDEN_DEBUG_COUNTS.get(key, 0) + 1
_EAGLE_ACCEPT_CP_DRAFT_HIDDEN_DEBUG_COUNTS[key] = count
if count <= 16 or count % 256 == 0:
logger.warning(
"[EAGLE_ACCEPT_DEBUG][cp_draft_hidden] " + message,
*args,
)
class DeepseekModelNextN(nn.Module):
@@ -150,10 +163,75 @@ class DeepseekModelNextN(nn.Module):
self._debug_cp_draft_shared_kv("fallback reason=missing_spec_hidden")
return None
padded_spec_hidden_shape = tuple(spec_hidden_states.shape)
hidden_states_backup = getattr(forward_batch, "hidden_states_backup", None)
if (
envs.SGLANG_CP_DRAFT_SHARED_KV.get()
and hidden_states_backup is not None
and hidden_states_backup.shape[0] != spec_hidden_states.shape[0]
):
# ForwardBatch.prepare_mlp_sync_batch pads EagleDraftInput.hidden_states
# to the global padded token count before model.forward() runs. For
# CP-local draft, the target side-channel is already CP-local; if we
# let the padded tensor look like a full-token tensor, the branch
# below will CP-split it a second time and corrupt draft features.
spec_hidden_states = hidden_states_backup
_log_eagle_accept_cp_draft_hidden_debug(
"use_backup",
"using pre-pad CP-local hidden backup. cp_rank=%s "
"padded_shape=%s backup_shape=%s full_tokens=%s local_tokens=%s",
get_attention_cp_rank(),
padded_spec_hidden_shape,
tuple(spec_hidden_states.shape),
full_num_tokens,
local_num_tokens,
)
if spec_hidden_states.shape[0] == local_num_tokens:
_log_eagle_accept_cp_draft_hidden_debug(
"local_direct",
"using CP-local hidden directly. cp_rank=%s shape=%s "
"full_tokens=%s local_tokens=%s",
get_attention_cp_rank(),
tuple(spec_hidden_states.shape),
full_num_tokens,
local_num_tokens,
)
return spec_hidden_states
if spec_hidden_states.shape[0] < local_num_tokens:
pad_rows = local_num_tokens - spec_hidden_states.shape[0]
if pad_rows <= max(get_attention_cp_size(), 1):
_log_eagle_accept_cp_draft_hidden_debug(
"local_pad",
"padding CP-local hidden to local token count. cp_rank=%s "
"shape=%s pad_rows=%s full_tokens=%s local_tokens=%s",
get_attention_cp_rank(),
tuple(spec_hidden_states.shape),
pad_rows,
full_num_tokens,
local_num_tokens,
)
return torch.cat(
(
spec_hidden_states,
spec_hidden_states.new_zeros(
pad_rows, *spec_hidden_states.shape[1:]
),
),
dim=0,
)
if spec_hidden_states.shape[0] == full_num_tokens:
_log_eagle_accept_cp_draft_hidden_debug(
"full_split",
"splitting full hidden for CP-local draft. cp_rank=%s shape=%s "
"full_tokens=%s local_tokens=%s",
get_attention_cp_rank(),
tuple(spec_hidden_states.shape),
full_num_tokens,
local_num_tokens,
)
return cp_split_and_rebuild_data(forward_batch, spec_hidden_states)
self._debug_cp_draft_shared_kv(