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

Reduce prefill EAGLE memory pressure under CP shared KV

Browse Source 
Prefill CP only needs the local hidden shard for DeepSeek NextN draft extend. The change adds a draft shared-KV path that captures target hidden locally, feeds only the CP-local slice into the draft model, and keeps draft KV writes/transfers on the same shared logical-to-physical page mapping as target KV.\n\nDebug logs are gated behind SGLANG_CP_DRAFT_SHARED_KV_DEBUG and cover scheduler pool selection, KV manager buffer registration, local physical writes, prefill sender filtering, transfer pages, and decode commit metadata so ETE runs can prove draft KV is sharded rather than full-concatenated on a prefill rank.\n\nConstraint: Prefill runs CP while decode remains DP, so prefill must avoid full hidden/KV materialization but decode still receives full logical KV pages.\nRejected: Keep draft extend on full hidden state | preserves correctness but wastes prefill memory and defeats CP shared-KV intent.\nRejected: Transfer draft KV with a separate mapping | target and draft pools share req_to_token logical indices, so duplicating mapping adds risk without benefit.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the debug logs until ETE evidence confirms draft MLA/index writes and transfer pages are CP-sharded on all ranks.\nTested: Remote compileall for changed CP draft, transfer, scheduler, NSA index, MLA write, and EAGLE files.\nNot-tested: Full GLM-5 EAGLE ETE with SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1 after this logging addition; local pytest intentionally not run.
This commit is contained in:
laoyao0822
2026-05-13 22:29:18 +08:00
parent 3fc7a5c18c
commit 99b669f8b9
16 changed files with 951 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+84
View File
@@ -77,6 +77,42 @@ if TYPE_CHECKING:
CLIP_MAX_NEW_TOKEN = envs.SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION.get()
def _cp_draft_shared_kv_debug(message: str, *args) -> None:
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
def _seq_summary(values) -> str:
if values is None:
return "None"
try:
size = len(values)
except TypeError:
return str(values)
if size == 0:
return "size=0"
try:
head = list(values[: min(8, size)])
except TypeError:
head = list(values)[: min(8, size)]
try:
min_val = min(values)
max_val = max(values)
return f"size={size} min={min_val} max={max_val} head={head}"
except (TypeError, ValueError):
return f"size={size} head={head}"
def _pool_summary(pool) -> str:
if pool is None:
return "None"
parts = [pool.__class__.__name__]
for attr in ("size", "page_size", "start_layer", "end_layer", "layer_num"):
if hasattr(pool, attr):
parts.append(f"{attr}={getattr(pool, attr)}")
return " ".join(parts)
def _kv_locs_to_page_indices_cpu(
kv_locs: torch.Tensor,
page_size: int,
@@ -321,16 +357,39 @@ class DecodePreallocQueue:
kv_data_ptrs, kv_data_lens, kv_item_lens = (
self.token_to_kv_pool.get_contiguous_buf_infos()
)
target_kv_buffer_count = len(kv_data_ptrs)
draft_kv_data_lens = []
draft_kv_item_lens = []
draft_kv_buffer_count = 0
if self.draft_token_to_kv_pool is not None:
# We should also transfer draft model kv cache. The indices are
# always shared with a target model.
draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = (
self.draft_token_to_kv_pool.get_contiguous_buf_infos()
)
draft_kv_buffer_count = len(draft_kv_data_ptrs)
kv_data_ptrs += draft_kv_data_ptrs
kv_data_lens += draft_kv_data_lens
kv_item_lens += draft_kv_item_lens
kv_args.draft_kv_buffer_start = target_kv_buffer_count
kv_args.draft_kv_buffer_count = draft_kv_buffer_count
_cp_draft_shared_kv_debug(
"decode_kv_manager cp_rank=%s target_pool=(%s) draft_pool=(%s) "
"target_bufs=%s draft_bufs=%s total_bufs=%s target_lens=%s "
"draft_lens=%s target_item_lens=%s draft_item_lens=%s",
self.tp_rank,
_pool_summary(self.token_to_kv_pool),
_pool_summary(self.draft_token_to_kv_pool),
target_kv_buffer_count,
draft_kv_buffer_count,
len(kv_data_ptrs),
_seq_summary(kv_data_lens[:target_kv_buffer_count]),
_seq_summary(draft_kv_data_lens),
_seq_summary(kv_item_lens[:target_kv_buffer_count]),
_seq_summary(draft_kv_item_lens),
)
kv_args.kv_data_ptrs = kv_data_ptrs
kv_args.kv_data_lens = kv_data_lens
kv_args.kv_item_lens = kv_item_lens
@@ -755,6 +814,19 @@ class DecodePreallocQueue:
self.req_to_metadata_buffer_idx_allocator.alloc()
)
assert decode_req.metadata_buffer_index is not None
_cp_draft_shared_kv_debug(
"decode_prealloc rid=%s room=%s origin_tokens=%s fill_tokens=%s "
"page_size=%s pages=%s state_pages=%s metadata_idx=%s has_draft_pool=%s",
decode_req.req.rid,
decode_req.req.bootstrap_room,
origin_input_len,
len(kv_loc),
page_size,
_seq_summary(page_indices),
_seq_summary(state_indices),
decode_req.metadata_buffer_index,
self.draft_token_to_kv_pool is not None,
)
decode_req.kv_receiver.init(
page_indices, decode_req.metadata_buffer_index, state_indices
)
@@ -1008,6 +1080,18 @@ class DecodeTransferQueue:
decode_req.req.output_topk_index = output_topk_index
decode_req.req.hidden_states_tensor = output_hidden_states
_cp_draft_shared_kv_debug(
"decode_transfer_commit rid=%s room=%s metadata_idx=%s cached_tokens=%s "
"topk_p_shape=%s topk_index_shape=%s hidden_shape=%s",
decode_req.req.rid,
decode_req.req.bootstrap_room,
idx,
decode_req.req.cached_tokens,
tuple(output_topk_p.shape) if output_topk_p is not None else None,
tuple(output_topk_index.shape) if output_topk_index is not None else None,
tuple(output_hidden_states.shape) if output_hidden_states is not None else None,
)
if decode_req.req.return_logprob:
decode_req.req.output_token_logprobs_val.append(
output_token_logprobs_val[0].item()
python/sglang/srt/disaggregation/mooncake/conn.py
+48
View File
@@ -51,6 +51,17 @@ def _cp_shared_debug_log(key: str, message: str, *args, limit: int = 64) -> None
logger.info("[CP_SHARED_KV_DEBUG] " + message, *args)
def _cp_draft_shared_kv_debug(message: str, *args, limit: int = 64) -> None:
if not envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
return
key = "draft:" + message.split(" ", 1)[0]
count = _CP_SHARED_DEBUG_COUNTS.get(key, 0)
if count >= limit:
return
_CP_SHARED_DEBUG_COUNTS[key] = count + 1
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
def _np_summary(arr) -> str:
if arr is None:
return "None"
@@ -246,6 +257,17 @@ class MooncakeKVManager(CommonKVManager):
def register_buffer_to_engine(self):
# Batch register KV data buffers
if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens:
_cp_draft_shared_kv_debug(
"register_buffers mode=%s cp_rank=%s total_kv_bufs=%s "
"draft_start=%s draft_count=%s kv_lens=%s kv_item_lens=%s",
self.disaggregation_mode,
self.attn_cp_rank,
len(self.kv_args.kv_data_ptrs),
getattr(self.kv_args, "draft_kv_buffer_start", None),
getattr(self.kv_args, "draft_kv_buffer_count", None),
_np_summary(self.kv_args.kv_data_lens),
_np_summary(self.kv_args.kv_item_lens),
)
self.engine.batch_register(
self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
)
@@ -847,6 +869,16 @@ class MooncakeKVManager(CommonKVManager):
chunked_dst_kv_indice = req.dst_kv_indices[
kv_chunk.index_slice
]
_cp_draft_shared_kv_debug(
"transfer_pages cp_rank=%s room=%s prefill_pages=%s "
"logical_positions=%s dst_pages=%s is_last=%s",
self.attn_cp_rank,
kv_chunk.room,
_np_summary(kv_chunk.prefill_kv_indices),
_np_summary(kv_chunk.logical_page_positions),
_np_summary(chunked_dst_kv_indice),
kv_chunk.is_last_chunk,
)
if envs.SGLANG_DEBUG_CP_SHARED_KV.get():
_cp_shared_debug_log(
"transfer_worker_kv",
@@ -1275,6 +1307,22 @@ class MooncakeKVSender(CommonKVSender):
_np_summary(state_logical_page_positions),
is_last_chunk,
)
_cp_draft_shared_kv_debug(
"sender_filter cp_rank=%s room=%s page_start=%s orig_kv_pages=%s "
"filtered_kv_pages=%s kv_positions=%s orig_state_pages=%s "
"filtered_state_pages=%s state_positions=%s is_last=%s draft_bufs=%s",
self.kv_mgr.attn_cp_rank,
self.bootstrap_room,
chunk_page_start,
_np_summary(orig_kv_indices),
_np_summary(kv_indices),
_np_summary(logical_page_positions),
_np_summary(orig_state_indices),
_np_summary(state_indices),
_np_summary(state_logical_page_positions),
is_last_chunk,
getattr(self.kv_mgr.kv_args, "draft_kv_buffer_count", None),
)
# Special handling for cp
elif self.kv_mgr.enable_all_cp_ranks_for_transfer:
kv_indices, index_slice = filter_kv_indices_for_cp_rank(
python/sglang/srt/disaggregation/prefill.py
+74
View File
@@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, List, Optional
import torch
from sglang.srt.disaggregation.base import KVPoll
from sglang.srt.environ import envs
from sglang.srt.disaggregation.common.conn import CommonKVManager
from sglang.srt.disaggregation.utils import (
FAKE_BOOTSTRAP_HOST,
@@ -61,6 +62,42 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _cp_draft_shared_kv_debug(message: str, *args) -> None:
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
def _seq_summary(values) -> str:
if values is None:
return "None"
try:
size = len(values)
except TypeError:
return str(values)
if size == 0:
return "size=0"
try:
head = list(values[: min(8, size)])
except TypeError:
head = list(values)[: min(8, size)]
try:
min_val = min(values)
max_val = max(values)
return f"size={size} min={min_val} max={max_val} head={head}"
except (TypeError, ValueError):
return f"size={size} head={head}"
def _pool_summary(pool) -> str:
if pool is None:
return "None"
parts = [pool.__class__.__name__]
for attr in ("size", "page_size", "start_layer", "end_layer", "layer_num"):
if hasattr(pool, attr):
parts.append(f"{attr}={getattr(pool, attr)}")
return " ".join(parts)
def _kv_locs_to_page_indices_cpu(
kv_locs: torch.Tensor,
page_size: int,
@@ -154,16 +191,39 @@ class PrefillBootstrapQueue:
self.token_to_kv_pool.get_contiguous_buf_infos()
)
target_kv_buffer_count = len(kv_data_ptrs)
draft_kv_data_lens = []
draft_kv_item_lens = []
draft_kv_buffer_count = 0
if self.draft_token_to_kv_pool is not None:
# We should also transfer draft model kv cache. The indices are
# always shared with a target model.
draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = (
self.draft_token_to_kv_pool.get_contiguous_buf_infos()
)
draft_kv_buffer_count = len(draft_kv_data_ptrs)
kv_data_ptrs += draft_kv_data_ptrs
kv_data_lens += draft_kv_data_lens
kv_item_lens += draft_kv_item_lens
kv_args.draft_kv_buffer_start = target_kv_buffer_count
kv_args.draft_kv_buffer_count = draft_kv_buffer_count
_cp_draft_shared_kv_debug(
"prefill_kv_manager cp_rank=%s target_pool=(%s) draft_pool=(%s) "
"target_bufs=%s draft_bufs=%s total_bufs=%s target_lens=%s "
"draft_lens=%s target_item_lens=%s draft_item_lens=%s",
self.tp_rank,
_pool_summary(self.token_to_kv_pool),
_pool_summary(self.draft_token_to_kv_pool),
target_kv_buffer_count,
draft_kv_buffer_count,
len(kv_data_ptrs),
_seq_summary(kv_data_lens[:target_kv_buffer_count]),
_seq_summary(draft_kv_data_lens),
_seq_summary(kv_item_lens[:target_kv_buffer_count]),
_seq_summary(draft_kv_item_lens),
)
kv_args.kv_data_ptrs = kv_data_ptrs
kv_args.kv_data_lens = kv_data_lens
kv_args.kv_item_lens = kv_item_lens
@@ -792,4 +852,18 @@ class SchedulerDisaggregationPrefillMixin:
f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
)
return
prefill_queue = getattr(self, "disagg_prefill_bootstrap_queue", None)
_cp_draft_shared_kv_debug(
"prefill_send_kv_chunk rid=%s room=%s start_idx=%s end_idx=%s "
"last_chunk=%s page_size=%s pages=%s state_pages=%s has_draft_pool=%s",
req.rid,
req.bootstrap_room,
start_idx,
end_idx,
last_chunk,
page_size,
_seq_summary(page_indices),
_seq_summary(state_indices),
getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None,
)
req.disagg_kv_sender.send(page_indices, state_indices)
python/sglang/srt/environ.py
+2
View File
@@ -213,6 +213,8 @@ class Envs:
SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)
SGLANG_CP_SHARED_KV_MATERIALIZE_NVTX = EnvBool(False)
SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1)
SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False)
SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False)
SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+10
View File
@@ -24,6 +24,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
filter_owned_logical_locs,
get_or_build_shared_paged_buffer_slot_remap,
is_current_only_extend_batch,
log_cp_draft_shared_kv_debug,
materialize_shared_paged_buffer,
tensor_debug_checksum,
tensor_debug_summary,
@@ -1497,6 +1498,15 @@ class Indexer(MultiPlatformOp):
physical_out_loc = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
if physical_out_loc is None:
return False
log_cp_draft_shared_kv_debug(
"index_write",
"index_write layer=%s tokens=%s physical_tokens=%s pool=%s key_shape=%s",
layer_id,
local_out_loc.numel(),
physical_out_loc.numel(),
forward_batch.token_to_kv_pool.__class__.__name__,
tuple(local_key.shape),
)
self._store_index_k_cache(
forward_batch=forward_batch,
layer_id=layer_id,
python/sglang/srt/layers/attention/nsa/utils.py
+29
View File
@@ -9,6 +9,7 @@ import torch.nn.functional as F
import triton
import triton.language as tl
from sglang.srt.environ import envs
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
@@ -29,6 +30,23 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_CP_DRAFT_SHARED_KV_DEBUG_COUNTS = {}
def log_cp_draft_shared_kv_debug(
key: str,
message: str,
*args,
limit: int = 128,
) -> None:
if not envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
return
count = _CP_DRAFT_SHARED_KV_DEBUG_COUNTS.get(key, 0)
if count >= limit:
return
_CP_DRAFT_SHARED_KV_DEBUG_COUNTS[key] = count + 1
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
def log_cp_shared_kv_direct_write_fallback(
reason: str,
@@ -565,6 +583,17 @@ def get_cp_shared_kv_local_physical_out_cache_loc(forward_batch: "ForwardBatch")
physical_out_cache_loc = layout.logical_locs_to_physical(
local_out_cache_loc
).contiguous()
log_cp_draft_shared_kv_debug(
"physical_out_loc",
"physical_out_loc cp_rank=%s cp_size=%s page_size=%s tokens=%s "
"physical_tokens=%s pool=%s",
layout.cp_rank,
layout.cp_size,
layout.page_size,
local_out_cache_loc.numel(),
physical_out_cache_loc.numel(),
getattr(forward_batch, "token_to_kv_pool", None).__class__.__name__,
)
forward_batch.cp_local_physical_out_cache_loc = physical_out_cache_loc
return physical_out_cache_loc
python/sglang/srt/layers/logits_processor.py
+9 -1
View File
@@ -71,6 +71,10 @@ class LogitsProcessorOutput:
# Used by speculative decoding (EAGLE)
# The last hidden layers
hidden_states: Optional[torch.Tensor] = None
# CP-local hidden states for draft prefill. This is intentionally separate
# from `hidden_states`: the logits path may use compact/narrow hidden while
# EAGLE draft still needs the local target hidden to build draft KV.
draft_hidden_states: Optional[torch.Tensor] = None
## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler
# he log probs of output tokens, if SGLANG_RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature.
@@ -314,7 +318,11 @@ class LogitsProcessor(nn.Module):
logits = self._get_logits(hidden_states, lm_head, logits_metadata)
return LogitsProcessorOutput(
next_token_logits=logits,
hidden_states=None,
hidden_states=(
hidden_states
if logits_metadata.capture_hidden_mode.is_last()
else None
),
mm_input_embeds=logits_metadata.mm_input_embeds,
)
python/sglang/srt/managers/schedule_batch.py
+2
View File
@@ -2317,6 +2317,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
else CaptureHiddenMode.NULL
)
),
capture_draft_hidden_states=False,
extend_input_logprob_token_ids=self.extend_input_logprob_token_ids,
is_prefill_only=self.is_prefill_only,
dimensions=self.dimensions,
@@ -2498,6 +2499,7 @@ class ModelWorkerBatch:
# If set, the output of the batch contains the hidden states of the run.
capture_hidden_mode: CaptureHiddenMode = None
capture_draft_hidden_states: bool = False
hicache_consumer_index: int = -1
# For matryoshka embeddings
python/sglang/srt/managers/scheduler.py
+26
View File
@@ -241,6 +241,22 @@ TEST_RETRACT = envs.SGLANG_TEST_RETRACT.get()
TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get()
TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get()
def _cp_draft_pool_summary(pool) -> str:
if pool is None:
return "None"
parts = [pool.__class__.__name__]
for attr in ("size", "page_size", "start_layer", "end_layer", "layer_num"):
if hasattr(pool, attr):
parts.append(f"{attr}={getattr(pool, attr)}")
return " ".join(parts)
def _cp_draft_shared_kv_debug(message: str, *args) -> None:
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
_is_npu = is_npu()
@@ -941,6 +957,16 @@ class Scheduler(
draft_token_to_kv_pool = self.draft_worker.model_runner.token_to_kv_pool
model_config = self.draft_worker.model_config
_cp_draft_shared_kv_debug(
"scheduler_disagg_init mode=%s spec_algorithm=%s draft_worker=%s "
"draft_pool=(%s) target_pool=(%s)",
self.disaggregation_mode,
self.spec_algorithm,
self.draft_worker is not None,
_cp_draft_pool_summary(draft_token_to_kv_pool),
_cp_draft_pool_summary(self.token_to_kv_pool_allocator.get_kvcache()),
)
if (
self.disaggregation_mode == DisaggregationMode.DECODE
): # *2 for the headroom.
python/sglang/srt/mem_cache/common.py
+10 -10
View File
@@ -369,16 +369,16 @@ def alloc_paged_token_slots_extend(
num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size
evict_result = evict_from_tree_cache(tree_cache, num_tokens)
logger.info(
"[MemCache-alloc] alloc_paged_token_slots_extend: extend_num_tokens=%d batch_size=%d num_tokens=%d page_size=%d "
"available_size=%d evicted=%d",
extend_num_tokens,
len(seq_lens_cpu),
num_tokens,
allocator.page_size,
allocator.available_size(),
getattr(evict_result, "num_tokens_evicted", 0),
)
# logger.info(
# "[MemCache-alloc] alloc_paged_token_slots_extend: extend_num_tokens=%d batch_size=%d num_tokens=%d page_size=%d "
# "available_size=%d evicted=%d",
# extend_num_tokens,
# len(seq_lens_cpu),
# num_tokens,
# allocator.page_size,
# allocator.available_size(),
# getattr(evict_result, "num_tokens_evicted", 0),
# )
alloc_extend_compute_owner = getattr(
allocator, "alloc_extend_compute_owner", None
python/sglang/srt/model_executor/forward_batch_info.py
+3
View File
@@ -399,6 +399,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
spec_algorithm: SpeculativeAlgorithm = None
mm_input_embeds: Optional[torch.Tensor] = None
capture_hidden_mode: CaptureHiddenMode = None
capture_draft_hidden_states: bool = False
draft_hidden_states: Optional[torch.Tensor] = None
# For padding
padded_static_len: int = -1 # -1 if not padded
@@ -484,6 +486,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
spec_algorithm=batch.spec_algorithm,
spec_info=batch.spec_info,
capture_hidden_mode=batch.capture_hidden_mode,
capture_draft_hidden_states=batch.capture_draft_hidden_states,
input_embeds=batch.input_embeds,
token_type_ids=batch.token_type_ids,
tbo_split_seq_index=batch.tbo_split_seq_index,
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+24
View File
@@ -9,6 +9,7 @@ from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.nsa.utils import (
get_cp_shared_kv_local_out_cache_loc,
get_cp_shared_kv_local_physical_out_cache_loc,
log_cp_draft_shared_kv_debug,
log_cp_shared_kv_direct_write_fallback,
nsa_use_prefill_cp,
)
@@ -623,6 +624,18 @@ class DeepseekMLAForwardMixin:
k_nope=k_nope,
k_rope=k_pe,
):
log_cp_draft_shared_kv_debug(
"mla_tai_write",
"mla_write path=tai_fused layer=%s cp_rank=%s cp_size=%s tokens=%s "
"pool=%s k_nope_shape=%s k_pe_shape=%s",
self.attn_mqa.layer_id,
layout.cp_rank,
layout.cp_size,
local_out_cache_loc.numel(),
forward_batch.token_to_kv_pool.__class__.__name__,
tuple(k_nope.shape),
tuple(k_pe.shape),
)
return True
physical_out_cache_loc = get_cp_shared_kv_local_physical_out_cache_loc(
@@ -630,6 +643,17 @@ class DeepseekMLAForwardMixin:
)
if physical_out_cache_loc is None:
return False
log_cp_draft_shared_kv_debug(
"mla_torch_write",
"mla_write path=torch layer=%s tokens=%s physical_tokens=%s pool=%s "
"k_nope_shape=%s k_pe_shape=%s",
self.attn_mqa.layer_id,
local_out_cache_loc.numel(),
physical_out_cache_loc.numel(),
forward_batch.token_to_kv_pool.__class__.__name__,
tuple(k_nope.shape),
tuple(k_pe.shape),
)
forward_batch.token_to_kv_pool.set_mla_kv_buffer(
self.attn_mqa,
physical_out_cache_loc,
python/sglang/srt/models/deepseek_nextn.py
+93 -14
View File
@@ -28,6 +28,8 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r
from sglang.srt.layers.attention.nsa.utils import (
can_cp_split,
cp_all_gather_rerange_output,
cp_collect_last_token_hidden,
cp_split_and_rebuild_1d,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
is_nsa_enable_prefill_cp,
@@ -130,6 +132,35 @@ class DeepseekModelNextN(nn.Module):
else:
self.cp_size = None
def _debug_cp_draft_shared_kv(self, message: str):
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
logger.info("[CP_DRAFT_SHARED_KV] %s", message)
def _get_cp_local_spec_hidden_states(
self,
forward_batch: ForwardBatch,
spec_hidden_states: torch.Tensor,
*,
full_num_tokens: int,
local_num_tokens: int,
) -> Optional[torch.Tensor]:
if spec_hidden_states is None:
self._debug_cp_draft_shared_kv("fallback reason=missing_spec_hidden")
return None
if spec_hidden_states.shape[0] == local_num_tokens:
return spec_hidden_states
if spec_hidden_states.shape[0] == full_num_tokens:
return cp_split_and_rebuild_data(forward_batch, spec_hidden_states)
self._debug_cp_draft_shared_kv(
"fallback reason=spec_hidden_shape_mismatch "
f"spec_tokens={spec_hidden_states.shape[0]} "
f"full_tokens={full_num_tokens} local_tokens={local_num_tokens}"
)
return None
def forward(
self,
input_ids: torch.Tensor,
@@ -145,25 +176,68 @@ class DeepseekModelNextN(nn.Module):
),
)
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
else:
hidden_states = input_embeds
use_cp = nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp)
use_cp_local_draft = use_cp and envs.SGLANG_CP_DRAFT_SHARED_KV.get()
if use_cp_local_draft:
local_input_ids = cp_split_and_rebuild_1d(forward_batch, input_ids)
local_num_tokens = local_input_ids.shape[0]
local_positions = cp_split_and_rebuild_position(forward_batch, positions)
spec_hidden_states = self._get_cp_local_spec_hidden_states(
forward_batch,
forward_batch.spec_info.hidden_states,
full_num_tokens=input_ids.shape[0],
local_num_tokens=local_num_tokens,
)
if spec_hidden_states is None:
use_cp_local_draft = False
else:
positions = local_positions
if input_embeds is None:
hidden_states = self.embed_tokens(local_input_ids)
elif input_embeds.shape[0] == local_num_tokens:
hidden_states = input_embeds
elif input_embeds.shape[0] == input_ids.shape[0]:
hidden_states = cp_split_and_rebuild_data(
forward_batch, input_embeds
)
else:
self._debug_cp_draft_shared_kv(
"fallback reason=input_embeds_shape_mismatch "
f"input_embed_tokens={input_embeds.shape[0]} "
f"full_tokens={input_ids.shape[0]} "
f"local_tokens={local_num_tokens}"
)
use_cp_local_draft = False
if not use_cp_local_draft:
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
else:
hidden_states = input_embeds
spec_hidden_states = forward_batch.spec_info.hidden_states
if hidden_states.shape[0] > 0:
hidden_states = self.eh_proj(
torch.cat(
(
self.enorm(hidden_states),
self.hnorm(forward_batch.spec_info.hidden_states),
self.hnorm(spec_hidden_states),
),
dim=-1,
)
)
if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp):
if use_cp and not use_cp_local_draft:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get() and use_cp_local_draft:
self._debug_cp_draft_shared_kv(
"local_path "
f"full_tokens={input_ids.shape[0]} "
f"local_tokens={hidden_states.shape[0]} "
f"capture_hidden_mode={forward_batch.capture_hidden_mode}"
)
residual = None
with get_global_expert_distribution_recorder().disable_this_region():
hidden_states, residual = self.decoder(
@@ -180,14 +254,19 @@ class DeepseekModelNextN(nn.Module):
else:
hidden_states = self.shared_head.norm(hidden_states)
if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp):
# allgather + rerrange
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if use_cp:
if use_cp_local_draft:
hidden_states = cp_collect_last_token_hidden(
hidden_states, forward_batch, self.cp_size
)
else:
# allgather + rerange
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
return hidden_states
python/sglang/srt/models/deepseek_v2.py
+8 -1
View File
@@ -2039,6 +2039,9 @@ class DeepseekV2Model(nn.Module):
else:
hidden_states, _ = self.norm(hidden_states, residual)
if getattr(forward_batch, "capture_draft_hidden_states", False):
forward_batch.draft_hidden_states = hidden_states
if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch):
if self._should_use_narrow_output_path(forward_batch):
hidden_states = cp_collect_last_token_hidden(
@@ -2221,9 +2224,13 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
hidden_states, aux_hidden_states = hidden_states
if self.pp_group.is_last_rank:
return self.logits_processor(
logits_output = self.logits_processor(
input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states
)
logits_output.draft_hidden_states = getattr(
forward_batch, "draft_hidden_states", None
)
return logits_output
else:
return hidden_states
python/sglang/srt/speculative/eagle_worker.py
+55 -5
View File
@@ -5,6 +5,7 @@ from typing import List, Optional, Tuple
import torch
from sglang.srt.distributed import get_tp_group
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import (
EAGLEDraftNpuGraphRunner,
)
@@ -276,6 +277,36 @@ class EAGLEWorker(TpModelWorker):
def draft_model_runner(self):
return self.model_runner
def _debug_cp_draft_shared_kv(self, message: str):
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
logger.info("[CP_DRAFT_SHARED_KV] %s", message)
def _can_use_cp_draft_shared_kv(self, batch: ScheduleBatch) -> bool:
if not envs.SGLANG_CP_DRAFT_SHARED_KV.get():
return False
if not (batch.forward_mode.is_extend() or batch.is_extend_in_batch):
self._debug_cp_draft_shared_kv("fallback reason=non_extend_mode")
return False
draft_architectures = getattr(
self.draft_model_runner.model_config.hf_config, "architectures", []
)
if "DeepseekV3ForCausalLMNextN" not in (draft_architectures or []):
self._debug_cp_draft_shared_kv(
f"fallback reason=unsupported_arch architectures={draft_architectures}"
)
return False
if not getattr(self.target_worker.model_runner, "uses_cp_shared_kv", False):
self._debug_cp_draft_shared_kv(
"fallback reason=target_cp_shared_kv_disabled"
)
return False
if not getattr(self.draft_model_runner, "uses_cp_shared_kv", False):
self._debug_cp_draft_shared_kv(
"fallback reason=draft_cp_shared_kv_disabled"
)
return False
return True
def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult:
"""Run speculative decoding forward.
@@ -298,9 +329,21 @@ class EAGLEWorker(TpModelWorker):
with self.draft_tp_context(
self.draft_model_runner.tp_group
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
draft_hidden_states = (
logits_output.draft_hidden_states
if logits_output.draft_hidden_states is not None
else logits_output.hidden_states
)
if (
envs.SGLANG_CP_DRAFT_SHARED_KV.get()
and draft_hidden_states is None
):
self._debug_cp_draft_shared_kv(
"fallback_failed reason=missing_target_hidden"
)
self.forward_draft_extend(
batch,
logits_output.hidden_states,
draft_hidden_states,
next_token_ids,
seq_lens_cpu,
logits_output.mm_input_embeds,
@@ -367,15 +410,22 @@ class EAGLEWorker(TpModelWorker):
batch: The batch to run. States could be modified.
Returns:
logits_output: The output of logits. It will contain the full hidden states.
logits_output: The output of logits. It contains full hidden states on the
legacy path, or CP-local draft hidden states in `draft_hidden_states`
when CP draft shared-KV is enabled.
next_token_ids: Next token ids generated.
seq_lens_cpu: CPU copy of sequence lengths for the draft prefill path.
can_run_cuda_graph: Whether the target prefill ran with cuda graph.
"""
# Forward with the target model and get hidden states.
# We need the full hidden states to prefill the KV cache of the draft model.
# Forward with the target model and get hidden states for draft prefill.
# CP draft shared-KV keeps this hidden side channel CP-local so target
# logits can still use the narrow output path.
model_worker_batch = batch.get_model_worker_batch()
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
if self._can_use_cp_draft_shared_kv(batch):
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.NULL
model_worker_batch.capture_draft_hidden_states = True
else:
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
batch_result = self.target_worker.forward_batch_generation(model_worker_batch)
logits_output, next_token_ids = (
batch_result.logits_output,