Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
92632a60ba | ||
|
|
4b3b367b63 | ||
|
|
a78da9b524 | ||
|
|
b081dd3d23 | ||
|
|
c74a4037fb | ||
|
|
b963295489 | ||
|
|
fa826e08b1 | ||
|
|
3810f531a8 | ||
|
|
2580c24d1b | ||
|
|
4f22146e51 | ||
|
|
12e3b82e52 | ||
|
|
8305f66fc8 | ||
|
|
21a4a16b4b | ||
|
|
fc954b7e08 | ||
|
|
c2059c4fb2 |
@@ -0,0 +1,50 @@
|
||||
# Gitea Actions 自动构建 sglang 镜像(海外节点,原版源)
|
||||
# 基底 lmsysorg/sglang:dev-dsv41(docker.io),依赖走 pypi.org 默认源。
|
||||
# 触发:push 到 dsv41-pd 分支。
|
||||
#
|
||||
# 前置条件(一次性,在 Gitea 实例上配置):
|
||||
# 1. 实例已注册 act_runner(Gitea Actions runner,标签含 ubuntu-latest)
|
||||
# 2. 仓库 Settings → Secrets 添加:
|
||||
# REGISTRY_USERNAME / REGISTRY_PASSWORD(推送镜像的账号,如 Gitea 访问令牌)
|
||||
# 3. 可选:Settings → Actions → Variables 添加 REGISTRY(默认 git.agentwithu.com,
|
||||
# 即 Gitea 自带容器 registry;也可填 docker.io 等)
|
||||
name: build-sglang-image
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [dsv41-pd]
|
||||
|
||||
env:
|
||||
# 直接写死 Gitea 自带 registry(vars context 在该实例上求值异常会导致回退 docker.io)
|
||||
REGISTRY: git.agentwithu.com
|
||||
IMAGE_NAME: minke.yu/sglang
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Resolve tag
|
||||
id: meta
|
||||
run: |
|
||||
SHA9=$(git rev-parse --short=9 HEAD)
|
||||
echo "tag=${{ github.ref_name }}-${SHA9}-$(date +%Y%m%d-%H%M)" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- uses: docker/setup-buildx-action@v3
|
||||
|
||||
- uses: docker/login-action@v3
|
||||
with:
|
||||
registry: ${{ env.REGISTRY }}
|
||||
username: ${{ secrets.REGISTRY_USERNAME }}
|
||||
password: ${{ secrets.REGISTRY_PASSWORD }}
|
||||
|
||||
- uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
file: Dockerfile.gitea
|
||||
push: true
|
||||
tags: |
|
||||
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ steps.meta.outputs.tag }}
|
||||
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ github.ref_name }}-latest
|
||||
# 注意:type=gha 缓存在本实例的 runner 上会 404(cache server 未配),勿加回
|
||||
@@ -0,0 +1,11 @@
|
||||
# Gitea Actions 用(海外节点):原版 docker.io 基底 + 原版 pypi
|
||||
# 与 b300 离线版(/data/ymk/build/Dockerfile)的区别:
|
||||
# - 基底直用 docker.io 的 lmsysorg/sglang:dev-dsv41(不走 umirror/DaoCloud)
|
||||
# - 不设 PIP_INDEX_URL,用默认 pypi.org
|
||||
# - 源码由 CI checkout 后经 COPY 进镜像(不用 in-image clone)
|
||||
FROM lmsysorg/sglang:dev-dsv41
|
||||
|
||||
# checkout 含 .git,setuptools-scm 可打戳(main 系分支 version 显示 dev 属正常)
|
||||
COPY . /sgl-workspace/sglang
|
||||
|
||||
RUN pip install --no-cache-dir -e /sgl-workspace/sglang/python
|
||||
@@ -249,23 +249,33 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
|
||||
if (
|
||||
read_ragged_verify_mode() is not RaggedVerifyMode.STATIC
|
||||
or cfg.disaggregation_transfer_backend != "mooncake"
|
||||
or cfg.dp_size != 1
|
||||
or cfg.enable_dp_attention
|
||||
or cfg.attn_cp_size != 1
|
||||
or cfg.dcp_size != 1
|
||||
):
|
||||
raise ValueError(
|
||||
"DeepSeek-V4.1 DSpark PD requires static verify, Mooncake, "
|
||||
"DP=1 and CP=1. Both servers must enable DSpark with the same "
|
||||
"block size and TP size."
|
||||
"and CP=1 on both servers. DP attention is supported when "
|
||||
"both servers use the same block size and target/draft KV layout."
|
||||
)
|
||||
|
||||
from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase
|
||||
|
||||
prefill_graph = cfg.cuda_graph_config.prefill
|
||||
if prefill_graph.backend != Backend.DISABLED and prefill_graph.max_seq_len is None:
|
||||
# The captured low-ratio indexer scores a static context width; 16k
|
||||
# keeps it inside the candidate window at under 1 ms per layer.
|
||||
cp_breakable_prefill = (
|
||||
cfg.enable_prefill_cp
|
||||
and cfg.cp_strategy == "interleave"
|
||||
and cfg.tp_size > 1
|
||||
and prefill_graph.backend == Backend.BREAKABLE
|
||||
)
|
||||
if (
|
||||
prefill_graph.backend != Backend.DISABLED
|
||||
and prefill_graph.max_seq_len is None
|
||||
and not cp_breakable_prefill
|
||||
):
|
||||
# The non-CP captured low-ratio indexer scores a static context width.
|
||||
# CP BCG runs these sources eagerly with live prefix metadata, so this
|
||||
# default would only force long-prefix CP batches back to eager.
|
||||
# Explicit max_seq_len values still constrain both paths.
|
||||
declare_resolution(
|
||||
server_args,
|
||||
"validate_deepseek_v41_features",
|
||||
|
||||
@@ -966,12 +966,14 @@ def handle_language_model_only(server_args: Any):
|
||||
):
|
||||
if flag:
|
||||
raise ValueError(f"--language-model-only cannot be combined with {name}")
|
||||
if cfg.disaggregation_mode != "null":
|
||||
hf_config = model_config_of(server_args).hf_config
|
||||
# V4.1 text-only workers use the standard PD KV transfer path.
|
||||
if cfg.disaggregation_mode != "null" and hf_config.model_type != "deepseek_v41":
|
||||
raise ValueError(
|
||||
"--language-model-only is incompatible with --disaggregation-mode "
|
||||
"prefill/decode"
|
||||
)
|
||||
architectures = model_config_of(server_args).hf_config.architectures
|
||||
architectures = hf_config.architectures
|
||||
if not any(
|
||||
a in server_args.LANGUAGE_MODEL_ONLY_ARCHITECTURES for a in architectures
|
||||
):
|
||||
|
||||
@@ -945,9 +945,33 @@ class CommonKVManager(BaseKVManager):
|
||||
"enable DSpark with the same block size and target/draft KV "
|
||||
"layout. Upgrade both servers together."
|
||||
)
|
||||
if info.attn_tp_size != self.attn_tp_size:
|
||||
same_tp_with_prefill_cp = (
|
||||
info.attn_cp_size > 1
|
||||
and (self.is_mla_backend or self.is_hybrid_mla_backend)
|
||||
and self.attn_cp_size == 1
|
||||
and info.attn_tp_size * info.attn_cp_size == self.attn_tp_size
|
||||
)
|
||||
# Combined branch (40323-series + 40177): prefill CP can also pair
|
||||
# with a DP-attention decode server. MLA KV is replicated across
|
||||
# prefill CP ranks, so per-rank layouts match when attn_tp matches.
|
||||
dp_decode_with_prefill_cp = (
|
||||
info.attn_cp_size > 1
|
||||
and self.attn_cp_size == 1
|
||||
and (self.is_mla_backend or self.is_hybrid_mla_backend)
|
||||
and info.attn_tp_size == self.attn_tp_size
|
||||
)
|
||||
non_cp_mla_layout = info.attn_cp_size == self.attn_cp_size == 1 and (
|
||||
self.is_mla_backend or self.is_hybrid_mla_backend
|
||||
)
|
||||
if info.attn_tp_size != self.attn_tp_size and not (
|
||||
same_tp_with_prefill_cp
|
||||
or dp_decode_with_prefill_cp
|
||||
or non_cp_mla_layout
|
||||
):
|
||||
raise RuntimeError(
|
||||
"DeepSeek-V4.1 DSpark PD requires the same TP size on both servers"
|
||||
"DeepSeek-V4.1 DSpark PD requires matching attention TP "
|
||||
"unless both servers use CP=1 with an MLA KV layout, "
|
||||
"or prefill runs CP with an MLA KV layout"
|
||||
)
|
||||
|
||||
if self.dcp_size > 1:
|
||||
|
||||
@@ -82,6 +82,115 @@ FAILED_SESSION_RECOVERIES = Counter(
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Intra-node NVLink transport helpers.
|
||||
#
|
||||
# Mooncake's IntraNodeNvlinkTransport can only register and reach *device*
|
||||
# memory (it IPC-opens the remote cudaMalloc segments). Host-resident regions
|
||||
# (aux buffers, some state components) cannot be registered: one host region
|
||||
# makes the whole registerLocalMemoryBatch fail, and the engine then rolls
|
||||
# back *every* region, leaving the segment descriptor empty and all KV
|
||||
# transfers failing with "Requested address ... not found". When the
|
||||
# intra-node NVLink transport is active we therefore
|
||||
# 1. register only device-memory regions, and
|
||||
# 2. route blocks whose source is host memory over the ordered zmq channel
|
||||
# (same ordering guarantee the aux TCP path relies on) instead of the
|
||||
# transfer engine.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
import ctypes as _ctypes
|
||||
|
||||
_CUDA_MEMORY_TYPE_DEVICE = 2
|
||||
|
||||
try:
|
||||
from cuda.bindings import runtime as _cudart
|
||||
except ImportError: # pragma: no cover - cuda-python is always present in images
|
||||
_cudart = None
|
||||
|
||||
|
||||
def _is_device_pointer(ptr: int) -> bool:
|
||||
"""Probe a *local* pointer with cudaPointerGetAttributes.
|
||||
|
||||
Only valid for pointers owned by this process (never probe remote
|
||||
segment addresses). Returns False on any error so the caller falls back
|
||||
to the safe host path.
|
||||
"""
|
||||
if _cudart is None:
|
||||
# Cannot tell; assume device so behavior stays unchanged.
|
||||
return True
|
||||
err, attr = _cudart.cudaPointerGetAttributes(int(ptr))
|
||||
if int(err) != 0:
|
||||
# Clear the error so subsequent CUDA calls are not poisoned.
|
||||
_cudart.cudaGetLastError()
|
||||
return False
|
||||
return int(attr.type) == _CUDA_MEMORY_TYPE_DEVICE
|
||||
|
||||
|
||||
def _read_bytes_from_address(addr: int, length: int) -> Optional[bytes]:
|
||||
if length <= 0:
|
||||
return b""
|
||||
if _is_device_pointer(addr):
|
||||
buf = bytearray(length)
|
||||
# cudaMemcpyDeviceToHost = 2; synchronous default-stream copy.
|
||||
err, = _cudart.cudaMemcpy(
|
||||
_ctypes.addressof((_ctypes.c_char * length).from_buffer(buf)),
|
||||
int(addr),
|
||||
length,
|
||||
2,
|
||||
)
|
||||
if int(err) != 0:
|
||||
logger.error(
|
||||
f"cudaMemcpy D2H failed (err={err}) for addr {hex(addr)} len {length}"
|
||||
)
|
||||
return None
|
||||
return bytes(buf)
|
||||
return _ctypes.string_at(int(addr), length)
|
||||
|
||||
|
||||
def _write_bytes_to_address(addr: int, data: bytes) -> bool:
|
||||
if not data:
|
||||
return True
|
||||
if _is_device_pointer(addr):
|
||||
buf = _ctypes.create_string_buffer(data, len(data))
|
||||
# cudaMemcpyHostToDevice = 1; synchronous default-stream copy.
|
||||
err, = _cudart.cudaMemcpy(
|
||||
int(addr), _ctypes.addressof(buf), len(data), 1
|
||||
)
|
||||
if int(err) != 0:
|
||||
logger.error(
|
||||
f"cudaMemcpy H2D failed (err={err}) for addr {hex(addr)} "
|
||||
f"len {len(data)}"
|
||||
)
|
||||
return False
|
||||
return True
|
||||
_ctypes.memmove(int(addr), data, len(data))
|
||||
return True
|
||||
|
||||
|
||||
_NVLINK_INTRA_ACTIVE = None
|
||||
|
||||
|
||||
def _nvlink_intra_transport_active() -> bool:
|
||||
"""Whether mooncake installed the intra-node NVLink transport.
|
||||
|
||||
Mirrors the env probing in mooncake's transfer_engine_impl.cpp: the
|
||||
transport is installed iff MC_INTRANODE_NVLINK is set (any value), or an
|
||||
equivalent protocol selection was made.
|
||||
"""
|
||||
global _NVLINK_INTRA_ACTIVE
|
||||
if _NVLINK_INTRA_ACTIVE is None:
|
||||
active = bool(
|
||||
os.environ.get("MC_INTRANODE_NVLINK")
|
||||
or os.environ.get("MC_INTRA_NVLINK")
|
||||
)
|
||||
if not active:
|
||||
proto = (os.environ.get("MOONCAKE_PROTOCOL") or "").strip().lower()
|
||||
active = proto in ("nvlink_intra", "nvlink-intra", "intra_nvlink")
|
||||
_NVLINK_INTRA_ACTIVE = active
|
||||
return _NVLINK_INTRA_ACTIVE
|
||||
|
||||
|
||||
# decode
|
||||
@dataclasses.dataclass
|
||||
class TransferInfo:
|
||||
@@ -214,6 +323,7 @@ class KVArgsRegisterInfo:
|
||||
|
||||
class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
AUX_DATA_HEADER = b"AUX_DATA"
|
||||
STATE_DATA_HEADER = b"STATE_DATA"
|
||||
# Implements teardown() below, so runtime PD role switching is supported.
|
||||
supports_role_switch = True
|
||||
|
||||
@@ -227,6 +337,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
|
||||
self.init_engine()
|
||||
self.register_buffer_to_engine()
|
||||
# session_id -> (endpoint, dst_port, room), used to route host-memory
|
||||
# transfer blocks over zmq when the intra-node NVLink transport is
|
||||
# active (it cannot reach host memory). Populated on bootstrap.
|
||||
self._session_endpoint_map = {}
|
||||
self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
|
||||
self.max_transfer_batch_indices = (
|
||||
envs.SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES.get()
|
||||
@@ -322,6 +436,13 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
Deduped because the unified memory pool reports one raw buffer as both
|
||||
its KV and its mamba state component, and double registration fails in
|
||||
the engine.
|
||||
|
||||
When the intra-node NVLink transport is active, host-memory regions
|
||||
(aux buffers, some state components) are skipped: the transport only
|
||||
accepts device memory, and a single host region fails the whole batch
|
||||
and triggers a full engine-side rollback that would unregister the KV
|
||||
pools too. Host-resident payloads are instead exchanged over the
|
||||
ordered zmq channel (see _transfer_data / send_aux).
|
||||
"""
|
||||
regions: List[Tuple[int, int]] = []
|
||||
seen: Set[Tuple[int, int]] = set()
|
||||
@@ -338,6 +459,24 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
|
||||
):
|
||||
add(ptrs, lens)
|
||||
|
||||
if _nvlink_intra_transport_active():
|
||||
device_regions = []
|
||||
skipped = []
|
||||
for ptr, length in regions:
|
||||
if _is_device_pointer(ptr):
|
||||
device_regions.append((ptr, length))
|
||||
else:
|
||||
skipped.append((ptr, length))
|
||||
if skipped:
|
||||
logger.info(
|
||||
"Intra-node NVLink transport: skipping %d host-memory "
|
||||
"regions from engine registration (they will be exchanged "
|
||||
"over the zmq channel instead): %s",
|
||||
len(skipped),
|
||||
[(hex(p), l) for p, l in skipped[:8]],
|
||||
)
|
||||
regions = device_regions
|
||||
return regions
|
||||
|
||||
def register_buffer_to_engine(self):
|
||||
@@ -748,10 +887,63 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
if not transfer_blocks:
|
||||
return 0
|
||||
|
||||
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
|
||||
return self.engine.batch_transfer_sync(
|
||||
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
|
||||
)
|
||||
if not _nvlink_intra_transport_active():
|
||||
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
|
||||
return self.engine.batch_transfer_sync(
|
||||
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
|
||||
)
|
||||
|
||||
# Intra-node NVLink transport can only move device memory. Partition
|
||||
# blocks by the *local source* pointer (probing a local pointer is
|
||||
# safe; the remote dst is never probed): device-sourced blocks go
|
||||
# through the engine as usual, host-sourced blocks are shipped over
|
||||
# the ordered zmq channel and written into the peer's buffer by the
|
||||
# receiver (see _handle_state_data). This mirrors the aux TCP path.
|
||||
device_blocks = []
|
||||
host_blocks = []
|
||||
for src, dst, length in transfer_blocks:
|
||||
if _is_device_pointer(src):
|
||||
device_blocks.append((src, dst, length))
|
||||
else:
|
||||
host_blocks.append((src, dst, length))
|
||||
|
||||
rc = 0
|
||||
if device_blocks:
|
||||
src_addrs, dst_addrs, lengths = zip(*device_blocks)
|
||||
rc = self.engine.batch_transfer_sync(
|
||||
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
|
||||
)
|
||||
if rc == 0 and host_blocks:
|
||||
rc = self._send_host_blocks_tcp(mooncake_session_id, host_blocks)
|
||||
return rc
|
||||
|
||||
def _send_host_blocks_tcp(self, mooncake_session_id, host_blocks):
|
||||
target = self._session_endpoint_map.get(mooncake_session_id)
|
||||
if target is None:
|
||||
logger.error(
|
||||
f"No zmq endpoint known for mooncake session "
|
||||
f"{mooncake_session_id}; cannot deliver {len(host_blocks)} "
|
||||
"host-memory transfer blocks"
|
||||
)
|
||||
return -1
|
||||
endpoint, dst_port, room = target
|
||||
na = NetworkAddress(endpoint, dst_port)
|
||||
for src, dst, length in host_blocks:
|
||||
data = _read_bytes_from_address(src, length)
|
||||
if data is None:
|
||||
return -1
|
||||
self._send_multipart_locked(
|
||||
na.to_tcp(),
|
||||
[
|
||||
MooncakeKVManager.STATE_DATA_HEADER,
|
||||
str(room).encode("ascii"),
|
||||
str(int(dst)).encode("ascii"),
|
||||
struct.pack(">I", len(data)),
|
||||
data,
|
||||
],
|
||||
is_ipv6=na.is_ipv6,
|
||||
)
|
||||
return 0
|
||||
|
||||
def _send_kvcache_generic(
|
||||
self,
|
||||
@@ -1482,8 +1674,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
):
|
||||
# TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free
|
||||
if (
|
||||
self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK"
|
||||
) or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get():
|
||||
(self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK")
|
||||
or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get()
|
||||
or _nvlink_intra_transport_active()
|
||||
):
|
||||
return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs)
|
||||
|
||||
transfer_blocks = []
|
||||
@@ -1566,6 +1760,71 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}"
|
||||
)
|
||||
|
||||
def _host_transfer_regions(self):
|
||||
"""Address ranges this process published as transfer targets.
|
||||
|
||||
Used to validate STATE_DATA writes. Built lazily because kv_args is
|
||||
fully populated only after registration.
|
||||
"""
|
||||
regions = getattr(self, "_host_transfer_regions_cache", None)
|
||||
if regions is None:
|
||||
regions = []
|
||||
for ptr, length in zip(
|
||||
self.kv_args.kv_data_ptrs or [], self.kv_args.kv_data_lens or []
|
||||
):
|
||||
regions.append((int(ptr), int(ptr) + int(length)))
|
||||
for ptr, length in zip(
|
||||
self.kv_args.aux_data_ptrs or [], self.kv_args.aux_data_lens or []
|
||||
):
|
||||
regions.append((int(ptr), int(ptr) + int(length)))
|
||||
for ptrs, lens in zip(
|
||||
self.kv_args.state_data_ptrs or [], self.kv_args.state_data_lens or []
|
||||
):
|
||||
for ptr, length in zip(ptrs or [], lens or []):
|
||||
regions.append((int(ptr), int(ptr) + int(length)))
|
||||
self._host_transfer_regions_cache = regions
|
||||
return regions
|
||||
|
||||
def _handle_state_data(self, msg: List[bytes]):
|
||||
"""Handle STATE_DATA messages received by the decode thread.
|
||||
|
||||
Carries one host-memory transfer block that could not go through the
|
||||
intra-node NVLink transport. Written directly into the local buffer at
|
||||
the destination address; ordering against the final status message is
|
||||
guaranteed by the shared per-endpoint zmq socket.
|
||||
"""
|
||||
room = int(msg[1].decode("ascii"))
|
||||
dst_addr = int(msg[2].decode("ascii"))
|
||||
data_length = struct.unpack(">I", msg[3])[0]
|
||||
data = msg[4]
|
||||
|
||||
if len(data) != data_length:
|
||||
logger.error(f"STATE_DATA length mismatch for bootstrap_room {room}")
|
||||
return
|
||||
|
||||
in_region = any(
|
||||
start <= dst_addr and dst_addr + len(data) <= end
|
||||
for start, end in self._host_transfer_regions()
|
||||
)
|
||||
if not in_region:
|
||||
logger.error(
|
||||
f"STATE_DATA for bootstrap_room {room} targets unknown region "
|
||||
f"{hex(dst_addr)}..{hex(dst_addr + len(data))}; dropping"
|
||||
)
|
||||
return
|
||||
|
||||
if not _write_bytes_to_address(dst_addr, data):
|
||||
logger.error(
|
||||
f"STATE_DATA write failed for bootstrap_room {room} at "
|
||||
f"{hex(dst_addr)} len {len(data)}"
|
||||
)
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
f"Received STATE_DATA for bootstrap_room {room} at {hex(dst_addr)} "
|
||||
f"with length:{len(data)}"
|
||||
)
|
||||
|
||||
def _get_dsa_cache_transfer_skip_flags(
|
||||
self, info: Optional[KVArgsRegisterInfo]
|
||||
) -> Tuple[bool, bool]:
|
||||
@@ -2412,6 +2671,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
):
|
||||
self._staging_outstanding.pop(kv_chunk.room, None)
|
||||
if kv_chunk.room in self.transfer_infos:
|
||||
for sid in self.transfer_infos[kv_chunk.room]:
|
||||
self._session_endpoint_map.pop(sid, None)
|
||||
self.transfer_infos.pop(kv_chunk.room)
|
||||
self.req_to_decode_prefix_len.pop(kv_chunk.room, None)
|
||||
if self.enable_staging:
|
||||
@@ -2557,6 +2818,11 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
self.transfer_infos[room][mooncake_session_id] = (
|
||||
TransferInfo.from_zmq(waiting_req_bytes)
|
||||
)
|
||||
self._session_endpoint_map[mooncake_session_id] = (
|
||||
self.transfer_infos[room][mooncake_session_id].endpoint,
|
||||
self.transfer_infos[room][mooncake_session_id].dst_port,
|
||||
room,
|
||||
)
|
||||
# NOTE: after bootstrapping we can mark the req as waiting for input
|
||||
if len(self.transfer_infos[room]) == required_dst_info_num:
|
||||
self.resolve_kv_replica_factor(self.transfer_infos[room])
|
||||
@@ -2585,6 +2851,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
|
||||
self._handle_aux_data(msg)
|
||||
continue
|
||||
if msg[0] == MooncakeKVManager.STATE_DATA_HEADER:
|
||||
self._handle_state_data(msg)
|
||||
continue
|
||||
|
||||
# Staging: prefill notifies a chunk written to staging buffer
|
||||
if msg[0] == b"CHUNK_READY":
|
||||
|
||||
@@ -736,7 +736,14 @@ class DSV4AttnMetadata:
|
||||
if src_val is None and dst_val is None:
|
||||
continue
|
||||
assert dst_val is not None, f"{field_name=} {src_val=} {dst_val=}"
|
||||
dst_val.copy_(src_val)
|
||||
shape_mismatch = dst_val.shape != src_val.shape
|
||||
assert not shape_mismatch or field_name in self._CP_GLOBAL_FIELDS, (
|
||||
f"Only CP-global replay metadata may use a shorter live prefix, "
|
||||
f"got {field_name=} {src_val.shape=} {dst_val.shape=}"
|
||||
)
|
||||
_copy_tensor_allowing_storage_alias(
|
||||
dst_val, src_val, pad_value=0 if shape_mismatch else None
|
||||
)
|
||||
|
||||
# These fields are safe to replace because captured kernels only need
|
||||
# the current per-replay objects, or the field is produced inside the
|
||||
@@ -988,6 +995,27 @@ def _prefill_graph_max_seq_len() -> Optional[int]:
|
||||
return get_exec().graph.cuda_graph_config.prefill.max_seq_len
|
||||
|
||||
|
||||
def _copy_tensor_allowing_storage_alias(
|
||||
dst: torch.Tensor, src: torch.Tensor, *, pad_value: Optional[int] = None
|
||||
) -> None:
|
||||
"""Copy replay metadata while preserving capture-stable destination addresses."""
|
||||
if dst is src:
|
||||
return
|
||||
if dst.untyped_storage().data_ptr() == src.untyped_storage().data_ptr():
|
||||
src = src.clone()
|
||||
if dst.shape == src.shape:
|
||||
dst.copy_(src)
|
||||
return
|
||||
assert (
|
||||
pad_value is not None
|
||||
and dst.ndim == src.ndim
|
||||
and dst.shape[0] >= src.shape[0]
|
||||
and dst.shape[1:] == src.shape[1:]
|
||||
), f"Cannot copy replay metadata from {src.shape=} to {dst.shape=}"
|
||||
dst.fill_(pad_value)
|
||||
dst[: src.shape[0]].copy_(src)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DSV4Metadata:
|
||||
core_attn_metadata: DSV4AttnMetadata
|
||||
@@ -1254,6 +1282,11 @@ class DeepseekV4AttnBackend(
|
||||
] = None
|
||||
self.online_c128_mtp = OnlineC128MTPController(self)
|
||||
self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device)
|
||||
# CP V4.1 consumers share compressed KV across layers. Separate ratio
|
||||
# workspaces keep those prefixes intact while each layer refreshes SWA.
|
||||
self.shared_compressed_prefill_workspaces = {
|
||||
ratio: SparsePrefillWorkspace(self.device) for ratio in (1, 2)
|
||||
}
|
||||
spec_alg = model_runner.spec_algorithm
|
||||
self.needs_cpu_seq_lens = not spec_alg.is_dspark() and (
|
||||
not _is_cuda or self.online_c128_mtp.enabled()
|
||||
@@ -1566,8 +1599,12 @@ class DeepseekV4AttnBackend(
|
||||
|
||||
@property
|
||||
def low_ratio_prefill_graph(self) -> bool:
|
||||
"""Whether ratio-1/2 sources use captured projections and indexer metadata."""
|
||||
return (
|
||||
bool(self.low_ratios) and _has_dense_fp4_indexer() and _is_sm100_or_newer()
|
||||
bool(self.low_ratios)
|
||||
and _has_dense_fp4_indexer()
|
||||
and _is_sm100_or_newer()
|
||||
and get_parallel().attn_cp_size == 1
|
||||
)
|
||||
|
||||
def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool:
|
||||
@@ -2870,7 +2907,7 @@ class DeepseekV4AttnBackend(
|
||||
q_lora[:num_local],
|
||||
positions[:num_local].to(torch.int64),
|
||||
forward_batch,
|
||||
torch.tensor(q_lens_cpu, dtype=torch.int32, device=x.device),
|
||||
self._move_to_device(q_lens_cpu),
|
||||
q_lens_cpu,
|
||||
)
|
||||
|
||||
@@ -3236,7 +3273,9 @@ class DeepseekV4AttnBackend(
|
||||
continue
|
||||
j = torch.arange(lc, device=device)
|
||||
slot_chunks.append(
|
||||
self.req_to_token[req_pool_indices[r], j * ratio].to(torch.int64)
|
||||
self.req_to_token[req_pool_indices[r : r + 1], j * ratio].to(
|
||||
torch.int64
|
||||
)
|
||||
// ratio
|
||||
)
|
||||
start += lc
|
||||
@@ -3261,7 +3300,7 @@ class DeepseekV4AttnBackend(
|
||||
weights = indexer.head_weights(x).float()
|
||||
compress_lens = ((pos + 1) // ratio).to(torch.int32)
|
||||
ks = torch.repeat_interleave(
|
||||
torch.tensor(starts, dtype=torch.int32, device=device),
|
||||
self._move_to_device(starts),
|
||||
q_lens.to(torch.int64),
|
||||
output_size=num_tokens,
|
||||
)
|
||||
@@ -3945,20 +3984,38 @@ class DeepseekV4AttnBackend(
|
||||
compress_ratio, core_attn_metadata, extra_page_size
|
||||
)
|
||||
n_compressed = flat_token_ids.shape[0]
|
||||
workspace = self.sparse_prefill_workspace.get(
|
||||
n_compressed + cache.swa_token_ids.shape[0]
|
||||
reuse_compressed = compress_ratio in (1, 2) and is_cp_active(forward_batch)
|
||||
workspace_pool = (
|
||||
self.shared_compressed_prefill_workspaces[compress_ratio]
|
||||
if reuse_compressed
|
||||
else self.sparse_prefill_workspace
|
||||
)
|
||||
workspace = workspace_pool.get(n_compressed + cache.swa_token_ids.shape[0])
|
||||
compressed_slice = workspace[:n_compressed]
|
||||
swa_slice = workspace[n_compressed:]
|
||||
|
||||
if compressed_slice is not None:
|
||||
dequantize_k_cache_paged(
|
||||
extra_k_cache,
|
||||
flat_token_ids,
|
||||
page_size=extra_page_size,
|
||||
out=compressed_slice,
|
||||
layout=token_to_kv_pool.get_extra_key_layout(layer_id),
|
||||
)
|
||||
source_key = None
|
||||
if reuse_compressed:
|
||||
source_layer = token_to_kv_pool.source_layer_of(layer_id)
|
||||
source_key = (source_layer, workspace.data_ptr())
|
||||
gather = cache.compressed[compress_ratio]
|
||||
# A source layer may have just updated its cache in place. Consumer
|
||||
# layers only reuse the compressed prefix; their top-k and SWA stay live.
|
||||
if (
|
||||
source_key is None
|
||||
or layer_id == source_key[0]
|
||||
or gather.dequantized_source != source_key
|
||||
):
|
||||
dequantize_k_cache_paged(
|
||||
extra_k_cache,
|
||||
flat_token_ids,
|
||||
page_size=extra_page_size,
|
||||
out=compressed_slice,
|
||||
layout=token_to_kv_pool.get_extra_key_layout(layer_id),
|
||||
)
|
||||
if source_key is not None:
|
||||
gather.dequantized_source = source_key
|
||||
dequantize_k_cache_paged(
|
||||
token_to_kv_pool.get_swa_key_buffer_radix(layer_id),
|
||||
cache.swa_token_ids,
|
||||
|
||||
@@ -78,10 +78,11 @@ def use_dsv4_q8kv8_sparse_prefill(dsv4_prefill_backend: str = "auto") -> bool:
|
||||
class SparsePrefillWorkspace:
|
||||
"""Backend-owned scratch storage for sparse prefill KV dequantization.
|
||||
|
||||
The workspace contents are fully overwritten before every attention call,
|
||||
so token buckets and compression ratios can safely share one buffer. Sparse
|
||||
prefill executes eagerly and serially on the supported paths, which makes it
|
||||
safe to replace the scratch allocation when a larger extent is needed.
|
||||
Callers normally overwrite the entire workspace. Shared compressed-KV
|
||||
callers keep separate workspaces per ratio and track prefix validity in the
|
||||
per-forward gather cache, including the allocation address. Sparse prefill
|
||||
executes eagerly and serially on the supported paths, so the allocation can
|
||||
be replaced when a larger extent is needed.
|
||||
"""
|
||||
|
||||
def __init__(self, device: torch.device):
|
||||
@@ -275,14 +276,18 @@ class CompressedGather:
|
||||
# chunk-invariant per request; subsequent layers only overwrite that prefix.
|
||||
combined_indices: Optional[torch.Tensor] = None
|
||||
combined_lens: Optional[torch.Tensor] = None
|
||||
# Valid only for this forward's gather layout. Each ratio has its own
|
||||
# workspace; its compressed prefix survives consumer layers' SWA writes.
|
||||
dequantized_source: Optional[tuple[int, int]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SparsePrefillChunkCache:
|
||||
"""Cache prefill-chunk metadata shared across layers.
|
||||
|
||||
Fields depend on request/token mappings and compressed page tables, not
|
||||
per-layer k_cache; per-layer top-k combinations are recomputed into reused
|
||||
Gather layouts depend on request/token mappings and compressed page tables.
|
||||
Shared-source dequantization keys live only for this forward; per-layer
|
||||
top-k combinations are recomputed into reused
|
||||
buffers.
|
||||
"""
|
||||
|
||||
|
||||
@@ -23,17 +23,22 @@ import torch
|
||||
|
||||
from sglang.srt.arg_groups.overrides import (
|
||||
attention_backends_of,
|
||||
model_config_of,
|
||||
resolved_view,
|
||||
resolving_view,
|
||||
)
|
||||
from sglang.srt.configs.model_config import is_deepseek_v4
|
||||
from sglang.srt.layers.cp.base import get_cp_strategy
|
||||
from sglang.srt.layers.cp.interleave import InterleaveCPStrategy
|
||||
from sglang.srt.layers.cp.padding import get_cp_padding_align_size
|
||||
from sglang.srt.layers.cp.utils import (
|
||||
cp_gather_after_forward,
|
||||
cp_shard_hidden_states,
|
||||
cp_split_before_forward,
|
||||
prepare_cp_forward,
|
||||
)
|
||||
from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
|
||||
from sglang.srt.layers.logits_processor import LogitsMetadata
|
||||
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -50,12 +55,18 @@ def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool:
|
||||
cfg = resolving_view(server_args)
|
||||
resolved = resolved_view(server_args)
|
||||
prefill_attention_backend, _ = attention_backends_of(resolved_view(server_args))
|
||||
supports_layout = (
|
||||
cfg.cp_strategy == "zigzag" and prefill_attention_backend == "trtllm_mha"
|
||||
) or (
|
||||
cfg.cp_strategy == "interleave"
|
||||
and prefill_attention_backend == "dsv4"
|
||||
and is_deepseek_v4(model_config_of(server_args).hf_config)
|
||||
)
|
||||
return (
|
||||
cfg.enable_prefill_cp
|
||||
and cfg.pp_size == 1
|
||||
and resolved.attn_cp_size == cfg.tp_size
|
||||
and cfg.cp_strategy == "zigzag"
|
||||
and prefill_attention_backend == "trtllm_mha"
|
||||
and supports_layout
|
||||
)
|
||||
|
||||
|
||||
@@ -67,8 +78,12 @@ def enable_cp_bcg_capture(server_args: ServerArgs) -> bool:
|
||||
def filter_prefill_cp_bcg_capture_num_tokens(
|
||||
capture_num_tokens: list[int], server_args: ServerArgs
|
||||
) -> list[int]:
|
||||
"""Keep only token buckets where the zigzag CP strategy can run."""
|
||||
min_num_tokens = resolved_view(server_args).attn_cp_size * 2
|
||||
"""Keep only token buckets where the configured CP strategy can run."""
|
||||
cfg = resolving_view(server_args)
|
||||
cp_segments_per_token_block = 2 if cfg.cp_strategy == "zigzag" else 1
|
||||
min_num_tokens = (
|
||||
resolved_view(server_args).attn_cp_size * cp_segments_per_token_block
|
||||
)
|
||||
filtered = [size for size in capture_num_tokens if size >= min_num_tokens]
|
||||
if not filtered:
|
||||
raise ValueError(
|
||||
@@ -96,6 +111,8 @@ class PrefillCPBCGInput:
|
||||
|
||||
input_embeds: torch.Tensor
|
||||
positions: torch.Tensor
|
||||
input_ids: Optional[torch.Tensor] = None
|
||||
num_token_non_padded: Optional[torch.Tensor] = None
|
||||
bucket_local_tokens: Dict[int, int] = field(default_factory=dict)
|
||||
live_local_tokens: int = 0
|
||||
|
||||
@@ -114,12 +131,22 @@ class PrefillCPBCGInput:
|
||||
(runner.max_num_tokens,),
|
||||
dtype=torch.int64,
|
||||
),
|
||||
input_ids=torch.zeros((runner.max_num_tokens,), dtype=torch.int64),
|
||||
num_token_non_padded=torch.zeros((), dtype=torch.int32),
|
||||
)
|
||||
|
||||
def required_local_tokens(self, extend_seq_lens: Any) -> Optional[int]:
|
||||
"""Return the aligned CP-local rows required by a live zigzag layout."""
|
||||
"""Return the aligned CP-local rows required by the active layout."""
|
||||
strategy = get_cp_strategy()
|
||||
if not isinstance(strategy, ZigzagCPStrategy) or extend_seq_lens is None:
|
||||
if extend_seq_lens is None:
|
||||
return None
|
||||
if isinstance(strategy, InterleaveCPStrategy):
|
||||
logical_tokens = (
|
||||
sum(int(length) for length in extend_seq_lens) + strategy.cp_size - 1
|
||||
) // strategy.cp_size
|
||||
align_size = get_cp_padding_align_size()
|
||||
return (logical_tokens + align_size - 1) // align_size * align_size
|
||||
if not isinstance(strategy, ZigzagCPStrategy):
|
||||
return None
|
||||
|
||||
cp_segment_num = strategy.cp_size * 2
|
||||
@@ -219,6 +246,7 @@ class PrefillCPBCGInput:
|
||||
raw_tokens = int(forward_batch.extend_num_tokens)
|
||||
global_input_ids = forward_batch.input_ids[:raw_tokens]
|
||||
global_positions = forward_batch.positions[:raw_tokens]
|
||||
local_input_ids = cp_shard_hidden_states(global_input_ids, forward_batch)
|
||||
global_input_embeds = runner.model_runner.model.get_input_embeddings()(
|
||||
global_input_ids
|
||||
)
|
||||
@@ -249,12 +277,31 @@ class PrefillCPBCGInput:
|
||||
|
||||
input_embeds = self.input_embeds[:captured_local_tokens]
|
||||
positions = self.positions[:captured_local_tokens]
|
||||
assert self.input_ids is not None
|
||||
input_ids = self.input_ids[:captured_local_tokens]
|
||||
input_embeds.zero_()
|
||||
positions.zero_()
|
||||
input_ids.zero_()
|
||||
input_embeds[:live_local_tokens].copy_(local_input_embeds)
|
||||
positions[:live_local_tokens].copy_(local_positions)
|
||||
input_ids[:live_local_tokens].copy_(local_input_ids)
|
||||
forward_batch.input_embeds = input_embeds
|
||||
forward_batch.positions = positions
|
||||
forward_batch._cp_positions = positions
|
||||
# Keep the global input_ids field intact: the runner uses its length to
|
||||
# select the global capture bucket. The DSV4 body consumes this fixed,
|
||||
# rank-local view for hash routing and MegaMoE.
|
||||
forward_batch._cp_input_ids = input_ids
|
||||
forward_batch.input_ids_global = input_ids
|
||||
if forward_batch.num_token_non_padded is not None:
|
||||
assert self.num_token_non_padded is not None
|
||||
metadata = forward_batch.attn_cp_metadata
|
||||
logical_tokens = (
|
||||
metadata.per_rank_logical_token or metadata.per_rank_actual_token
|
||||
)
|
||||
strategy = get_cp_strategy()
|
||||
assert strategy is not None
|
||||
self.num_token_non_padded.fill_(logical_tokens[strategy.cp_rank])
|
||||
forward_batch.num_token_non_padded = self.num_token_non_padded
|
||||
self.live_local_tokens = live_local_tokens
|
||||
|
||||
|
||||
@@ -307,10 +354,50 @@ def execute_prefill_cp_bcg(
|
||||
static_forward_batch,
|
||||
torch.cuda.current_stream(),
|
||||
)
|
||||
return model.logits_processor(
|
||||
forward_batch.input_ids,
|
||||
if aux_hidden_states is not None:
|
||||
if torch.is_tensor(aux_hidden_states):
|
||||
aux_hidden_states = cp_gather_after_forward(
|
||||
aux_hidden_states, static_forward_batch, torch.cuda.current_stream()
|
||||
)
|
||||
else:
|
||||
aux_hidden_states = [
|
||||
cp_gather_after_forward(
|
||||
aux, static_forward_batch, torch.cuda.current_stream()
|
||||
)
|
||||
for aux in aux_hidden_states
|
||||
]
|
||||
hidden_states_before_norm = None
|
||||
if isinstance(hidden_states, tuple):
|
||||
assert len(hidden_states) == 2
|
||||
hidden_states, hidden_states_before_norm = hidden_states
|
||||
|
||||
input_ids = forward_batch.input_ids
|
||||
logits_metadata = forward_batch
|
||||
tail = None
|
||||
language_model = getattr(model, "model", None)
|
||||
if (
|
||||
capture_aux_hidden_states
|
||||
and getattr(language_model, "late_layer_start", None) is not None
|
||||
and forward_batch.forward_mode.is_extend_without_speculative()
|
||||
):
|
||||
tail_metadata = runner.model_runner.attn_backend.tail_forward_metadata
|
||||
tail = tail_metadata.late_layer_tail
|
||||
input_ids = tail.rows(input_ids)
|
||||
logits_metadata = LogitsMetadata.from_forward_batch(forward_batch)
|
||||
logits_metadata.extend_seq_lens = tail.extend_seq_lens
|
||||
logits_metadata.extend_seq_lens_cpu = tail.extend_seq_lens_cpu
|
||||
logits_metadata.extend_logprob_start_lens_cpu = tail.extend_seq_lens_cpu
|
||||
|
||||
output = model.logits_processor(
|
||||
input_ids,
|
||||
hidden_states,
|
||||
model.lm_head,
|
||||
forward_batch,
|
||||
logits_metadata,
|
||||
aux_hidden_states,
|
||||
hidden_states_before_norm=(
|
||||
None if aux_hidden_states is not None else hidden_states_before_norm
|
||||
),
|
||||
)
|
||||
if tail is not None:
|
||||
output.hidden_states_token_indices = tail.token_indices
|
||||
return output
|
||||
|
||||
@@ -216,20 +216,33 @@ def _run_mega_routed(
|
||||
|
||||
if num_tokens > 0:
|
||||
router_logits = moe.gate(hidden_states, forward_batch=forward_batch)
|
||||
topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {}
|
||||
topk_output = moe.topk(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
num_token_non_padded=(
|
||||
forward_batch.num_token_non_padded
|
||||
if forward_batch is not None
|
||||
else None
|
||||
),
|
||||
expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
|
||||
layer_id=moe.layer_id,
|
||||
),
|
||||
**topk_kwargs,
|
||||
num_token_non_padded = (
|
||||
forward_batch.num_token_non_padded if forward_batch is not None else None
|
||||
)
|
||||
if isinstance(
|
||||
getattr(moe.gate, "e_score_correction_bias_vl", None), torch.Tensor
|
||||
):
|
||||
# V4.1 uses a different correction bias for image-token rows. The
|
||||
# MegaMoE transport consumes the same routed ids/weights as TopK.
|
||||
from sglang.srt.multimodal.dsv41.vl_routing import vision_topk
|
||||
|
||||
topk_output = vision_topk(
|
||||
moe,
|
||||
router_logits,
|
||||
input_ids_global,
|
||||
num_token_non_padded=num_token_non_padded,
|
||||
)
|
||||
else:
|
||||
topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {}
|
||||
topk_output = moe.topk(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
num_token_non_padded=num_token_non_padded,
|
||||
expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
|
||||
layer_id=moe.layer_id,
|
||||
),
|
||||
**topk_kwargs,
|
||||
)
|
||||
topk_ids = topk_output.topk_ids
|
||||
topk_weights = topk_output.topk_weights
|
||||
else:
|
||||
|
||||
@@ -155,6 +155,10 @@ def free_kv_row_segments(
|
||||
|
||||
def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
|
||||
if getattr(req, "skip_radix_cache_insert", False):
|
||||
kv_indices = tree_cache.req_to_token_pool.req_to_token[
|
||||
req.kv.req_pool_idx, : len(req.get_fill_ids())
|
||||
]
|
||||
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
|
||||
return
|
||||
|
||||
tree_cache.cache_unfinished_req(req, **kwargs)
|
||||
|
||||
@@ -385,14 +385,22 @@ class EagerRunner(BaseRunner):
|
||||
"""
|
||||
model = self.model_runner.model
|
||||
|
||||
input_ids = forward_batch.input_ids
|
||||
input_embeds = kwargs.get("input_embeds")
|
||||
# Multimodal spans must be embedded in global token order, before CP
|
||||
# slicing. The model may also normalize image hash IDs for its router.
|
||||
prepare_inputs = getattr(model, "prepare_language_model_inputs", None)
|
||||
if prepare_inputs is not None:
|
||||
input_ids, input_embeds = prepare_inputs(
|
||||
input_ids, forward_batch, input_embeds
|
||||
)
|
||||
if input_embeds is None:
|
||||
input_embeds = model.get_input_embeddings()(forward_batch.input_ids)
|
||||
input_embeds = model.get_input_embeddings()(input_ids)
|
||||
with cp_shard_model_inputs(
|
||||
input_embeds,
|
||||
forward_batch.positions,
|
||||
forward_batch,
|
||||
forward_batch.input_ids,
|
||||
input_ids,
|
||||
) as (sharded_input_embeds, sharded_positions, model_input_ids):
|
||||
model_kwargs = {"input_embeds": sharded_input_embeds}
|
||||
if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None:
|
||||
@@ -437,7 +445,7 @@ class EagerRunner(BaseRunner):
|
||||
if aux_hidden_states is None:
|
||||
logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
|
||||
return model.logits_processor(
|
||||
forward_batch.input_ids,
|
||||
input_ids,
|
||||
hidden_states,
|
||||
model.lm_head,
|
||||
forward_batch,
|
||||
|
||||
@@ -701,6 +701,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
||||
|
||||
def _get_layer_model_positions(self, forward_batch: ForwardBatch) -> torch.Tensor:
|
||||
"""Mirror outer multimodal wrappers when BCG captures layer_model directly."""
|
||||
cp_positions = getattr(forward_batch, "_cp_positions", None)
|
||||
if cp_positions is not None:
|
||||
return cp_positions
|
||||
if forward_batch.mrope_positions is None:
|
||||
return forward_batch.positions
|
||||
|
||||
@@ -782,7 +785,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
||||
if self._uses_eager_prefill_tail():
|
||||
# BCG / Full: capture the transformer body only.
|
||||
positions = self._get_layer_model_positions(forward_batch)
|
||||
input_ids = forward_batch.input_ids
|
||||
input_ids = getattr(
|
||||
forward_batch, "_cp_input_ids", forward_batch.input_ids
|
||||
)
|
||||
kwargs = _build_layer_model_forward_kwargs(
|
||||
self.layer_model, forward_batch, pp_proxy_tensors
|
||||
)
|
||||
@@ -1336,9 +1341,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
||||
batch_max_context_len=batch_max_context_len,
|
||||
):
|
||||
return False
|
||||
if getattr(self, "enable_cp_bcg_capture", False) and is_cp_active(
|
||||
forward_batch
|
||||
):
|
||||
if getattr(self, "enable_cp_bcg_capture", False):
|
||||
if not is_cp_active(forward_batch):
|
||||
return False
|
||||
assert self.prefill_cp_bcg_input is not None
|
||||
if (
|
||||
self.prefill_cp_bcg_input.select_replay_bucket_for_batch(
|
||||
|
||||
@@ -182,6 +182,7 @@ from sglang.srt.multimodal.deepseek_v41_image_processing import (
|
||||
)
|
||||
from sglang.srt.runtime_context import (
|
||||
get_device,
|
||||
get_disagg,
|
||||
get_exec,
|
||||
get_forward,
|
||||
get_parallel,
|
||||
@@ -2039,7 +2040,10 @@ class MQALayer(MqaAttentionBase):
|
||||
if (
|
||||
forward_batch.forward_mode.is_extend()
|
||||
and is_in_breakable_cuda_graph()
|
||||
and not getattr(attn_backend, "low_ratio_prefill_graph", False)
|
||||
and (
|
||||
dsa_use_prefill_cp(forward_batch)
|
||||
or not getattr(attn_backend, "low_ratio_prefill_graph", False)
|
||||
)
|
||||
):
|
||||
bcg_deepseek_v4_low_ratio_sources(self, x, q_lora, positions)
|
||||
else:
|
||||
@@ -2650,7 +2654,8 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
is_nextn=is_nextn,
|
||||
is_deepseek_v4=True,
|
||||
vl_correction_bias=config.model_type == "deepseek_v41"
|
||||
and config.vision_n_layers > 0,
|
||||
and config.vision_n_layers > 0
|
||||
and not getattr(config, "language_model_only", False),
|
||||
)
|
||||
|
||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
@@ -3872,7 +3877,15 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
finally:
|
||||
forward_batch.num_token_non_padded = saved_num_token_non_padded
|
||||
if _use_cp and get_moe_a2a_backend().is_none():
|
||||
hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states)
|
||||
if self.config.model_type == "deepseek_v41":
|
||||
parallel = get_parallel()
|
||||
hidden_states = parallel.tp_group.all_reduce(hidden_states)
|
||||
parallel = get_parallel()
|
||||
hidden_states = hidden_states.tensor_split(parallel.attn_cp_size)[
|
||||
parallel.attn_cp_rank
|
||||
].contiguous()
|
||||
else:
|
||||
hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states)
|
||||
elif _use_tp_moe_gather:
|
||||
hidden_states, global_hidden_states = (
|
||||
get_local_dp_buffer(get_parallel().tp_group),
|
||||
@@ -4402,11 +4415,18 @@ class DeepseekV4Model(nn.Module):
|
||||
)
|
||||
if self.engram_hasher is not None:
|
||||
if cp_extend:
|
||||
# n-gram hashing needs each token's predecessors: hash the whole prompt
|
||||
# N-gram hashing needs each token's predecessors, so hash the
|
||||
# whole prompt before selecting this CP rank's interleaved rows.
|
||||
# The hasher builds request-to-token indices dynamically; keep
|
||||
# that work at an eager break during breakable graph capture.
|
||||
total = int(forward_batch.attn_cp_metadata.total_seq_lens)
|
||||
hash_ids = self.engram_hasher(
|
||||
forward_batch.input_ids[:total], forward_batch
|
||||
)
|
||||
global_input_ids = forward_batch.input_ids[:total]
|
||||
if is_in_breakable_cuda_graph():
|
||||
hash_ids = bcg_deepseek_v4_engram_hash_ids(
|
||||
self.engram_hasher, global_input_ids
|
||||
)
|
||||
else:
|
||||
hash_ids = self.engram_hasher(global_input_ids, forward_batch)
|
||||
parallel = get_parallel()
|
||||
hash_ids = hash_ids[parallel.attn_cp_rank :: parallel.attn_cp_size]
|
||||
pad_rows = hidden_states.shape[0] - hash_ids.shape[0]
|
||||
@@ -4839,6 +4859,13 @@ class DeepseekV4Model(nn.Module):
|
||||
return hidden_states, pre_hc_head
|
||||
|
||||
|
||||
def _v41_vision_a2a_supported() -> bool:
|
||||
backend = get_moe_a2a_backend()
|
||||
return backend.is_none() or (
|
||||
backend.is_megamoe() and get_disagg().disaggregation_mode == "decode"
|
||||
)
|
||||
|
||||
|
||||
class DeepseekV4ForCausalLM(nn.Module):
|
||||
supports_cuda_vmm_feature_transport = True
|
||||
|
||||
@@ -4864,14 +4891,19 @@ class DeepseekV4ForCausalLM(nn.Module):
|
||||
self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config)
|
||||
self.determine_num_fused_shared_experts()
|
||||
self.vision = None
|
||||
if config.model_type == "deepseek_v41" and config.vision_n_layers > 0:
|
||||
if (
|
||||
config.model_type == "deepseek_v41"
|
||||
and config.vision_n_layers > 0
|
||||
and not getattr(config, "language_model_only", False)
|
||||
):
|
||||
if (
|
||||
get_parallel().attn_cp_size != 1
|
||||
or get_parallel().pp_group.world_size != 1
|
||||
or not get_moe_a2a_backend().is_none()
|
||||
or get_pp_group().world_size != 1
|
||||
or not _v41_vision_a2a_supported()
|
||||
):
|
||||
raise ValueError(
|
||||
"V4.1 vision currently supports TP/EP/DP without CP, PP or MoE A2A"
|
||||
"V4.1 vision supports TP/EP/DP without CP or PP; "
|
||||
"MoE A2A is supported only with MegaMoE on a PD decode node"
|
||||
)
|
||||
|
||||
args = SimpleNamespace(**vars(config), dim=config.hidden_size)
|
||||
@@ -5078,16 +5110,19 @@ class DeepseekV4ForCausalLM(nn.Module):
|
||||
0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
|
||||
)
|
||||
|
||||
def forward(
|
||||
def prepare_language_model_inputs(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
input_embeds: Optional[torch.Tensor] = None,
|
||||
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||
) -> torch.Tensor:
|
||||
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
|
||||
"""Prepare full-sequence image embeddings and model IDs before CP splits.
|
||||
|
||||
Scheduler hash IDs stay intact for multimodal cache keys; the language
|
||||
model uses image_token_id for Engram masking and visual MoE routing.
|
||||
"""
|
||||
if (
|
||||
self.vision is not None
|
||||
getattr(self, "vision", None) is not None
|
||||
and not forward_batch.forward_mode.is_decode()
|
||||
and not forward_batch.forward_mode.is_target_verify()
|
||||
and forward_batch.mm_inputs is not None
|
||||
@@ -5096,7 +5131,7 @@ class DeepseekV4ForCausalLM(nn.Module):
|
||||
if input_embeds is not None:
|
||||
raise ValueError("Cannot combine input_embeds and image inputs")
|
||||
input_embeds = self._prepare_mm_embeddings(input_ids, forward_batch)
|
||||
if self.vision is not None and not (
|
||||
if getattr(self, "vision", None) is not None and not (
|
||||
forward_batch.forward_mode.is_decode_or_idle()
|
||||
or forward_batch.forward_mode.is_target_verify()
|
||||
):
|
||||
@@ -5106,6 +5141,19 @@ class DeepseekV4ForCausalLM(nn.Module):
|
||||
input_ids >= MM_PAD_SHIFT_VALUE, self.config.image_token_id
|
||||
)
|
||||
|
||||
return input_ids, input_embeds
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
input_embeds: Optional[torch.Tensor] = None,
|
||||
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||
) -> torch.Tensor:
|
||||
input_ids, input_embeds = self.prepare_language_model_inputs(
|
||||
input_ids, forward_batch, input_embeds
|
||||
)
|
||||
with get_attn_tp_context().maybe_input_scattered(forward_batch):
|
||||
hidden_states = self.model.forward(
|
||||
input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
|
||||
|
||||
@@ -376,6 +376,7 @@ class ServerArgs:
|
||||
# ===== END TO BE REFACTORED ====
|
||||
|
||||
LANGUAGE_MODEL_ONLY_ARCHITECTURES = (
|
||||
"DeepseekV4ForCausalLM",
|
||||
"MuseGlimmerForConditionalGeneration",
|
||||
"Cosmos3ForConditionalGeneration",
|
||||
"Cosmos3EdgeForConditionalGeneration",
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Small CPU tensors; production CP slicing/gather, mocked collective transport."""
|
||||
|
||||
from contextlib import ExitStack, contextmanager, nullcontext
|
||||
from types import SimpleNamespace as NS
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.cp.interleave import InterleaveCPStrategy
|
||||
from sglang.srt.layers.cp.padding import pad_logical_token_to_physical
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
||||
|
||||
CP = "sglang.srt.layers.cp"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def cp_context(size, rank, lengths=(3, 6), prefix_lengths=(7, 13)):
|
||||
"""Keep real interleave indexing/padding; replace only runtime context."""
|
||||
strategy = InterleaveCPStrategy(size)
|
||||
parallel = NS(attn_cp_size=size, attn_cp_rank=rank, attn_cp_group=None)
|
||||
batch = NS(
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
input_ids=torch.arange(1, sum(lengths) + 1),
|
||||
positions=torch.cat(
|
||||
[
|
||||
torch.arange(prefix, prefix + length)
|
||||
for prefix, length in zip(prefix_lengths, lengths)
|
||||
]
|
||||
),
|
||||
extend_seq_lens_cpu=list(lengths),
|
||||
extend_prefix_lens_cpu=list(prefix_lengths),
|
||||
mm_inputs=None,
|
||||
spec_info=None,
|
||||
)
|
||||
batch.attn_cp_metadata = strategy.build_metadata(
|
||||
sum(lengths), [p + n for p, n in zip(prefix_lengths, lengths)], list(lengths)
|
||||
)
|
||||
with ExitStack() as stack:
|
||||
for module in ("base", "utils", "padding", "interleave"):
|
||||
stack.enter_context(
|
||||
patch(CP + "." + module + ".get_parallel", return_value=parallel)
|
||||
)
|
||||
stack.enter_context(patch(CP + ".utils.get_cp_strategy", return_value=strategy))
|
||||
stack.enter_context(
|
||||
patch(CP + ".padding.get_cp_padding_align_size", return_value=size)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch(
|
||||
CP + ".utils.get_moe_a2a_backend", return_value=NS(is_none=lambda: True)
|
||||
)
|
||||
)
|
||||
pad_logical_token_to_physical(batch.attn_cp_metadata)
|
||||
yield strategy, batch
|
||||
|
||||
|
||||
@contextmanager
|
||||
def simulated_collective(strategy, batch, global_tensor):
|
||||
"""Inject peer buffers into all-gather; retain production unpadding/reordering."""
|
||||
physical = max(batch.attn_cp_metadata.per_rank_actual_token)
|
||||
buffers = []
|
||||
for rank in range(strategy.cp_size):
|
||||
buf = global_tensor.new_zeros((physical, *global_tensor.shape[1:]))
|
||||
local = global_tensor[rank :: strategy.cp_size]
|
||||
buf[: len(local)] = local
|
||||
buffers.append(buf)
|
||||
|
||||
def gather(output, local):
|
||||
torch.testing.assert_close(local, buffers[strategy.cp_rank], rtol=0, atol=0)
|
||||
output.copy_(torch.cat(buffers))
|
||||
|
||||
with (
|
||||
patch(
|
||||
CP + ".interleave.use_symmetric_memory",
|
||||
side_effect=lambda *a, **k: nullcontext(),
|
||||
),
|
||||
patch(CP + ".interleave.is_allocation_symmetric", return_value=False),
|
||||
patch(CP + ".interleave.attn_cp_all_gather_into_tensor", side_effect=gather),
|
||||
):
|
||||
yield
|
||||
@@ -0,0 +1,151 @@
|
||||
"""Shared compressed-KV workspace validity across layers and forwards."""
|
||||
import sys
|
||||
from types import SimpleNamespace as NS
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention import deepseek_v4_backend as backend
|
||||
from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
|
||||
CompressedGather,
|
||||
SparsePrefillWorkspace,
|
||||
WORKSPACE_DIM,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
def test_shared_compressed_dequant_lifetime():
|
||||
device = "cuda"
|
||||
obj = object.__new__(backend.DeepseekV4AttnBackend)
|
||||
obj.sparse_prefill_workspace = SparsePrefillWorkspace(device)
|
||||
obj.shared_compressed_prefill_workspaces = {
|
||||
ratio: SparsePrefillWorkspace(device) for ratio in (1, 2)
|
||||
}
|
||||
obj.softmax_scale = 0.1
|
||||
obj.head_dim_v = WORKSPACE_DIM
|
||||
sources = {0: 0, 1: 1, 2: 0, 3: 1, 4: 4, 5: 4, 6: 6, 7: 7}
|
||||
ratios = {0: 1, 1: 2, 2: 1, 3: 2, 4: 1, 5: 1, 6: 0, 7: 4}
|
||||
compressed = {
|
||||
source: torch.full(
|
||||
(128, 1, WORKSPACE_DIM),
|
||||
float(source + 1),
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
for source in (0, 1, 4, 7)
|
||||
}
|
||||
swa = torch.empty((16, 1, WORKSPACE_DIM), device=device, dtype=torch.bfloat16)
|
||||
pool = NS(
|
||||
source_layer_of=lambda layer: sources[layer],
|
||||
get_extra_key_page_size=lambda layer: 1,
|
||||
get_extra_key_buffer=lambda layer: compressed[sources[layer]],
|
||||
get_extra_key_layout=lambda layer: None,
|
||||
get_swa_key_buffer_radix=lambda layer: swa,
|
||||
get_swa_key_layout=lambda: None,
|
||||
)
|
||||
calls = []
|
||||
active = [True]
|
||||
|
||||
def dequant(src, indices, *, out, **kwargs):
|
||||
calls.append("swa" if src is swa else sources_by_ptr[src.data_ptr()])
|
||||
out.copy_(src.index_select(0, indices.long()))
|
||||
|
||||
sources_by_ptr = {v.data_ptr(): k for k, v in compressed.items()}
|
||||
|
||||
def make_cache(n):
|
||||
gathers = {
|
||||
ratio: CompressedGather(
|
||||
flat_token_ids=torch.arange(n // ratio, device=device, dtype=torch.int32),
|
||||
compressed_base=torch.zeros(1, device=device, dtype=torch.int32),
|
||||
swa_base=torch.zeros(1, device=device, dtype=torch.int32),
|
||||
)
|
||||
for ratio in (1, 2, 4)
|
||||
}
|
||||
indices = torch.zeros((4, 128), device=device, dtype=torch.int32)
|
||||
lengths = torch.full((4,), 3, device=device, dtype=torch.int32)
|
||||
cache = NS(
|
||||
compressed=gathers,
|
||||
swa_token_ids=torch.arange(3, device=device),
|
||||
swa_page_size=1,
|
||||
c0_combined_indices=indices,
|
||||
c0_combined_lens=lengths,
|
||||
)
|
||||
cache.layer_inputs = lambda ratio, core, page: (
|
||||
gathers[ratio].flat_token_ids,
|
||||
indices,
|
||||
lengths,
|
||||
)
|
||||
return cache
|
||||
|
||||
def forward(layer):
|
||||
return obj._forward_prefill_sparse(
|
||||
torch.empty((4, 1, 1, WORKSPACE_DIM), device=device, dtype=torch.bfloat16),
|
||||
layer,
|
||||
ratios[layer],
|
||||
NS(),
|
||||
pool,
|
||||
NS(),
|
||||
torch.zeros(1, device=device),
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(backend, "is_cp_active", side_effect=lambda _: active[0]),
|
||||
patch.object(backend, "dequantize_k_cache_paged", side_effect=dequant),
|
||||
patch(
|
||||
"sgl_kernel.flash_mla.flash_mla_sparse_fwd",
|
||||
side_effect=lambda **kw: (kw["kv"].clone(), None, None),
|
||||
),
|
||||
):
|
||||
# Same-size replay, then growth and shrink must all read freshly written KV.
|
||||
for step, n in enumerate((8, 8, 40, 4)):
|
||||
obj.forward_metadata = NS(sparse_prefill_cache=make_cache(n))
|
||||
for source, tensor in compressed.items():
|
||||
tensor.fill_(source + 1 + step * 10)
|
||||
for layer, should_dequant in (
|
||||
(0, True),
|
||||
(1, True),
|
||||
(2, False),
|
||||
(3, False),
|
||||
(4, True),
|
||||
(5, False),
|
||||
(2, True),
|
||||
(3, False),
|
||||
(6, False),
|
||||
(7, True),
|
||||
(7, True),
|
||||
):
|
||||
swa.fill_(100 + layer + step)
|
||||
before = len(calls)
|
||||
active[0] = True
|
||||
actual = forward(layer)
|
||||
actual_calls = calls[before:]
|
||||
assert actual_calls.count("swa") == 1
|
||||
assert len(actual_calls) == 1 + int(should_dequant)
|
||||
active[0] = False
|
||||
expected = forward(layer)
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
# Workspace replacement invalidates even an unchanged source identity.
|
||||
obj.shared_compressed_prefill_workspaces[1].get(256 + step * 256)
|
||||
active[0] = True
|
||||
before = len(calls)
|
||||
actual = forward(2)
|
||||
assert calls[before:] == [0, "swa"]
|
||||
active[0] = False
|
||||
torch.testing.assert_close(actual, forward(2), rtol=0, atol=0)
|
||||
|
||||
# Re-executing a producer may mutate the same cache address in place.
|
||||
compressed[0].add_(1)
|
||||
active[0] = True
|
||||
before = len(calls)
|
||||
actual = forward(0)
|
||||
assert calls[before:] == [0, "swa"]
|
||||
active[0] = False
|
||||
torch.testing.assert_close(actual, forward(0), rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, *sys.argv[1:]]))
|
||||
@@ -0,0 +1,286 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from sglang.srt.disaggregation.base.conn import StateType
|
||||
from sglang.srt.disaggregation.common.conn import (
|
||||
CommonKVBootstrapServer,
|
||||
CommonKVManager,
|
||||
)
|
||||
from sglang.srt.disaggregation.decode import DecodePreallocQueue
|
||||
from sglang.srt.disaggregation.utils import get_dsv41_spec_layout
|
||||
from sglang.srt.mem_cache.common import retraction_backup
|
||||
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def make_layout():
|
||||
args = SimpleNamespace(
|
||||
mla_compression_ratios=[0, 2, 1],
|
||||
kv_layer_ids=[1, 2],
|
||||
kv_item_lens=[512, 1024],
|
||||
state_types=[StateType.SWA, StateType.DSV4_REQUEST_STATE, StateType.SWA],
|
||||
state_item_lens=[[512], [32768], [512]],
|
||||
)
|
||||
with get_context().override_server_args(
|
||||
speculative_algorithm="DSPARK", speculative_num_draft_tokens=6
|
||||
):
|
||||
return get_dsv41_spec_layout(args)
|
||||
|
||||
|
||||
class TestDSV41DSparkPD(CustomTestCase):
|
||||
def test_bootstrap_validates_before_caching(self):
|
||||
layout = make_layout()
|
||||
cases = [("matching", layout, layout, 4, True), ("legacy", None, None, 2, True)]
|
||||
for key, value in (
|
||||
("num_draft_tokens", 5),
|
||||
("kv_layer_ids", [2, 1]),
|
||||
("kv_item_lens", [256, 1024]),
|
||||
("state_types", ["swa", "c128_state"]),
|
||||
("state_item_lens", [[512], [8192], [512]]),
|
||||
):
|
||||
different = copy.deepcopy(layout)
|
||||
different[key] = value
|
||||
cases.append((key, layout, different, 4, False))
|
||||
cases += [
|
||||
("prefill_only", None, layout, 4, False),
|
||||
("decode_only_or_old_prefill", layout, None, 4, False),
|
||||
("tp_mismatch", layout, layout, 2, False),
|
||||
]
|
||||
for name, local, peer, tp_size, supported in cases:
|
||||
with self.subTest(name=name):
|
||||
manager = object.__new__(CommonKVManager)
|
||||
manager.prefill_info_table = {}
|
||||
manager.kv_args = SimpleNamespace(page_size=256)
|
||||
manager.kv_cache_dtype_str = "fp8_e4m3"
|
||||
manager.dsv41_spec_layout = local
|
||||
manager.attn_tp_size = 4
|
||||
manager.dcp_size = 1
|
||||
manager._resolve_rank_mapping = Mock()
|
||||
response = Mock(status_code=200)
|
||||
response.json.return_value = dict(
|
||||
attn_tp_size=tp_size,
|
||||
attn_cp_size=1,
|
||||
dp_size=1,
|
||||
pp_size=1,
|
||||
page_size=256,
|
||||
kv_cache_dtype="fp8_e4m3",
|
||||
follow_bootstrap_room=True,
|
||||
dsv41_spec_layout=peer,
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.disaggregation.common.conn.requests.get",
|
||||
return_value=response,
|
||||
) as fetch:
|
||||
if supported:
|
||||
self.assertTrue(
|
||||
manager.try_ensure_parallel_info("prefill:8998")
|
||||
)
|
||||
self.assertTrue(
|
||||
manager.try_ensure_parallel_info("prefill:8998")
|
||||
)
|
||||
fetch.assert_called_once()
|
||||
else:
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError, "DeepSeek-V4.1 DSpark PD"
|
||||
):
|
||||
manager.try_ensure_parallel_info("prefill:8998")
|
||||
self.assertFalse(manager.prefill_info_table)
|
||||
manager._resolve_rank_mapping.assert_not_called()
|
||||
|
||||
def test_python_bootstrap_preserves_layout_and_rejects_mixed_ranks(self):
|
||||
with patch.object(CommonKVBootstrapServer, "run"):
|
||||
server = CommonKVBootstrapServer("127.0.0.1", 8998)
|
||||
layout = make_layout()
|
||||
payload = dict(
|
||||
attn_tp_size=1,
|
||||
attn_tp_rank=0,
|
||||
attn_cp_size=1,
|
||||
attn_cp_rank=0,
|
||||
attn_dp_size=1,
|
||||
attn_dp_rank=0,
|
||||
pp_size=1,
|
||||
pp_rank=0,
|
||||
system_dp_size=1,
|
||||
system_dp_rank=0,
|
||||
rank_ip="127.0.0.1",
|
||||
rank_port=1234,
|
||||
page_size=256,
|
||||
kv_cache_dtype="fp8_e4m3",
|
||||
dsv41_spec_layout=layout,
|
||||
)
|
||||
request = Mock(json=AsyncMock(return_value=payload))
|
||||
self.assertEqual(asyncio.run(server._handle_route_put(request)).status, 200)
|
||||
query = Mock(
|
||||
query={
|
||||
key: "-1"
|
||||
for key in (
|
||||
"prefill_dp_rank",
|
||||
"prefill_cp_rank",
|
||||
"target_tp_rank",
|
||||
"target_pp_rank",
|
||||
)
|
||||
}
|
||||
)
|
||||
response = asyncio.run(server._handle_route_get(query))
|
||||
self.assertEqual(json.loads(response.text)["dsv41_spec_layout"], layout)
|
||||
payload["dsv41_spec_layout"] = None
|
||||
self.assertEqual(asyncio.run(server._handle_route_put(request)).status, 400)
|
||||
self.assertEqual(server._registered_count, 1)
|
||||
self.assertEqual(server.dsv41_spec_layout, layout)
|
||||
|
||||
def test_retraction_recomputes_from_prefill_and_replays_boundary_token(self):
|
||||
pool = object.__new__(DeepSeekV4TokenToKVPool)
|
||||
pool.compression_ratios = [0, 2, 1]
|
||||
pool.device = "cuda"
|
||||
allocator = Mock(get_kvcache=Mock(return_value=pool))
|
||||
for algorithm in (None, "DSPARK"):
|
||||
with (
|
||||
self.subTest(algorithm=algorithm),
|
||||
get_context().override_server_args(speculative_algorithm=algorithm),
|
||||
patch("torch.get_device_module") as device_module,
|
||||
):
|
||||
req = SimpleNamespace(
|
||||
output_ids=[7, 8],
|
||||
bootstrap_host="prefill",
|
||||
time_stats=Mock(),
|
||||
offload_kv_cache=Mock(),
|
||||
)
|
||||
request_pool = Mock()
|
||||
self.assertTrue(
|
||||
retraction_backup(
|
||||
req, Mock(), request_pool, allocator, "cpu_tensor"
|
||||
)
|
||||
)
|
||||
queue = SimpleNamespace(
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
_check_if_req_exceed_kv_capacity=Mock(return_value=False),
|
||||
_create_receiver_and_enqueue=Mock(),
|
||||
_resolve_prefill_dp_rank=Mock(return_value=0),
|
||||
retracted_queue=[],
|
||||
pending_reqs=[],
|
||||
)
|
||||
DecodePreallocQueue.add(queue, req, is_retracted=True)
|
||||
if algorithm == "DSPARK":
|
||||
req.offload_kv_cache.assert_not_called()
|
||||
device_module.return_value.synchronize.assert_called_once_with(
|
||||
"cuda"
|
||||
)
|
||||
self.assertEqual(req.output_ids, [7])
|
||||
self.assertEqual(req.pd_rebootstrap_forced_output_id, 8)
|
||||
self.assertTrue(req.pd_rebootstrap_in_progress)
|
||||
queue._create_receiver_and_enqueue.assert_called_once_with(
|
||||
req, is_rebootstrap=True
|
||||
)
|
||||
self.assertFalse(queue.retracted_queue)
|
||||
else:
|
||||
req.offload_kv_cache.assert_called_once_with(
|
||||
request_pool, allocator
|
||||
)
|
||||
device_module.assert_not_called()
|
||||
self.assertEqual(req.output_ids, [7, 8])
|
||||
self.assertEqual(queue.retracted_queue, [req])
|
||||
|
||||
|
||||
class TestDSV41CPPDHandshake(CustomTestCase):
|
||||
def make(self, rank=0, hybrid=True):
|
||||
m = object.__new__(CommonKVManager)
|
||||
m.prefill_info_table = {}
|
||||
m.kv_args = SimpleNamespace(page_size=256, engine_rank=rank)
|
||||
m.kv_cache_dtype_str = "fp8_e4m3"
|
||||
m.dsv41_spec_layout = {"kv_item_lens": [512], "state_item_lens": [[32768]]}
|
||||
m.attn_tp_size = 4
|
||||
m.attn_cp_size = 1
|
||||
m.attn_cp_rank = 0
|
||||
m.dcp_size = 1
|
||||
m.is_mla_backend = False
|
||||
m.is_hybrid_mla_backend = hybrid
|
||||
m.enable_all_cp_ranks_for_transfer = True
|
||||
m.pp_size = 1
|
||||
m.pp_rank = 0
|
||||
return m
|
||||
|
||||
def fetch(self, m, tp, cp, layout=None):
|
||||
response = Mock(status_code=200)
|
||||
response.json.return_value = dict(
|
||||
attn_tp_size=tp,
|
||||
attn_cp_size=cp,
|
||||
dp_size=1,
|
||||
pp_size=1,
|
||||
page_size=256,
|
||||
kv_cache_dtype="fp8_e4m3",
|
||||
follow_bootstrap_room=True,
|
||||
dsv41_spec_layout=layout or m.dsv41_spec_layout,
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.disaggregation.common.conn.requests.get", return_value=response
|
||||
):
|
||||
return m.try_ensure_parallel_info("prefill:8761")
|
||||
|
||||
def test_cp4_maps_all_shards_to_each_decode_rank(self):
|
||||
for rank in range(4):
|
||||
m = self.make(rank)
|
||||
self.assertTrue(self.fetch(m, 1, 4))
|
||||
info = m.prefill_info_table["prefill:8761"]
|
||||
self.assertEqual(info.target_tp_ranks, [0])
|
||||
self.assertEqual(info.target_cp_ranks, [0, 1, 2, 3])
|
||||
self.assertEqual(info.required_prefill_response_num, 4)
|
||||
self.assertEqual(info.required_dst_info_num, 4)
|
||||
|
||||
def test_dsv4_pool_is_classified_as_mla(self):
|
||||
from sglang.srt.disaggregation.utils import is_mla_backend
|
||||
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
|
||||
|
||||
pool = object.__new__(DeepSeekV4TokenToKVPool)
|
||||
self.assertTrue(is_mla_backend(pool))
|
||||
m = self.make(hybrid=False)
|
||||
m.is_mla_backend = is_mla_backend(pool)
|
||||
self.assertTrue(self.fetch(m, 1, 4))
|
||||
self.assertEqual(
|
||||
m.prefill_info_table["prefill:8761"].required_prefill_response_num, 4
|
||||
)
|
||||
|
||||
def test_cp2_tp2_maps_corresponding_tp_and_both_cp_ranks(self):
|
||||
for rank in range(4):
|
||||
m = self.make(rank)
|
||||
self.assertTrue(self.fetch(m, 2, 2))
|
||||
info = m.prefill_info_table["prefill:8761"]
|
||||
self.assertEqual(info.target_tp_ranks, [rank // 2])
|
||||
self.assertEqual(info.target_cp_ranks, [0, 1])
|
||||
self.assertEqual(info.required_prefill_response_num, 2)
|
||||
|
||||
def test_plain_tp4_unchanged(self):
|
||||
m = self.make(3)
|
||||
self.assertTrue(self.fetch(m, 4, 1))
|
||||
info = m.prefill_info_table["prefill:8761"]
|
||||
self.assertEqual(info.target_tp_ranks, [3])
|
||||
self.assertEqual(info.target_cp_ranks, [0])
|
||||
|
||||
def test_unequal_model_tp_rejected(self):
|
||||
for tp, cp in [(2, 1), (1, 2), (1, 8)]:
|
||||
m = self.make()
|
||||
with self.assertRaisesRegex(RuntimeError, "same TP size"):
|
||||
self.fetch(m, tp, cp)
|
||||
self.assertFalse(m.prefill_info_table)
|
||||
|
||||
def test_nonhybrid_cp_mismatch_rejected(self):
|
||||
m = self.make(hybrid=False)
|
||||
with self.assertRaisesRegex(RuntimeError, "same TP size"):
|
||||
self.fetch(m, 1, 4)
|
||||
|
||||
def test_layout_mismatch_still_rejected(self):
|
||||
m = self.make()
|
||||
with self.assertRaisesRegex(RuntimeError, "layout mismatch"):
|
||||
self.fetch(m, 1, 4, {"kv_item_lens": [1024], "state_item_lens": [[32768]]})
|
||||
self.assertFalse(m.prefill_info_table)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Chunk continuation when fake PD transfer skips shared radix insertion."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace as NS
|
||||
from unittest.mock import Mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.common import maybe_cache_unfinished_req
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestFakeTransferChunkProgress(CustomTestCase):
|
||||
def test_chunk_progress_without_shared_insert(self):
|
||||
slots = torch.arange(16385, dtype=torch.int32).reshape(1, -1)
|
||||
cache = NS(
|
||||
req_to_token_pool=NS(req_to_token=slots), cache_unfinished_req=Mock()
|
||||
)
|
||||
req = NS(
|
||||
skip_radix_cache_insert=True,
|
||||
kv=NS(req_pool_idx=0, cache_protected_len=0),
|
||||
get_fill_ids=lambda: range(16384),
|
||||
)
|
||||
maybe_cache_unfinished_req(req, cache, chunked=True)
|
||||
self.assertEqual(16385 - len(req.prefix_indices), 1)
|
||||
self.assertEqual(req.kv.cache_protected_len, 0)
|
||||
cache.cache_unfinished_req.assert_not_called()
|
||||
self.assertEqual(req.prefix_indices.dtype, torch.int64)
|
||||
slots[0, 0] = -1
|
||||
self.assertEqual(req.prefix_indices[0].item(), 0)
|
||||
req.get_fill_ids = lambda: range(16385)
|
||||
maybe_cache_unfinished_req(req, cache, chunked=True)
|
||||
self.assertEqual(len(req.prefix_indices), 16385)
|
||||
|
||||
def test_real_transfer_preserves_cache_path(self):
|
||||
req = NS(skip_radix_cache_insert=False)
|
||||
cache = NS(cache_unfinished_req=Mock())
|
||||
maybe_cache_unfinished_req(req, cache, chunked=True)
|
||||
cache.cache_unfinished_req.assert_called_once_with(req, chunked=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,316 @@
|
||||
"""V4.1 image/text CP input contracts; vision and model compute are mocked."""
|
||||
|
||||
import unittest
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace as NS
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import MultimodalInputs
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
||||
from sglang.srt.model_executor.runner.eager_runner import EagerRunner
|
||||
from sglang.srt.models.deepseek_v4 import MM_PAD_SHIFT_VALUE, DeepseekV4ForCausalLM
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.dsv41_cp_test_utils import cp_context, simulated_collective
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
MODEL = "sglang.srt.models.deepseek_v4"
|
||||
RUNNER = "sglang.srt.model_executor.runner.eager_runner"
|
||||
IMAGE_ID = 129264
|
||||
|
||||
|
||||
class TestDSV41MultimodalCP(CustomTestCase):
|
||||
def test_image_spans_cross_ranks_before_shard_and_gather(self):
|
||||
# Two image spans with distinct cache hashes; first request is text-only.
|
||||
original = torch.tensor(
|
||||
[
|
||||
7,
|
||||
8,
|
||||
9,
|
||||
MM_PAD_SHIFT_VALUE + 11,
|
||||
MM_PAD_SHIFT_VALUE + 11,
|
||||
10,
|
||||
MM_PAD_SHIFT_VALUE + 23,
|
||||
MM_PAD_SHIFT_VALUE + 23,
|
||||
12,
|
||||
]
|
||||
)
|
||||
normalized = torch.tensor(
|
||||
[7, 8, 9, IMAGE_ID, IMAGE_ID, 10, IMAGE_ID, IMAGE_ID, 12]
|
||||
)
|
||||
full = torch.arange(36, dtype=torch.float32).reshape(9, 4)
|
||||
# Distinct image features expose using text embeddings or wrong row order.
|
||||
full[3:5] += 1000
|
||||
full[6:8] += 2000
|
||||
for size in (2, 4):
|
||||
for rank in range(size):
|
||||
with (
|
||||
self.subTest(size=size, rank=rank),
|
||||
cp_context(size, rank) as (strategy, batch),
|
||||
):
|
||||
batch.input_ids = original.clone()
|
||||
batch.mm_inputs = [None, MultimodalInputs(mm_items=[])]
|
||||
model = NS(
|
||||
vision=object(),
|
||||
config=NS(image_token_id=IMAGE_ID),
|
||||
get_input_embeddings=Mock(
|
||||
side_effect=AssertionError(
|
||||
"Raw image hashes entered text embeddings"
|
||||
)
|
||||
),
|
||||
_prepare_mm_embeddings=Mock(return_value=full),
|
||||
capture_aux_hidden_states=False,
|
||||
pp_group=NS(is_last_rank=True),
|
||||
lm_head=object(),
|
||||
logits_processor=Mock(return_value="ok"),
|
||||
)
|
||||
model.prepare_language_model_inputs = lambda ids, fb, emb: (
|
||||
DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
model, ids, fb, emb
|
||||
)
|
||||
)
|
||||
|
||||
def body(ids, positions, fb, input_embeds):
|
||||
model._prepare_mm_embeddings.assert_called_once_with(
|
||||
batch.input_ids, batch
|
||||
)
|
||||
n = len(normalized[rank::size])
|
||||
torch.testing.assert_close(ids[:n], normalized[rank::size])
|
||||
torch.testing.assert_close(input_embeds[:n], full[rank::size])
|
||||
torch.testing.assert_close(
|
||||
positions[:n], batch.positions[rank::size]
|
||||
)
|
||||
self.assertFalse(
|
||||
(fb.input_ids_global >= MM_PAD_SHIFT_VALUE).any().item()
|
||||
)
|
||||
return input_embeds
|
||||
|
||||
model.model = body
|
||||
with (
|
||||
simulated_collective(strategy, batch, full),
|
||||
patch(RUNNER + ".torch.cuda.current_stream", return_value=None),
|
||||
):
|
||||
result = EagerRunner._execute_extend_cp(
|
||||
NS(model_runner=NS(model=model)), batch, {}
|
||||
)
|
||||
self.assertEqual(result, "ok")
|
||||
torch.testing.assert_close(
|
||||
model.logits_processor.call_args.args[0], normalized
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
model.logits_processor.call_args.args[1], full
|
||||
)
|
||||
torch.testing.assert_close(batch.input_ids, original)
|
||||
self.assertFalse(hasattr(batch, "input_ids_global"))
|
||||
|
||||
def test_chunk_prefix_metadata_and_scheduler_hashes_survive_embedder(self):
|
||||
for prefixes, lengths in (([0, 0], [3, 6]), ([16384, 127], [3, 6])):
|
||||
with self.subTest(prefixes=prefixes):
|
||||
ids = torch.tensor([7, 8, 9] + [MM_PAD_SHIFT_VALUE + 17] * 6)
|
||||
original = ids.clone()
|
||||
image = MultimodalInputs(mm_items=[])
|
||||
batch = NS(
|
||||
mm_inputs=[None, image],
|
||||
extend_prefix_lens_cpu=prefixes,
|
||||
extend_seq_lens_cpu=lengths,
|
||||
)
|
||||
full = torch.arange(27, dtype=torch.float32).reshape(9, 3)
|
||||
embedding = object()
|
||||
model = NS(get_input_embeddings=lambda: embedding)
|
||||
|
||||
def embed(**kwargs):
|
||||
self.assertEqual(kwargs["extend_prefix_lens"], prefixes)
|
||||
self.assertEqual(kwargs["extend_seq_lens"], lengths)
|
||||
self.assertIs(kwargs["mm_inputs_list"][1], image)
|
||||
self.assertEqual(kwargs["mm_inputs_list"][0].mm_items, [])
|
||||
self.assertIs(kwargs["input_embedding"], embedding)
|
||||
self.assertNotEqual(kwargs["input_ids"].data_ptr(), ids.data_ptr())
|
||||
kwargs["input_ids"].zero_()
|
||||
return full, {}
|
||||
|
||||
with patch(MODEL + ".embed_mm_inputs", side_effect=embed) as mocked:
|
||||
result = DeepseekV4ForCausalLM._prepare_mm_embeddings(
|
||||
model, ids, batch
|
||||
)
|
||||
mocked.assert_called_once()
|
||||
self.assertIs(result, full)
|
||||
self.assertIs(batch.mm_input_embeds, full)
|
||||
torch.testing.assert_close(ids, original)
|
||||
|
||||
def test_vision_enabled_text_batch_skips_image_encoder(self):
|
||||
for mm_inputs in (None, [None, None], []):
|
||||
with self.subTest(mm_inputs=mm_inputs):
|
||||
ids = torch.tensor([4, 5, 6])
|
||||
model = NS(
|
||||
vision=object(),
|
||||
config=NS(image_token_id=IMAGE_ID),
|
||||
_prepare_mm_embeddings=Mock(),
|
||||
)
|
||||
batch = NS(forward_mode=ForwardMode.EXTEND, mm_inputs=mm_inputs)
|
||||
result, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
model, ids, batch
|
||||
)
|
||||
torch.testing.assert_close(result, ids)
|
||||
self.assertIsNone(embeds)
|
||||
model._prepare_mm_embeddings.assert_not_called()
|
||||
|
||||
def test_decode_idle_and_verify_preserve_vocab_ids(self):
|
||||
for mode in (ForwardMode.DECODE, ForwardMode.IDLE, ForwardMode.TARGET_VERIFY):
|
||||
with self.subTest(mode=mode):
|
||||
ids = torch.tensor([4, IMAGE_ID, 6])
|
||||
model = NS(
|
||||
vision=object(),
|
||||
config=NS(image_token_id=IMAGE_ID),
|
||||
_prepare_mm_embeddings=Mock(),
|
||||
)
|
||||
batch = NS(forward_mode=mode, mm_inputs=None)
|
||||
result, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
model, ids, batch
|
||||
)
|
||||
self.assertIs(result, ids)
|
||||
self.assertIsNone(embeds)
|
||||
model._prepare_mm_embeddings.assert_not_called()
|
||||
|
||||
def test_image_embedding_failure_does_not_mutate_scheduler_ids(self):
|
||||
ids = torch.tensor([7, MM_PAD_SHIFT_VALUE + 12, 8])
|
||||
original = ids.clone()
|
||||
batch = NS(
|
||||
mm_inputs=[MultimodalInputs(mm_items=[])],
|
||||
extend_prefix_lens_cpu=[0],
|
||||
extend_seq_lens_cpu=[3],
|
||||
)
|
||||
model = NS(get_input_embeddings=lambda: object())
|
||||
|
||||
def fail(**kwargs):
|
||||
kwargs["input_ids"].zero_()
|
||||
raise RuntimeError("vision failure")
|
||||
|
||||
with patch(MODEL + ".embed_mm_inputs", side_effect=fail):
|
||||
with self.assertRaisesRegex(RuntimeError, "vision failure"):
|
||||
DeepseekV4ForCausalLM._prepare_mm_embeddings(model, ids, batch)
|
||||
torch.testing.assert_close(ids, original)
|
||||
self.assertFalse(hasattr(batch, "mm_input_embeds"))
|
||||
|
||||
|
||||
class TestDSV41MultimodalInputs(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.ids = torch.tensor(
|
||||
[7, MM_PAD_SHIFT_VALUE + 12, MM_PAD_SHIFT_VALUE + 12, 9, 10]
|
||||
)
|
||||
self.original = self.ids.clone()
|
||||
self.embeds = torch.arange(15, dtype=torch.float32).reshape(5, 3)
|
||||
self.model = NS(
|
||||
vision=object(),
|
||||
config=NS(image_token_id=129264),
|
||||
_prepare_mm_embeddings=Mock(return_value=self.embeds),
|
||||
)
|
||||
self.batch = NS(
|
||||
input_ids=self.ids,
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
mm_inputs=[MultimodalInputs(mm_items=[])],
|
||||
)
|
||||
|
||||
def test_prepare_global_embeddings_and_normalized_ids(self):
|
||||
ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
self.model, self.ids, self.batch
|
||||
)
|
||||
self.assertEqual(ids.tolist(), [7, 129264, 129264, 9, 10])
|
||||
self.assertIs(embeds, self.embeds)
|
||||
self.model._prepare_mm_embeddings.assert_called_once_with(self.ids, self.batch)
|
||||
self.assertTrue(torch.equal(self.ids, self.original))
|
||||
|
||||
def test_reject_preembedded_images(self):
|
||||
with self.assertRaisesRegex(ValueError, "Cannot combine"):
|
||||
DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
self.model, self.ids, self.batch, self.embeds
|
||||
)
|
||||
|
||||
def test_text_only_model_keeps_existing_embeddings(self):
|
||||
self.model.vision = None
|
||||
ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
self.model, self.ids, self.batch, self.embeds
|
||||
)
|
||||
self.assertIs(ids, self.ids)
|
||||
self.assertIs(embeds, self.embeds)
|
||||
self.model._prepare_mm_embeddings.assert_not_called()
|
||||
|
||||
def test_text_subclass_without_vision_module(self):
|
||||
del self.model.vision
|
||||
ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
self.model, self.ids, self.batch, self.embeds
|
||||
)
|
||||
self.assertIs(ids, self.ids)
|
||||
self.assertIs(embeds, self.embeds)
|
||||
|
||||
def test_decode_keeps_vocabulary_ids(self):
|
||||
self.batch.forward_mode = ForwardMode.DECODE
|
||||
ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
self.model, self.ids, self.batch
|
||||
)
|
||||
self.assertIs(ids, self.ids)
|
||||
self.assertIsNone(embeds)
|
||||
self.model._prepare_mm_embeddings.assert_not_called()
|
||||
|
||||
def test_embedding_does_not_mutate_scheduler_hashes(self):
|
||||
self.batch.extend_prefix_lens_cpu = [0]
|
||||
self.batch.extend_seq_lens_cpu = [5]
|
||||
self.model.get_input_embeddings = lambda: None
|
||||
|
||||
def embed(**kwargs):
|
||||
kwargs["input_ids"].zero_()
|
||||
return (self.embeds, {})
|
||||
|
||||
with patch("sglang.srt.models.deepseek_v4.embed_mm_inputs", side_effect=embed):
|
||||
result = DeepseekV4ForCausalLM._prepare_mm_embeddings(
|
||||
self.model, self.ids, self.batch
|
||||
)
|
||||
self.assertTrue(torch.equal(self.ids, self.original))
|
||||
self.assertIs(result, self.batch.mm_input_embeds)
|
||||
|
||||
def test_cp_runner_prepares_before_sharding_and_uses_model_ids_for_logits(self):
|
||||
normalized = torch.tensor([7, 129264, 129264, 9, 10])
|
||||
self.batch.positions = torch.arange(5)
|
||||
self.model.prepare_language_model_inputs = lambda ids, batch, emb: (
|
||||
DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
self.model, ids, batch, emb
|
||||
)
|
||||
)
|
||||
self.model.get_input_embeddings = Mock(
|
||||
side_effect=AssertionError("Raw hashes must not enter text embedding")
|
||||
)
|
||||
self.model.model = Mock(return_value=self.embeds[1::4])
|
||||
self.model.capture_aux_hidden_states = False
|
||||
self.model.pp_group = NS(is_last_rank=True)
|
||||
self.model.lm_head = object()
|
||||
self.model.logits_processor = Mock(return_value="ok")
|
||||
|
||||
@contextmanager
|
||||
def shard(embeds, positions, batch, ids):
|
||||
self.assertIs(embeds, self.embeds)
|
||||
self.assertTrue(torch.equal(ids, normalized))
|
||||
yield (embeds[1::4], positions[1::4], ids[1::4])
|
||||
|
||||
runner = NS(model_runner=NS(model=self.model))
|
||||
module = "sglang.srt.model_executor.runner.eager_runner"
|
||||
with (
|
||||
patch(module + ".cp_shard_model_inputs", side_effect=shard),
|
||||
patch(module + ".cp_gather_after_forward", return_value=self.embeds),
|
||||
patch(module + ".torch.cuda.current_stream", return_value=None),
|
||||
):
|
||||
result = EagerRunner._execute_extend_cp(runner, self.batch, {})
|
||||
self.assertEqual(result, "ok")
|
||||
args, kwargs = self.model.model.call_args
|
||||
self.assertEqual(args[0].tolist(), [129264])
|
||||
self.assertTrue(torch.equal(kwargs["input_embeds"], self.embeds[1::4]))
|
||||
self.assertTrue(
|
||||
torch.equal(self.model.logits_processor.call_args.args[0], normalized)
|
||||
)
|
||||
self.assertTrue(torch.equal(self.batch.input_ids, self.original))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,219 @@
|
||||
"""Pure-language V4.1 CP input, padding and DSpark state regressions."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace as NS
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.cp.utils import (
|
||||
cp_gather_after_forward,
|
||||
cp_shard_model_inputs,
|
||||
is_cp_active,
|
||||
)
|
||||
from sglang.srt.model_executor.runner.eager_runner import EagerRunner
|
||||
from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.dsv41_cp_test_utils import cp_context, simulated_collective
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
RUNNER = "sglang.srt.model_executor.runner.eager_runner"
|
||||
|
||||
|
||||
class TestDSV41TextCP(CustomTestCase):
|
||||
def test_interleave_roundtrip_mixed_lengths_prefix_and_padding(self):
|
||||
for size in (2, 4):
|
||||
for length in (4, 5, 9, 127, 128, 129):
|
||||
for rank in range(size):
|
||||
with (
|
||||
self.subTest(size=size, length=length, rank=rank),
|
||||
cp_context(size, rank, (1, length - 1), (0, 16384)) as (
|
||||
strategy,
|
||||
batch,
|
||||
),
|
||||
):
|
||||
embeddings = torch.arange(
|
||||
length * 3, dtype=torch.float32
|
||||
).reshape(length, 3)
|
||||
original_ids = batch.input_ids.clone()
|
||||
with cp_shard_model_inputs(
|
||||
embeddings, batch.positions, batch, batch.input_ids
|
||||
) as (local, positions, ids):
|
||||
count = len(embeddings[rank::size])
|
||||
torch.testing.assert_close(
|
||||
local[:count], embeddings[rank::size]
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
positions[:count], batch.positions[rank::size]
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
ids[:count], batch.input_ids[rank::size]
|
||||
)
|
||||
self.assertEqual(
|
||||
torch.count_nonzero(local[count:]).item(), 0
|
||||
)
|
||||
self.assertEqual(torch.count_nonzero(ids[count:]).item(), 0)
|
||||
with simulated_collective(strategy, batch, embeddings):
|
||||
restored = cp_gather_after_forward(local, batch)
|
||||
torch.testing.assert_close(
|
||||
restored, embeddings, rtol=0, atol=0
|
||||
)
|
||||
torch.testing.assert_close(batch.input_ids, original_ids)
|
||||
self.assertFalse(hasattr(batch, "input_ids_global"))
|
||||
|
||||
def test_speculative_state_and_global_ids_restored_on_exception(self):
|
||||
for size in (2, 4):
|
||||
for rank in range(size):
|
||||
for had_global in (False, True):
|
||||
with (
|
||||
self.subTest(size=size, rank=rank, had_global=had_global),
|
||||
cp_context(size, rank) as (_, batch),
|
||||
):
|
||||
full = torch.arange(36, dtype=torch.float32).reshape(9, 4)
|
||||
batch.spec_info = NS(hidden_states=full)
|
||||
previous = object()
|
||||
if had_global:
|
||||
batch.input_ids_global = previous
|
||||
with self.assertRaisesRegex(RuntimeError, "injected"):
|
||||
with cp_shard_model_inputs(
|
||||
full, batch.positions, batch, batch.input_ids
|
||||
):
|
||||
n = len(full[rank::size])
|
||||
torch.testing.assert_close(
|
||||
batch.spec_info.hidden_states[:n], full[rank::size]
|
||||
)
|
||||
# Global MoE IDs are in rank-major order, with padding.
|
||||
physical = sum(
|
||||
batch.attn_cp_metadata.per_rank_actual_token
|
||||
)
|
||||
padded = batch.input_ids.new_zeros(physical)
|
||||
padded[:9] = batch.input_ids
|
||||
expected = torch.cat(
|
||||
[padded[r::size] for r in range(size)]
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
batch.input_ids_global, expected
|
||||
)
|
||||
raise RuntimeError("injected")
|
||||
self.assertIs(batch.spec_info.hidden_states, full)
|
||||
if had_global:
|
||||
self.assertIs(batch.input_ids_global, previous)
|
||||
else:
|
||||
self.assertFalse(hasattr(batch, "input_ids_global"))
|
||||
|
||||
def test_short_prompt_falls_back_from_cp(self):
|
||||
with cp_context(4, 0, (1, 2), (0, 0)) as (_, batch):
|
||||
self.assertFalse(is_cp_active(batch))
|
||||
|
||||
def test_runner_text_embedding_and_preembedded_paths(self):
|
||||
for preembedded in (False, True):
|
||||
for rank in range(4):
|
||||
with (
|
||||
self.subTest(preembedded=preembedded, rank=rank),
|
||||
cp_context(4, rank) as (strategy, batch),
|
||||
):
|
||||
full = torch.arange(27, dtype=torch.float32).reshape(9, 3)
|
||||
embedding = Mock(return_value=full)
|
||||
model = NS(
|
||||
vision=None,
|
||||
_prepare_mm_embeddings=Mock(
|
||||
side_effect=AssertionError("Text must not invoke vision")
|
||||
),
|
||||
get_input_embeddings=Mock(return_value=embedding),
|
||||
pp_group=NS(is_last_rank=True),
|
||||
lm_head=object(),
|
||||
capture_aux_hidden_states=False,
|
||||
logits_processor=Mock(return_value="ok"),
|
||||
)
|
||||
model.prepare_language_model_inputs = lambda ids, fb, emb: (
|
||||
DeepseekV4ForCausalLM.prepare_language_model_inputs(
|
||||
model, ids, fb, emb
|
||||
)
|
||||
)
|
||||
|
||||
def body(ids, positions, fb, input_embeds):
|
||||
n = len(full[rank::4])
|
||||
torch.testing.assert_close(ids[:n], batch.input_ids[rank::4])
|
||||
torch.testing.assert_close(
|
||||
positions[:n], batch.positions[rank::4]
|
||||
)
|
||||
torch.testing.assert_close(input_embeds[:n], full[rank::4])
|
||||
return input_embeds
|
||||
|
||||
model.model = body
|
||||
with (
|
||||
simulated_collective(strategy, batch, full),
|
||||
patch(RUNNER + ".torch.cuda.current_stream", return_value=None),
|
||||
):
|
||||
result = EagerRunner._execute_extend_cp(
|
||||
NS(model_runner=NS(model=model)),
|
||||
batch,
|
||||
{"input_embeds": full} if preembedded else {},
|
||||
)
|
||||
self.assertEqual(result, "ok")
|
||||
if preembedded:
|
||||
model.get_input_embeddings.assert_not_called()
|
||||
else:
|
||||
embedding.assert_called_once_with(batch.input_ids)
|
||||
model._prepare_mm_embeddings.assert_not_called()
|
||||
args = model.logits_processor.call_args.args
|
||||
torch.testing.assert_close(args[0], batch.input_ids)
|
||||
torch.testing.assert_close(args[1], full)
|
||||
|
||||
def test_dspark_aux_tensor_and_list_gathered_without_pre_norm_override(self):
|
||||
for as_list in (False, True):
|
||||
with self.subTest(as_list=as_list), cp_context(4, 2) as (strategy, batch):
|
||||
full = torch.arange(27, dtype=torch.float32).reshape(9, 3)
|
||||
local = strategy.shard_hidden_states(full, batch)
|
||||
aux = [local.clone(), local.clone()] if as_list else local.clone()
|
||||
model = NS(
|
||||
get_input_embeddings=lambda: lambda ids: full,
|
||||
model=Mock(return_value=((local, local.clone()), aux)),
|
||||
capture_aux_hidden_states=True,
|
||||
pp_group=NS(is_last_rank=True),
|
||||
lm_head=object(),
|
||||
logits_processor=Mock(return_value="ok"),
|
||||
)
|
||||
with (
|
||||
simulated_collective(strategy, batch, full),
|
||||
patch(RUNNER + ".torch.cuda.current_stream", return_value=None),
|
||||
):
|
||||
EagerRunner._execute_extend_cp(
|
||||
NS(model_runner=NS(model=model)), batch, {}
|
||||
)
|
||||
args, kwargs = model.logits_processor.call_args
|
||||
torch.testing.assert_close(args[1], full)
|
||||
for tensor in args[4] if as_list else [args[4]]:
|
||||
torch.testing.assert_close(tensor, full)
|
||||
self.assertNotIn("hidden_states_before_norm", kwargs)
|
||||
|
||||
def test_target_hidden_states_before_norm_preserved_without_dspark_aux(self):
|
||||
with cp_context(4, 1) as (strategy, batch):
|
||||
full = torch.arange(27, dtype=torch.float32).reshape(9, 3)
|
||||
local = strategy.shard_hidden_states(full, batch)
|
||||
model = NS(
|
||||
get_input_embeddings=lambda: lambda ids: full,
|
||||
model=Mock(return_value=(local, local.clone())),
|
||||
capture_aux_hidden_states=False,
|
||||
pp_group=NS(is_last_rank=True),
|
||||
lm_head=object(),
|
||||
logits_processor=Mock(return_value="ok"),
|
||||
)
|
||||
with (
|
||||
simulated_collective(strategy, batch, full),
|
||||
patch(RUNNER + ".torch.cuda.current_stream", return_value=None),
|
||||
):
|
||||
EagerRunner._execute_extend_cp(
|
||||
NS(model_runner=NS(model=model)), batch, {}
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
model.logits_processor.call_args.kwargs["hidden_states_before_norm"],
|
||||
full,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,73 @@
|
||||
"""V4.1 language-model-only PD configuration validation."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace as NS
|
||||
from unittest.mock import patch
|
||||
|
||||
from sglang.srt.arg_groups import model_hook
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestDSV41TextOnlyPDPolicy(CustomTestCase):
|
||||
def validate(
|
||||
self,
|
||||
model_type="deepseek_v41",
|
||||
mode="prefill",
|
||||
arch="DeepseekV4ForCausalLM",
|
||||
**flags,
|
||||
):
|
||||
cfg = NS(
|
||||
language_model_only=True,
|
||||
encoder_only=False,
|
||||
language_only=False,
|
||||
enable_prefix_mm_cache=False,
|
||||
enable_broadcast_mm_inputs_process=False,
|
||||
mm_enable_dp_encoder=False,
|
||||
disaggregation_mode=mode,
|
||||
)
|
||||
for name, value in flags.items():
|
||||
setattr(cfg, name, value)
|
||||
model = NS(hf_config=NS(model_type=model_type, architectures=[arch]))
|
||||
args = NS(
|
||||
LANGUAGE_MODEL_ONLY_ARCHITECTURES=ServerArgs.LANGUAGE_MODEL_ONLY_ARCHITECTURES
|
||||
)
|
||||
with (
|
||||
patch.object(model_hook, "resolving_view", return_value=cfg),
|
||||
patch.object(model_hook, "model_config_of", return_value=model),
|
||||
):
|
||||
model_hook.handle_language_model_only(args)
|
||||
|
||||
def test_v41_modes(self):
|
||||
for mode in ("null", "prefill", "decode"):
|
||||
with self.subTest(mode=mode):
|
||||
self.validate(mode=mode)
|
||||
|
||||
def test_other_models_still_reject_pd(self):
|
||||
with self.assertRaisesRegex(ValueError, "incompatible"):
|
||||
self.validate(model_type="cosmos3", arch="Cosmos3ForConditionalGeneration")
|
||||
|
||||
def test_encoder_options_still_rejected(self):
|
||||
for flag in (
|
||||
"encoder_only",
|
||||
"language_only",
|
||||
"enable_prefix_mm_cache",
|
||||
"enable_broadcast_mm_inputs_process",
|
||||
"mm_enable_dp_encoder",
|
||||
):
|
||||
with (
|
||||
self.subTest(flag=flag),
|
||||
self.assertRaisesRegex(ValueError, "cannot be combined"),
|
||||
):
|
||||
self.validate(**{flag: True})
|
||||
|
||||
def test_unknown_arch_rejected(self):
|
||||
with self.assertRaisesRegex(ValueError, "does not support"):
|
||||
self.validate(arch="UnknownArchitecture")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user