minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[kimi k3][pd disagg] support pp prefill + dcp decode with dspark (#40045)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-09-20 00:15:28 -07:00
committed by GitHub
parent 22f02cc339
commit f4c256354c
21 changed files with 724 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+3 -1
View File
@@ -601,8 +601,10 @@ def _handle_dspark(server_args: ServerArgs) -> None:
)
if cfg.pp_size != 1:
if cfg.disaggregation_mode != "prefill":
raise ValueError(
"Currently DSpark speculative decoding only supports pp_size == 1."
"DSpark pipeline parallelism requires PD prefill; "
"decode and non-disaggregated serving require pp_size == 1."
)
if cfg.speculative_draft_model_path is None:
python/sglang/srt/arg_groups/validation_hook.py
+8 -1
View File
@@ -57,7 +57,14 @@ def check_pipeline_parallel_compat(
assert cfg.disable_overlap_schedule, (
"Pipeline parallelism is not compatible with overlap schedule"
)
if cfg.speculative_algorithm is not None:
if cfg.speculative_algorithm == "DSPARK":
assert cfg.disaggregation_mode == "prefill", (
"Pipeline parallel DSPARK requires disaggregation-mode=prefill"
)
assert not envs.SGLANG_ENABLE_PP_SPEC.get(), (
"SGLANG_ENABLE_PP_SPEC does not support DSPARK PD prefill"
)
elif cfg.speculative_algorithm is not None:
assert (
cfg.speculative_algorithm.upper() == "EAGLE"
and not cfg.enable_multi_layer_eagle
python/sglang/srt/disaggregation/mooncake/conn.py
+105 -7
View File
@@ -154,6 +154,7 @@ class KVArgsRegisterInfo:
dst_dcp_rank: int = 0
requires_dcp_relayout: bool = False
dcp_token_item_lens: Optional[List[int]] = None
dst_kv_item_lens: List[int] = dataclasses.field(default_factory=list)
staging_base_ptr: int = 0
staging_total_size: int = 0
staging: Optional[StagingRegisterInfo] = None
@@ -201,6 +202,11 @@ class KVArgsRegisterInfo:
dst_dcp_rank=(
int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0
),
dst_kv_item_lens=(
list(struct.unpack(f"{len(msg[19]) // 8}Q", msg[19]))
if len(msg) > 19 and msg[19]
else []
),
# Note: always put the staging field at the final
staging=StagingRegisterInfo.from_zmq_fields(msg, 14, slot_ids_index=18),
)
@@ -1090,11 +1096,16 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
executor: concurrent.futures.ThreadPoolExecutor,
dst_layer_ids: List[int],
pack_buffer=None,
dst_kv_item_lens: Optional[List[int]] = None,
dst_tp_rank: int = 0,
dst_attn_tp_size: Optional[int] = None,
) -> int:
if num_kv_tokens is None:
raise ValueError("PD DCP transfer requires num_kv_tokens")
physical_page_size = self.kv_args.page_size
if dst_kv_item_lens and len(dst_kv_item_lens) != len(dst_kv_ptrs):
raise ValueError("PD DCP destination KV lengths must match its buffers")
src_layer_ids = self.kv_args.kv_layer_ids
if src_layer_ids or dst_layer_ids:
dst_indices = resolve_dcp_dst_entry_indices(
@@ -1105,11 +1116,17 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
)
src_kv_ptrs = self.kv_args.kv_data_ptrs
dst_kv_ptrs = [dst_kv_ptrs[j] for j in dst_indices]
if dst_kv_item_lens:
dst_kv_item_lens = [dst_kv_item_lens[j] for j in dst_indices]
else:
src_kv_ptrs, dst_kv_ptrs, _ = self.get_mla_kv_ptrs_with_pp(
self.kv_args.kv_data_ptrs,
dst_kv_ptrs,
)
if dst_kv_item_lens:
_, dst_kv_item_lens, _ = self.get_mla_kv_ptrs_with_pp(
self.kv_args.kv_item_lens, dst_kv_item_lens
)
num_draft = self.kv_args.num_draft_entries
num_target = len(src_kv_ptrs) - num_draft
@@ -1155,20 +1172,87 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
)
for entry in range(num_target)
]
sliced_draft_params = []
if num_draft > 0 and plan.draft_src_token_indices.size:
if not dst_kv_item_lens and dst_attn_tp_size not in (
None,
self.attn_tp_size,
):
raise ValueError(
"PD DCP with different draft TP sizes requires destination KV lengths"
)
draft_groups = group_concurrent_contiguous(
plan.draft_src_token_indices,
plan.draft_dst_token_indices,
)
layers_params += [
for entry in range(num_target, num_target + num_draft):
src_width = dcp_token_item_lens[entry]
dst_width = src_width
if dst_kv_item_lens:
dst_width, remainder = divmod(
dst_kv_item_lens[entry], physical_page_size * dst_dcp_size
)
if remainder or dst_width <= 0:
raise ValueError("Invalid PD DCP draft destination token width")
if src_width == dst_width:
layers_params.append(
(
src_kv_ptrs[num_target + entry],
dst_kv_ptrs[num_target + entry],
dcp_token_item_lens[num_target + entry],
src_kv_ptrs[entry],
dst_kv_ptrs[entry],
src_width,
draft_groups,
)
for entry in range(num_draft)
]
)
continue
if self.is_mla_backend:
raise ValueError(
"PD DCP draft head slicing is unsupported for pure MLA: "
"dummy prefill senders may omit draft head shards"
)
copy_width = min(src_width, dst_width)
if max(src_width, dst_width) % copy_width:
raise ValueError("PD DCP draft KV head shards must divide evenly")
if dst_attn_tp_size is None:
raise ValueError(
"PD DCP draft head slicing requires destination TP size"
)
src_span = src_width * self.attn_tp_size
dst_span = dst_width * dst_attn_tp_size
src_rank = (self.kv_args.engine_rank % self.attn_tp_size) // max(
1, src_span // dst_span
)
dst_rank = dst_tp_rank // max(1, dst_span // src_span)
src_offset = (dst_rank * dst_width) % src_width
dst_offset = (src_rank * src_width) % dst_width
sliced_draft_params.append(
(
src_kv_ptrs[entry] + src_offset,
dst_kv_ptrs[entry] + dst_offset,
src_width,
dst_width,
copy_width,
)
)
def process_sliced_draft(params) -> int:
batch_size = self.max_transfer_batch_indices
if batch_size <= 0:
batch_size = 4096
for start in range(0, plan.draft_src_token_indices.size, batch_size):
src_indices = plan.draft_src_token_indices[start : start + batch_size]
dst_indices = plan.draft_dst_token_indices[start : start + batch_size]
blocks = []
for src_ptr, dst_ptr, src_width, dst_width, copy_width in params:
src_addrs = src_ptr + src_indices * src_width
dst_addrs = dst_ptr + dst_indices * dst_width
blocks.extend(
(int(src), int(dst), copy_width)
for src, dst in zip(src_addrs, dst_addrs)
)
ret = self._transfer_data(mooncake_session_id, blocks)
if ret != 0:
return ret
return 0
def set_transfer_blocks(
src_ptr: int, dst_ptr: int, token_item_len: int, groups
@@ -1196,12 +1280,19 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
executor.submit(process_layer, *layer_params)
for layer_params in layers_params
]
futures.extend(
executor.submit(process_sliced_draft, [params])
for params in sliced_draft_params
)
return self._await_transfer_futures(futures)
transfer_blocks = []
for layer_params in layers_params:
transfer_blocks.extend(set_transfer_blocks(*layer_params))
return self._transfer_data(mooncake_session_id, transfer_blocks)
ret = self._transfer_data(mooncake_session_id, transfer_blocks)
if ret != 0 or not sliced_draft_params:
return ret
return process_sliced_draft(sliced_draft_params)
def send_kvcache_slice(
self,
@@ -2114,6 +2205,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
target_rank_registration_info.dst_kv_layer_ids
),
pack_buffer=pack_buffer,
dst_kv_item_lens=target_rank_registration_info.dst_kv_item_lens,
dst_tp_rank=target_rank_registration_info.dst_tp_rank,
dst_attn_tp_size=target_rank_registration_info.dst_attn_tp_size,
)
elif (
self.is_mla_backend
@@ -2810,6 +2904,10 @@ class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
dst_dcp_size,
dst_dcp_rank,
packed_staging_slot_layer_ids,
struct.pack(
f"{len(self.kv_mgr.kv_args.kv_item_lens)}Q",
*self.kv_mgr.kv_args.kv_item_lens,
),
]
)
except zmq.ZMQError:
python/sglang/srt/managers/scheduler_pp_mixin.py
+15 -1
View File
@@ -886,7 +886,11 @@ class SchedulerPPMixin:
# Draft extend runs only on the last stage, but every rank needs its relayed
# output to fill PD auxiliary buffers.
draft_input = result.next_draft_input
if draft_input is not None and draft_input.topk_p is not None:
if (
draft_input is not None
and not batch.spec_algorithm.is_dspark()
and draft_input.topk_p is not None
):
tensor_dict["draft_topk_p"] = draft_input.topk_p.contiguous()
tensor_dict["draft_topk_index"] = draft_input.topk_index.contiguous()
tensor_dict["draft_hidden_states"] = draft_input.hidden_states.contiguous()
@@ -1138,6 +1142,16 @@ class SchedulerPPMixin:
dsa_topk_indices=pp_outputs.tensors.get("draft_dsa_topk_indices"),
)
batch.spec_info = next_draft_input
elif batch.spec_algorithm.is_dspark():
from sglang.srt.speculative.dspark_components.dspark_draft import (
make_next_draft_input,
)
next_draft_input = make_next_draft_input(
bonus_tokens=next_token_ids,
new_seq_lens=batch.seq_lens,
)
batch.spec_info = next_draft_input
if self._pp_spec_relay:
# Gated single-instance PP+spec: the sampled first token roots
python/sglang/srt/model_executor/model_runner.py
+7
View File
@@ -799,6 +799,13 @@ class ModelRunner:
enable_batch_invariant_mode()
def get_pp_proxy_dspark_hidden_size(self) -> int:
return misc_utils.resolve_pp_proxy_dspark_hidden_size(
model=self.model,
pp_size=self.ps.pp_size,
pp_rank=self.ps.pp_rank,
)
def get_pp_proxy_topk_size(self) -> Optional[int]:
return misc_utils.resolve_pp_proxy_topk_size(
model_config=self.model_config,
python/sglang/srt/model_executor/model_runner_components/misc_utils.py
+16 -1
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable
from sglang.srt.configs.model_config import (
dsa_layer_skips_topk,
@@ -90,3 +90,18 @@ def resolve_pp_proxy_residual_num_blocks(
if block_size is None:
return None
return (start_layer + block_size - 1) // block_size
@runtime_checkable
class _SupportsDSparkPPProxy(Protocol):
def get_pp_proxy_dspark_hidden_size(self) -> int: ...
def resolve_pp_proxy_dspark_hidden_size(
*, model: Any, pp_size: int, pp_rank: int
) -> int:
if pp_size <= 1 or pp_rank == 0:
return 0
if isinstance(model, _SupportsDSparkPPProxy):
return model.get_pp_proxy_dspark_hidden_size()
return 0
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+3
View File
@@ -418,6 +418,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
pp_proxy_residual_num_blocks=(
self.model_runner.get_pp_proxy_residual_num_blocks()
),
pp_proxy_dspark_hidden_size=(
self.model_runner.get_pp_proxy_dspark_hidden_size()
),
)
self.buffers.share_buffers()
# FB-shared slot registry adopting DecodeInputBuffers storage (same
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+16 -3
View File
@@ -389,6 +389,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
pp_proxy_residual_num_blocks=(
self.model_runner.get_pp_proxy_residual_num_blocks()
),
pp_proxy_dspark_hidden_size=(
self.model_runner.get_pp_proxy_dspark_hidden_size()
),
)
self.buffers.share_buffers()
# Token-axis FB-shared slot registry adopting PrefillInputBuffers
@@ -599,8 +602,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
f"unsupported for this model architecture."
) from exc
params = list(inspect.signature(self.layer_model.forward).parameters)
self._input_embeds_arg_idx = (
params.index("input_embeds") if "input_embeds" in params else None
self._input_embeds_arg_idx = next(
(
params.index(name)
for name in ("input_embeds", "inputs_embeds")
if name in params
),
None,
)
# --- aiter chip info pre-warming (AMD) -------------------------
@@ -1930,6 +1938,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
"""A text-only batch would otherwise replay the captured input_embeds."""
ie_idx = self._input_embeds_arg_idx
ie = layer_kwargs.get("input_embeds")
if ie is None:
ie = layer_kwargs.get("inputs_embeds")
if ie is None and ie_idx is not None and len(args) > ie_idx:
ie = args[ie_idx]
if ie is None:
@@ -1968,7 +1978,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# text-only batches they are get_input_embeddings()(input_ids).
# Copy them into the slot before replay so the graph sees the
# current request's embeddings (mirrors main's BCG closure).
if self.buffer_registry.has_slot("input_embeds"):
if (
self.model_runner.pp_group.is_first_rank
and self.buffer_registry.has_slot("input_embeds")
):
self._fill_input_embeds_slot(args, layer_kwargs, static_num_tokens)
hs = self.backend.replay(shape_key, static_forward_batch, **kwargs)
return _slice_output_rows(hs, raw_num_tokens) if full_path else hs
python/sglang/srt/model_executor/runner_utils/buffers.py
+9
View File
@@ -67,6 +67,7 @@ def _allocate_pp_proxy_tensors(
hc_hidden_size: Optional[int] = None,
pp_proxy_topk_size: Optional[int] = None,
pp_proxy_residual_num_blocks: Optional[int] = None,
pp_proxy_dspark_hidden_size: int = 0,
) -> Dict[str, torch.Tensor]:
"""Allocate the stable buffers consumed by an incoming PP proxy."""
is_mhc = hc_hidden_size is not None
@@ -87,6 +88,10 @@ def _allocate_pp_proxy_tensors(
pp_proxy_tensors["topk_indices"] = torch.zeros(
(max_num_tokens, pp_proxy_topk_size), dtype=torch.int32
)
if pp_proxy_dspark_hidden_size:
pp_proxy_tensors["dspark_hidden_states"] = torch.zeros(
(max_num_tokens, pp_proxy_dspark_hidden_size), dtype=dtype
)
return pp_proxy_tensors
@@ -136,6 +141,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
hc_hidden_size: Optional[int] = None,
pp_proxy_topk_size: Optional[int] = None,
pp_proxy_residual_num_blocks: Optional[int] = None,
pp_proxy_dspark_hidden_size: int = 0,
) -> DecodeInputBuffers:
with torch.device(device):
input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
@@ -173,6 +179,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
hc_hidden_size=hc_hidden_size,
pp_proxy_topk_size=pp_proxy_topk_size,
pp_proxy_residual_num_blocks=pp_proxy_residual_num_blocks,
pp_proxy_dspark_hidden_size=pp_proxy_dspark_hidden_size,
)
if pp_size > 1
else None
@@ -275,6 +282,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
hc_hidden_size: Optional[int] = None,
pp_proxy_topk_size: Optional[int] = None,
pp_proxy_residual_num_blocks: Optional[int] = None,
pp_proxy_dspark_hidden_size: int = 0,
) -> PrefillInputBuffers:
with torch.device(device):
input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
@@ -311,6 +319,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
hc_hidden_size=hc_hidden_size,
pp_proxy_topk_size=pp_proxy_topk_size,
pp_proxy_residual_num_blocks=pp_proxy_residual_num_blocks,
pp_proxy_dspark_hidden_size=pp_proxy_dspark_hidden_size,
)
if pp_size > 1 and not is_first_pp_rank
else None
python/sglang/srt/models/kimi_k3.py
+29 -8
View File
@@ -3047,6 +3047,18 @@ class KimiK3LinearModel(nn.Module):
)
sp_sharded = False
aux_hidden_states = []
if (
self.dspark_layers_to_capture is not None
and not self.pp_group.is_first_rank
):
if "dspark_hidden_states" in pp_proxy_tensors.tensors:
aux_hidden_states.append(pp_proxy_tensors["dspark_hidden_states"])
if self.start_layer - 1 in self.dspark_layers_to_capture:
aux_hidden_states.append(
self._dspark_capture_stream(
self.start_layer - 1, hidden_states, residual, attn_res
)
)
for i in range(self.start_layer, self.end_layer):
if sp_sharded and not self.layers[i]._sp_moe:
hidden_states = _sp_all_gather_rows(hidden_states)
@@ -3065,6 +3077,7 @@ class KimiK3LinearModel(nn.Module):
if (
self.dspark_layers_to_capture is not None
and i in self.dspark_layers_to_capture
and (i + 1 < self.end_layer or self.pp_group.is_last_rank)
):
aux_hidden_states.append(
self._dspark_capture_stream(i, hidden_states, residual, attn_res)
@@ -3078,9 +3091,12 @@ class KimiK3LinearModel(nn.Module):
# full stream head (bit-identical to the fused fold).
hidden_states = residual + hidden_states
residual = attn_res.block_residual # raw bank across ranks
return PPProxyTensors(
{"hidden_states": hidden_states, "residual": residual}
proxy_tensors = {"hidden_states": hidden_states, "residual": residual}
if aux_hidden_states:
proxy_tensors["dspark_hidden_states"] = torch.cat(
aux_hidden_states, dim=-1
)
return PPProxyTensors(proxy_tensors)
if hidden_states.shape[0] != 0:
if attn_res is not None:
@@ -3204,13 +3220,13 @@ class KimiK3LinearForCausalLM(nn.Module):
def get_input_embeddings(self):
return self.model.embed_tokens
def get_pp_proxy_dspark_hidden_size(self) -> int:
layers = self.model.dspark_layers_to_capture or []
return self.config.hidden_size * sum(
layer < self.model.start_layer - 1 for layer in layers
)
def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None:
if self.pp_group.world_size > 1:
# Capture layers living on non-last PP ranks would be silently
# skipped (the flag is only set on the last rank).
raise NotImplementedError("DSPARK aux hidden capture requires PP=1.")
if not self.pp_group.is_last_rank:
return
if layer_ids is None:
raise ValueError(
"DSPARK requires explicit layer_ids for aux hidden capture."
@@ -3667,6 +3683,11 @@ class KimiK3ForConditionalGeneration(nn.Module):
raise AttributeError("lm_head is not available in encoder-only mode")
return self.language_model.lm_head
def get_pp_proxy_dspark_hidden_size(self) -> int:
if self.language_model is None:
return 0
return self.language_model.get_pp_proxy_dspark_hidden_size()
def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None:
if self.language_model is None:
raise AttributeError(
python/sglang/srt/models/kimi_linear.py
+17 -9
View File
@@ -717,6 +717,12 @@ class KimiLinearModel(nn.Module):
device=device,
)
aux_hidden_states = []
if (
self.dspark_layers_to_capture is not None
and not self.pp_group.is_first_rank
and "dspark_hidden_states" in pp_proxy_tensors.tensors
):
aux_hidden_states.append(pp_proxy_tensors["dspark_hidden_states"])
for i in range(self.start_layer, self.end_layer):
ctx = get_global_expert_distribution_recorder().with_current_layer(i)
with ctx:
@@ -737,12 +743,12 @@ class KimiLinearModel(nn.Module):
)
if not self.pp_group.is_last_rank:
return PPProxyTensors(
{
"hidden_states": hidden_states,
"residual": residual,
}
proxy_tensors = {"hidden_states": hidden_states, "residual": residual}
if aux_hidden_states:
proxy_tensors["dspark_hidden_states"] = torch.cat(
aux_hidden_states, dim=-1
)
return PPProxyTensors(proxy_tensors)
else:
if hidden_states.shape[0] != 0:
if residual is None:
@@ -787,11 +793,13 @@ class KimiLinearForCausalLM(nn.Module):
def get_input_embeddings(self):
return self.model.embed_tokens
def get_pp_proxy_dspark_hidden_size(self) -> int:
layers = self.model.dspark_layers_to_capture or []
return self.config.hidden_size * sum(
layer < self.model.start_layer for layer in layers
)
def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None:
if self.pp_group.world_size > 1:
raise NotImplementedError("DSPARK aux hidden capture requires PP=1.")
if not self.pp_group.is_last_rank:
return
if layer_ids is None:
raise ValueError(
"DSPARK requires explicit layer_ids for aux hidden capture."
python/sglang/srt/speculative/draft_worker_common.py
+2
View File
@@ -69,6 +69,7 @@ def build_draft_tp_worker(
algo_label: str,
attention_backend_override: Optional[str] = None,
draft_worker_cls: type[TpModelWorker] = TpModelWorker,
random_seed: Optional[int] = None,
) -> DraftWorkerBundle:
# An override names a draft-specific backend the caller has already
# validated (e.g. a self-drafting architecture); it skips the generic
@@ -90,6 +91,7 @@ def build_draft_tp_worker(
ps=ps,
nccl_port=nccl_port,
is_draft_worker=True,
random_seed=random_seed,
# The draft runs at absolute target positions.
context_length=target_model_config.context_len,
draft_attention_backend=draft_backend,
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
+49 -7
View File
@@ -9,6 +9,7 @@ from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
is_unified_kv_triton,
)
from sglang.srt.configs.hybrid_arch import mambaish_config
from sglang.srt.distributed import get_pp_group
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.layers.logprob_processor import compute_spec_logprobs
@@ -20,6 +21,7 @@ from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardMode,
PPProxyTensors,
compute_position,
)
from sglang.srt.runtime_context import (
@@ -126,6 +128,8 @@ def _configure_target_hidden_projection(
class DSparkWorkerV2(BaseSpecWorker):
"""Non-last PP stages run only the target; draft state belongs to the last stage."""
def __init__(
self,
server_args: ServerArgs,
@@ -145,6 +149,10 @@ class DSparkWorkerV2(BaseSpecWorker):
self.model_runner = target_worker.model_runner
self.page_size = get_schedule().page_size
self.device = target_worker.device
self._draft_worker = None
self._hosts_draft = get_pp_group().is_last_rank
if not self._hosts_draft:
return
self._draft_is_moe = draft_is_deepseek_v4()
self._draft_dp_context_enabled = (
@@ -178,6 +186,7 @@ class DSparkWorkerV2(BaseSpecWorker):
DSV4_DRAFT_ATTENTION_BACKEND if self._draft_is_moe else None
),
draft_worker_cls=draft_worker_cls,
random_seed=target_worker.random_seed,
)
self._draft_worker = bundle.draft_worker
self.draft_model_runner = bundle.draft_model_runner
@@ -397,10 +406,12 @@ class DSparkWorkerV2(BaseSpecWorker):
@property
def carries_confidence(self) -> bool:
return self._verify_planner.carries_confidence
return self._hosts_draft and self._verify_planner.carries_confidence
@property
def spec_v2_attn_backends(self) -> tuple:
if not self._hosts_draft:
return super().spec_v2_attn_backends
return (
self._target_worker.model_runner.attn_backend,
self.draft_model_runner.attn_backend,
@@ -422,6 +433,8 @@ class DSparkWorkerV2(BaseSpecWorker):
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
if not self._hosts_draft:
return
self._draft_worker.alloc_memory_pool(
memory_pool_config=memory_pool_config,
req_to_token_pool=req_to_token_pool,
@@ -429,6 +442,8 @@ class DSparkWorkerV2(BaseSpecWorker):
)
def init_attention_backends(self):
if not self._hosts_draft:
return
with draft_pp_context(), self._draft_context():
self._draft_worker.init_attention_backends()
self._target_hidden_projection_enabled = _configure_target_hidden_projection(
@@ -449,6 +464,8 @@ class DSparkWorkerV2(BaseSpecWorker):
)
def init_cuda_graphs(self):
if not self._hosts_draft:
return
capture_decode_cuda_graph = self._decode_graph_allowed
available_mem = self._tp_sync.available_memory_gb(
SpecTpSyncSite.DSPARK_MEM,
@@ -510,19 +527,29 @@ class DSparkWorkerV2(BaseSpecWorker):
pass
def set_dspark_forced_budget_frac(self, frac: Optional[float]) -> None:
if not self._hosts_draft:
return
self._forced_budget_frac = frac
self._verify_planner.set_forced_budget_frac(frac)
def dump_info_records(self) -> Optional[dict]:
if not self._hosts_draft:
return None
return self._observers.dump_info_records()
def clear_info_records(self) -> None:
if not self._hosts_draft:
return
self._observers.clear_info_records()
def block_accept_estimate_log_suffix(self) -> Optional[str]:
if not self._hosts_draft:
return None
return self._observers.block_accept_estimate_log_suffix()
def note_request_finished(self, *, rid: str, natural_stop: bool) -> None:
if not self._hosts_draft:
return
self._observers.note_request_finished(rid=rid, natural_stop=natural_stop)
def forward_batch_generation(
@@ -531,19 +558,30 @@ class DSparkWorkerV2(BaseSpecWorker):
on_publish=None,
grammar_barrier=None,
*,
pp_proxy_tensors=None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> GenerationBatchResult:
# The non-overlap scheduler passes this keyword even when PP=1.
assert pp_proxy_tensors is None, "DSpark does not support pipeline parallelism"
if not self._hosts_draft:
batch_output = self.target_worker.forward_batch_generation(
batch,
pp_proxy_tensors=pp_proxy_tensors,
capture_hidden_mode=CaptureHiddenMode.FULL,
)
batch_output.new_seq_lens = batch.seq_lens
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
return batch_output
if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
self._verify_planner.note_non_decode_step()
self._observers.note_prefill_step()
return self._forward_prefill(batch, on_publish)
return self._forward_prefill(batch, on_publish, pp_proxy_tensors)
return self._forward_decode(batch, on_publish, grammar_barrier)
def _forward_prefill(
self, batch: ScheduleBatch, on_publish
self,
batch: ScheduleBatch,
on_publish,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> GenerationBatchResult:
if batch.forward_mode.is_idle():
if get_parallel().enable_dp_attention:
@@ -553,7 +591,9 @@ class DSparkWorkerV2(BaseSpecWorker):
return self._decode_idle_result(on_publish=on_publish)
batch_output = self.target_worker.forward_batch_generation(
batch, capture_hidden_mode=CaptureHiddenMode.FULL
batch,
pp_proxy_tensors=pp_proxy_tensors,
capture_hidden_mode=CaptureHiddenMode.FULL,
)
# BCG replay skips model-side Python, so re-evaluate the same pure predicate.
target_hidden_is_projected = (
@@ -1010,4 +1050,6 @@ class DSparkWorkerV2(BaseSpecWorker):
)
def get_confidence_budget_prepare(self):
if not self._hosts_draft:
return None
return self._verify_planner.confidence_budget_prepare()
test/registered/kernel/attention/test_kda_gate_beta_cumsum.py → test/registered/kernels/ops/attention/test_kda_gate_beta_cumsum.py
View File
test/registered/unit/disaggregation/test_disaggregation_wire.py
+3
View File
@@ -96,6 +96,9 @@ class TestDisaggregationWire(unittest.TestCase):
self.assertEqual(info.staging_total_size, 4096)
self.assertEqual(info.dst_dcp_size, 4)
self.assertEqual(info.dst_dcp_rank, 2)
self.assertEqual(info.dst_kv_item_lens, [])
info = KVArgsRegisterInfo.from_zmq(msg + [b"", struct.pack("Q", 128)])
self.assertEqual(info.dst_kv_item_lens, [128])
def test_int_lists_roundtrip(self):
cases = [
test/registered/unit/disaggregation/test_mooncake_transfer_batching.py
+156
View File
@@ -130,5 +130,161 @@ class TestMooncakeTransferBatching(unittest.TestCase):
)
class TestDcpDraftHeadTransfer(unittest.TestCase):
def test_transfers_draft_heads_to_logical_destination_rows(self):
for src_tp, dst_tp in ((4, 8), (8, 4), (8, 8), (4, 32), (32, 4)):
for custom_pool in (False, True):
for batch_size in (0, 37):
with self.subTest(
src_tp=src_tp,
dst_tp=dst_tp,
custom_pool=custom_pool,
batch_size=batch_size,
):
self._check_transfer(src_tp, dst_tp, custom_pool, batch_size)
def test_rejects_pure_mla_with_unequal_draft_head_widths(self):
for src_tp, dst_tp in ((4, 8), (8, 4)):
with self.subTest(src_tp=src_tp, dst_tp=dst_tp):
with self.assertRaisesRegex(ValueError, "dummy prefill senders"):
self._check_transfer(src_tp, dst_tp, False, 37, pure_mla=True)
def test_sliced_draft_stops_after_failed_batch(self):
self._check_transfer(4, 8, False, 37, fail_draft=True)
def _check_transfer(
self, src_tp, dst_tp, custom_pool, batch_size, fail_draft=False, pure_mla=False
):
page_size, tokens, heads, head_bytes = 64, 249, 16, 4
src_width, dst_width = (
max(1, heads // src_tp) * head_bytes,
max(1, heads // dst_tp) * head_bytes,
)
src_pages = np.array([1, 3, 4, 7], dtype=np.int32)
logical = np.arange(tokens)
src_rows = src_pages[logical // page_size] * page_size + logical % page_size
expected = (
np.arange(tokens * heads * head_bytes, dtype=np.int64)
.reshape(tokens, heads, head_bytes)
.astype(np.uint8)
)
for dst_rank in range(dst_tp):
dst_buffers = {
base: np.zeros(16384 * max(8, dst_width), dtype=np.uint8)
for base in (1000000, 2000000, 3000000, 4000000)
}
source_ranks = (
range(dst_rank * src_tp // dst_tp, (dst_rank + 1) * src_tp // dst_tp)
if src_tp >= dst_tp
else [dst_rank * src_tp // dst_tp]
)
for src_rank in source_ranks:
src_head_start = (src_rank // max(1, src_tp // heads)) * max(
1, heads // src_tp
)
source = np.zeros(1024 * src_width, dtype=np.uint8)
source.reshape(-1, src_width)[src_rows] = expected[
:, src_head_start : src_head_start + max(1, heads // src_tp)
].reshape(tokens, src_width)
target = np.zeros(1024 * 8, dtype=np.uint8)
target.reshape(-1, 8)[src_rows] = (
np.arange(tokens * 8).reshape(tokens, 8).astype(np.uint8)
)
src_buffers = {10000: target, 100000: source, 200000: source}
failed_batches = []
def transfer(
session, blocks, src_buffers=src_buffers, dst_buffers=dst_buffers
):
draft_blocks = [block for block in blocks if block[1] >= 3000000]
if fail_draft and draft_blocks:
failed_batches.append(draft_blocks)
return 17
if batch_size and src_width != dst_width:
self.assertLessEqual(
len(draft_blocks), batch_size * (1 if custom_pool else 2)
)
for src, dst, size in blocks:
src_base = max(base for base in src_buffers if base <= src)
dst_base = max(base for base in dst_buffers if base <= dst)
dst_buffers[dst_base][
dst - dst_base : dst - dst_base + size
] = src_buffers[src_base][
src - src_base : src - src_base + size
]
return 0
manager = SimpleNamespace(
is_mla_backend=pure_mla,
kv_args=SimpleNamespace(
page_size=page_size,
kv_layer_ids=[47, 93, 93],
kv_data_ptrs=[10000, 100000, 200000],
num_draft_entries=2,
engine_rank=src_rank + 2 * src_tp,
),
attn_tp_size=src_tp,
max_transfer_batch_indices=batch_size,
enable_custom_mem_pool=custom_pool,
_transfer_data=transfer,
_await_transfer_futures=lambda futures: max(
f.result() for f in futures
),
)
with concurrent.futures.ThreadPoolExecutor() as executor:
result = MooncakeKVManager.send_kvcache_dcp(
manager,
"session",
src_pages,
[1000000, 2000000, 3000000, 4000000],
np.array([2], dtype=np.int32),
dcp_token_item_lens=[8, src_width, src_width],
dst_dcp_size=dst_tp,
dst_dcp_rank=dst_rank,
src_page_offset=0,
decode_prefix_len=0,
num_kv_tokens=tokens,
executor=executor,
dst_layer_ids=[3, 47, 93, 93],
dst_kv_item_lens=[
page_size * 8,
page_size * 8,
page_size * dst_tp * dst_width,
page_size * dst_tp * dst_width,
],
dst_tp_rank=dst_rank,
dst_attn_tp_size=dst_tp,
)
if fail_draft:
self.assertEqual(result, 17)
self.assertEqual(len(failed_batches), 1)
return
self.assertEqual(result, 0)
dst_head_start = (dst_rank // max(1, dst_tp // heads)) * max(
1, heads // dst_tp
)
for base in (3000000, 4000000):
actual = dst_buffers[base].reshape(-1, dst_width)[
2 * page_size * dst_tp + logical
]
np.testing.assert_array_equal(
actual,
expected[
:,
dst_head_start : dst_head_start + max(1, heads // dst_tp),
].reshape(tokens, dst_width),
)
owned = np.arange(dst_rank, tokens, dst_tp)
actual_target = dst_buffers[2000000].reshape(-1, 8)[
2 * page_size + owned // dst_tp
]
np.testing.assert_array_equal(
actual_target,
np.arange(tokens * 8).reshape(tokens, 8).astype(np.uint8)[owned],
)
self.assertFalse(dst_buffers[1000000].any())
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_auxiliary_output.py
+3
View File
@@ -660,6 +660,7 @@ def test_pipeline_parallel_auxiliary_output_round_trip():
next_token_ids=torch.tensor([7]),
)
batch = SimpleNamespace(
spec_algorithm=SpeculativeAlgorithm.NONE,
return_logprob=False,
req_pool_indices=torch.tensor([3]),
input_ids=torch.tensor([5]),
@@ -705,6 +706,7 @@ def test_pipeline_parallel_dsa_seed_round_trip(dsa_topk_indices):
next_draft_input=draft_input,
)
batch = SimpleNamespace(
spec_algorithm=SpeculativeAlgorithm.EAGLE3,
return_logprob=False,
req_pool_indices=torch.tensor([3]),
input_ids=torch.tensor([5]),
@@ -745,6 +747,7 @@ def test_pipeline_parallel_auxiliary_output_stays_packed_before_first_rank():
next_token_ids=torch.tensor([7]),
)
batch = SimpleNamespace(
spec_algorithm=SpeculativeAlgorithm.NONE,
return_logprob=False,
req_pool_indices=torch.tensor([3]),
input_ids=torch.tensor([5]),
test/registered/unit/managers/test_pp_cp_rank_offsets.py
+54
View File
@@ -2,8 +2,11 @@ import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import (
CustomTestCase,
enter_scope,
maybe_stub_sgl_kernel,
published_topology,
@@ -239,5 +242,56 @@ class TestPPCPRankOffsets(unittest.TestCase):
)
class TestDSparkPPOutput(CustomTestCase):
def test_output_ring_rebinds_dspark_state_on_each_stage(self):
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
from sglang.srt.speculative.dspark_components.dspark_draft import (
make_next_draft_input,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
payloads = []
scheduler = SimpleNamespace(
_pp_spec_relay=False,
pp_group=SimpleNamespace(is_first_rank=False),
future_map=SimpleNamespace(
stash=lambda indices, value: payloads.append(value)
),
)
tokens = torch.tensor([13, 29])
batch = SimpleNamespace(
return_logprob=False,
req_pool_indices=torch.tensor([0, 1]),
seq_lens=torch.tensor([8, 15]),
spec_algorithm=SpeculativeAlgorithm.DSPARK,
spec_info=object(),
)
wire = SchedulerPPMixin._pp_prepare_tensor_dict(
scheduler,
SimpleNamespace(
next_token_ids=tokens,
next_draft_input=make_next_draft_input(
bonus_tokens=tokens, new_seq_lens=batch.seq_lens
),
logits_output=None,
),
batch,
)
self.assertNotIn("draft_topk_p", wire)
result = SchedulerPPMixin._pp_prep_batch_result(
scheduler,
batch,
SimpleNamespace(can_run_cuda_graph=False),
PPProxyTensors(wire),
)
self.assertIsInstance(result.next_draft_input, DFlashDraftInputV2)
self.assertIs(batch.spec_info, result.next_draft_input)
torch.testing.assert_close(batch.spec_info.bonus_tokens, tokens)
torch.testing.assert_close(batch.spec_info.new_seq_lens, batch.seq_lens)
torch.testing.assert_close(payloads[0].bonus_tokens, tokens)
self.assertEqual(payloads[0].hidden_states.numel(), 0)
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py
+166 -1
View File
@@ -11,6 +11,9 @@ from sglang.srt.model_executor.cuda_graph_buffer_registry import (
build_prefill_registry,
)
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.model_executor.model_runner_components.misc_utils import (
resolve_pp_proxy_dspark_hidden_size,
)
from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
PrefillCudaGraphRunner,
_build_layer_model_forward_kwargs,
@@ -52,6 +55,25 @@ def _make_pp_buffers_and_registry():
class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
def test_dspark_proxy_width_requires_receiving_stage_and_model_support(self):
class Model:
def get_pp_proxy_dspark_hidden_size(self):
return 16
for model, pp_size, pp_rank, expected in (
(Model(), 1, 0, 0),
(Model(), 2, 0, 0),
(Model(), 2, 1, 16),
(object(), 2, 1, 0),
):
with self.subTest(pp_size=pp_size, pp_rank=pp_rank, expected=expected):
self.assertEqual(
resolve_pp_proxy_dspark_hidden_size(
model=model, pp_size=pp_size, pp_rank=pp_rank
),
expected,
)
def test_pp_proxy_stable_buffers_accept_full_and_hidden_only_contracts(self):
buffers, registry = _make_pp_buffers_and_registry()
full_proxy = PPProxyTensors(
@@ -170,6 +192,7 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
pp_size=2,
is_first_pp_rank=False,
pp_proxy_residual_num_blocks=3,
pp_proxy_dspark_hidden_size=16,
)
self.assertEqual(
@@ -177,9 +200,151 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
key: tuple(value.shape)
for key, value in buffers.pp_proxy_tensors.items()
},
{"hidden_states": (16, 8), "residual": (16, 3, 8)},
{
"hidden_states": (16, 8),
"residual": (16, 3, 8),
"dspark_hidden_states": (16, 16),
},
)
def test_dspark_proxy_width_respects_deferred_k3_boundary_capture(self):
from sglang.srt.models.kimi_k3 import (
KimiK3ForConditionalGeneration,
KimiK3LinearForCausalLM,
)
from sglang.srt.models.kimi_linear import KimiLinearForCausalLM
for start, k3_count, linear_count in [
(0, 0, 0),
(8, 0, 1),
(24, 1, 2),
(52, 2, 3),
]:
model = SimpleNamespace(
config=SimpleNamespace(hidden_size=8),
model=SimpleNamespace(
start_layer=start, dspark_layers_to_capture=[7, 23, 51]
),
)
self.assertEqual(
KimiK3LinearForCausalLM.get_pp_proxy_dspark_hidden_size(model),
k3_count * 8,
)
self.assertEqual(
KimiLinearForCausalLM.get_pp_proxy_dspark_hidden_size(model),
linear_count * 8,
)
model.get_pp_proxy_dspark_hidden_size = lambda: (
KimiK3LinearForCausalLM.get_pp_proxy_dspark_hidden_size(model)
)
self.assertEqual(
KimiK3ForConditionalGeneration.get_pp_proxy_dspark_hidden_size(
SimpleNamespace(language_model=model)
),
k3_count * 8,
)
def test_body_replay_uses_plural_embeds_and_skips_embedding_on_later_pp_stage(self):
for first_rank in (True, False):
for positional in (True, False):
with self.subTest(first_rank=first_rank, positional=positional):
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
runner._is_full_backend = False
runner._input_embeds_arg_idx = 3
backing = torch.zeros(4, 8)
supplied = torch.full((4, 8), 7.0) if first_rank else None
runner.buffer_registry = SimpleNamespace(
has_slot=lambda name: name == "input_embeds",
get_slot=lambda _: SimpleNamespace(
slice_for=lambda *args: backing
),
)
layer = SimpleNamespace(forward=lambda *args, **kwargs: None)
original = layer.forward
runner.layer_model = layer
sentinel = object()
runner.backend = SimpleNamespace(
replay=lambda *args, **kwargs: sentinel
)
runner._prefill_forward_context = lambda *args, **kwargs: (
nullcontext()
)
def outer_forward(ids, positions, batch):
if positional:
return layer.forward(None, positions, batch, supplied)
return layer.forward(
None, positions, batch, inputs_embeds=supplied
)
runner.model_runner = SimpleNamespace(
pp_group=SimpleNamespace(is_first_rank=first_rank),
model=SimpleNamespace(forward=outer_forward),
)
batch = SimpleNamespace(
input_ids=torch.arange(4),
positions=torch.arange(4),
mm_input_embeds=None,
)
result = runner._execute_body_capture(batch, batch, 4, 4, None)
self.assertIs(result, sentinel)
self.assertIs(layer.forward, original)
if first_rank:
torch.testing.assert_close(backing, supplied)
else:
self.assertEqual(torch.count_nonzero(backing).item(), 0)
def test_dspark_proxy_replay_updates_features_and_clears_padding(self):
buffers = PrefillInputBuffers.create(
device=torch.device("cpu"),
max_bs=1,
max_num_tokens=8,
cache_loc_dtype=torch.int64,
is_multimodal=False,
hidden_size=4,
dtype=torch.float32,
enable_mamba_track=False,
pp_size=2,
pp_proxy_dspark_hidden_size=8,
)
registry = build_prefill_registry(
device=torch.device("cpu"),
max_bs=1,
max_num_token=8,
cache_loc_dtype=torch.int64,
share_pool=False,
source=buffers,
)
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
runner.buffers = buffers
runner.model_runner = SimpleNamespace(
pp_group=SimpleNamespace(is_first_rank=False)
)
captured = runner._capture_pp_proxy_tensors(8)["dspark_hidden_states"]
ptr = captured.data_ptr()
for count, value in [(7, 2.0), (3, 5.0)]:
tokens = torch.arange(count)
proxy = PPProxyTensors(
{
"hidden_states": torch.zeros(count, 4),
"residual": torch.zeros(count, 4),
"dspark_hidden_states": torch.full((count, 8), value),
}
)
registry.fill_from(
SimpleNamespace(
input_ids=tokens, positions=tokens, out_cache_loc=tokens
),
raw_bs=1,
padded_bs=1,
raw_num_tokens=count,
padded_num_tokens=8,
pp_proxy_tensors=proxy,
)
self.assertEqual(captured.data_ptr(), ptr)
torch.testing.assert_close(captured[:count], proxy["dspark_hidden_states"])
self.assertEqual(torch.count_nonzero(captured[count:]).item(), 0)
def test_pipeline_proxy_output_is_supported(self):
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
runner.raw_num_tokens = 3
test/registered/unit/server_args/test_server_args.py
+22
View File
@@ -2501,6 +2501,28 @@ class TestPipelineParallelCompat(CustomTestCase):
def test_no_speculative_decoding_is_fine(self):
check_pipeline_parallel_compat(self._cfg())
def test_dspark_pd_prefill_does_not_require_eagle_architecture(self):
check_pipeline_parallel_compat(self._cfg(speculative_algorithm="DSPARK"))
def test_dspark_is_rejected_outside_pd_prefill(self):
for mode in ("decode", "null"):
with self.subTest(mode=mode):
with self.assertRaisesRegex(AssertionError, "DSPARK.*prefill"):
check_pipeline_parallel_compat(
self._cfg(
speculative_algorithm="DSPARK", disaggregation_mode=mode
)
)
def test_dspark_rejects_eagle_pp_relay(self):
with patch.object(
validation_hook.envs.SGLANG_ENABLE_PP_SPEC, "get", return_value=True
):
with self.assertRaisesRegex(AssertionError, "SGLANG_ENABLE_PP_SPEC"):
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="DSPARK")
)
def test_eagle_is_allowed_on_prefill(self):
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="EAGLE"),
test/registered/unit/spec/test_dspark_target_hidden_projection.py
+35
View File
@@ -9,6 +9,7 @@ from sglang.srt.models.dspark import DSparkDraftMixin
from sglang.srt.speculative.dspark_components.dspark_kv_inject import (
TargetHiddenKvInjector,
)
from sglang.srt.speculative.dspark_components.dspark_worker_v2 import DSparkWorkerV2
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -37,6 +38,40 @@ class _Attention:
class DSparkTargetHiddenProjectionTest(CustomTestCase):
def test_nonfinal_prefill_stage_only_forwards_target_proxies(self) -> None:
proxy = object()
target = SimpleNamespace(
model_runner=SimpleNamespace(attn_backend=object(), spec_algorithm=None),
device="cpu",
forward_batch_generation=lambda batch, *, pp_proxy_tensors, capture_hidden_mode: (
SimpleNamespace(pp_hidden_states_proxy_tensors=pp_proxy_tensors)
),
)
with (
mock.patch(
"sglang.srt.speculative.dspark_components.dspark_worker_v2.get_pp_group",
return_value=SimpleNamespace(is_last_rank=False),
),
mock.patch(
"sglang.srt.speculative.dspark_components.dspark_worker_v2.get_schedule",
return_value=SimpleNamespace(page_size=1),
),
):
worker = DSparkWorkerV2(None, 0, None, 0, target)
worker.alloc_memory_pool()
worker.init_attention_backends()
worker.init_cuda_graphs()
batch = SimpleNamespace(seq_lens=torch.tensor([8]))
result = worker.forward_batch_generation(batch, pp_proxy_tensors=proxy)
self.assertIs(result.pp_hidden_states_proxy_tensors, proxy)
self.assertIs(result.new_seq_lens, batch.seq_lens)
self.assertIsNone(worker.get_confidence_budget_prepare())
self.assertIsNone(worker.primary_draft_kv_pool)
self.assertEqual(worker.preloaded_weights_bytes, 0)
self.assertEqual(
worker.spec_v2_attn_backends, (target.model_runner.attn_backend,)
)
def test_single_aux_hidden_state_is_returned_without_copy(self) -> None:
hidden_states = torch.empty(2, 3)