Compare commits

...
16 Commits
Author SHA1 Message Date
minke.yu 104218d9ed ci: tag builds with -ci-<UTC datetime> suffix (naming rule: <branch>-<sha9>-<ci|local>-<UTC>)
build-sglang-image / build (push) In progress
2026-09-22 20:37:57 +08:00
minke.yu 92632a60ba ci: drop gha cache backend (runner cache server 404s on this instance)
build-sglang-image / build (push) Successful in 50m51s
2026-09-22 18:49:20 +08:00
minke.yu 4b3b367b63 ci: hardcode Gitea registry (vars context falls back to docker.io on this instance)
build-sglang-image / build (push) Failing after 24m13s
2026-09-22 18:23:35 +08:00
minke.yu a78da9b524 ci: Gitea Actions workflow to build sglang image on push to dsv41-pd
build-sglang-image / build (push) Failing after 27m21s
Overseas node: original docker.io base (lmsysorg/sglang:dev-dsv41) and
default pypi. Tags follow the ymkymx convention <branch>-<sha9>-<datetime>
plus a moving <branch>-latest. Registry defaults to the Gitea instance's
container registry, override with the REGISTRY variable.
2026-09-22 16:59:13 +08:00
minke.yu b081dd3d23 Merge branch 'main' into dsv41-pd 2026-09-22 14:56:35 +08:00
minke.yu c74a4037fb mooncake: probe pointers via cuda.bindings instead of ctypes libcudart
The ctypes-based probing loaded a second libcudart copy into the process,
which corrupted CUDA/runtime state and segfaulted the scheduler processes
shortly after registration (reproducible on prefill ranks, intermittent on
decode). cuda.bindings.runtime.cudaPointerGetAttributes is a properly
typed binding and boots cleanly on both roles.
2026-09-21 12:46:25 +08:00
minke.yu b963295489 mooncake: support host-memory payloads with intra-node NVLink transport
IntraNodeNvlinkTransport can only register/reach device memory. sglang's
register list includes host regions (aux buffers, some state components);
one host region failed the whole registerLocalMemoryBatch and the engine
rolled back every region, leaving the segment descriptor empty so all KV
transfers failed with 'Requested address ... not found'.

When the intra-node NVLink transport is active (MC_INTRANODE_NVLINK etc.):
- register only device-memory regions (probe via cudaPointerGetAttributes)
- route transfer blocks whose local source is host memory over the ordered
  per-endpoint zmq channel (new STATE_DATA message) instead of the engine;
  the receiver validates and writes the block into its buffer, relying on
  the same ordering guarantee as the aux TCP path
- force aux over the TCP socket path automatically
2026-09-21 11:53:26 +08:00
abing fa826e08b1 Reuse shared compressed KV dequantization in DeepSeek V4.1 CP prefill
(cherry picked from commit 44eb378e94610e4ccc6b990bf09c105f8c8f9ee9)
2026-09-20 22:12:47 +08:00
SYChen123 3810f531a8 Support V4.1 decode vision MegaMoE
(cherry picked from commit 98aa79972d6d31f30a3c5bc83b9d7e3009ae638a)
2026-09-20 22:12:43 +08:00
SYChen123 2580c24d1b Support V4.1 DP attention in DP-only DSpark PD
(cherry picked from commit 6ef0ebba730b3fffde4b59b953c5c9b58a35c4ed)
2026-09-20 22:11:48 +08:00
abing 4f22146e51 Avoid host synchronization in DeepSeek V4.1 CP prefill
(cherry picked from commit 510596019fd5782d494c9b5227bf784d46661f6f)
2026-09-20 22:07:30 +08:00
abing 12e3b82e52 run pass cp+megamoe+bcg
(cherry picked from commit 85b8a150ef69c35cbd8c3fe940e41b414b903757)
2026-09-20 22:07:29 +08:00
abing 8305f66fc8 fix bug
(cherry picked from commit 60d31cb0c268c4b733ab290e133e3763b65f090e)
2026-09-20 22:07:24 +08:00
abing 21a4a16b4b update code
(cherry picked from commit 1453150a0bdf2cce425e9e2da83ef16b8e8f2859)
2026-09-20 22:07:24 +08:00
abing fc954b7e08 add test
(cherry picked from commit 094abd5e90af77507acd7c91513f69b5b29b685b)
2026-09-20 22:07:24 +08:00
abing c2059c4fb2 run pass llm cp
(cherry picked from commit 1d85394563d96cf22d5d84f15c0e9043582419e4)
2026-09-20 22:07:23 +08:00
22 changed files with 1848 additions and 84 deletions
+50
View File
@@ -0,0 +1,50 @@
# Gitea Actions 自动构建 sglang 镜像(海外节点,原版源)
# 基底 lmsysorg/sglang:dev-dsv41docker.io),依赖走 pypi.org 默认源。
# 触发:push 到 dsv41-pd 分支。
#
# 前置条件(一次性,在 Gitea 实例上配置):
# 1. 实例已注册 act_runnerGitea 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 自带 registryvars 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}-ci-$(date -u +%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 上会 404cache server 未配),勿加回
+11
View File
@@ -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 含 .gitsetuptools-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",
+4 -2
View File
@@ -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.
"""
+97 -10
View File
@@ -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
+26 -13
View File
@@ -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:
+4
View File
@@ -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(
+65 -17
View File
@@ -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
+1
View File
@@ -376,6 +376,7 @@ class ServerArgs:
# ===== END TO BE REFACTORED ====
LANGUAGE_MODEL_ONLY_ARCHITECTURES = (
"DeepseekV4ForCausalLM",
"MuseGlimmerForConditionalGeneration",
"Cosmos3ForConditionalGeneration",
"Cosmos3EdgeForConditionalGeneration",
+79
View File
@@ -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()