Compare commits
28
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4f0ca95477 | ||
|
|
b91137ab98 | ||
|
|
67d8368a84 | ||
|
|
db7d2cb7db | ||
|
|
f23179ce99 | ||
|
|
8266769b2d | ||
|
|
3b671d6086 | ||
|
|
7925735a3e | ||
|
|
6833498646 | ||
|
|
b48e2cb1eb | ||
|
|
bfeb7cd9b2 | ||
|
|
ddf5207630 | ||
|
|
104218d9ed | ||
|
|
92632a60ba | ||
|
|
4b3b367b63 | ||
|
|
a78da9b524 | ||
|
|
b081dd3d23 | ||
|
|
c74a4037fb | ||
|
|
b963295489 | ||
|
|
fa826e08b1 | ||
|
|
3810f531a8 | ||
|
|
2580c24d1b | ||
|
|
4f22146e51 | ||
|
|
12e3b82e52 | ||
|
|
8305f66fc8 | ||
|
|
21a4a16b4b | ||
|
|
fc954b7e08 | ||
|
|
c2059c4fb2 |
@@ -0,0 +1,57 @@
|
|||||||
|
# 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, dsv41-pd-visioncp]
|
||||||
|
# compose/部署配置 与 workflow 自身的改动不触发镜像构建(省 runner 与推送带宽)。
|
||||||
|
# 改了 workflow 想重建镜像时,需伴随任意代码改动或手动触发。
|
||||||
|
paths-ignore:
|
||||||
|
- 'deploy/**'
|
||||||
|
- '.gitea/**'
|
||||||
|
|
||||||
|
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}-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
|
||||||
|
# cache-from: type=registry
|
||||||
|
# cache-to: type=registry,mode=max
|
||||||
|
tags: |
|
||||||
|
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ steps.meta.outputs.tag }}
|
||||||
|
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ github.ref_name }}-latest
|
||||||
|
|
||||||
@@ -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
|
||||||
@@ -0,0 +1,21 @@
|
|||||||
|
# 优化版(2026-09-24):解决清华源限流下每次构建重下 torch 的问题。
|
||||||
|
# 要点:
|
||||||
|
# 1. pyproject build-system.requires 含 torch==2.13.0,pip 隔离构建环境每次都重下 ~6GB。
|
||||||
|
# 基底镜像已装 torch 2.13 → 用 --no-build-isolation 复用,构建环境零下载。
|
||||||
|
# 2. 基底缺 setuptools-scm(版本打戳要用),单层预装 + pip cache mount,只下载一次。
|
||||||
|
# 3. RUN 的 pip cache mount 持久化在宿主机,依赖解析命中的包不再重复下载。
|
||||||
|
FROM uhub.service.ucloud.cn/umirror/sglang:dev-dsv41
|
||||||
|
|
||||||
|
ARG PIP_INDEX=https://pypi.tuna.tsinghua.edu.cn/simple
|
||||||
|
ENV PIP_INDEX_URL=${PIP_INDEX}
|
||||||
|
|
||||||
|
# 构建后端(pyproject build-system.requires,torch 除外——基底已有)
|
||||||
|
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||||
|
pip install "setuptools>=61" "setuptools-scm>=8" "setuptools-rust>=1.11" wheel
|
||||||
|
|
||||||
|
# 整个源码树(含 .git,用于 setuptools-scm 版本打戳)
|
||||||
|
COPY sglang/ /sgl-workspace/sglang/
|
||||||
|
|
||||||
|
# editable 安装,全量依赖解析(pyproject.toml 变化时走这个)
|
||||||
|
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||||
|
pip install --no-build-isolation -e /sgl-workspace/sglang/python
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
# fast 路径:python/pyproject.toml 未变化时使用(build.sh 自动判断)。
|
||||||
|
# --no-deps 跳过全部依赖解析(连 index 元数据请求都省掉),纯源码迭代秒级完成。
|
||||||
|
# 依赖有变化时必须走 Dockerfile 全量路径(build.sh 按 pyproject sha256 自动切换)。
|
||||||
|
FROM uhub.service.ucloud.cn/umirror/sglang:dev-dsv41
|
||||||
|
|
||||||
|
ARG PIP_INDEX=https://pypi.tuna.tsinghua.edu.cn/simple
|
||||||
|
ENV PIP_INDEX_URL=${PIP_INDEX}
|
||||||
|
|
||||||
|
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||||
|
pip install "setuptools>=61" "setuptools-scm>=8" "setuptools-rust>=1.11" wheel
|
||||||
|
|
||||||
|
COPY sglang/ /sgl-workspace/sglang/
|
||||||
|
|
||||||
|
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||||
|
pip install --no-deps --no-build-isolation -e /sgl-workspace/sglang/python
|
||||||
@@ -0,0 +1,33 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
# 用法: /data/ymk/build/build.sh [tag]
|
||||||
|
# 默认 tag: ymkymx/sglang:<分支>-<sha9>-local-<UTC构建日期时间>(命名规则见 AGENTS.md)
|
||||||
|
# 另打 <分支>-latest-local 别名。
|
||||||
|
#
|
||||||
|
# 依赖路径自动选择:
|
||||||
|
# python/pyproject.toml 的 sha256 与上次成功构建一致 → Dockerfile.fast(--no-deps,秒级)
|
||||||
|
# 不一致(或 FULL=1)→ Dockerfile 全量依赖解析,成功后记录新 hash
|
||||||
|
set -euo pipefail
|
||||||
|
cd /data/ymk
|
||||||
|
SHA=$(git -C sglang rev-parse --short=9 HEAD)
|
||||||
|
BR=$(git -C sglang rev-parse --abbrev-ref HEAD)
|
||||||
|
DT=$(date -u +%Y%m%d-%H%M)
|
||||||
|
TAG=${1:-ymkymx/sglang:$BR-$SHA-local-$DT}
|
||||||
|
|
||||||
|
HASH_FILE=build/.last_pyproject_sha256
|
||||||
|
CUR=$(sha256sum sglang/python/pyproject.toml | cut -d' ' -f1)
|
||||||
|
PREV=$(cat "$HASH_FILE" 2>/dev/null || echo none)
|
||||||
|
|
||||||
|
DOCKERFILE=build/Dockerfile
|
||||||
|
if [ "$CUR" = "$PREV" ] && [ "${FULL:-0}" != "1" ]; then
|
||||||
|
DOCKERFILE=build/Dockerfile.fast
|
||||||
|
echo "== pyproject 未变,走 fast 路径(--no-deps)"
|
||||||
|
else
|
||||||
|
echo "== pyproject 有变化或 FULL=1,走全量依赖解析"
|
||||||
|
fi
|
||||||
|
|
||||||
|
echo "== building $TAG from sglang@$(git -C sglang log --oneline -1) [$DOCKERFILE]"
|
||||||
|
docker build -f "$DOCKERFILE" -t "$TAG" .
|
||||||
|
# 构建成功才记录 hash / 打别名
|
||||||
|
echo "$CUR" > "$HASH_FILE"
|
||||||
|
docker tag "$TAG" "ymkymx/sglang:$BR-latest-local"
|
||||||
|
echo "== done: $TAG (别名 $BR-latest-local)"
|
||||||
@@ -0,0 +1,33 @@
|
|||||||
|
# B300 ds41 部署 compose 档案
|
||||||
|
|
||||||
|
b300-01 上 DeepSeek-V4.1-Flash 各部署方案的 docker compose 存档。
|
||||||
|
**本目录是 source of truth**;b300 上的运行目录 `/data/ymk/ds41/` 是工作副本。
|
||||||
|
|
||||||
|
方案说明、参数矩阵、特殊 env、已知坑见飞书部署手册(组内共享):
|
||||||
|
《B300 sglang 部署手册》 https://u04wb5irxz.feishu.cn/docx/FOi8dbnybodHr0xIfkPcQ6gtnRx
|
||||||
|
|
||||||
|
## 同步流程
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# 本机(D:\B300\sglang)
|
||||||
|
git add deploy/ && git commit -m "deploy: ..."
|
||||||
|
git push origin dsv41-pd-visioncp # Gitea(paths-ignore,不触发镜像构建)
|
||||||
|
git push b300 dsv41-pd-visioncp # b300 裸仓库
|
||||||
|
|
||||||
|
# b300-01
|
||||||
|
git -C /data/ymk/sglang pull # deploy/ 出现在 /data/ymk/sglang/deploy/b300-ds41/
|
||||||
|
```
|
||||||
|
|
||||||
|
## CI 说明
|
||||||
|
|
||||||
|
`.gitea/workflows/build-image.yaml` 的 push 触发器带 `paths-ignore: ['deploy/**', '.gitea/**']`,
|
||||||
|
改本目录不会触发镜像构建。注意:workflow 自身的改动也不再自动触发,需要重建镜像时
|
||||||
|
伴随代码改动 push 或手动触发。
|
||||||
|
|
||||||
|
## 目录规则
|
||||||
|
|
||||||
|
- 命名:`dockerserve-<拓扑>-<特性>-<卡位>.yml`(如 `-b4` = 后 4 卡)
|
||||||
|
- 退役文件不要留在主目录:移入 `archive-YYYYMMDD/`(b300 侧同样执行)
|
||||||
|
- 禁止提交 `.bak` 备份文件
|
||||||
|
- `mooncake-store.json` 是 hicache L3 mooncake store 配置,被
|
||||||
|
`dockerserve-pd2-p-r1/r2.yml` 通过 `SGLANG_HICACHE_MOONCAKE_CONFIG_PATH` 引用
|
||||||
@@ -0,0 +1,171 @@
|
|||||||
|
# 3x cp2-P + 1x dp2-D PD 分离 + L3 mooncake 互联缓存 + engram host table
|
||||||
|
# 2026-09-23 实验;基线见 archive-20260923/ 与 D:\B300\experiments\3cp2-pd\baseline-20260923.md
|
||||||
|
# GPU: P1=0,1 P2=2,3 P3=4,5 D=6,7
|
||||||
|
# 端口: P1=30000 P2=30010 P3=30020 D=30001 router=30002
|
||||||
|
# bootstrap: P1=8998 P2=8999 P3=9000 D=9001
|
||||||
|
x-sglang-common: &sglang-common
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
|
||||||
|
x-p-environment: &p-environment
|
||||||
|
SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE: "1"
|
||||||
|
SGLANG_RAGGED_VERIFY_MODE: static
|
||||||
|
MC_INTRANODE_NVLINK: "true"
|
||||||
|
MC_INTRA_NVLINK: "true"
|
||||||
|
SGLANG_MOONCAKE_SEND_AUX_TCP: "1"
|
||||||
|
SGLANG_HICACHE_MOONCAKE_CONFIG_PATH: /data/ymk/ds41/mooncake-store.json
|
||||||
|
SGLANG_DISAGGREGATION_QUEUE_SIZE: "16"
|
||||||
|
SGLANG_DISAGGREGATION_THREAD_POOL_SIZE: "32"
|
||||||
|
|
||||||
|
services:
|
||||||
|
p1:
|
||||||
|
<<: *sglang-common
|
||||||
|
container_name: ds41-cp2-p1
|
||||||
|
environment:
|
||||||
|
<<: *p-environment
|
||||||
|
CUDA_VISIBLE_DEVICES: "0,1"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30000
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
|
--disaggregation-bootstrap-port 8998
|
||||||
|
|
||||||
|
p2:
|
||||||
|
<<: *sglang-common
|
||||||
|
container_name: ds41-cp2-p2
|
||||||
|
environment:
|
||||||
|
<<: *p-environment
|
||||||
|
CUDA_VISIBLE_DEVICES: "2,3"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30010
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
|
--disaggregation-bootstrap-port 8999
|
||||||
|
|
||||||
|
p3:
|
||||||
|
<<: *sglang-common
|
||||||
|
container_name: ds41-cp2-p3
|
||||||
|
environment:
|
||||||
|
<<: *p-environment
|
||||||
|
CUDA_VISIBLE_DEVICES: "4,5"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30020
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
|
--disaggregation-bootstrap-port 9000
|
||||||
|
|
||||||
|
d:
|
||||||
|
<<: *sglang-common
|
||||||
|
container_name: ds41-cp2-d
|
||||||
|
environment:
|
||||||
|
SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE: "1"
|
||||||
|
SGLANG_RAGGED_VERIFY_MODE: static
|
||||||
|
MC_INTRANODE_NVLINK: "true"
|
||||||
|
MC_INTRA_NVLINK: "true"
|
||||||
|
SGLANG_MOONCAKE_SEND_AUX_TCP: "1"
|
||||||
|
SGLANG_DISAGGREGATION_QUEUE_SIZE: "16"
|
||||||
|
SGLANG_DISAGGREGATION_THREAD_POOL_SIZE: "32"
|
||||||
|
CUDA_VISIBLE_DEVICES: "6,7"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2 --dp-size 2
|
||||||
|
--enable-dp-attention --enable-dp-lm-head
|
||||||
|
--mem-fraction-static 0.80
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30001
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--disaggregation-mode decode
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
|
--disaggregation-bootstrap-port 9001
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:main-ee5fcdf0d-20260920-0142
|
||||||
|
container_name: ds41-sglang-b4
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30001:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||||
|
- PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.60
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--dp-size 4
|
||||||
|
--enable-dp-attention
|
||||||
|
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:main-ee5fcdf0d-20260920-0142
|
||||||
|
container_name: ds41-sglang-b4
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30001:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||||
|
- PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.60
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--enable-decoder-swa-bounded-replay
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
|
||||||
|
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,42 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:main-ee5fcdf0d-20260920-0142
|
||||||
|
container_name: ds41-sglang-b4
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30001:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--enable-decoder-swa-bounded-replay
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,160 @@
|
|||||||
|
# 4×cp2 独立实例(非 PD)+ hicache L3(mooncake) + dspark + engram host table 卸载(2026-09-23 晚)
|
||||||
|
# 用户指定组合:4 cp2 hicache l3 dspark roundrobin loadbalance engram offload
|
||||||
|
# 每实例:tp2 ep2 + interleave CP2 + dspark b5 + hicache L3 write_through(size 0) + engram host table
|
||||||
|
# (tp2 权重必须靠 engram 卸载才放得下,见 goals-track G1.9 反转)
|
||||||
|
# GPU: a=0,1 b=2,3 c=4,5 d=6,7;端口: 30000/30010/30020/30030;rr router=30002
|
||||||
|
# 参考:dockerserve-cpdspark-l3.yml(hicache 参数)、dockerserve-3cp2-pd.yml(cp2/engram 参数)
|
||||||
|
x-common: &common
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
# 热补丁(2026-09-23):CP 对称 idle 检查,修 health-check × hicache drain 死锁
|
||||||
|
# 源文件:/data/ymk/sglang(dsv41-pd-visioncp 分支工作区已改);补丁脚本见
|
||||||
|
# D:\B300\experiments\cp2x4\patch_healthcheck_cp_idle.py
|
||||||
|
- /data/ymk/sglang/python/sglang/srt/managers/scheduler.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py
|
||||||
|
environment: &env
|
||||||
|
SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE: "1"
|
||||||
|
SGLANG_RAGGED_VERIFY_MODE: static
|
||||||
|
MC_MS_AUTO_DISC: "0"
|
||||||
|
MOONCAKE_MASTER: 127.0.0.1:50051
|
||||||
|
MOONCAKE_TE_META_DATA_SERVER: P2PHANDSHAKE
|
||||||
|
MOONCAKE_PROTOCOL: tcp
|
||||||
|
MOONCAKE_GLOBAL_SEGMENT_SIZE: 300gb
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
|
||||||
|
services:
|
||||||
|
a:
|
||||||
|
<<: *common
|
||||||
|
container_name: ds41-cp2-a
|
||||||
|
environment:
|
||||||
|
<<: *env
|
||||||
|
CUDA_VISIBLE_DEVICES: "0,1"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30000
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
b:
|
||||||
|
<<: *common
|
||||||
|
container_name: ds41-cp2-b
|
||||||
|
environment:
|
||||||
|
<<: *env
|
||||||
|
CUDA_VISIBLE_DEVICES: "2,3"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30010
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
c:
|
||||||
|
<<: *common
|
||||||
|
container_name: ds41-cp2-c
|
||||||
|
environment:
|
||||||
|
<<: *env
|
||||||
|
CUDA_VISIBLE_DEVICES: "4,5"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30020
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
d:
|
||||||
|
<<: *common
|
||||||
|
container_name: ds41-cp2-d2
|
||||||
|
environment:
|
||||||
|
<<: *env
|
||||||
|
CUDA_VISIBLE_DEVICES: "6,7"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30030
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
# cp8 单机非 PD 测试(2026-09-23 晚)
|
||||||
|
# 背景:dp8 单机 dspark 健康但长上下文 prefill 结构性慢(单请求只落 1 rank),
|
||||||
|
# 换 cp8(tp8 + interleave prefill CP8,单请求 prefill 切到 8 卡)对照。
|
||||||
|
# 关 engram、无 hicache L3、无 PD。cp+dspark 无 replay 是已知健康组合
|
||||||
|
# (cp4+dspark+replay 三元组必崩,本配置不开 replay)。
|
||||||
|
# GPU 0-7,端口 30000,容器 ds41-cp8,project cp8
|
||||||
|
services:
|
||||||
|
cp8:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-cp8
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
environment:
|
||||||
|
SGLANG_RAGGED_VERIFY_MODE: "static"
|
||||||
|
CUDA_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7"
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 8 --ep-size 8
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--mem-fraction-static 0.80
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30000
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-ddf520763-local-20260923-0536
|
||||||
|
container_name: ds41-cpdspark-l3-b
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||||
|
- MC_MS_AUTO_DISC=0
|
||||||
|
- MOONCAKE_MASTER=127.0.0.1:50051
|
||||||
|
- MOONCAKE_TE_META_DATA_SERVER=P2PHANDSHAKE
|
||||||
|
- MOONCAKE_PROTOCOL=tcp
|
||||||
|
- MOONCAKE_GLOBAL_SEGMENT_SIZE=400gb
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4 --ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30001
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,51 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-ddf520763-local-20260923-0536
|
||||||
|
container_name: ds41-cpdspark-l3-b
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||||
|
- MC_MS_AUTO_DISC=0
|
||||||
|
- MOONCAKE_MASTER=127.0.0.1:50051
|
||||||
|
- MOONCAKE_TE_META_DATA_SERVER=P2PHANDSHAKE
|
||||||
|
- MOONCAKE_PROTOCOL=tcp
|
||||||
|
- MOONCAKE_GLOBAL_SEGMENT_SIZE=400gb
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4 --ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30001
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-ddf520763-local-20260923-0536
|
||||||
|
container_name: ds41-cpdspark-l3-a
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
- MC_MS_AUTO_DISC=0
|
||||||
|
- MOONCAKE_MASTER=127.0.0.1:50051
|
||||||
|
- MOONCAKE_TE_META_DATA_SERVER=P2PHANDSHAKE
|
||||||
|
- MOONCAKE_PROTOCOL=tcp
|
||||||
|
- MOONCAKE_GLOBAL_SEGMENT_SIZE=400gb
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4 --ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30000
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,51 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-ddf520763-local-20260923-0536
|
||||||
|
container_name: ds41-cpdspark-l3-a
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
- MC_MS_AUTO_DISC=0
|
||||||
|
- MOONCAKE_MASTER=127.0.0.1:50051
|
||||||
|
- MOONCAKE_TE_META_DATA_SERVER=P2PHANDSHAKE
|
||||||
|
- MOONCAKE_PROTOCOL=tcp
|
||||||
|
- MOONCAKE_GLOBAL_SEGMENT_SIZE=400gb
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4 --ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30000
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-prefill-cp --cp-strategy interleave
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-92632a60b-20260922-2100
|
||||||
|
container_name: ds41-cpdspark-nr-b4
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30001:30001"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-92632a60b-20260922-2100
|
||||||
|
container_name: ds41-cpdspark-nr
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30000:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:main-ee5fcdf0d-20260920-0142
|
||||||
|
container_name: ds41-cpdspark
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30000:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--enable-decoder-swa-bounded-replay
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,62 @@
|
|||||||
|
# dp2 对照(2026-09-24 上午):验证 dp2(dp-attention)下 dspark 是否正常。
|
||||||
|
# 与 dockerserve-tp2.yml 逐参数对齐,唯一差异 = 加 --dp-size 2 --enable-dp-attention --enable-dp-lm-head。
|
||||||
|
# GPU 6-7,端口 30030 直连。镜像/补丁挂载/L3/engram/tok8 全部相同。
|
||||||
|
# 注:dp+dspark 不能开 --enable-decoder-swa-bounded-replay(本配置本来就没开)。
|
||||||
|
x-common: &common
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
- /data/ymk/sglang/python/sglang/srt/managers/scheduler.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py
|
||||||
|
environment: &env
|
||||||
|
SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE: "1"
|
||||||
|
SGLANG_RAGGED_VERIFY_MODE: static
|
||||||
|
MC_MS_AUTO_DISC: "0"
|
||||||
|
MOONCAKE_MASTER: 127.0.0.1:50051
|
||||||
|
MOONCAKE_TE_META_DATA_SERVER: P2PHANDSHAKE
|
||||||
|
MOONCAKE_PROTOCOL: tcp
|
||||||
|
MOONCAKE_GLOBAL_SEGMENT_SIZE: 300gb
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
|
||||||
|
services:
|
||||||
|
dp2:
|
||||||
|
<<: *common
|
||||||
|
container_name: ds41-dp2
|
||||||
|
environment:
|
||||||
|
<<: *env
|
||||||
|
CUDA_VISIBLE_DEVICES: "6,7"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--dp-size 2 --enable-dp-attention --enable-dp-lm-head
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30030
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
@@ -0,0 +1,42 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:main-ee5fcdf0d-20260920-0142
|
||||||
|
container_name: ds41-dp4test
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30000:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--enable-decoder-swa-bounded-replay
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
# dp8+tp8 单机非 PD 对照测试(2026-09-23 晚)
|
||||||
|
# 背景:3cp2+dp2 PD 部署 dspark accept 回归(32k ctx accept len 3.14→1.31,
|
||||||
|
# 见 experiments/3cp2-pd/bench-20260923.md「未决问题」)。本配置退回单机验证
|
||||||
|
# dspark 健康度:engram host table 关闭、无 hicache L3、无 PD、无 mooncake。
|
||||||
|
# 参数基准:dockerserve-3cp2-pd.yml 的 d 服务(dp 路径),去掉 PD/L3/engram 相关。
|
||||||
|
# GPU 0-7,端口 30000,容器 ds41-dp8,project dp8
|
||||||
|
services:
|
||||||
|
dp8:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-dp8
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
environment:
|
||||||
|
SGLANG_RAGGED_VERIFY_MODE: "static"
|
||||||
|
CUDA_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7"
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 8 --ep-size 8 --dp-size 8
|
||||||
|
--enable-dp-attention --enable-dp-lm-head
|
||||||
|
--mem-fraction-static 0.80
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30000
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
@@ -0,0 +1,47 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:main-ee5fcdf0d-20260920-0142
|
||||||
|
container_name: ds41-prefill
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30000:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-decoder-swa-bounded-replay
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2.5
|
||||||
|
--hicache-write-policy write_back
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--optimistic-prefill-attempts 4
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,51 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-pd-decode
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--dp-size 4
|
||||||
|
--enable-dp-attention
|
||||||
|
--enable-dp-lm-head
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30001
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--disaggregation-mode decode
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-ddf520763-local-20260923-0536
|
||||||
|
container_name: ds41-pd-decode
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--dp-size 4
|
||||||
|
--enable-dp-attention
|
||||||
|
--enable-dp-lm-head
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30001
|
||||||
|
--enable-cache-report
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--disaggregation-mode decode
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,55 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-c74a4037f-20260921-1300
|
||||||
|
container_name: ds41-pd-prefill-dp
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--dp-size 4
|
||||||
|
--enable-dp-attention
|
||||||
|
--enable-dp-lm-head
|
||||||
|
--load-balance-method total_tokens
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30010
|
||||||
|
--enable-cache-report
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2.5
|
||||||
|
--hicache-write-policy write_back
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-bootstrap-port 8918
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,53 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-pd-prefill
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2.5
|
||||||
|
--hicache-write-policy write_back
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
@@ -0,0 +1,52 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-ddf520763-local-20260923-0536
|
||||||
|
container_name: ds41-pd-prefill
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2.5
|
||||||
|
--hicache-write-policy write_back
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,57 @@
|
|||||||
|
# PD2 二分 Round 2 D 侧(2026-09-24):干净 D + 嫌疑项②
|
||||||
|
# SGLANG_DISAGGREGATION_QUEUE_SIZE=16 + SGLANG_DISAGGREGATION_THREAD_POOL_SIZE=32
|
||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-pd2-decode
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=6,7
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
- SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1
|
||||||
|
- SGLANG_DISAGGREGATION_QUEUE_SIZE=16
|
||||||
|
- SGLANG_DISAGGREGATION_THREAD_POOL_SIZE=32
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
- /data/ymk/sglang/python/sglang/srt/managers/scheduler.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2
|
||||||
|
--ep-size 2
|
||||||
|
--dp-size 2
|
||||||
|
--enable-dp-attention
|
||||||
|
--enable-dp-lm-head
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30030
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--disaggregation-mode decode
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
@@ -0,0 +1,54 @@
|
|||||||
|
# PD 干净对照实验 D 节点(2026-09-24 上午):dp2 decode,配 dockerserve-pd2-p.yml 使用。
|
||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-pd2-decode
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=6,7
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
- SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
- /data/ymk/sglang/python/sglang/srt/managers/scheduler.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2
|
||||||
|
--ep-size 2
|
||||||
|
--dp-size 2
|
||||||
|
--enable-dp-attention
|
||||||
|
--enable-dp-lm-head
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30030
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--disaggregation-mode decode
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
@@ -0,0 +1,67 @@
|
|||||||
|
# PD2 二分 Round 1(2026-09-24):干净基线 + 嫌疑项①「P 侧 L3 hicache 块 + mooncake-store.json」
|
||||||
|
# 相对 dockerserve-pd2-p.yml 的差异:
|
||||||
|
# env 加 SGLANG_HICACHE_MOONCAKE_CONFIG_PATH=/data/ymk/ds41/mooncake-store.json
|
||||||
|
# hicache 参数块换成坏部署同款:page_first_direct + direct + write_through + mooncake +
|
||||||
|
# wait_complete + size 0(替代 write_back L2)
|
||||||
|
# D 侧不变(坏部署 D 本无 L3)。探针:bs16 low_entropy decode 看 accept len。
|
||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-pd2-prefill
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
- SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1
|
||||||
|
- SGLANG_HICACHE_MOONCAKE_CONFIG_PATH=/data/ymk/ds41/mooncake-store.json
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
- /data/ymk/sglang/python/sglang/srt/managers/scheduler.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2
|
||||||
|
--ep-size 2
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30020
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
@@ -0,0 +1,65 @@
|
|||||||
|
# PD2 二分 Round 2 P 侧(2026-09-24):R1(L3 块)之上再加嫌疑项②
|
||||||
|
# SGLANG_DISAGGREGATION_QUEUE_SIZE=16 + SGLANG_DISAGGREGATION_THREAD_POOL_SIZE=32
|
||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-pd2-prefill
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
- SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1
|
||||||
|
- SGLANG_HICACHE_MOONCAKE_CONFIG_PATH=/data/ymk/ds41/mooncake-store.json
|
||||||
|
- SGLANG_DISAGGREGATION_QUEUE_SIZE=16
|
||||||
|
- SGLANG_DISAGGREGATION_THREAD_POOL_SIZE=32
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
- /data/ymk/sglang/python/sglang/srt/managers/scheduler.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2
|
||||||
|
--ep-size 2
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30020
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
@@ -0,0 +1,59 @@
|
|||||||
|
# PD 干净对照实验(2026-09-24 上午):验证「cp2-P → dp2-D」传输路径下 dspark 是否正常。
|
||||||
|
# 用户疑问:单机六形态健康不等于 PD 链路健康,可能两边传输没对齐。
|
||||||
|
# 以已验证可用的 dockerserve-pd-{p,d}-vision.yml 为底,缩到 2+2 卡;P 加 engram host table
|
||||||
|
# (2 卡权重放不下,见 G1.9)。P=GPU4-5 端口 30020,D=GPU6-7 端口 30030,mini_lb=30004。
|
||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
container_name: ds41-pd2-prefill
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
pid: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=4,5
|
||||||
|
- SGLANG_RAGGED_VERIFY_MODE=static
|
||||||
|
- MC_INTRANODE_NVLINK=true
|
||||||
|
- MC_INTRA_NVLINK=true
|
||||||
|
- SGLANG_MOONCAKE_SEND_AUX_TCP=1
|
||||||
|
- SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
- /data/ymk/sglang/python/sglang/srt/managers/scheduler.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2
|
||||||
|
--ep-size 2
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30020
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2.5
|
||||||
|
--hicache-write-policy write_back
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--disaggregation-mode prefill
|
||||||
|
--disaggregation-transfer-backend mooncake
|
||||||
@@ -0,0 +1,42 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:main-ee5fcdf0d-20260920-0142
|
||||||
|
container_name: ds41-replay-off
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30000:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,42 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: ymkymx/sglang:main-ee5fcdf0d-20260920-0142
|
||||||
|
container_name: ds41-replay-on
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
ports:
|
||||||
|
- "30000:30000"
|
||||||
|
environment:
|
||||||
|
- CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--speculative-algorithm DSPARK
|
||||||
|
--speculative-dspark-block-size 5
|
||||||
|
--enable-decoder-swa-bounded-replay
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
@@ -0,0 +1,92 @@
|
|||||||
|
# 纯 tp2 单实例对照(2026-09-24 早):与 dockerserve-cp2x4.yml 逐参数对齐,唯一差异 = 去掉
|
||||||
|
# --enable-prefill-cp --cp-strategy interleave(即无 CP),用于测「tp2 vs cp2」的 prefill 效率差。
|
||||||
|
# GPU 4-5(a)/6-7(b),端口 30020/30030;2×tp2 聚合经 rr router 30003。
|
||||||
|
# 镜像/补丁挂载/L3/engram/tok8 全部与 cp2x4 相同。
|
||||||
|
x-common: &common
|
||||||
|
image: ymkymx/sglang:dsv41-pd-visioncp-db7d2cb7d-fix
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
network_mode: host
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
- /data/ymk/cache/sglang:/root/.cache/sglang
|
||||||
|
# 与 cp2x4 保持同一份 scheduler.py(含未 commit 的 CP idle 补丁),保证唯一变量是 CP 开关
|
||||||
|
- /data/ymk/sglang/python/sglang/srt/managers/scheduler.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py
|
||||||
|
environment: &env
|
||||||
|
SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE: "1"
|
||||||
|
SGLANG_RAGGED_VERIFY_MODE: static
|
||||||
|
MC_MS_AUTO_DISC: "0"
|
||||||
|
MOONCAKE_MASTER: 127.0.0.1:50051
|
||||||
|
MOONCAKE_TE_META_DATA_SERVER: P2PHANDSHAKE
|
||||||
|
MOONCAKE_PROTOCOL: tcp
|
||||||
|
MOONCAKE_GLOBAL_SEGMENT_SIZE: 300gb
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
|
||||||
|
services:
|
||||||
|
tp2:
|
||||||
|
<<: *common
|
||||||
|
container_name: ds41-tp2
|
||||||
|
environment:
|
||||||
|
<<: *env
|
||||||
|
CUDA_VISIBLE_DEVICES: "4,5"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30020
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
|
b:
|
||||||
|
<<: *common
|
||||||
|
container_name: ds41-tp2-b
|
||||||
|
environment:
|
||||||
|
<<: *env
|
||||||
|
CUDA_VISIBLE_DEVICES: "6,7"
|
||||||
|
command: >
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 2 --ep-size 2
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--cuda-graph-max-bs-decode 32
|
||||||
|
--reasoning-parser auto --tool-call-parser auto
|
||||||
|
--default-chat-template-kwargs '{"thinking": true, "reasoning_effort": "high"}'
|
||||||
|
--max-running-requests 64
|
||||||
|
--tokenizer-worker-num 8
|
||||||
|
--host 0.0.0.0 --port 30030
|
||||||
|
--enable-cache-report --enable-metrics
|
||||||
|
--speculative-algorithm DSPARK --speculative-dspark-block-size 5
|
||||||
|
--enable-hierarchical-cache
|
||||||
|
--hicache-ratio 2
|
||||||
|
--hicache-mem-layout page_first_direct
|
||||||
|
--hicache-io-backend direct
|
||||||
|
--hicache-write-policy write_through
|
||||||
|
--hicache-storage-backend mooncake
|
||||||
|
--hicache-storage-prefetch-policy wait_complete
|
||||||
|
--hicache-size 0
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
services:
|
||||||
|
sglang:
|
||||||
|
image: uhub.service.ucloud.cn/umirror/sglang:dev-dsv41
|
||||||
|
shm_size: "32gb"
|
||||||
|
ipc: host
|
||||||
|
privileged: true
|
||||||
|
|
||||||
|
ports:
|
||||||
|
- "30000:30000"
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
- /data:/data
|
||||||
|
|
||||||
|
deploy:
|
||||||
|
resources:
|
||||||
|
reservations:
|
||||||
|
devices:
|
||||||
|
- driver: nvidia
|
||||||
|
count: all
|
||||||
|
capabilities: [gpu]
|
||||||
|
|
||||||
|
command: >-
|
||||||
|
sglang serve
|
||||||
|
--trust-remote-code
|
||||||
|
--model-path /data/models/DeepSeek-V4.1-Flash
|
||||||
|
--tp 4
|
||||||
|
--ep-size 4
|
||||||
|
--mem-fraction-static 0.75
|
||||||
|
--attention-backend dsv4
|
||||||
|
--moe-runner-backend flashinfer_mxfp4
|
||||||
|
--enable-decoder-swa-bounded-replay
|
||||||
|
--cuda-graph-max-bs-decode 64
|
||||||
|
--reasoning-parser auto
|
||||||
|
--tool-call-parser auto
|
||||||
|
--max-running-requests 64
|
||||||
|
--host 0.0.0.0
|
||||||
|
--port 30000
|
||||||
|
--enable-cache-report
|
||||||
|
--enable-metrics
|
||||||
|
--enable-prefill-cp
|
||||||
|
--cp-strategy interleave
|
||||||
|
--json-model-override-args '{"vision_n_layers": 0}'
|
||||||
|
|
||||||
@@ -0,0 +1,6 @@
|
|||||||
|
{
|
||||||
|
"master_server_address": "127.0.0.1:50051",
|
||||||
|
"metadata_server": "P2PHANDSHAKE",
|
||||||
|
"global_segment_size": "300gb",
|
||||||
|
"protocol": "tcp"
|
||||||
|
}
|
||||||
@@ -181,12 +181,15 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
cfg = resolving_view(server_args)
|
cfg = resolving_view(server_args)
|
||||||
if model_config_of(server_args).hf_config.model_type != "deepseek_v41":
|
hf_config = model_config_of(server_args).hf_config
|
||||||
|
if hf_config.model_type != "deepseek_v41":
|
||||||
if cfg.enable_encoder_swa_bounded_replay:
|
if cfg.enable_encoder_swa_bounded_replay:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"--enable-encoder-swa-bounded-replay requires DeepSeek-V4.1"
|
"--enable-encoder-swa-bounded-replay requires DeepSeek-V4.1"
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
if hf_config.vision_n_layers > 0 and cfg.enable_prefill_cp:
|
||||||
|
_validate_deepseek_v41_vision_prefill_cp(server_args)
|
||||||
if cfg.enable_encoder_swa_bounded_replay:
|
if cfg.enable_encoder_swa_bounded_replay:
|
||||||
from sglang.srt.model_executor.cuda_graph_config import Backend
|
from sglang.srt.model_executor.cuda_graph_config import Backend
|
||||||
|
|
||||||
@@ -197,7 +200,8 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
|
|||||||
cfg.cuda_graph_config.prefill.backend != Backend.DISABLED,
|
cfg.cuda_graph_config.prefill.backend != Backend.DISABLED,
|
||||||
),
|
),
|
||||||
("DP attention", cfg.enable_dp_attention),
|
("DP attention", cfg.enable_dp_attention),
|
||||||
("context parallelism", cfg.attn_cp_size > 1),
|
# Prefill CP declares attn_cp_size and DP attention only later.
|
||||||
|
("context parallelism", cfg.attn_cp_size > 1 or cfg.enable_prefill_cp),
|
||||||
("external cache linker", cfg.enable_unified_cache_external_linker),
|
("external cache linker", cfg.enable_unified_cache_external_linker),
|
||||||
("unified memory", cfg.enable_unified_memory),
|
("unified memory", cfg.enable_unified_memory),
|
||||||
("PD disaggregation", cfg.disaggregation_mode != "null"),
|
("PD disaggregation", cfg.disaggregation_mode != "null"),
|
||||||
@@ -249,23 +253,33 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
|
|||||||
if (
|
if (
|
||||||
read_ragged_verify_mode() is not RaggedVerifyMode.STATIC
|
read_ragged_verify_mode() is not RaggedVerifyMode.STATIC
|
||||||
or cfg.disaggregation_transfer_backend != "mooncake"
|
or cfg.disaggregation_transfer_backend != "mooncake"
|
||||||
or cfg.dp_size != 1
|
|
||||||
or cfg.enable_dp_attention
|
|
||||||
or cfg.attn_cp_size != 1
|
or cfg.attn_cp_size != 1
|
||||||
or cfg.dcp_size != 1
|
or cfg.dcp_size != 1
|
||||||
):
|
):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"DeepSeek-V4.1 DSpark PD requires static verify, Mooncake, "
|
"DeepSeek-V4.1 DSpark PD requires static verify, Mooncake, "
|
||||||
"DP=1 and CP=1. Both servers must enable DSpark with the same "
|
"and CP=1 on both servers. DP attention is supported when "
|
||||||
"block size and TP size."
|
"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
|
from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase
|
||||||
|
|
||||||
prefill_graph = cfg.cuda_graph_config.prefill
|
prefill_graph = cfg.cuda_graph_config.prefill
|
||||||
if prefill_graph.backend != Backend.DISABLED and prefill_graph.max_seq_len is None:
|
cp_breakable_prefill = (
|
||||||
# The captured low-ratio indexer scores a static context width; 16k
|
cfg.enable_prefill_cp
|
||||||
# keeps it inside the candidate window at under 1 ms per layer.
|
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(
|
declare_resolution(
|
||||||
server_args,
|
server_args,
|
||||||
"validate_deepseek_v41_features",
|
"validate_deepseek_v41_features",
|
||||||
@@ -296,3 +310,40 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
|
|||||||
"--enable-decoder-swa-bounded-replay cannot be combined with "
|
"--enable-decoder-swa-bounded-replay cannot be combined with "
|
||||||
f"{feature} yet; disable one of them."
|
f"{feature} yet; disable one of them."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_deepseek_v41_vision_prefill_cp(server_args: ServerArgs) -> None:
|
||||||
|
from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase
|
||||||
|
|
||||||
|
cfg = resolving_view(server_args)
|
||||||
|
if cfg.cp_strategy != "interleave":
|
||||||
|
raise ValueError(
|
||||||
|
"DeepSeek-V4.1 vision with prefill CP requires --cp-strategy "
|
||||||
|
f"interleave; got {cfg.cp_strategy!r}."
|
||||||
|
)
|
||||||
|
if cfg.cuda_graph_config.prefill.backend != Backend.DISABLED:
|
||||||
|
# The CP runner merges image features eagerly; no capture path replays it.
|
||||||
|
locked = getattr(server_args, "_cuda_graph_config_locked", set())
|
||||||
|
if (Phase.PREFILL, "backend") in locked:
|
||||||
|
raise ValueError(
|
||||||
|
"DeepSeek-V4.1 vision with prefill CP runs eager prefill; remove "
|
||||||
|
"the explicit prefill CUDA graph backend."
|
||||||
|
)
|
||||||
|
declare_resolution(
|
||||||
|
server_args,
|
||||||
|
"validate_deepseek_v41_features",
|
||||||
|
cuda_graph_config=with_phase(
|
||||||
|
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
|
||||||
|
),
|
||||||
|
)
|
||||||
|
logger.warning(
|
||||||
|
"Disabling the prefill CUDA graph for DeepSeek-V4.1 vision with prefill CP."
|
||||||
|
)
|
||||||
|
if (
|
||||||
|
str(cfg.speculative_algorithm).upper() == "DSPARK"
|
||||||
|
and cfg.enable_decoder_swa_bounded_replay
|
||||||
|
):
|
||||||
|
raise ValueError(
|
||||||
|
"DeepSeek-V4.1 vision with prefill CP does not support DSpark together "
|
||||||
|
"with --enable-decoder-swa-bounded-replay yet."
|
||||||
|
)
|
||||||
|
|||||||
@@ -966,12 +966,14 @@ def handle_language_model_only(server_args: Any):
|
|||||||
):
|
):
|
||||||
if flag:
|
if flag:
|
||||||
raise ValueError(f"--language-model-only cannot be combined with {name}")
|
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(
|
raise ValueError(
|
||||||
"--language-model-only is incompatible with --disaggregation-mode "
|
"--language-model-only is incompatible with --disaggregation-mode "
|
||||||
"prefill/decode"
|
"prefill/decode"
|
||||||
)
|
)
|
||||||
architectures = model_config_of(server_args).hf_config.architectures
|
architectures = hf_config.architectures
|
||||||
if not any(
|
if not any(
|
||||||
a in server_args.LANGUAGE_MODEL_ONLY_ARCHITECTURES for a in architectures
|
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 "
|
"enable DSpark with the same block size and target/draft KV "
|
||||||
"layout. Upgrade both servers together."
|
"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(
|
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:
|
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
|
# decode
|
||||||
@dataclasses.dataclass
|
@dataclasses.dataclass
|
||||||
class TransferInfo:
|
class TransferInfo:
|
||||||
@@ -214,6 +323,7 @@ class KVArgsRegisterInfo:
|
|||||||
|
|
||||||
class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
||||||
AUX_DATA_HEADER = b"AUX_DATA"
|
AUX_DATA_HEADER = b"AUX_DATA"
|
||||||
|
STATE_DATA_HEADER = b"STATE_DATA"
|
||||||
# Implements teardown() below, so runtime PD role switching is supported.
|
# Implements teardown() below, so runtime PD role switching is supported.
|
||||||
supports_role_switch = True
|
supports_role_switch = True
|
||||||
|
|
||||||
@@ -227,6 +337,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
|||||||
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
|
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
|
||||||
self.init_engine()
|
self.init_engine()
|
||||||
self.register_buffer_to_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.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
|
||||||
self.max_transfer_batch_indices = (
|
self.max_transfer_batch_indices = (
|
||||||
envs.SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES.get()
|
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
|
Deduped because the unified memory pool reports one raw buffer as both
|
||||||
its KV and its mamba state component, and double registration fails in
|
its KV and its mamba state component, and double registration fails in
|
||||||
the engine.
|
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]] = []
|
regions: List[Tuple[int, int]] = []
|
||||||
seen: Set[Tuple[int, int]] = set()
|
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
|
self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
|
||||||
):
|
):
|
||||||
add(ptrs, 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
|
return regions
|
||||||
|
|
||||||
def register_buffer_to_engine(self):
|
def register_buffer_to_engine(self):
|
||||||
@@ -748,10 +887,63 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
|||||||
if not transfer_blocks:
|
if not transfer_blocks:
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
|
if not _nvlink_intra_transport_active():
|
||||||
return self.engine.batch_transfer_sync(
|
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
|
||||||
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
|
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(
|
def _send_kvcache_generic(
|
||||||
self,
|
self,
|
||||||
@@ -1482,8 +1674,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
|||||||
):
|
):
|
||||||
# TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free
|
# TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free
|
||||||
if (
|
if (
|
||||||
self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK"
|
(self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK")
|
||||||
) or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get():
|
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)
|
return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs)
|
||||||
|
|
||||||
transfer_blocks = []
|
transfer_blocks = []
|
||||||
@@ -1566,6 +1760,71 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
|||||||
f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}"
|
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(
|
def _get_dsa_cache_transfer_skip_flags(
|
||||||
self, info: Optional[KVArgsRegisterInfo]
|
self, info: Optional[KVArgsRegisterInfo]
|
||||||
) -> Tuple[bool, bool]:
|
) -> Tuple[bool, bool]:
|
||||||
@@ -2412,6 +2671,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
|||||||
):
|
):
|
||||||
self._staging_outstanding.pop(kv_chunk.room, None)
|
self._staging_outstanding.pop(kv_chunk.room, None)
|
||||||
if kv_chunk.room in self.transfer_infos:
|
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.transfer_infos.pop(kv_chunk.room)
|
||||||
self.req_to_decode_prefix_len.pop(kv_chunk.room, None)
|
self.req_to_decode_prefix_len.pop(kv_chunk.room, None)
|
||||||
if self.enable_staging:
|
if self.enable_staging:
|
||||||
@@ -2557,6 +2818,11 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
|||||||
self.transfer_infos[room][mooncake_session_id] = (
|
self.transfer_infos[room][mooncake_session_id] = (
|
||||||
TransferInfo.from_zmq(waiting_req_bytes)
|
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
|
# NOTE: after bootstrapping we can mark the req as waiting for input
|
||||||
if len(self.transfer_infos[room]) == required_dst_info_num:
|
if len(self.transfer_infos[room]) == required_dst_info_num:
|
||||||
self.resolve_kv_replica_factor(self.transfer_infos[room])
|
self.resolve_kv_replica_factor(self.transfer_infos[room])
|
||||||
@@ -2585,6 +2851,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
|
|||||||
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
|
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
|
||||||
self._handle_aux_data(msg)
|
self._handle_aux_data(msg)
|
||||||
continue
|
continue
|
||||||
|
if msg[0] == MooncakeKVManager.STATE_DATA_HEADER:
|
||||||
|
self._handle_state_data(msg)
|
||||||
|
continue
|
||||||
|
|
||||||
# Staging: prefill notifies a chunk written to staging buffer
|
# Staging: prefill notifies a chunk written to staging buffer
|
||||||
if msg[0] == b"CHUNK_READY":
|
if msg[0] == b"CHUNK_READY":
|
||||||
|
|||||||
@@ -736,7 +736,14 @@ class DSV4AttnMetadata:
|
|||||||
if src_val is None and dst_val is None:
|
if src_val is None and dst_val is None:
|
||||||
continue
|
continue
|
||||||
assert dst_val is not None, f"{field_name=} {src_val=} {dst_val=}"
|
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
|
# These fields are safe to replace because captured kernels only need
|
||||||
# the current per-replay objects, or the field is produced inside the
|
# 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
|
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
|
@dataclass
|
||||||
class DSV4Metadata:
|
class DSV4Metadata:
|
||||||
core_attn_metadata: DSV4AttnMetadata
|
core_attn_metadata: DSV4AttnMetadata
|
||||||
@@ -1254,6 +1282,11 @@ class DeepseekV4AttnBackend(
|
|||||||
] = None
|
] = None
|
||||||
self.online_c128_mtp = OnlineC128MTPController(self)
|
self.online_c128_mtp = OnlineC128MTPController(self)
|
||||||
self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device)
|
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
|
spec_alg = model_runner.spec_algorithm
|
||||||
self.needs_cpu_seq_lens = not spec_alg.is_dspark() and (
|
self.needs_cpu_seq_lens = not spec_alg.is_dspark() and (
|
||||||
not _is_cuda or self.online_c128_mtp.enabled()
|
not _is_cuda or self.online_c128_mtp.enabled()
|
||||||
@@ -1566,8 +1599,12 @@ class DeepseekV4AttnBackend(
|
|||||||
|
|
||||||
@property
|
@property
|
||||||
def low_ratio_prefill_graph(self) -> bool:
|
def low_ratio_prefill_graph(self) -> bool:
|
||||||
|
"""Whether ratio-1/2 sources use captured projections and indexer metadata."""
|
||||||
return (
|
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:
|
def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool:
|
||||||
@@ -2870,7 +2907,7 @@ class DeepseekV4AttnBackend(
|
|||||||
q_lora[:num_local],
|
q_lora[:num_local],
|
||||||
positions[:num_local].to(torch.int64),
|
positions[:num_local].to(torch.int64),
|
||||||
forward_batch,
|
forward_batch,
|
||||||
torch.tensor(q_lens_cpu, dtype=torch.int32, device=x.device),
|
self._move_to_device(q_lens_cpu),
|
||||||
q_lens_cpu,
|
q_lens_cpu,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -3236,7 +3273,9 @@ class DeepseekV4AttnBackend(
|
|||||||
continue
|
continue
|
||||||
j = torch.arange(lc, device=device)
|
j = torch.arange(lc, device=device)
|
||||||
slot_chunks.append(
|
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
|
// ratio
|
||||||
)
|
)
|
||||||
start += lc
|
start += lc
|
||||||
@@ -3261,7 +3300,7 @@ class DeepseekV4AttnBackend(
|
|||||||
weights = indexer.head_weights(x).float()
|
weights = indexer.head_weights(x).float()
|
||||||
compress_lens = ((pos + 1) // ratio).to(torch.int32)
|
compress_lens = ((pos + 1) // ratio).to(torch.int32)
|
||||||
ks = torch.repeat_interleave(
|
ks = torch.repeat_interleave(
|
||||||
torch.tensor(starts, dtype=torch.int32, device=device),
|
self._move_to_device(starts),
|
||||||
q_lens.to(torch.int64),
|
q_lens.to(torch.int64),
|
||||||
output_size=num_tokens,
|
output_size=num_tokens,
|
||||||
)
|
)
|
||||||
@@ -3447,17 +3486,44 @@ class DeepseekV4AttnBackend(
|
|||||||
self.candidate_indexer.publish_decode(inputs, page_indices, raw_indices)
|
self.candidate_indexer.publish_decode(inputs, page_indices, raw_indices)
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
logits = deep_gemm_fp4_paged_mqa_logits(
|
if isinstance(metadata.deep_gemm_metadata, list):
|
||||||
(q_fp4, q_sf),
|
topk_plans = metadata.topk_metadata_chunks
|
||||||
k_cache,
|
assert not metadata.use_topk_v2 or topk_plans is not None
|
||||||
weights,
|
for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()):
|
||||||
metadata.compressed_seq_lens,
|
logits = deep_gemm_fp4_paged_mqa_logits(
|
||||||
metadata.page_table,
|
(q_fp4[rows], q_sf[rows]),
|
||||||
metadata.deep_gemm_metadata,
|
k_cache,
|
||||||
metadata.max_compressed_seq_len,
|
weights[rows],
|
||||||
)
|
metadata.compressed_seq_lens[rows],
|
||||||
# TODO(dark): add bf16 topk
|
metadata.page_table[rows],
|
||||||
topk_transform_paged_from_metadata(logits, metadata, page_indices, raw_indices)
|
plan,
|
||||||
|
metadata.max_compressed_seq_len,
|
||||||
|
)
|
||||||
|
# TODO(dark): add bf16 topk
|
||||||
|
topk_transform_paged_from_metadata(
|
||||||
|
logits,
|
||||||
|
metadata,
|
||||||
|
page_indices,
|
||||||
|
raw_indices,
|
||||||
|
rows=rows,
|
||||||
|
topk_metadata=(
|
||||||
|
topk_plans[chunk_idx] if topk_plans is not None else None
|
||||||
|
),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logits = deep_gemm_fp4_paged_mqa_logits(
|
||||||
|
(q_fp4, q_sf),
|
||||||
|
k_cache,
|
||||||
|
weights,
|
||||||
|
metadata.compressed_seq_lens,
|
||||||
|
metadata.page_table,
|
||||||
|
metadata.deep_gemm_metadata,
|
||||||
|
metadata.max_compressed_seq_len,
|
||||||
|
)
|
||||||
|
# TODO(dark): add bf16 topk
|
||||||
|
topk_transform_paged_from_metadata(
|
||||||
|
logits, metadata, page_indices, raw_indices
|
||||||
|
)
|
||||||
|
|
||||||
# TODO(candidate): Hopper decode still publishes / consumes masks inline (torch
|
# TODO(candidate): Hopper decode still publishes / consumes masks inline (torch
|
||||||
# top-k); move into the candidate indexer with the prefill paths.
|
# top-k); move into the candidate indexer with the prefill paths.
|
||||||
@@ -3945,20 +4011,38 @@ class DeepseekV4AttnBackend(
|
|||||||
compress_ratio, core_attn_metadata, extra_page_size
|
compress_ratio, core_attn_metadata, extra_page_size
|
||||||
)
|
)
|
||||||
n_compressed = flat_token_ids.shape[0]
|
n_compressed = flat_token_ids.shape[0]
|
||||||
workspace = self.sparse_prefill_workspace.get(
|
reuse_compressed = compress_ratio in (1, 2) and is_cp_active(forward_batch)
|
||||||
n_compressed + cache.swa_token_ids.shape[0]
|
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]
|
compressed_slice = workspace[:n_compressed]
|
||||||
swa_slice = workspace[n_compressed:]
|
swa_slice = workspace[n_compressed:]
|
||||||
|
|
||||||
if compressed_slice is not None:
|
if compressed_slice is not None:
|
||||||
dequantize_k_cache_paged(
|
source_key = None
|
||||||
extra_k_cache,
|
if reuse_compressed:
|
||||||
flat_token_ids,
|
source_layer = token_to_kv_pool.source_layer_of(layer_id)
|
||||||
page_size=extra_page_size,
|
source_key = (source_layer, workspace.data_ptr())
|
||||||
out=compressed_slice,
|
gather = cache.compressed[compress_ratio]
|
||||||
layout=token_to_kv_pool.get_extra_key_layout(layer_id),
|
# 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(
|
dequantize_k_cache_paged(
|
||||||
token_to_kv_pool.get_swa_key_buffer_radix(layer_id),
|
token_to_kv_pool.get_swa_key_buffer_radix(layer_id),
|
||||||
cache.swa_token_ids,
|
cache.swa_token_ids,
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ from sglang.srt.layers.attention.dsv4.candidate_indexer import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.layers.attention.dsv4.indexer import (
|
from sglang.srt.layers.attention.dsv4.indexer import (
|
||||||
deep_gemm_fp4_paged_mqa_logits,
|
deep_gemm_fp4_paged_mqa_logits,
|
||||||
|
topk_transform_paged_from_metadata,
|
||||||
)
|
)
|
||||||
|
|
||||||
CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16
|
CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16
|
||||||
@@ -176,6 +177,10 @@ class DeepGemmCandidateIndexer:
|
|||||||
metadata."""
|
metadata."""
|
||||||
metadata = inputs.metadata
|
metadata = inputs.metadata
|
||||||
seq_lens = metadata.compressed_seq_lens.reshape(-1)
|
seq_lens = metadata.compressed_seq_lens.reshape(-1)
|
||||||
|
if isinstance(metadata.deep_gemm_metadata, list):
|
||||||
|
return self._publish_decode_chunked(
|
||||||
|
inputs, page_indices, raw_indices, seq_lens
|
||||||
|
)
|
||||||
logits = deep_gemm_fp4_paged_mqa_logits(
|
logits = deep_gemm_fp4_paged_mqa_logits(
|
||||||
(inputs.q_fp4, inputs.q_sf),
|
(inputs.q_fp4, inputs.q_sf),
|
||||||
inputs.k_cache,
|
inputs.k_cache,
|
||||||
@@ -231,6 +236,83 @@ class DeepGemmCandidateIndexer:
|
|||||||
ready=ready,
|
ready=ready,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _publish_decode_chunked(
|
||||||
|
self,
|
||||||
|
inputs: IndexerInputs,
|
||||||
|
page_indices: torch.Tensor,
|
||||||
|
raw_indices: Optional[torch.Tensor],
|
||||||
|
seq_lens: torch.Tensor,
|
||||||
|
) -> SparseBlockTable:
|
||||||
|
"""Publish an eager forward whose dense logits are bounded by row chunks.
|
||||||
|
|
||||||
|
CUDA-graph metadata always carries one tensor schedule and keeps using the
|
||||||
|
asynchronous fast path above. The exceptional eager path stays on the
|
||||||
|
current stream so each chunk's full logits can be released before the next.
|
||||||
|
"""
|
||||||
|
metadata = inputs.metadata
|
||||||
|
block_chunks = []
|
||||||
|
phys_block_chunks = []
|
||||||
|
valid_len_chunks = []
|
||||||
|
topk_plans = metadata.topk_metadata_chunks
|
||||||
|
assert not metadata.use_topk_v2 or topk_plans is not None
|
||||||
|
|
||||||
|
for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()):
|
||||||
|
logits = deep_gemm_fp4_paged_mqa_logits(
|
||||||
|
(inputs.q_fp4[rows], inputs.q_sf[rows]),
|
||||||
|
inputs.k_cache,
|
||||||
|
inputs.weights[rows],
|
||||||
|
metadata.compressed_seq_lens[rows],
|
||||||
|
metadata.page_table[rows],
|
||||||
|
plan,
|
||||||
|
metadata.max_compressed_seq_len,
|
||||||
|
)
|
||||||
|
topk_transform_paged_from_metadata(
|
||||||
|
logits,
|
||||||
|
metadata,
|
||||||
|
page_indices,
|
||||||
|
raw_indices,
|
||||||
|
rows=rows,
|
||||||
|
topk_metadata=(
|
||||||
|
topk_plans[chunk_idx] if topk_plans is not None else None
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
chunk_seq_lens = seq_lens[rows]
|
||||||
|
nblocks, row_valid_lens = candidate_row_lens(
|
||||||
|
chunk_seq_lens, self.topk_blocks
|
||||||
|
)
|
||||||
|
blocks = amax_topk_blocks(logits, chunk_seq_lens, nblocks, self.topk_blocks)
|
||||||
|
phys_blocks = sort_candidate_blocks(
|
||||||
|
blocks,
|
||||||
|
chunk_seq_lens,
|
||||||
|
metadata.page_table[rows],
|
||||||
|
metadata.compressed_page_size,
|
||||||
|
)
|
||||||
|
block_chunks.append(blocks)
|
||||||
|
phys_block_chunks.append(phys_blocks)
|
||||||
|
valid_len_chunks.append(row_valid_lens)
|
||||||
|
|
||||||
|
blocks = torch.cat(block_chunks)
|
||||||
|
phys_blocks = torch.cat(phys_block_chunks)
|
||||||
|
row_valid_lens = torch.cat(valid_len_chunks)
|
||||||
|
schedule = build_sparse_indexer_schedule(
|
||||||
|
blocks,
|
||||||
|
seq_lens,
|
||||||
|
metadata.page_table,
|
||||||
|
metadata.compressed_page_size,
|
||||||
|
inputs.q_fp4.dtype,
|
||||||
|
self._request_ids(inputs.request_ids, inputs.num_rows, blocks.device),
|
||||||
|
)
|
||||||
|
ready = torch.cuda.Event()
|
||||||
|
ready.record(torch.cuda.current_stream())
|
||||||
|
return SparseBlockTable(
|
||||||
|
blocks=blocks,
|
||||||
|
schedule=schedule,
|
||||||
|
phys_blocks=phys_blocks,
|
||||||
|
valid_lens=row_valid_lens,
|
||||||
|
ready=ready,
|
||||||
|
)
|
||||||
|
|
||||||
def _scores(self, table: SparseBlockTable, inputs: IndexerInputs) -> torch.Tensor:
|
def _scores(self, table: SparseBlockTable, inputs: IndexerInputs) -> torch.Tensor:
|
||||||
return sparse_logits(
|
return sparse_logits(
|
||||||
inputs.q_fp4,
|
inputs.q_fp4,
|
||||||
|
|||||||
@@ -475,27 +475,40 @@ def topk_transform_paged_from_metadata(
|
|||||||
metadata,
|
metadata,
|
||||||
page_indices: torch.Tensor,
|
page_indices: torch.Tensor,
|
||||||
raw_indices: Optional[torch.Tensor] = None,
|
raw_indices: Optional[torch.Tensor] = None,
|
||||||
|
*,
|
||||||
|
rows: Optional[slice] = None,
|
||||||
|
topk_metadata: Optional[torch.Tensor] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Pool slots into ``page_indices`` (``-1`` past the valid count) and, when given,
|
"""Pool slots into ``page_indices`` (``-1`` past the valid count) and, when given,
|
||||||
positions into ``raw_indices``; ``metadata`` is a ``PagedIndexerMetadata``."""
|
positions into ``raw_indices``; ``metadata`` is a ``PagedIndexerMetadata``."""
|
||||||
|
if rows is None:
|
||||||
|
seq_lens = metadata.compressed_seq_lens
|
||||||
|
page_table = metadata.page_table
|
||||||
|
out_page_indices = page_indices
|
||||||
|
out_raw_indices = raw_indices
|
||||||
|
else:
|
||||||
|
seq_lens = metadata.compressed_seq_lens[rows]
|
||||||
|
page_table = metadata.page_table[rows]
|
||||||
|
out_page_indices = page_indices[rows]
|
||||||
|
out_raw_indices = raw_indices[rows] if raw_indices is not None else None
|
||||||
if metadata.use_topk_v2:
|
if metadata.use_topk_v2:
|
||||||
topk_transform_paged_v2(
|
topk_transform_paged_v2(
|
||||||
logits,
|
logits,
|
||||||
metadata.compressed_seq_lens,
|
seq_lens,
|
||||||
metadata.page_table,
|
page_table,
|
||||||
page_indices,
|
out_page_indices,
|
||||||
metadata.compressed_page_size,
|
metadata.compressed_page_size,
|
||||||
metadata.topk_metadata,
|
metadata.topk_metadata if topk_metadata is None else topk_metadata,
|
||||||
raw_indices,
|
out_raw_indices,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
topk_transform_paged(
|
topk_transform_paged(
|
||||||
logits,
|
logits,
|
||||||
metadata.compressed_seq_lens,
|
seq_lens,
|
||||||
metadata.page_table,
|
page_table,
|
||||||
page_indices,
|
out_page_indices,
|
||||||
metadata.compressed_page_size,
|
metadata.compressed_page_size,
|
||||||
raw_indices,
|
out_raw_indices,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context
|
|||||||
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
|
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
|
||||||
is_in_tc_piecewise_cuda_graph,
|
is_in_tc_piecewise_cuda_graph,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
|
||||||
from sglang.srt.utils import is_hip, is_sm120_supported, is_xpu
|
from sglang.srt.utils import is_hip, is_sm120_supported, is_xpu
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -306,7 +307,8 @@ class PagedIndexerMetadata:
|
|||||||
):
|
):
|
||||||
return None
|
return None
|
||||||
if (
|
if (
|
||||||
torch.cuda.is_current_stream_capturing()
|
get_is_capture_mode()
|
||||||
|
or torch.cuda.is_current_stream_capturing()
|
||||||
or is_in_breakable_cuda_graph()
|
or is_in_breakable_cuda_graph()
|
||||||
or is_in_tc_piecewise_cuda_graph()
|
or is_in_tc_piecewise_cuda_graph()
|
||||||
):
|
):
|
||||||
@@ -325,14 +327,27 @@ class PagedIndexerMetadata:
|
|||||||
|
|
||||||
def row_chunks(self):
|
def row_chunks(self):
|
||||||
num_rows = self.compressed_seq_lens.shape[0]
|
num_rows = self.compressed_seq_lens.shape[0]
|
||||||
if self.row_chunk <= 0:
|
if self.row_chunk > 0:
|
||||||
|
rows_per_chunk = self.row_chunk
|
||||||
|
elif isinstance(self.deep_gemm_metadata, list):
|
||||||
|
assert self.rows_per_chunk is not None, (
|
||||||
|
"chunked DeepGEMM metadata requires rows_per_chunk"
|
||||||
|
)
|
||||||
|
rows_per_chunk = self.rows_per_chunk
|
||||||
|
else:
|
||||||
return [(slice(0, num_rows), self.deep_gemm_metadata)]
|
return [(slice(0, num_rows), self.deep_gemm_metadata)]
|
||||||
return [
|
|
||||||
(slice(start, min(start + self.row_chunk, num_rows)), plan)
|
chunks = [
|
||||||
|
(slice(start, min(start + rows_per_chunk, num_rows)), plan)
|
||||||
for start, plan in zip(
|
for start, plan in zip(
|
||||||
range(0, num_rows, self.row_chunk), self.deep_gemm_metadata
|
range(0, num_rows, rows_per_chunk), self.deep_gemm_metadata
|
||||||
)
|
)
|
||||||
]
|
]
|
||||||
|
assert chunks and chunks[-1][0].stop == num_rows, (
|
||||||
|
f"chunk schedules do not cover all rows: {num_rows=} {rows_per_chunk=} "
|
||||||
|
f"{len(chunks)=}"
|
||||||
|
)
|
||||||
|
return chunks
|
||||||
|
|
||||||
def copy_(self, other: PagedIndexerMetadata):
|
def copy_(self, other: PagedIndexerMetadata):
|
||||||
# A chunked schedule list has no in-place copy; rebind it instead.
|
# A chunked schedule list has no in-place copy; rebind it instead.
|
||||||
|
|||||||
@@ -78,10 +78,11 @@ def use_dsv4_q8kv8_sparse_prefill(dsv4_prefill_backend: str = "auto") -> bool:
|
|||||||
class SparsePrefillWorkspace:
|
class SparsePrefillWorkspace:
|
||||||
"""Backend-owned scratch storage for sparse prefill KV dequantization.
|
"""Backend-owned scratch storage for sparse prefill KV dequantization.
|
||||||
|
|
||||||
The workspace contents are fully overwritten before every attention call,
|
Callers normally overwrite the entire workspace. Shared compressed-KV
|
||||||
so token buckets and compression ratios can safely share one buffer. Sparse
|
callers keep separate workspaces per ratio and track prefix validity in the
|
||||||
prefill executes eagerly and serially on the supported paths, which makes it
|
per-forward gather cache, including the allocation address. Sparse prefill
|
||||||
safe to replace the scratch allocation when a larger extent is needed.
|
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):
|
def __init__(self, device: torch.device):
|
||||||
@@ -275,14 +276,18 @@ class CompressedGather:
|
|||||||
# chunk-invariant per request; subsequent layers only overwrite that prefix.
|
# chunk-invariant per request; subsequent layers only overwrite that prefix.
|
||||||
combined_indices: Optional[torch.Tensor] = None
|
combined_indices: Optional[torch.Tensor] = None
|
||||||
combined_lens: 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
|
@dataclass
|
||||||
class SparsePrefillChunkCache:
|
class SparsePrefillChunkCache:
|
||||||
"""Cache prefill-chunk metadata shared across layers.
|
"""Cache prefill-chunk metadata shared across layers.
|
||||||
|
|
||||||
Fields depend on request/token mappings and compressed page tables, not
|
Gather layouts depend on request/token mappings and compressed page tables.
|
||||||
per-layer k_cache; per-layer top-k combinations are recomputed into reused
|
Shared-source dequantization keys live only for this forward; per-layer
|
||||||
|
top-k combinations are recomputed into reused
|
||||||
buffers.
|
buffers.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|||||||
@@ -23,17 +23,22 @@ import torch
|
|||||||
|
|
||||||
from sglang.srt.arg_groups.overrides import (
|
from sglang.srt.arg_groups.overrides import (
|
||||||
attention_backends_of,
|
attention_backends_of,
|
||||||
|
model_config_of,
|
||||||
resolved_view,
|
resolved_view,
|
||||||
resolving_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.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.padding import get_cp_padding_align_size
|
||||||
from sglang.srt.layers.cp.utils import (
|
from sglang.srt.layers.cp.utils import (
|
||||||
cp_gather_after_forward,
|
cp_gather_after_forward,
|
||||||
|
cp_shard_hidden_states,
|
||||||
cp_split_before_forward,
|
cp_split_before_forward,
|
||||||
prepare_cp_forward,
|
prepare_cp_forward,
|
||||||
)
|
)
|
||||||
from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
|
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
|
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -50,12 +55,18 @@ def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool:
|
|||||||
cfg = resolving_view(server_args)
|
cfg = resolving_view(server_args)
|
||||||
resolved = resolved_view(server_args)
|
resolved = resolved_view(server_args)
|
||||||
prefill_attention_backend, _ = attention_backends_of(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 (
|
return (
|
||||||
cfg.enable_prefill_cp
|
cfg.enable_prefill_cp
|
||||||
and cfg.pp_size == 1
|
and cfg.pp_size == 1
|
||||||
and resolved.attn_cp_size == cfg.tp_size
|
and resolved.attn_cp_size == cfg.tp_size
|
||||||
and cfg.cp_strategy == "zigzag"
|
and supports_layout
|
||||||
and prefill_attention_backend == "trtllm_mha"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -67,8 +78,12 @@ def enable_cp_bcg_capture(server_args: ServerArgs) -> bool:
|
|||||||
def filter_prefill_cp_bcg_capture_num_tokens(
|
def filter_prefill_cp_bcg_capture_num_tokens(
|
||||||
capture_num_tokens: list[int], server_args: ServerArgs
|
capture_num_tokens: list[int], server_args: ServerArgs
|
||||||
) -> list[int]:
|
) -> list[int]:
|
||||||
"""Keep only token buckets where the zigzag CP strategy can run."""
|
"""Keep only token buckets where the configured CP strategy can run."""
|
||||||
min_num_tokens = resolved_view(server_args).attn_cp_size * 2
|
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]
|
filtered = [size for size in capture_num_tokens if size >= min_num_tokens]
|
||||||
if not filtered:
|
if not filtered:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
@@ -96,6 +111,8 @@ class PrefillCPBCGInput:
|
|||||||
|
|
||||||
input_embeds: torch.Tensor
|
input_embeds: torch.Tensor
|
||||||
positions: 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)
|
bucket_local_tokens: Dict[int, int] = field(default_factory=dict)
|
||||||
live_local_tokens: int = 0
|
live_local_tokens: int = 0
|
||||||
|
|
||||||
@@ -114,12 +131,22 @@ class PrefillCPBCGInput:
|
|||||||
(runner.max_num_tokens,),
|
(runner.max_num_tokens,),
|
||||||
dtype=torch.int64,
|
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]:
|
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()
|
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
|
return None
|
||||||
|
|
||||||
cp_segment_num = strategy.cp_size * 2
|
cp_segment_num = strategy.cp_size * 2
|
||||||
@@ -219,6 +246,7 @@ class PrefillCPBCGInput:
|
|||||||
raw_tokens = int(forward_batch.extend_num_tokens)
|
raw_tokens = int(forward_batch.extend_num_tokens)
|
||||||
global_input_ids = forward_batch.input_ids[:raw_tokens]
|
global_input_ids = forward_batch.input_ids[:raw_tokens]
|
||||||
global_positions = forward_batch.positions[: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_embeds = runner.model_runner.model.get_input_embeddings()(
|
||||||
global_input_ids
|
global_input_ids
|
||||||
)
|
)
|
||||||
@@ -249,12 +277,31 @@ class PrefillCPBCGInput:
|
|||||||
|
|
||||||
input_embeds = self.input_embeds[:captured_local_tokens]
|
input_embeds = self.input_embeds[:captured_local_tokens]
|
||||||
positions = self.positions[: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_()
|
input_embeds.zero_()
|
||||||
positions.zero_()
|
positions.zero_()
|
||||||
|
input_ids.zero_()
|
||||||
input_embeds[:live_local_tokens].copy_(local_input_embeds)
|
input_embeds[:live_local_tokens].copy_(local_input_embeds)
|
||||||
positions[:live_local_tokens].copy_(local_positions)
|
positions[:live_local_tokens].copy_(local_positions)
|
||||||
|
input_ids[:live_local_tokens].copy_(local_input_ids)
|
||||||
forward_batch.input_embeds = input_embeds
|
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
|
self.live_local_tokens = live_local_tokens
|
||||||
|
|
||||||
|
|
||||||
@@ -307,10 +354,50 @@ def execute_prefill_cp_bcg(
|
|||||||
static_forward_batch,
|
static_forward_batch,
|
||||||
torch.cuda.current_stream(),
|
torch.cuda.current_stream(),
|
||||||
)
|
)
|
||||||
return model.logits_processor(
|
if aux_hidden_states is not None:
|
||||||
forward_batch.input_ids,
|
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,
|
hidden_states,
|
||||||
model.lm_head,
|
model.lm_head,
|
||||||
forward_batch,
|
logits_metadata,
|
||||||
aux_hidden_states,
|
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:
|
if num_tokens > 0:
|
||||||
router_logits = moe.gate(hidden_states, forward_batch=forward_batch)
|
router_logits = moe.gate(hidden_states, forward_batch=forward_batch)
|
||||||
topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {}
|
num_token_non_padded = (
|
||||||
topk_output = moe.topk(
|
forward_batch.num_token_non_padded if forward_batch is not None else None
|
||||||
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,
|
|
||||||
)
|
)
|
||||||
|
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_ids = topk_output.topk_ids
|
||||||
topk_weights = topk_output.topk_weights
|
topk_weights = topk_output.topk_weights
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -0,0 +1,486 @@
|
|||||||
|
"""One owner rank encodes each image span and broadcasts it to the ranks that
|
||||||
|
run the same prefill chunk; every agreement precedes the payload it guards."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from contextlib import contextmanager
|
||||||
|
from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple
|
||||||
|
|
||||||
|
import msgspec
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
|
||||||
|
disable_symmetric_memory_context,
|
||||||
|
restore_symmetric_memory_context,
|
||||||
|
)
|
||||||
|
from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
SpanKey = Tuple[Optional[int], int]
|
||||||
|
SpanEncoder = Callable[[List[Any]], torch.Tensor | List[torch.Tensor]]
|
||||||
|
SpanSignature = Callable[[Any, int], Tuple[Any, ...]]
|
||||||
|
|
||||||
|
LOCAL_HIT = 0
|
||||||
|
OWNER_CACHE_BROADCAST = 1
|
||||||
|
OWNER_ENCODE_BROADCAST = 2
|
||||||
|
|
||||||
|
PHASE_PREPARE = "prepare"
|
||||||
|
PHASE_FEATURES = "features"
|
||||||
|
PHASE_FINALIZE = "finalize"
|
||||||
|
|
||||||
|
|
||||||
|
class MmOwnerProtocolError(RuntimeError):
|
||||||
|
"""Raised with identical text on every group member after a group-agreed failure."""
|
||||||
|
|
||||||
|
|
||||||
|
class ImageSpanRequest(msgspec.Struct, frozen=True):
|
||||||
|
hash: Optional[int]
|
||||||
|
span_len: int
|
||||||
|
item: Any
|
||||||
|
inside_chunk: bool
|
||||||
|
duplicates: List[Any] = []
|
||||||
|
|
||||||
|
|
||||||
|
class ImageSpanKey(msgspec.Struct, frozen=True):
|
||||||
|
hash: Optional[int]
|
||||||
|
span_len: int
|
||||||
|
geometry: Optional[Tuple[Any, ...]]
|
||||||
|
|
||||||
|
|
||||||
|
class RankManifest(msgspec.Struct, frozen=True):
|
||||||
|
rank: int
|
||||||
|
keys: List[ImageSpanKey]
|
||||||
|
cached: List[bool]
|
||||||
|
dtype: str
|
||||||
|
width: int
|
||||||
|
rids: List[str]
|
||||||
|
error: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
class OwnerPlan(msgspec.Struct, frozen=True):
|
||||||
|
actions: List[int]
|
||||||
|
owners: List[int]
|
||||||
|
error: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
class RankStatus(msgspec.Struct, frozen=True):
|
||||||
|
rank: int
|
||||||
|
error: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
def select_owner_group(parallel) -> Optional[Any]:
|
||||||
|
"""The group whose members all execute the same requests, or None when a
|
||||||
|
single rank already encodes every image it sees."""
|
||||||
|
replication = parallel.tp_size // parallel.attn_dp_size
|
||||||
|
if replication <= 1:
|
||||||
|
return None
|
||||||
|
if parallel.attn_cp_size == 1:
|
||||||
|
group = parallel.attn_tp_group
|
||||||
|
elif parallel.attn_dp_size == 1 and parallel.attn_cp_size == parallel.tp_size:
|
||||||
|
group = parallel.attn_cp_group
|
||||||
|
else:
|
||||||
|
return None
|
||||||
|
return group if group.world_size == replication else None
|
||||||
|
|
||||||
|
|
||||||
|
def has_owner_span_work(
|
||||||
|
mm_inputs: Sequence[Any],
|
||||||
|
extend_prefix_lens: Sequence[int],
|
||||||
|
extend_seq_lens: Sequence[int],
|
||||||
|
) -> bool:
|
||||||
|
"""Host-side mirror of the per-image scheduling path: does any raw
|
||||||
|
single-span image overlap the chunk on every rank of the group."""
|
||||||
|
for mm_input, prefix_len, extend_len in zip(
|
||||||
|
mm_inputs, extend_prefix_lens, extend_seq_lens
|
||||||
|
):
|
||||||
|
if mm_input is None or extend_len <= 0:
|
||||||
|
continue
|
||||||
|
items = [item for item in mm_input.mm_items if item is not None]
|
||||||
|
if not items or any(
|
||||||
|
item.precomputed_embeddings is not None or len(item.offsets) != 1
|
||||||
|
for item in items
|
||||||
|
):
|
||||||
|
continue
|
||||||
|
for item in items:
|
||||||
|
start, end = item.offsets[0]
|
||||||
|
if end >= prefix_len and start < prefix_len + extend_len:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
class MmOwnerSession(msgspec.Struct):
|
||||||
|
group: Any
|
||||||
|
device: Any
|
||||||
|
dtype: Any
|
||||||
|
width: int
|
||||||
|
rids: List[str]
|
||||||
|
signature: Any
|
||||||
|
engaged: bool
|
||||||
|
phase: str = PHASE_PREPARE
|
||||||
|
in_collective: bool = False
|
||||||
|
|
||||||
|
def resolve(
|
||||||
|
self,
|
||||||
|
requests: Sequence[ImageSpanRequest],
|
||||||
|
cache: MultiModalStaticCache,
|
||||||
|
encode: SpanEncoder,
|
||||||
|
) -> Dict[SpanKey, torch.Tensor]:
|
||||||
|
if not self.engaged:
|
||||||
|
raise RuntimeError(
|
||||||
|
"owner protocol reached for a chunk whose host metadata has no image span"
|
||||||
|
)
|
||||||
|
# Owners allocate different amounts than receivers, so none of these
|
||||||
|
# buffers may come out of a symmetric pool.
|
||||||
|
saved_context = disable_symmetric_memory_context()
|
||||||
|
try:
|
||||||
|
return _resolve_owner_features(self, requests, cache, encode)
|
||||||
|
finally:
|
||||||
|
restore_symmetric_memory_context(saved_context)
|
||||||
|
|
||||||
|
def features_ready(self) -> None:
|
||||||
|
self._complete()
|
||||||
|
self.phase = PHASE_FINALIZE
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def uncaptured(self) -> Iterator[None]:
|
||||||
|
# A failure inside a collective leaves the group in an unknown state;
|
||||||
|
# no later exchange may try to agree on it.
|
||||||
|
self.in_collective = True
|
||||||
|
yield
|
||||||
|
self.in_collective = False
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def fence(self) -> Iterator[None]:
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
except Exception as exc:
|
||||||
|
self._fail(exc)
|
||||||
|
raise
|
||||||
|
self._complete()
|
||||||
|
|
||||||
|
def _fail(self, exc: BaseException) -> None:
|
||||||
|
if (
|
||||||
|
not self.engaged
|
||||||
|
or self.in_collective
|
||||||
|
or isinstance(exc, MmOwnerProtocolError)
|
||||||
|
):
|
||||||
|
raise exc
|
||||||
|
text = _describe(self, self.phase, exc)
|
||||||
|
if self.phase == PHASE_PREPARE:
|
||||||
|
try:
|
||||||
|
_exchange_manifest(self, _manifest(self, [], [], error=text))
|
||||||
|
except MmOwnerProtocolError as agreed:
|
||||||
|
raise agreed from exc
|
||||||
|
_exchange_status(self, text, exc)
|
||||||
|
|
||||||
|
def _complete(self) -> None:
|
||||||
|
if not self.engaged:
|
||||||
|
return
|
||||||
|
if self.phase == PHASE_PREPARE:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"owner protocol {self.phase} completed without a manifest exchange"
|
||||||
|
)
|
||||||
|
error = None
|
||||||
|
cause = None
|
||||||
|
try:
|
||||||
|
_synchronize(self.device)
|
||||||
|
except Exception as exc:
|
||||||
|
cause = exc
|
||||||
|
error = _describe(self, self.phase, exc)
|
||||||
|
_exchange_status(self, error, cause)
|
||||||
|
|
||||||
|
|
||||||
|
def _manifest(
|
||||||
|
session: MmOwnerSession,
|
||||||
|
keys: List[ImageSpanKey],
|
||||||
|
cached: List[bool],
|
||||||
|
error: Optional[str] = None,
|
||||||
|
) -> RankManifest:
|
||||||
|
return RankManifest(
|
||||||
|
rank=session.group.rank_in_group,
|
||||||
|
keys=keys,
|
||||||
|
cached=cached,
|
||||||
|
dtype=str(session.dtype),
|
||||||
|
width=session.width,
|
||||||
|
rids=list(session.rids),
|
||||||
|
error=error,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _exchange_manifest(session: MmOwnerSession, manifest: RankManifest) -> OwnerPlan:
|
||||||
|
group = session.group
|
||||||
|
with session.uncaptured():
|
||||||
|
manifests = group.all_gather_object(manifest)
|
||||||
|
plan = _plan_or_error(session, manifests) if group.rank_in_group == 0 else None
|
||||||
|
plan = group.broadcast_object(plan, src=0)
|
||||||
|
session.phase = PHASE_FEATURES
|
||||||
|
if plan.error is not None:
|
||||||
|
raise MmOwnerProtocolError(plan.error)
|
||||||
|
return plan
|
||||||
|
|
||||||
|
|
||||||
|
def _plan_or_error(session: MmOwnerSession, manifests: List[RankManifest]) -> OwnerPlan:
|
||||||
|
try:
|
||||||
|
return _make_plan(manifests)
|
||||||
|
except Exception as exc:
|
||||||
|
return OwnerPlan(actions=[], owners=[], error=_describe(session, "plan", exc))
|
||||||
|
|
||||||
|
|
||||||
|
def _exchange_status(
|
||||||
|
session: MmOwnerSession, error: Optional[str], cause: Optional[BaseException]
|
||||||
|
) -> None:
|
||||||
|
with session.uncaptured():
|
||||||
|
statuses = session.group.all_gather_object(
|
||||||
|
RankStatus(rank=session.group.rank_in_group, error=error)
|
||||||
|
)
|
||||||
|
_raise_first_error(statuses, cause)
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_owner_features(
|
||||||
|
session: MmOwnerSession,
|
||||||
|
requests: Sequence[ImageSpanRequest],
|
||||||
|
cache: MultiModalStaticCache,
|
||||||
|
encode: SpanEncoder,
|
||||||
|
) -> Dict[SpanKey, torch.Tensor]:
|
||||||
|
group = session.group
|
||||||
|
features: Dict[SpanKey, torch.Tensor] = {}
|
||||||
|
keys: List[ImageSpanKey] = []
|
||||||
|
cached: List[bool] = []
|
||||||
|
error = None
|
||||||
|
try:
|
||||||
|
keys, cached = _pin_local_cache(session, requests, cache, features)
|
||||||
|
except Exception as exc:
|
||||||
|
error = _describe(session, "manifest", exc)
|
||||||
|
plan = _exchange_manifest(session, _manifest(session, keys, cached, error))
|
||||||
|
|
||||||
|
if all(action == LOCAL_HIT for action in plan.actions):
|
||||||
|
return features
|
||||||
|
|
||||||
|
buffers: Dict[int, torch.Tensor] = {}
|
||||||
|
error = None
|
||||||
|
try:
|
||||||
|
buffers = _prepare_transfers(session, requests, keys, plan, features, encode)
|
||||||
|
_synchronize(session.device)
|
||||||
|
except Exception as exc:
|
||||||
|
error = _describe(session, "encode", exc)
|
||||||
|
_exchange_status(session, error, None)
|
||||||
|
|
||||||
|
with session.uncaptured():
|
||||||
|
for index, (action, owner) in enumerate(zip(plan.actions, plan.owners)):
|
||||||
|
if action != LOCAL_HIT:
|
||||||
|
group.broadcast(buffers[index], src=owner)
|
||||||
|
|
||||||
|
for index, key in enumerate(keys):
|
||||||
|
if plan.actions[index] == LOCAL_HIT:
|
||||||
|
continue
|
||||||
|
span = buffers[index]
|
||||||
|
features[(key.hash, key.span_len)] = span
|
||||||
|
cache.set(key.hash, EmbeddingResult(embedding=span))
|
||||||
|
return features
|
||||||
|
|
||||||
|
|
||||||
|
def _pin_local_cache(
|
||||||
|
session: MmOwnerSession,
|
||||||
|
requests: Sequence[ImageSpanRequest],
|
||||||
|
cache: MultiModalStaticCache,
|
||||||
|
features: Dict[SpanKey, torch.Tensor],
|
||||||
|
) -> Tuple[List[ImageSpanKey], List[bool]]:
|
||||||
|
keys: List[ImageSpanKey] = []
|
||||||
|
cached: List[bool] = []
|
||||||
|
for request in requests:
|
||||||
|
if request.hash is None:
|
||||||
|
raise ValueError(
|
||||||
|
f"image span of {request.span_len} tokens has no content hash"
|
||||||
|
)
|
||||||
|
geometry = session.signature(request.item, request.span_len)
|
||||||
|
for duplicate in request.duplicates:
|
||||||
|
other = session.signature(duplicate, request.span_len)
|
||||||
|
if other != geometry:
|
||||||
|
raise ValueError(
|
||||||
|
f"image hash {request.hash} ({request.span_len} tokens) occurs "
|
||||||
|
f"with different geometry: {geometry} vs {other}"
|
||||||
|
)
|
||||||
|
keys.append(
|
||||||
|
ImageSpanKey(
|
||||||
|
hash=request.hash, span_len=request.span_len, geometry=geometry
|
||||||
|
)
|
||||||
|
)
|
||||||
|
span = _valid_cached_span(session, cache, request)
|
||||||
|
if span is not None:
|
||||||
|
features[(request.hash, request.span_len)] = span
|
||||||
|
cached.append(span is not None)
|
||||||
|
return keys, cached
|
||||||
|
|
||||||
|
|
||||||
|
def _valid_cached_span(
|
||||||
|
session: MmOwnerSession,
|
||||||
|
cache: MultiModalStaticCache,
|
||||||
|
request: ImageSpanRequest,
|
||||||
|
) -> Optional[torch.Tensor]:
|
||||||
|
entry = cache.get_single(request.hash)
|
||||||
|
if entry is None:
|
||||||
|
return None
|
||||||
|
span = entry.embedding
|
||||||
|
if (
|
||||||
|
span.dim() == 2
|
||||||
|
and span.shape[0] == request.span_len
|
||||||
|
and span.shape[1] == session.width
|
||||||
|
and span.dtype == session.dtype
|
||||||
|
and span.device == session.device
|
||||||
|
):
|
||||||
|
return span
|
||||||
|
logger.warning(
|
||||||
|
"Discarding cached multimodal embedding that cannot serve the current "
|
||||||
|
"image span: cache_key=%s expected=(%d, %d, %s) cached=(%s, %s).",
|
||||||
|
request.hash,
|
||||||
|
request.span_len,
|
||||||
|
session.width,
|
||||||
|
session.dtype,
|
||||||
|
tuple(span.shape),
|
||||||
|
span.dtype,
|
||||||
|
)
|
||||||
|
cache.free(request.hash, None)
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _make_plan(manifests: List[RankManifest]) -> OwnerPlan:
|
||||||
|
for manifest in manifests:
|
||||||
|
if manifest.error is not None:
|
||||||
|
return OwnerPlan(actions=[], owners=[], error=manifest.error)
|
||||||
|
lead = manifests[0]
|
||||||
|
for manifest in manifests[1:]:
|
||||||
|
if (manifest.keys, manifest.dtype, manifest.width, manifest.rids) != (
|
||||||
|
lead.keys,
|
||||||
|
lead.dtype,
|
||||||
|
lead.width,
|
||||||
|
lead.rids,
|
||||||
|
):
|
||||||
|
return OwnerPlan(
|
||||||
|
actions=[],
|
||||||
|
owners=[],
|
||||||
|
error=(
|
||||||
|
"image manifest mismatch between group ranks 0 and "
|
||||||
|
f"{manifest.rank}: rids={lead.rids} vs {manifest.rids}, "
|
||||||
|
f"keys={lead.keys} vs {manifest.keys}, "
|
||||||
|
f"dtype={lead.dtype} vs {manifest.dtype}, "
|
||||||
|
f"width={lead.width} vs {manifest.width}"
|
||||||
|
),
|
||||||
|
)
|
||||||
|
replication = len(manifests)
|
||||||
|
actions: List[int] = []
|
||||||
|
owners: List[int] = []
|
||||||
|
for index, key in enumerate(lead.keys):
|
||||||
|
owner = key.hash % replication
|
||||||
|
if all(manifest.cached[index] for manifest in manifests):
|
||||||
|
action = LOCAL_HIT
|
||||||
|
elif manifests[owner].cached[index]:
|
||||||
|
action = OWNER_CACHE_BROADCAST
|
||||||
|
else:
|
||||||
|
action = OWNER_ENCODE_BROADCAST
|
||||||
|
actions.append(action)
|
||||||
|
owners.append(owner)
|
||||||
|
return OwnerPlan(actions=actions, owners=owners)
|
||||||
|
|
||||||
|
|
||||||
|
def _prepare_transfers(
|
||||||
|
session: MmOwnerSession,
|
||||||
|
requests: Sequence[ImageSpanRequest],
|
||||||
|
keys: List[ImageSpanKey],
|
||||||
|
plan: OwnerPlan,
|
||||||
|
features: Dict[SpanKey, torch.Tensor],
|
||||||
|
encode: SpanEncoder,
|
||||||
|
) -> Dict[int, torch.Tensor]:
|
||||||
|
rank = session.group.rank_in_group
|
||||||
|
buffers: Dict[int, torch.Tensor] = {}
|
||||||
|
owned: List[int] = []
|
||||||
|
for index, (action, owner) in enumerate(zip(plan.actions, plan.owners)):
|
||||||
|
if action == LOCAL_HIT:
|
||||||
|
continue
|
||||||
|
if owner != rank:
|
||||||
|
try:
|
||||||
|
buffers[index] = _new_span_buffer(session, keys[index])
|
||||||
|
except Exception as exc:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"receive buffer for image hash {keys[index].hash} shape "
|
||||||
|
f"{(keys[index].span_len, session.width)} {session.dtype} "
|
||||||
|
f"failed: {type(exc).__name__}: {exc}"
|
||||||
|
) from exc
|
||||||
|
elif action == OWNER_CACHE_BROADCAST:
|
||||||
|
key = (keys[index].hash, keys[index].span_len)
|
||||||
|
buffers[index] = features[key].contiguous()
|
||||||
|
else:
|
||||||
|
owned.append(index)
|
||||||
|
if owned:
|
||||||
|
owned_hashes = [keys[index].hash for index in owned]
|
||||||
|
try:
|
||||||
|
encoded = encode([requests[index].item for index in owned])
|
||||||
|
except Exception as exc:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"owner encode of image hashes {owned_hashes} failed: "
|
||||||
|
f"{type(exc).__name__}: {exc}"
|
||||||
|
) from exc
|
||||||
|
spans = _split_spans(encoded, [keys[index].span_len for index in owned])
|
||||||
|
for index, span in zip(owned, spans):
|
||||||
|
buffers[index] = _validated_span(session, keys[index], span)
|
||||||
|
return buffers
|
||||||
|
|
||||||
|
|
||||||
|
def _new_span_buffer(session: MmOwnerSession, key: ImageSpanKey) -> torch.Tensor:
|
||||||
|
return torch.empty(
|
||||||
|
(key.span_len, session.width), device=session.device, dtype=session.dtype
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _split_spans(
|
||||||
|
encoded: torch.Tensor | List[torch.Tensor], span_lens: List[int]
|
||||||
|
) -> List[torch.Tensor]:
|
||||||
|
if isinstance(encoded, list):
|
||||||
|
if len(encoded) != len(span_lens):
|
||||||
|
raise ValueError(
|
||||||
|
f"encoder returned {len(encoded)} spans for {len(span_lens)} images"
|
||||||
|
)
|
||||||
|
return [span.reshape(-1, span.shape[-1]) for span in encoded]
|
||||||
|
encoded = encoded.reshape(-1, encoded.shape[-1])
|
||||||
|
if encoded.shape[0] != sum(span_lens):
|
||||||
|
raise ValueError(
|
||||||
|
f"encoder returned {encoded.shape[0]} rows for spans of {span_lens}"
|
||||||
|
)
|
||||||
|
return list(torch.split(encoded, span_lens, dim=0))
|
||||||
|
|
||||||
|
|
||||||
|
def _validated_span(
|
||||||
|
session: MmOwnerSession, key: ImageSpanKey, span: torch.Tensor
|
||||||
|
) -> torch.Tensor:
|
||||||
|
expected = (key.span_len, session.width)
|
||||||
|
if tuple(span.shape) != expected or span.dtype != session.dtype:
|
||||||
|
raise ValueError(
|
||||||
|
f"encoded span for hash={key.hash} has shape {tuple(span.shape)} "
|
||||||
|
f"dtype {span.dtype}; expected {expected} {session.dtype}"
|
||||||
|
)
|
||||||
|
if span.device != session.device:
|
||||||
|
span = span.to(session.device)
|
||||||
|
return span.contiguous()
|
||||||
|
|
||||||
|
|
||||||
|
def _synchronize(device) -> None:
|
||||||
|
if device.type == "cuda":
|
||||||
|
torch.cuda.current_stream(device).synchronize()
|
||||||
|
|
||||||
|
|
||||||
|
def _describe(session: MmOwnerSession, stage: str, exc: BaseException) -> str:
|
||||||
|
return (
|
||||||
|
f"multimodal owner protocol failed during {stage} on group rank "
|
||||||
|
f"{session.group.rank_in_group} (global rank "
|
||||||
|
f"{session.group.ranks[session.group.rank_in_group]}, rids={list(session.rids)}): "
|
||||||
|
f"{type(exc).__name__}: {exc}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _raise_first_error(
|
||||||
|
statuses: List[RankStatus], cause: Optional[BaseException]
|
||||||
|
) -> None:
|
||||||
|
for status in statuses:
|
||||||
|
if status.error is not None:
|
||||||
|
raise MmOwnerProtocolError(status.error) from cause
|
||||||
@@ -5,6 +5,7 @@ from typing import Callable, Dict, List, Optional, Tuple
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.managers.mm_owner_embedding import ImageSpanRequest, MmOwnerSession
|
||||||
from sglang.srt.managers.schedule_batch import MultimodalDataItem
|
from sglang.srt.managers.schedule_batch import MultimodalDataItem
|
||||||
from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
|
from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
|
||||||
from sglang.srt.multimodal.evs import EVSEmbeddingResult
|
from sglang.srt.multimodal.evs import EVSEmbeddingResult
|
||||||
@@ -339,43 +340,24 @@ def _batch_encode_per_image_misses(
|
|||||||
unique_misses: Dict[Tuple[Optional[int], int], Tuple[MultimodalDataItem, int]] = {}
|
unique_misses: Dict[Tuple[Optional[int], int], Tuple[MultimodalDataItem, int]] = {}
|
||||||
hash_to_embedding: Dict[Tuple[Optional[int], int], torch.Tensor] = {}
|
hash_to_embedding: Dict[Tuple[Optional[int], int], torch.Tensor] = {}
|
||||||
|
|
||||||
# Phase 1a: find overlapping items per request and collect cache misses
|
# Phase 1a: collect cache misses over the unique overlapping spans
|
||||||
for req_info in per_image_requests:
|
for span in _collect_image_span_requests(per_image_requests):
|
||||||
chunk_start = req_info.extend_prefix_len
|
cache_key = (span.hash, span.span_len)
|
||||||
chunk_end = chunk_start + req_info.extend_seq_len # exclusive
|
cached = embedding_cache.get_single(span.hash)
|
||||||
overlapping = []
|
if cached is not None:
|
||||||
if req_info.extend_seq_len > 0:
|
cached_embedding = cached.embedding
|
||||||
for idx, (item, (start, end)) in enumerate(
|
cached_token_count = _embedding_token_count(cached_embedding)
|
||||||
zip(req_info.items, req_info.items_offset)
|
if cached_token_count == span.span_len:
|
||||||
):
|
hash_to_embedding[cache_key] = cached_embedding
|
||||||
if end >= chunk_start and start < chunk_end:
|
|
||||||
overlapping.append((idx, item, start, end))
|
|
||||||
req_info.overlapping = overlapping
|
|
||||||
|
|
||||||
for _idx, item, start, end in overlapping:
|
|
||||||
expected_token_count = end - start + 1
|
|
||||||
cache_key = (item.hash, expected_token_count)
|
|
||||||
if cache_key in hash_to_embedding:
|
|
||||||
continue
|
continue
|
||||||
cached = embedding_cache.get_single(item.hash)
|
_discard_mismatched_cached_embedding(
|
||||||
if cached is not None:
|
span.hash, span.span_len, cached_token_count
|
||||||
cached_embedding = cached.embedding
|
)
|
||||||
cached_token_count = _embedding_token_count(cached_embedding)
|
elif (
|
||||||
if cached_token_count == expected_token_count:
|
span.inside_chunk and span.item.can_defer_cuda_ipc_feature_reconstruction()
|
||||||
hash_to_embedding[cache_key] = cached_embedding
|
):
|
||||||
else:
|
span.item.model_specific_data[BORROW_CUDA_IPC_FEATURE_KEY] = True
|
||||||
_discard_mismatched_cached_embedding(
|
unique_misses[cache_key] = (span.item, span.span_len)
|
||||||
item.hash, expected_token_count, cached_token_count
|
|
||||||
)
|
|
||||||
unique_misses[cache_key] = (item, expected_token_count)
|
|
||||||
elif cache_key not in unique_misses:
|
|
||||||
if (
|
|
||||||
start >= chunk_start
|
|
||||||
and end < chunk_end
|
|
||||||
and item.can_defer_cuda_ipc_feature_reconstruction()
|
|
||||||
):
|
|
||||||
item.model_specific_data[BORROW_CUDA_IPC_FEATURE_KEY] = True
|
|
||||||
unique_misses[cache_key] = (item, expected_token_count)
|
|
||||||
|
|
||||||
# Phase 1b: single ViT call for all unique cache misses
|
# Phase 1b: single ViT call for all unique cache misses
|
||||||
if unique_misses:
|
if unique_misses:
|
||||||
@@ -412,6 +394,52 @@ def _batch_encode_per_image_misses(
|
|||||||
return hash_to_embedding
|
return hash_to_embedding
|
||||||
|
|
||||||
|
|
||||||
|
def _collect_image_span_requests(
|
||||||
|
per_image_requests: List[PerImageRequestInfo],
|
||||||
|
) -> List[ImageSpanRequest]:
|
||||||
|
spans: Dict[
|
||||||
|
Tuple[Optional[int], int],
|
||||||
|
Tuple[MultimodalDataItem, bool, List[MultimodalDataItem]],
|
||||||
|
] = {}
|
||||||
|
for req_info in per_image_requests:
|
||||||
|
chunk_start = req_info.extend_prefix_len
|
||||||
|
chunk_end = chunk_start + req_info.extend_seq_len # exclusive
|
||||||
|
overlapping = []
|
||||||
|
if req_info.extend_seq_len > 0:
|
||||||
|
for idx, (item, (start, end)) in enumerate(
|
||||||
|
zip(req_info.items, req_info.items_offset)
|
||||||
|
):
|
||||||
|
if end >= chunk_start and start < chunk_end:
|
||||||
|
overlapping.append((idx, item, start, end))
|
||||||
|
req_info.overlapping = overlapping
|
||||||
|
|
||||||
|
for _idx, item, start, end in overlapping:
|
||||||
|
cache_key = (item.hash, end - start + 1)
|
||||||
|
if cache_key in spans:
|
||||||
|
spans[cache_key][2].append(item)
|
||||||
|
continue
|
||||||
|
spans[cache_key] = (item, start >= chunk_start and end < chunk_end, [])
|
||||||
|
return [
|
||||||
|
ImageSpanRequest(
|
||||||
|
hash=item_hash,
|
||||||
|
span_len=span_len,
|
||||||
|
item=item,
|
||||||
|
inside_chunk=inside_chunk,
|
||||||
|
duplicates=duplicates,
|
||||||
|
)
|
||||||
|
for (item_hash, span_len), (item, inside_chunk, duplicates) in spans.items()
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _owner_span_encoder(data_embedding_func: DataEmbeddingFunc, device: torch.device):
|
||||||
|
def encode(items: List[MultimodalDataItem]):
|
||||||
|
if not _can_skip_pre_embed_feature_move(data_embedding_func):
|
||||||
|
_move_items_to_device(items, device)
|
||||||
|
return data_embedding_func(items)
|
||||||
|
|
||||||
|
return encode
|
||||||
|
|
||||||
|
|
||||||
def _get_chunked_embedding_by_item(
|
def _get_chunked_embedding_by_item(
|
||||||
data_embedding_func: DataEmbeddingFunc,
|
data_embedding_func: DataEmbeddingFunc,
|
||||||
embedding_items_per_req: List[MultimodalDataItem],
|
embedding_items_per_req: List[MultimodalDataItem],
|
||||||
@@ -537,6 +565,7 @@ def _get_chunked_prefill_embedding(
|
|||||||
extend_length: List[int],
|
extend_length: List[int],
|
||||||
items_offset_list: List[List[Tuple[int, int]]],
|
items_offset_list: List[List[Tuple[int, int]]],
|
||||||
input_ids: torch.Tensor,
|
input_ids: torch.Tensor,
|
||||||
|
mm_owner: Optional[MmOwnerSession] = None,
|
||||||
) -> tuple[torch.Tensor | None, torch.Tensor]:
|
) -> tuple[torch.Tensor | None, torch.Tensor]:
|
||||||
"""
|
"""
|
||||||
Chunked prefill embedding: encode items across all requests and extract
|
Chunked prefill embedding: encode items across all requests and extract
|
||||||
@@ -598,7 +627,22 @@ def _get_chunked_prefill_embedding(
|
|||||||
|
|
||||||
# Phase 1: batch encode all per-image cache misses in ONE ViT call
|
# Phase 1: batch encode all per-image cache misses in ONE ViT call
|
||||||
hash_to_embedding: Dict[Tuple[Optional[int], int], torch.Tensor] = {}
|
hash_to_embedding: Dict[Tuple[Optional[int], int], torch.Tensor] = {}
|
||||||
if per_image_requests:
|
if per_image_requests and mm_owner is not None:
|
||||||
|
# The owner protocol must see every overlapping span before any local
|
||||||
|
# cache filtering: a rank-local hit can never skip a group collective.
|
||||||
|
span_requests = _collect_image_span_requests(per_image_requests)
|
||||||
|
if mm_owner.engaged:
|
||||||
|
hash_to_embedding = mm_owner.resolve(
|
||||||
|
span_requests,
|
||||||
|
cache=embedding_cache,
|
||||||
|
encode=_owner_span_encoder(data_embedding_func, device),
|
||||||
|
)
|
||||||
|
elif span_requests:
|
||||||
|
raise RuntimeError(
|
||||||
|
"owner eligibility saw no image span in this chunk, but "
|
||||||
|
f"scheduling found {len(span_requests)}"
|
||||||
|
)
|
||||||
|
elif per_image_requests:
|
||||||
hash_to_embedding = _batch_encode_per_image_misses(
|
hash_to_embedding = _batch_encode_per_image_misses(
|
||||||
data_embedding_func, per_image_requests, device
|
data_embedding_func, per_image_requests, device
|
||||||
)
|
)
|
||||||
@@ -701,6 +745,7 @@ def get_embedding_and_mask(
|
|||||||
prefix_length: List[int],
|
prefix_length: List[int],
|
||||||
extend_length: List[int],
|
extend_length: List[int],
|
||||||
items_offset_list: List[List[Tuple[int, int]]],
|
items_offset_list: List[List[Tuple[int, int]]],
|
||||||
|
mm_owner: Optional[MmOwnerSession] = None,
|
||||||
) -> Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]:
|
) -> Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]:
|
||||||
"""
|
"""
|
||||||
Generate multimodal embeddings and create a mask for identifying their positions in the input sequence.
|
Generate multimodal embeddings and create a mask for identifying their positions in the input sequence.
|
||||||
@@ -741,6 +786,7 @@ def get_embedding_and_mask(
|
|||||||
extend_length,
|
extend_length,
|
||||||
items_offset_list,
|
items_offset_list,
|
||||||
input_ids,
|
input_ids,
|
||||||
|
mm_owner=mm_owner,
|
||||||
)
|
)
|
||||||
if embedding is None:
|
if embedding is None:
|
||||||
return None, None, input_ids
|
return None, None, input_ids
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ import pickle
|
|||||||
import sys
|
import sys
|
||||||
from abc import abstractmethod
|
from abc import abstractmethod
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
|
from contextlib import nullcontext
|
||||||
from multiprocessing import shared_memory
|
from multiprocessing import shared_memory
|
||||||
from typing import Any, Dict, List, Optional, Tuple
|
from typing import Any, Dict, List, Optional, Tuple
|
||||||
|
|
||||||
@@ -24,6 +25,7 @@ from sglang.srt.managers.io_struct import (
|
|||||||
TokenizedEmbeddingReqInput,
|
TokenizedEmbeddingReqInput,
|
||||||
TokenizedGenerateReqInput,
|
TokenizedGenerateReqInput,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.managers.mm_owner_embedding import MmOwnerSession
|
||||||
|
|
||||||
# Preserve the existing initialization import for downstream callers.
|
# Preserve the existing initialization import for downstream callers.
|
||||||
from sglang.srt.managers.mm_schedule import (
|
from sglang.srt.managers.mm_schedule import (
|
||||||
@@ -397,6 +399,7 @@ def embed_mm_inputs(
|
|||||||
data_embedding_func_mapping: Dict[Modality, DataEmbeddingFunc] = None,
|
data_embedding_func_mapping: Dict[Modality, DataEmbeddingFunc] = None,
|
||||||
placeholder_tokens: dict[Modality, List[int]] = None,
|
placeholder_tokens: dict[Modality, List[int]] = None,
|
||||||
use_deepstack: Dict[Modality, bool] = {},
|
use_deepstack: Dict[Modality, bool] = {},
|
||||||
|
mm_owner: Optional[MmOwnerSession] = None,
|
||||||
) -> Optional[torch.Tensor]:
|
) -> Optional[torch.Tensor]:
|
||||||
"""
|
"""
|
||||||
Embed multimodal inputs and integrate them with text token embeddings.
|
Embed multimodal inputs and integrate them with text token embeddings.
|
||||||
@@ -478,6 +481,7 @@ def embed_mm_inputs(
|
|||||||
prefix_length=extend_prefix_lens,
|
prefix_length=extend_prefix_lens,
|
||||||
extend_length=extend_seq_lens,
|
extend_length=extend_seq_lens,
|
||||||
items_offset_list=items_offsets,
|
items_offset_list=items_offsets,
|
||||||
|
mm_owner=mm_owner,
|
||||||
)
|
)
|
||||||
|
|
||||||
if use_deepstack.get(modality, None) and embedding is not None:
|
if use_deepstack.get(modality, None) and embedding is not None:
|
||||||
@@ -498,7 +502,12 @@ def embed_mm_inputs(
|
|||||||
# filled with the hash values of the multimodal for the prefix matching in the radix attention.
|
# filled with the hash values of the multimodal for the prefix matching in the radix attention.
|
||||||
# There values are useless because their embeddings will be replaced by vision embeddings anyway.
|
# There values are useless because their embeddings will be replaced by vision embeddings anyway.
|
||||||
input_ids.clamp_(min=0, max=vocab_size - 1)
|
input_ids.clamp_(min=0, max=vocab_size - 1)
|
||||||
input_embeds = input_embedding(input_ids)
|
if mm_owner is not None:
|
||||||
|
# The text embedding may all-reduce across TP; a rank-local failure in
|
||||||
|
# feature preparation has to be agreed on before any rank enters it.
|
||||||
|
mm_owner.features_ready()
|
||||||
|
with mm_owner.uncaptured() if mm_owner is not None else nullcontext():
|
||||||
|
input_embeds = input_embedding(input_ids)
|
||||||
|
|
||||||
# deepstack embedding
|
# deepstack embedding
|
||||||
if use_deepstack:
|
if use_deepstack:
|
||||||
@@ -525,7 +534,9 @@ def embed_mm_inputs(
|
|||||||
_scatter_mm_embedding(dest=input_embeds, mask=mask, src=embedding)
|
_scatter_mm_embedding(dest=input_embeds, mask=mask, src=embedding)
|
||||||
if use_deepstack.get(modality, None):
|
if use_deepstack.get(modality, None):
|
||||||
_scatter_mm_embedding(
|
_scatter_mm_embedding(
|
||||||
dest=input_deepstack_embeds, mask=mask, src=deepstack_embeddings[i]
|
dest=input_deepstack_embeds,
|
||||||
|
mask=mask,
|
||||||
|
src=deepstack_embeddings[i],
|
||||||
)
|
)
|
||||||
|
|
||||||
return input_embeds, other_info
|
return input_embeds, other_info
|
||||||
|
|||||||
@@ -1019,6 +1019,15 @@ class PrefillAdder:
|
|||||||
else AddReqResult.CONTINUE
|
else AddReqResult.CONTINUE
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def can_share_extend_batch(self, req: Req) -> bool:
|
||||||
|
# Token embedding overrides embed the batch's raw input_ids before the
|
||||||
|
# model runs, and that lookup cannot index multimodal placeholder hash IDs.
|
||||||
|
if req.positional_embed_overrides is not None:
|
||||||
|
return all(r.multimodal_inputs is None for r in self.can_run_list)
|
||||||
|
if req.multimodal_inputs is not None:
|
||||||
|
return all(r.positional_embed_overrides is None for r in self.can_run_list)
|
||||||
|
return True
|
||||||
|
|
||||||
def add_chunked_req(self, req: Req):
|
def add_chunked_req(self, req: Req):
|
||||||
if self.dllm_config is not None:
|
if self.dllm_config is not None:
|
||||||
_rem_tokens = self._get_dllm_remain_tokens()
|
_rem_tokens = self._get_dllm_remain_tokens()
|
||||||
|
|||||||
@@ -2097,9 +2097,13 @@ class Scheduler(
|
|||||||
vmm_errors = self._materialize_cuda_vmm_inputs(recv_req)
|
vmm_errors = self._materialize_cuda_vmm_inputs(recv_req)
|
||||||
|
|
||||||
# Skip health check when server is busy — ongoing requests already carry health info.
|
# Skip health check when server is busy — ongoing requests already carry health info.
|
||||||
if is_health_check_generate_req(recv_req) and not self.is_fully_idle(
|
# NOTE: the admit/skip decision must be identical on every CP/TP rank.
|
||||||
for_health_check=True
|
# is_fully_idle() includes rank-local hicache drain queues, which diverge
|
||||||
):
|
# across ranks right after activity; a divergent decision lets one rank
|
||||||
|
# dispatch the health-check generate while others piggyback-skip, breaking
|
||||||
|
# collective ordering (deadlock: one rank blocks in the hicache drain
|
||||||
|
# all_reduce while another waits in the CP request broadcast).
|
||||||
|
if is_health_check_generate_req(recv_req) and not self.is_sched_idle_cp_symmetric():
|
||||||
self.return_health_check_ipcs.append(
|
self.return_health_check_ipcs.append(
|
||||||
getattr(recv_req, "http_worker_ipc", None)
|
getattr(recv_req, "http_worker_ipc", None)
|
||||||
)
|
)
|
||||||
@@ -3940,6 +3944,8 @@ class Scheduler(
|
|||||||
for req in self.waiting_queue:
|
for req in self.waiting_queue:
|
||||||
if self.enable_lora and not self.can_schedule_lora_req(req, running_loras):
|
if self.enable_lora and not self.can_schedule_lora_req(req, running_loras):
|
||||||
continue
|
continue
|
||||||
|
if not adder.can_share_extend_batch(req):
|
||||||
|
break
|
||||||
|
|
||||||
running_bs = len(running_batch.reqs)
|
running_bs = len(running_batch.reqs)
|
||||||
candidate_beam_width = (
|
candidate_beam_width = (
|
||||||
@@ -4980,6 +4986,31 @@ class Scheduler(
|
|||||||
else:
|
else:
|
||||||
self.metrics_reporter.record_scheduler_active()
|
self.metrics_reporter.record_scheduler_active()
|
||||||
|
|
||||||
|
def is_sched_idle_cp_symmetric(self) -> bool:
|
||||||
|
"""Idle check using only state that is identical across CP/TP ranks.
|
||||||
|
|
||||||
|
Request/batch/queue state is collectively maintained (requests arrive
|
||||||
|
via broadcast, batches are collectively scheduled), so every rank
|
||||||
|
computes the same result. Rank-local hicache drain and disagg transfer
|
||||||
|
queues are deliberately excluded: those are exactly the terms that
|
||||||
|
diverge across ranks and caused the CP health-check deadlock
|
||||||
|
(hicache drain all_reduce vs CP request broadcast cross-collective
|
||||||
|
wait, seen on cp2/cp4 + hicache L3 right after router health checks).
|
||||||
|
|
||||||
|
Used only for health-check admission; all other idle logic keeps using
|
||||||
|
is_fully_idle().
|
||||||
|
"""
|
||||||
|
return (
|
||||||
|
self.running_batch.is_empty()
|
||||||
|
and self.chunked_req is None
|
||||||
|
and not self.dllm_manager.any_staging_reqs()
|
||||||
|
and (self.last_batch is None or self.last_batch.is_empty())
|
||||||
|
and (not self.enable_overlap or len(self.result_queue) == 0)
|
||||||
|
and self._pp_microbatches_drained()
|
||||||
|
and len(self.waiting_queue) == 0
|
||||||
|
and len(self.grammar_manager.grammar_queue) == 0
|
||||||
|
)
|
||||||
|
|
||||||
def is_fully_idle(self, for_health_check=False) -> bool:
|
def is_fully_idle(self, for_health_check=False) -> bool:
|
||||||
# Health check piggybacks on running requests in process_output.
|
# Health check piggybacks on running requests in process_output.
|
||||||
# Only running_batch + waiting_queue guarantee active GPU processing;
|
# Only running_batch + waiting_queue guarantee active GPU processing;
|
||||||
|
|||||||
@@ -1279,6 +1279,16 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
raise ValueError(
|
raise ValueError(
|
||||||
"encoder SWA replay cannot return cached prompt logprobs"
|
"encoder SWA replay cannot return cached prompt logprobs"
|
||||||
)
|
)
|
||||||
|
requests_embed_overrides = obj.positional_embed_overrides is not None or (
|
||||||
|
isinstance(obj, EmbeddingReqInput)
|
||||||
|
and obj.embed_overrides is not None
|
||||||
|
and obj.embed_override_token_id is not None
|
||||||
|
)
|
||||||
|
if requests_embed_overrides and obj.contains_mm_input():
|
||||||
|
raise ValueError(
|
||||||
|
"embedding overrides cannot be combined with image, video, or audio "
|
||||||
|
"inputs"
|
||||||
|
)
|
||||||
_max_req_len = self.context_len
|
_max_req_len = self.context_len
|
||||||
input_token_num = len(input_ids) if input_ids is not None else 0
|
input_token_num = len(input_ids) if input_ids is not None else 0
|
||||||
input_token_num += self.num_reserved_tokens
|
input_token_num += self.num_reserved_tokens
|
||||||
|
|||||||
@@ -155,6 +155,10 @@ def free_kv_row_segments(
|
|||||||
|
|
||||||
def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
|
def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
|
||||||
if getattr(req, "skip_radix_cache_insert", False):
|
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
|
return
|
||||||
|
|
||||||
tree_cache.cache_unfinished_req(req, **kwargs)
|
tree_cache.cache_unfinished_req(req, **kwargs)
|
||||||
|
|||||||
@@ -1646,6 +1646,7 @@ class ModelRunner:
|
|||||||
forward_batch.replace_embeds is not None
|
forward_batch.replace_embeds is not None
|
||||||
and forward_batch.replace_positions is not None
|
and forward_batch.replace_positions is not None
|
||||||
):
|
):
|
||||||
|
misc_utils.validate_replace_embeds_batch(forward_batch)
|
||||||
# Token embedding overrides: get base embeddings, scatter replacements
|
# Token embedding overrides: get base embeddings, scatter replacements
|
||||||
if "input_embeds" not in kwargs:
|
if "input_embeds" not in kwargs:
|
||||||
embed_layer = self.model.get_input_embeddings()
|
embed_layer = self.model.get_input_embeddings()
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ from sglang.srt.server_args import CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACK
|
|||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.configs.model_config import ModelConfig
|
from sglang.srt.configs.model_config import ModelConfig
|
||||||
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -105,3 +106,24 @@ def resolve_pp_proxy_dspark_hidden_size(
|
|||||||
if isinstance(model, _SupportsDSparkPPProxy):
|
if isinstance(model, _SupportsDSparkPPProxy):
|
||||||
return model.get_pp_proxy_dspark_hidden_size()
|
return model.get_pp_proxy_dspark_hidden_size()
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
def validate_replace_embeds_batch(forward_batch: ForwardBatch) -> None:
|
||||||
|
if forward_batch.mm_inputs is None:
|
||||||
|
return
|
||||||
|
for mm_inputs, prefix_len, extend_len in zip(
|
||||||
|
forward_batch.mm_inputs,
|
||||||
|
forward_batch.extend_prefix_lens_cpu,
|
||||||
|
forward_batch.extend_seq_lens_cpu,
|
||||||
|
):
|
||||||
|
if mm_inputs is None:
|
||||||
|
continue
|
||||||
|
chunk_end = prefix_len + extend_len
|
||||||
|
for item in mm_inputs.mm_items:
|
||||||
|
for start, end in item.offsets or ():
|
||||||
|
if start < chunk_end and end >= prefix_len:
|
||||||
|
# Placeholder rows carry hash IDs the base embedding lookup cannot index.
|
||||||
|
raise ValueError(
|
||||||
|
"Token embedding overrides cannot share an extend batch with "
|
||||||
|
"multimodal placeholders"
|
||||||
|
)
|
||||||
|
|||||||
@@ -385,14 +385,23 @@ class EagerRunner(BaseRunner):
|
|||||||
"""
|
"""
|
||||||
model = self.model_runner.model
|
model = self.model_runner.model
|
||||||
|
|
||||||
|
input_ids = forward_batch.input_ids
|
||||||
input_embeds = kwargs.get("input_embeds")
|
input_embeds = kwargs.get("input_embeds")
|
||||||
|
if hasattr(model, "prepare_model_inputs"):
|
||||||
|
# Multimodal offsets are request-global, so the merge and the
|
||||||
|
# placeholder-ID remap must see the full extend layout first.
|
||||||
|
input_ids, input_embeds = model.prepare_model_inputs(
|
||||||
|
input_ids=input_ids,
|
||||||
|
forward_batch=forward_batch,
|
||||||
|
input_embeds=input_embeds,
|
||||||
|
)
|
||||||
if input_embeds is None:
|
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(
|
with cp_shard_model_inputs(
|
||||||
input_embeds,
|
input_embeds,
|
||||||
forward_batch.positions,
|
forward_batch.positions,
|
||||||
forward_batch,
|
forward_batch,
|
||||||
forward_batch.input_ids,
|
input_ids,
|
||||||
) as (sharded_input_embeds, sharded_positions, model_input_ids):
|
) as (sharded_input_embeds, sharded_positions, model_input_ids):
|
||||||
model_kwargs = {"input_embeds": sharded_input_embeds}
|
model_kwargs = {"input_embeds": sharded_input_embeds}
|
||||||
if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None:
|
if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None:
|
||||||
@@ -437,7 +446,7 @@ class EagerRunner(BaseRunner):
|
|||||||
if aux_hidden_states is None:
|
if aux_hidden_states is None:
|
||||||
logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
|
logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
|
||||||
return model.logits_processor(
|
return model.logits_processor(
|
||||||
forward_batch.input_ids,
|
input_ids,
|
||||||
hidden_states,
|
hidden_states,
|
||||||
model.lm_head,
|
model.lm_head,
|
||||||
forward_batch,
|
forward_batch,
|
||||||
|
|||||||
@@ -701,6 +701,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
|
|
||||||
def _get_layer_model_positions(self, forward_batch: ForwardBatch) -> torch.Tensor:
|
def _get_layer_model_positions(self, forward_batch: ForwardBatch) -> torch.Tensor:
|
||||||
"""Mirror outer multimodal wrappers when BCG captures layer_model directly."""
|
"""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:
|
if forward_batch.mrope_positions is None:
|
||||||
return forward_batch.positions
|
return forward_batch.positions
|
||||||
|
|
||||||
@@ -782,7 +785,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
if self._uses_eager_prefill_tail():
|
if self._uses_eager_prefill_tail():
|
||||||
# BCG / Full: capture the transformer body only.
|
# BCG / Full: capture the transformer body only.
|
||||||
positions = self._get_layer_model_positions(forward_batch)
|
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(
|
kwargs = _build_layer_model_forward_kwargs(
|
||||||
self.layer_model, forward_batch, pp_proxy_tensors
|
self.layer_model, forward_batch, pp_proxy_tensors
|
||||||
)
|
)
|
||||||
@@ -1336,9 +1341,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
batch_max_context_len=batch_max_context_len,
|
batch_max_context_len=batch_max_context_len,
|
||||||
):
|
):
|
||||||
return False
|
return False
|
||||||
if getattr(self, "enable_cp_bcg_capture", False) and is_cp_active(
|
if getattr(self, "enable_cp_bcg_capture", False):
|
||||||
forward_batch
|
if not is_cp_active(forward_batch):
|
||||||
):
|
return False
|
||||||
assert self.prefill_cp_bcg_input is not None
|
assert self.prefill_cp_bcg_input is not None
|
||||||
if (
|
if (
|
||||||
self.prefill_cp_bcg_input.select_replay_bucket_for_batch(
|
self.prefill_cp_bcg_input.select_replay_bucket_for_batch(
|
||||||
|
|||||||
@@ -74,6 +74,7 @@ from sglang.srt.layers.communicator_dsa_cp import (
|
|||||||
dsa_cp_gather_hidden_states,
|
dsa_cp_gather_hidden_states,
|
||||||
dsa_cp_reduce_scatter_hidden_states,
|
dsa_cp_reduce_scatter_hidden_states,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.layers.cp.base import is_zigzag
|
||||||
from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
|
from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
|
||||||
from sglang.srt.layers.cp.utils import (
|
from sglang.srt.layers.cp.utils import (
|
||||||
cp_gather_full_sequence_states,
|
cp_gather_full_sequence_states,
|
||||||
@@ -120,6 +121,11 @@ from sglang.srt.layers.quantization.mxfp8_input import Mxfp8SwizzledInput
|
|||||||
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
|
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
|
||||||
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
|
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
|
||||||
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
||||||
|
from sglang.srt.managers.mm_owner_embedding import (
|
||||||
|
MmOwnerSession,
|
||||||
|
has_owner_span_work,
|
||||||
|
select_owner_group,
|
||||||
|
)
|
||||||
from sglang.srt.managers.mm_utils import (
|
from sglang.srt.managers.mm_utils import (
|
||||||
MultiModalityDataPaddingPatternMultimodalTokens,
|
MultiModalityDataPaddingPatternMultimodalTokens,
|
||||||
embed_mm_inputs,
|
embed_mm_inputs,
|
||||||
@@ -182,6 +188,7 @@ from sglang.srt.multimodal.deepseek_v41_image_processing import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.runtime_context import (
|
from sglang.srt.runtime_context import (
|
||||||
get_device,
|
get_device,
|
||||||
|
get_disagg,
|
||||||
get_exec,
|
get_exec,
|
||||||
get_forward,
|
get_forward,
|
||||||
get_parallel,
|
get_parallel,
|
||||||
@@ -2039,7 +2046,10 @@ class MQALayer(MqaAttentionBase):
|
|||||||
if (
|
if (
|
||||||
forward_batch.forward_mode.is_extend()
|
forward_batch.forward_mode.is_extend()
|
||||||
and is_in_breakable_cuda_graph()
|
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)
|
bcg_deepseek_v4_low_ratio_sources(self, x, q_lora, positions)
|
||||||
else:
|
else:
|
||||||
@@ -2650,7 +2660,8 @@ class DeepseekV4DecoderLayer(nn.Module):
|
|||||||
is_nextn=is_nextn,
|
is_nextn=is_nextn,
|
||||||
is_deepseek_v4=True,
|
is_deepseek_v4=True,
|
||||||
vl_correction_bias=config.model_type == "deepseek_v41"
|
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)
|
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||||
@@ -3872,7 +3883,15 @@ class DeepseekV4DecoderLayer(nn.Module):
|
|||||||
finally:
|
finally:
|
||||||
forward_batch.num_token_non_padded = saved_num_token_non_padded
|
forward_batch.num_token_non_padded = saved_num_token_non_padded
|
||||||
if _use_cp and get_moe_a2a_backend().is_none():
|
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:
|
elif _use_tp_moe_gather:
|
||||||
hidden_states, global_hidden_states = (
|
hidden_states, global_hidden_states = (
|
||||||
get_local_dp_buffer(get_parallel().tp_group),
|
get_local_dp_buffer(get_parallel().tp_group),
|
||||||
@@ -4402,11 +4421,18 @@ class DeepseekV4Model(nn.Module):
|
|||||||
)
|
)
|
||||||
if self.engram_hasher is not None:
|
if self.engram_hasher is not None:
|
||||||
if cp_extend:
|
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)
|
total = int(forward_batch.attn_cp_metadata.total_seq_lens)
|
||||||
hash_ids = self.engram_hasher(
|
global_input_ids = forward_batch.input_ids[:total]
|
||||||
forward_batch.input_ids[:total], forward_batch
|
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()
|
parallel = get_parallel()
|
||||||
hash_ids = hash_ids[parallel.attn_cp_rank :: parallel.attn_cp_size]
|
hash_ids = hash_ids[parallel.attn_cp_rank :: parallel.attn_cp_size]
|
||||||
pad_rows = hidden_states.shape[0] - hash_ids.shape[0]
|
pad_rows = hidden_states.shape[0] - hash_ids.shape[0]
|
||||||
@@ -4839,6 +4865,13 @@ class DeepseekV4Model(nn.Module):
|
|||||||
return hidden_states, pre_hc_head
|
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):
|
class DeepseekV4ForCausalLM(nn.Module):
|
||||||
supports_cuda_vmm_feature_transport = True
|
supports_cuda_vmm_feature_transport = True
|
||||||
|
|
||||||
@@ -4864,14 +4897,23 @@ class DeepseekV4ForCausalLM(nn.Module):
|
|||||||
self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config)
|
self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config)
|
||||||
self.determine_num_fused_shared_experts()
|
self.determine_num_fused_shared_experts()
|
||||||
self.vision = None
|
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 (
|
if (
|
||||||
get_parallel().attn_cp_size != 1
|
get_parallel().pp_group.world_size != 1
|
||||||
or get_parallel().pp_group.world_size != 1
|
or not _v41_vision_a2a_supported()
|
||||||
or not get_moe_a2a_backend().is_none()
|
|
||||||
):
|
):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"V4.1 vision currently supports TP/EP/DP without CP, PP or MoE A2A"
|
"V4.1 vision supports TP/EP/DP without PP; "
|
||||||
|
"MoE A2A is supported only with MegaMoE on a PD decode node"
|
||||||
|
)
|
||||||
|
if get_parallel().attn_cp_size != 1 and (_is_npu or is_zigzag()):
|
||||||
|
raise ValueError(
|
||||||
|
"V4.1 vision context parallelism requires the CUDA interleave "
|
||||||
|
"strategy; NPU and zigzag CP are not supported yet"
|
||||||
)
|
)
|
||||||
|
|
||||||
args = SimpleNamespace(**vars(config), dim=config.hidden_size)
|
args = SimpleNamespace(**vars(config), dim=config.hidden_size)
|
||||||
@@ -4880,6 +4922,11 @@ class DeepseekV4ForCausalLM(nn.Module):
|
|||||||
self.image_start = nn.Parameter(torch.empty(config.hidden_size))
|
self.image_start = nn.Parameter(torch.empty(config.hidden_size))
|
||||||
self.image_end = nn.Parameter(torch.empty(config.hidden_size))
|
self.image_end = nn.Parameter(torch.empty(config.hidden_size))
|
||||||
self.image_newline = nn.Parameter(torch.empty(config.hidden_size))
|
self.image_newline = nn.Parameter(torch.empty(config.hidden_size))
|
||||||
|
self.mm_owner_group = (
|
||||||
|
select_owner_group(get_parallel())
|
||||||
|
if self.vision is not None and _is_cuda
|
||||||
|
else None
|
||||||
|
)
|
||||||
self.model = DeepseekV4Model(
|
self.model = DeepseekV4Model(
|
||||||
config, quant_config, prefix=add_prefix("model", prefix)
|
config, quant_config, prefix=add_prefix("model", prefix)
|
||||||
)
|
)
|
||||||
@@ -5014,7 +5061,42 @@ class DeepseekV4ForCausalLM(nn.Module):
|
|||||||
spans.append(span)
|
spans.append(span)
|
||||||
return spans
|
return spans
|
||||||
|
|
||||||
def _prepare_mm_embeddings(self, input_ids, forward_batch):
|
def _image_span_signature(self, item, span_len: int):
|
||||||
|
h, w = int(item.n_vit_h), int(item.n_vit_w)
|
||||||
|
r = self.config.vision_downsample_ratio
|
||||||
|
expected = len(image_token_types((h + r - 1) // r, (w + r - 1) // r))
|
||||||
|
if expected != span_len:
|
||||||
|
raise ValueError(
|
||||||
|
f"image grid {(h, w)} yields {expected} span tokens, "
|
||||||
|
f"placeholder has {span_len}"
|
||||||
|
)
|
||||||
|
plan = item.model_specific_data.get(GPU_PLAN_KEY)
|
||||||
|
feature = item.feature
|
||||||
|
return (
|
||||||
|
h,
|
||||||
|
w,
|
||||||
|
tuple(feature.shape) if isinstance(feature, torch.Tensor) else None,
|
||||||
|
None if plan is None else tuple(sorted(plan.items())),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _mm_owner_session(self, forward_batch) -> Optional[MmOwnerSession]:
|
||||||
|
if self.mm_owner_group is None:
|
||||||
|
return None
|
||||||
|
return MmOwnerSession(
|
||||||
|
group=self.mm_owner_group,
|
||||||
|
device=self.image_start.device,
|
||||||
|
dtype=self.image_start.dtype,
|
||||||
|
width=self.config.hidden_size,
|
||||||
|
rids=list(forward_batch.rids or ()),
|
||||||
|
signature=self._image_span_signature,
|
||||||
|
engaged=has_owner_span_work(
|
||||||
|
forward_batch.mm_inputs,
|
||||||
|
forward_batch.extend_prefix_lens_cpu,
|
||||||
|
forward_batch.extend_seq_lens_cpu,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _prepare_mm_embeddings(self, input_ids, forward_batch, mm_owner):
|
||||||
# Keep scheduler hash IDs intact: the shared embedder clamps its input in place.
|
# Keep scheduler hash IDs intact: the shared embedder clamps its input in place.
|
||||||
input_embeds, _ = embed_mm_inputs(
|
input_embeds, _ = embed_mm_inputs(
|
||||||
mm_inputs_list=[
|
mm_inputs_list=[
|
||||||
@@ -5026,6 +5108,7 @@ class DeepseekV4ForCausalLM(nn.Module):
|
|||||||
input_ids=input_ids.clone(),
|
input_ids=input_ids.clone(),
|
||||||
input_embedding=self.get_input_embeddings(),
|
input_embedding=self.get_input_embeddings(),
|
||||||
multimodal_model=self,
|
multimodal_model=self,
|
||||||
|
mm_owner=mm_owner,
|
||||||
)
|
)
|
||||||
forward_batch.mm_input_embeds = input_embeds
|
forward_batch.mm_input_embeds = input_embeds
|
||||||
return input_embeds
|
return input_embeds
|
||||||
@@ -5033,6 +5116,41 @@ class DeepseekV4ForCausalLM(nn.Module):
|
|||||||
def get_input_embeddings(self) -> nn.Module:
|
def get_input_embeddings(self) -> nn.Module:
|
||||||
return self.model.get_input_embeddings()
|
return self.model.get_input_embeddings()
|
||||||
|
|
||||||
|
def prepare_model_inputs(
|
||||||
|
self,
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
input_embeds: Optional[torch.Tensor],
|
||||||
|
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
|
||||||
|
if self.vision is None:
|
||||||
|
return input_ids, input_embeds
|
||||||
|
has_images = (
|
||||||
|
not forward_batch.forward_mode.is_decode()
|
||||||
|
and not forward_batch.forward_mode.is_target_verify()
|
||||||
|
and forward_batch.mm_inputs is not None
|
||||||
|
and any(x is not None for x in forward_batch.mm_inputs)
|
||||||
|
)
|
||||||
|
if has_images and input_embeds is not None:
|
||||||
|
raise ValueError("Cannot combine input_embeds and image inputs")
|
||||||
|
mm_owner = self._mm_owner_session(forward_batch) if has_images else None
|
||||||
|
# Peers may only enter the body or the CP shard once every rank has
|
||||||
|
# finished all of its fallible input preparation, the remap included.
|
||||||
|
with mm_owner.fence() if mm_owner is not None else nullcontext():
|
||||||
|
if has_images:
|
||||||
|
input_embeds = self._prepare_mm_embeddings(
|
||||||
|
input_ids, forward_batch, mm_owner
|
||||||
|
)
|
||||||
|
if not (
|
||||||
|
forward_batch.forward_mode.is_decode_or_idle()
|
||||||
|
or forward_batch.forward_mode.is_target_verify()
|
||||||
|
):
|
||||||
|
# Decode/verify IDs are already vocabulary IDs; remap prompt image
|
||||||
|
# hashes for Engram and routing.
|
||||||
|
input_ids = input_ids.masked_fill(
|
||||||
|
input_ids >= MM_PAD_SHIFT_VALUE, self.config.image_token_id
|
||||||
|
)
|
||||||
|
return input_ids, input_embeds
|
||||||
|
|
||||||
def set_dspark_layers_to_capture(self, layer_ids: List[int]) -> None:
|
def set_dspark_layers_to_capture(self, layer_ids: List[int]) -> None:
|
||||||
if not self.pp_group.is_last_rank:
|
if not self.pp_group.is_last_rank:
|
||||||
return
|
return
|
||||||
@@ -5078,6 +5196,19 @@ class DeepseekV4ForCausalLM(nn.Module):
|
|||||||
0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
|
0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def prepare_language_model_inputs(
|
||||||
|
self,
|
||||||
|
input_ids: 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_model_inputs(
|
||||||
|
input_ids=input_ids, forward_batch=forward_batch, input_embeds=input_embeds
|
||||||
|
)
|
||||||
|
|
||||||
|
return input_ids, input_embeds
|
||||||
|
|
||||||
def forward(
|
def forward(
|
||||||
self,
|
self,
|
||||||
input_ids: torch.Tensor,
|
input_ids: torch.Tensor,
|
||||||
@@ -5086,26 +5217,9 @@ class DeepseekV4ForCausalLM(nn.Module):
|
|||||||
input_embeds: Optional[torch.Tensor] = None,
|
input_embeds: Optional[torch.Tensor] = None,
|
||||||
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
if (
|
input_ids, input_embeds = self.prepare_language_model_inputs(
|
||||||
self.vision is not None
|
input_ids, forward_batch, input_embeds
|
||||||
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
|
|
||||||
and any(x is not None for x in forward_batch.mm_inputs)
|
|
||||||
):
|
|
||||||
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 (
|
|
||||||
forward_batch.forward_mode.is_decode_or_idle()
|
|
||||||
or forward_batch.forward_mode.is_target_verify()
|
|
||||||
):
|
|
||||||
# Decode/verify IDs are already vocabulary IDs; remap prompt image
|
|
||||||
# hashes for Engram and routing.
|
|
||||||
input_ids = input_ids.masked_fill(
|
|
||||||
input_ids >= MM_PAD_SHIFT_VALUE, self.config.image_token_id
|
|
||||||
)
|
|
||||||
|
|
||||||
with get_attn_tp_context().maybe_input_scattered(forward_batch):
|
with get_attn_tp_context().maybe_input_scattered(forward_batch):
|
||||||
hidden_states = self.model.forward(
|
hidden_states = self.model.forward(
|
||||||
input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
|
input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
|
||||||
|
|||||||
@@ -222,6 +222,7 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
|
|||||||
self.quant_config = quant_config
|
self.quant_config = quant_config
|
||||||
self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config)
|
self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config)
|
||||||
self.determine_num_fused_shared_experts()
|
self.determine_num_fused_shared_experts()
|
||||||
|
self.vision = None
|
||||||
|
|
||||||
self.model = DeepseekV4ModelNextN(
|
self.model = DeepseekV4ModelNextN(
|
||||||
config, quant_config, prefix=add_prefix("model", prefix)
|
config, quant_config, prefix=add_prefix("model", prefix)
|
||||||
|
|||||||
@@ -376,6 +376,7 @@ class ServerArgs:
|
|||||||
# ===== END TO BE REFACTORED ====
|
# ===== END TO BE REFACTORED ====
|
||||||
|
|
||||||
LANGUAGE_MODEL_ONLY_ARCHITECTURES = (
|
LANGUAGE_MODEL_ONLY_ARCHITECTURES = (
|
||||||
|
"DeepseekV4ForCausalLM",
|
||||||
"MuseGlimmerForConditionalGeneration",
|
"MuseGlimmerForConditionalGeneration",
|
||||||
"Cosmos3ForConditionalGeneration",
|
"Cosmos3ForConditionalGeneration",
|
||||||
"Cosmos3EdgeForConditionalGeneration",
|
"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()
|
||||||
@@ -752,6 +752,33 @@ class TestPagedIndexerMetadataChunking(CustomTestCase):
|
|||||||
row chunks the indexer loops over; a mismatch would silently score rows
|
row chunks the indexer loops over; a mismatch would silently score rows
|
||||||
with another chunk's schedule."""
|
with another chunk's schedule."""
|
||||||
|
|
||||||
|
def test_capture_warmup_skips_dynamic_budget_but_eager_forward_uses_it(self):
|
||||||
|
metadata = SimpleNamespace(
|
||||||
|
use_prefill_cuda_graph=False,
|
||||||
|
compressed_seq_lens=SimpleNamespace(
|
||||||
|
is_cuda=True, device=SimpleNamespace(index=0)
|
||||||
|
),
|
||||||
|
max_compressed_seq_len=65536,
|
||||||
|
)
|
||||||
|
for capture_mode in (True, False):
|
||||||
|
with (
|
||||||
|
self.subTest(capture_mode=capture_mode),
|
||||||
|
patch(f"{_METADATA}.get_is_capture_mode", return_value=capture_mode),
|
||||||
|
patch("torch.cuda.is_current_stream_capturing", return_value=False),
|
||||||
|
patch(f"{_METADATA}.is_in_breakable_cuda_graph", return_value=False),
|
||||||
|
patch(f"{_METADATA}.is_in_tc_piecewise_cuda_graph", return_value=False),
|
||||||
|
patch(
|
||||||
|
f"{_METADATA}.mqa_logits_budget_bytes", return_value=4096
|
||||||
|
) as budget,
|
||||||
|
):
|
||||||
|
result = PagedIndexerMetadata._mqa_logits_budget(metadata, num_rows=256)
|
||||||
|
if capture_mode:
|
||||||
|
self.assertIsNone(result)
|
||||||
|
budget.assert_not_called()
|
||||||
|
else:
|
||||||
|
self.assertEqual(result, 4096)
|
||||||
|
budget.assert_called_once_with(device_index=0, allow_sync=True)
|
||||||
|
|
||||||
def _build(self, *, num_rows: int, budget, use_topk_v2: bool):
|
def _build(self, *, num_rows: int, budget, use_topk_v2: bool):
|
||||||
deep_gemm = SimpleNamespace(
|
deep_gemm = SimpleNamespace(
|
||||||
get_num_sms=MagicMock(return_value=1),
|
get_num_sms=MagicMock(return_value=1),
|
||||||
@@ -806,6 +833,12 @@ class TestPagedIndexerMetadataChunking(CustomTestCase):
|
|||||||
|
|
||||||
self.assertIsInstance(metadata.deep_gemm_metadata, list)
|
self.assertIsInstance(metadata.deep_gemm_metadata, list)
|
||||||
self.assertEqual(len(metadata.deep_gemm_metadata), len(chunks))
|
self.assertEqual(len(metadata.deep_gemm_metadata), len(chunks))
|
||||||
|
metadata_chunks = metadata.row_chunks()
|
||||||
|
self.assertEqual([rows for rows, _ in metadata_chunks], chunks)
|
||||||
|
for (_, actual_plan), expected_plan in zip(
|
||||||
|
metadata_chunks, metadata.deep_gemm_metadata
|
||||||
|
):
|
||||||
|
self.assertIs(actual_plan, expected_plan)
|
||||||
schedule_rows = [
|
schedule_rows = [
|
||||||
call.args[0]
|
call.args[0]
|
||||||
for call in deep_gemm.get_paged_mqa_logits_metadata.call_args_list
|
for call in deep_gemm.get_paged_mqa_logits_metadata.call_args_list
|
||||||
@@ -875,6 +908,92 @@ class TestChunkedTopKMatchesUnchunked(CustomTestCase):
|
|||||||
self.assertTrue(torch.equal(run(rows_per_chunk), expected))
|
self.assertTrue(torch.equal(run(rows_per_chunk), expected))
|
||||||
|
|
||||||
|
|
||||||
|
class TestChunkedCandidatePublisher(CustomTestCase):
|
||||||
|
def test_each_deep_gemm_call_receives_one_tensor_schedule(self):
|
||||||
|
from sglang.srt.layers.attention.dsv4 import candidate_indexer_deep_gemm as mod
|
||||||
|
|
||||||
|
num_rows, width = 5, 16
|
||||||
|
chunks = [slice(0, 2), slice(2, 4), slice(4, 5)]
|
||||||
|
plans = [torch.tensor([i], dtype=torch.uint8) for i in range(len(chunks))]
|
||||||
|
topk_plans = [torch.tensor([i], dtype=torch.int32) for i in range(len(chunks))]
|
||||||
|
metadata = SimpleNamespace(
|
||||||
|
compressed_seq_lens=torch.full((num_rows, 1), width, dtype=torch.int32),
|
||||||
|
page_table=torch.zeros((num_rows, 1), dtype=torch.int32),
|
||||||
|
deep_gemm_metadata=plans,
|
||||||
|
max_compressed_seq_len=width,
|
||||||
|
compressed_page_size=64,
|
||||||
|
topk_metadata_chunks=topk_plans,
|
||||||
|
use_topk_v2=True,
|
||||||
|
row_chunks=lambda: list(zip(chunks, plans)),
|
||||||
|
)
|
||||||
|
inputs = SimpleNamespace(
|
||||||
|
q_fp4=torch.zeros((num_rows, 1, 2, 64), dtype=torch.int8),
|
||||||
|
q_sf=torch.zeros((num_rows, 1, 2), dtype=torch.int32),
|
||||||
|
k_cache=torch.zeros((1, 64, 1, 68), dtype=torch.uint8),
|
||||||
|
weights=torch.zeros((num_rows, 2), dtype=torch.float32),
|
||||||
|
metadata=metadata,
|
||||||
|
request_ids=torch.arange(num_rows),
|
||||||
|
num_rows=num_rows,
|
||||||
|
)
|
||||||
|
page_indices = torch.full((num_rows, 4), -1, dtype=torch.int32)
|
||||||
|
raw_indices = torch.full_like(page_indices, -1)
|
||||||
|
indexer = object.__new__(mod.DeepGemmCandidateIndexer)
|
||||||
|
indexer.topk_blocks = 2
|
||||||
|
|
||||||
|
deep_gemm = MagicMock(
|
||||||
|
side_effect=lambda q, *_args: torch.zeros(
|
||||||
|
(q[0].shape[0], width), dtype=torch.float32
|
||||||
|
)
|
||||||
|
)
|
||||||
|
topk = MagicMock()
|
||||||
|
event = MagicMock()
|
||||||
|
stream = MagicMock()
|
||||||
|
with (
|
||||||
|
patch.object(mod, "deep_gemm_fp4_paged_mqa_logits", deep_gemm),
|
||||||
|
patch.object(mod, "topk_transform_paged_from_metadata", topk),
|
||||||
|
patch.object(
|
||||||
|
mod,
|
||||||
|
"candidate_row_lens",
|
||||||
|
side_effect=lambda lens, _topk: (
|
||||||
|
torch.ones_like(lens),
|
||||||
|
lens.clone(),
|
||||||
|
),
|
||||||
|
),
|
||||||
|
patch.object(
|
||||||
|
mod,
|
||||||
|
"amax_topk_blocks",
|
||||||
|
side_effect=lambda _logits, lens, _nblocks, topk_blocks: torch.zeros(
|
||||||
|
(lens.shape[0], topk_blocks), dtype=torch.int32
|
||||||
|
),
|
||||||
|
),
|
||||||
|
patch.object(
|
||||||
|
mod,
|
||||||
|
"sort_candidate_blocks",
|
||||||
|
side_effect=lambda blocks, *_args: blocks + 1,
|
||||||
|
),
|
||||||
|
patch.object(
|
||||||
|
mod,
|
||||||
|
"build_sparse_indexer_schedule",
|
||||||
|
return_value=torch.tensor([7], dtype=torch.uint8),
|
||||||
|
),
|
||||||
|
patch.object(mod.torch.cuda, "Event", return_value=event),
|
||||||
|
patch.object(mod.torch.cuda, "current_stream", return_value=stream),
|
||||||
|
):
|
||||||
|
table = indexer.publish_decode(inputs, page_indices, raw_indices)
|
||||||
|
|
||||||
|
self.assertEqual(deep_gemm.call_count, len(chunks))
|
||||||
|
for call, plan in zip(deep_gemm.call_args_list, plans):
|
||||||
|
self.assertIs(call.args[5], plan)
|
||||||
|
self.assertIsInstance(call.args[5], torch.Tensor)
|
||||||
|
self.assertEqual([call.kwargs["rows"] for call in topk.call_args_list], chunks)
|
||||||
|
for call, plan in zip(topk.call_args_list, topk_plans):
|
||||||
|
self.assertIs(call.kwargs["topk_metadata"], plan)
|
||||||
|
self.assertEqual(table.blocks.shape, (num_rows, indexer.topk_blocks))
|
||||||
|
self.assertEqual(table.phys_blocks.shape, table.blocks.shape)
|
||||||
|
self.assertEqual(table.valid_lens.shape, (num_rows,))
|
||||||
|
event.record.assert_called_once_with(stream)
|
||||||
|
|
||||||
|
|
||||||
class TestCandidateIndexerGating(CustomTestCase):
|
class TestCandidateIndexerGating(CustomTestCase):
|
||||||
def test_candidate_indexer_gating(self):
|
def test_candidate_indexer_gating(self):
|
||||||
from sglang.srt.layers.attention.dsv4 import candidate_indexer
|
from sglang.srt.layers.attention.dsv4 import candidate_indexer
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ Covers:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import unittest
|
import unittest
|
||||||
|
from types import SimpleNamespace
|
||||||
from unittest.mock import AsyncMock, MagicMock
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
@@ -17,10 +18,16 @@ from sglang.srt.constants import MIS_DELIMITER_TOKEN_ID
|
|||||||
from sglang.srt.entrypoints.openai.utils import convert_embeds_to_tensors
|
from sglang.srt.entrypoints.openai.utils import convert_embeds_to_tensors
|
||||||
from sglang.srt.managers.embed_types import PositionalEmbeds
|
from sglang.srt.managers.embed_types import PositionalEmbeds
|
||||||
from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
|
from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
|
||||||
|
from sglang.srt.managers.schedule_batch import (
|
||||||
|
Modality,
|
||||||
|
MultimodalDataItem,
|
||||||
|
MultimodalInputs,
|
||||||
|
)
|
||||||
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
||||||
from sglang.srt.managers.tokenizer_manager_score_mixin import (
|
from sglang.srt.managers.tokenizer_manager_score_mixin import (
|
||||||
TokenizerManagerScoreMixin,
|
TokenizerManagerScoreMixin,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||||
from sglang.srt.runtime_context import publish, reset_context
|
from sglang.srt.runtime_context import publish, reset_context
|
||||||
from sglang.srt.server_args import ServerArgs
|
from sglang.srt.server_args import ServerArgs
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
@@ -642,5 +649,87 @@ class TestScoreRequestValidation(CustomTestCase):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestEmbedOverridesRejectMultimodal(CustomTestCase):
|
||||||
|
def setUp(self):
|
||||||
|
reset_context()
|
||||||
|
self.addCleanup(reset_context)
|
||||||
|
publish(ServerArgs(model_path="dummy"), role="tokenizer")
|
||||||
|
self.manager = TokenizerManager.__new__(TokenizerManager)
|
||||||
|
self.manager.context_len = 128
|
||||||
|
self.manager.num_reserved_tokens = 0
|
||||||
|
self.manager.allow_auto_truncate = False
|
||||||
|
self.manager.validate_total_tokens = False
|
||||||
|
self.manager.is_generation = True
|
||||||
|
|
||||||
|
def _request(self, **fields):
|
||||||
|
return GenerateReqInput(
|
||||||
|
input_ids=[10, 50, 20],
|
||||||
|
sampling_params={},
|
||||||
|
positional_embed_overrides=PositionalEmbeds(embeds=[_vec()], positions=[1]),
|
||||||
|
**fields,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_request_with_image_is_rejected(self):
|
||||||
|
req = self._request(image_data=["image.png"])
|
||||||
|
with self.assertRaisesRegex(ValueError, "overrides cannot be combined"):
|
||||||
|
self.manager._validate_one_request(req, req.input_ids)
|
||||||
|
text_only = self._request()
|
||||||
|
self.manager._validate_one_request(text_only, text_only.input_ids)
|
||||||
|
|
||||||
|
def test_unresolved_embedding_overrides_with_image_are_rejected(self):
|
||||||
|
"""EmbeddingReqInput resolves embed_overrides only after validation, so
|
||||||
|
the unresolved form must be caught at admission too."""
|
||||||
|
self.manager.is_generation = False
|
||||||
|
req = EmbeddingReqInput(
|
||||||
|
input_ids=[10, 50, 20],
|
||||||
|
sampling_params={},
|
||||||
|
embed_override_token_id=50,
|
||||||
|
embed_overrides=[_vec()],
|
||||||
|
image_data=["image.png"],
|
||||||
|
)
|
||||||
|
with self.assertRaisesRegex(ValueError, "overrides cannot be combined"):
|
||||||
|
self.manager._validate_one_request(req, req.input_ids)
|
||||||
|
req.image_data = None
|
||||||
|
self.manager._validate_one_request(req, req.input_ids)
|
||||||
|
|
||||||
|
def test_mixed_extend_batch_is_rejected_before_embedding_lookup(self):
|
||||||
|
"""Placeholder rows hold hash IDs, so the base lookup must never run
|
||||||
|
on a batch whose chunk also covers multimodal placeholders."""
|
||||||
|
embed_layer = MagicMock(
|
||||||
|
side_effect=AssertionError("embedding lookup must not run")
|
||||||
|
)
|
||||||
|
runner = SimpleNamespace(
|
||||||
|
_pp_kwargs=lambda pp_proxy_tensors: {},
|
||||||
|
model=SimpleNamespace(get_input_embeddings=lambda: embed_layer),
|
||||||
|
is_generation=True,
|
||||||
|
)
|
||||||
|
image = MultimodalDataItem(
|
||||||
|
modality=Modality.IMAGE, feature=torch.zeros(1), offsets=[(0, 1)]
|
||||||
|
)
|
||||||
|
image.set_hash(1234)
|
||||||
|
forward_batch = SimpleNamespace(
|
||||||
|
input_embeds=None,
|
||||||
|
input_ids=torch.tensor([1, 2, image.pad_value, image.pad_value]),
|
||||||
|
replace_embeds=torch.full((1, HIDDEN_DIM), 5.0),
|
||||||
|
replace_positions=torch.tensor([0]),
|
||||||
|
mm_inputs=[None, MultimodalInputs(mm_items=[image])],
|
||||||
|
extend_prefix_lens_cpu=[0, 0],
|
||||||
|
extend_seq_lens_cpu=[2, 2],
|
||||||
|
)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(ValueError, "cannot share an extend batch"):
|
||||||
|
ModelRunner._extend_forward_kwargs(runner, forward_batch, None)
|
||||||
|
embed_layer.assert_not_called()
|
||||||
|
|
||||||
|
# A decoding image request in a mixed chunk has no placeholder rows here.
|
||||||
|
forward_batch.input_ids = torch.tensor([1, 2, 3])
|
||||||
|
forward_batch.extend_prefix_lens_cpu = [0, 5]
|
||||||
|
forward_batch.extend_seq_lens_cpu = [2, 1]
|
||||||
|
embed_layer.side_effect = None
|
||||||
|
embed_layer.return_value = torch.zeros(3, HIDDEN_DIM)
|
||||||
|
kwargs = ModelRunner._extend_forward_kwargs(runner, forward_batch, None)
|
||||||
|
self.assertTrue(torch.equal(kwargs["input_embeds"][0], _vec(5.0)))
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -28,7 +28,13 @@ from sglang.srt.runtime_context import get_context
|
|||||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||||
from sglang.srt.utils.common import Range
|
from sglang.srt.utils.common import Range
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||||
|
|
||||||
|
maybe_stub_sgl_kernel()
|
||||||
|
|
||||||
|
import sglang.srt.managers.scheduler as scheduler_module
|
||||||
|
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||||
|
from sglang.srt.managers.scheduler import Scheduler
|
||||||
|
|
||||||
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
||||||
|
|
||||||
@@ -296,6 +302,139 @@ class TestPrefillAdder(CustomTestCase):
|
|||||||
)
|
)
|
||||||
self.assertEqual(adder.can_run_list, [first])
|
self.assertEqual(adder.can_run_list, [first])
|
||||||
|
|
||||||
|
def test_embed_override_and_multimodal_requests_never_share_a_batch(self):
|
||||||
|
def tagged(rid, *, multimodal=False, overrides=False):
|
||||||
|
req = self.create_shared_req(rid)
|
||||||
|
req.multimodal_inputs = object() if multimodal else None
|
||||||
|
req.positional_embed_overrides = object() if overrides else None
|
||||||
|
return req
|
||||||
|
|
||||||
|
for first, second in (
|
||||||
|
(tagged("image", multimodal=True), tagged("override", overrides=True)),
|
||||||
|
(tagged("override", overrides=True), tagged("image", multimodal=True)),
|
||||||
|
):
|
||||||
|
with self.subTest(first=first.rid):
|
||||||
|
adder = self.create_shared_adder()
|
||||||
|
self.assertTrue(adder.can_share_extend_batch(first))
|
||||||
|
adder.add_one_req(
|
||||||
|
first, has_chunked_req=False, truncation_align_size=None
|
||||||
|
)
|
||||||
|
self.assertEqual(adder.can_run_list, [first])
|
||||||
|
self.assertFalse(adder.can_share_extend_batch(second))
|
||||||
|
self.assertTrue(adder.can_share_extend_batch(tagged("text")))
|
||||||
|
|
||||||
|
adder = self.create_shared_adder()
|
||||||
|
chunked = tagged("chunked-image", multimodal=True)
|
||||||
|
chunked.full_untruncated_fill_ids = list(range(64))
|
||||||
|
self.assertIs(adder.add_chunked_req(chunked), chunked)
|
||||||
|
self.assertFalse(
|
||||||
|
adder.can_share_extend_batch(tagged("override", overrides=True))
|
||||||
|
)
|
||||||
|
|
||||||
|
def create_admission_scheduler(self, *, chunked_req) -> Scheduler:
|
||||||
|
allocator = self.create_token_allocator(available_size=4096)
|
||||||
|
allocator.page_size = 1
|
||||||
|
self.mock_tree_cache.supports_mamba.return_value = False
|
||||||
|
self.mock_tree_cache.is_tree_cache.return_value = False
|
||||||
|
self.mock_tree_cache.supports_fast_match_prefix.return_value = False
|
||||||
|
self.mock_tree_cache.storage_prefetch_retries = None
|
||||||
|
scheduler = Scheduler.__new__(Scheduler)
|
||||||
|
scheduler.grammar_manager = SimpleNamespace(has_waiting_grammars=lambda: False)
|
||||||
|
scheduler.enable_priority_preemption = False
|
||||||
|
scheduler.enable_priority_scheduling = False
|
||||||
|
scheduler.is_hybrid_swa = False
|
||||||
|
scheduler.min_free_slots_delayer = None
|
||||||
|
scheduler.get_num_allocatable_reqs = lambda *args, **kwargs: 64
|
||||||
|
scheduler.policy = SchedulePolicy(
|
||||||
|
policy="fcfs",
|
||||||
|
tree_cache=self.mock_tree_cache,
|
||||||
|
enable_hierarchical_cache=False,
|
||||||
|
enable_priority_scheduling=False,
|
||||||
|
schedule_low_priority_values_first=False,
|
||||||
|
)
|
||||||
|
scheduler.processed_tokens_counter = 0
|
||||||
|
scheduler.chunked_prefill_size = 16
|
||||||
|
scheduler.dynamic_chunk_sizer = None
|
||||||
|
scheduler.tp_worker = SimpleNamespace(
|
||||||
|
model_runner=SimpleNamespace(attn_backend=object(), prefill_aware_swa=False)
|
||||||
|
)
|
||||||
|
scheduler.page_size = 1
|
||||||
|
scheduler.tree_cache = self.mock_tree_cache
|
||||||
|
scheduler.token_to_kv_pool_allocator = allocator
|
||||||
|
scheduler.new_token_ratio_tracker = SimpleNamespace(current=1.0)
|
||||||
|
scheduler.max_prefill_tokens = 16384
|
||||||
|
scheduler.is_mixed_chunk = False
|
||||||
|
scheduler.priority_scheduling_preemption_threshold = 0
|
||||||
|
scheduler.max_prefill_bs = 64
|
||||||
|
scheduler.max_running_requests = 64
|
||||||
|
scheduler.dllm_config = None
|
||||||
|
scheduler.enable_lora = False
|
||||||
|
scheduler.req_to_token_pool = SimpleNamespace()
|
||||||
|
scheduler.disaggregation_mode = DisaggregationMode.NULL
|
||||||
|
scheduler.enable_hicache_storage = False
|
||||||
|
scheduler.enable_hierarchical_cache = False
|
||||||
|
scheduler.enable_unified_cache_external_linker = False
|
||||||
|
scheduler.truncation_align_size = None
|
||||||
|
scheduler.model_config = None
|
||||||
|
scheduler.enable_overlap = False
|
||||||
|
scheduler.spec_algorithm = None
|
||||||
|
scheduler.load_inquirer = MagicMock()
|
||||||
|
scheduler.chunked_req = chunked_req
|
||||||
|
scheduler.waiting_queue = []
|
||||||
|
return scheduler
|
||||||
|
|
||||||
|
def run_admission_pass(self, scheduler: Scheduler) -> list:
|
||||||
|
running_batch = self.create_running_batch()
|
||||||
|
running_batch.batch_is_full = False
|
||||||
|
with (
|
||||||
|
patch.object(scheduler_module, "ScheduleBatch") as schedule_batch,
|
||||||
|
patch.object(scheduler_module, "PrefillStats"),
|
||||||
|
patch.object(scheduler_module, "set_time_batch"),
|
||||||
|
):
|
||||||
|
new_batch, _ = scheduler._get_new_batch_prefill_raw(None, running_batch)
|
||||||
|
if new_batch is None:
|
||||||
|
return []
|
||||||
|
admitted = list(schedule_batch.init_new.call_args.args[0])
|
||||||
|
for req in admitted:
|
||||||
|
req.prefix_indices = list(range(req.extend_range.end))
|
||||||
|
return admitted
|
||||||
|
|
||||||
|
def test_fcfs_admits_override_request_once_image_continuation_drains(self):
|
||||||
|
"""An override request at the queue head must be admitted once the image
|
||||||
|
chunk ahead of it drains, even while more image requests keep arriving."""
|
||||||
|
|
||||||
|
def tagged(rid, length, *, multimodal=False, overrides=False):
|
||||||
|
req = self.create_shared_req(rid)
|
||||||
|
req.origin_input_ids = list(range(length))
|
||||||
|
req.full_untruncated_fill_ids = list(range(length))
|
||||||
|
req.multimodal_inputs = object() if multimodal else None
|
||||||
|
req.positional_embed_overrides = object() if overrides else None
|
||||||
|
req.beam_group = None
|
||||||
|
req.inflight_middle_chunks = 0
|
||||||
|
return req
|
||||||
|
|
||||||
|
continuation = tagged("image-continuation", 20, multimodal=True)
|
||||||
|
continuation.prefix_indices = list(range(16))
|
||||||
|
scheduler = self.create_admission_scheduler(chunked_req=continuation)
|
||||||
|
override = tagged("override", 4, overrides=True)
|
||||||
|
scheduler.waiting_queue = [override]
|
||||||
|
|
||||||
|
admitted_at = None
|
||||||
|
for pass_index in range(6):
|
||||||
|
scheduler.waiting_queue.append(
|
||||||
|
tagged(f"image-{pass_index}", 16, multimodal=True)
|
||||||
|
)
|
||||||
|
admitted = self.run_admission_pass(scheduler)
|
||||||
|
self.assertFalse(
|
||||||
|
any(r.multimodal_inputs is not None for r in admitted)
|
||||||
|
and any(r.positional_embed_overrides is not None for r in admitted)
|
||||||
|
)
|
||||||
|
if any(r is override for r in admitted):
|
||||||
|
admitted_at = pass_index
|
||||||
|
break
|
||||||
|
self.assertIsNotNone(admitted_at)
|
||||||
|
self.assertNotIn(override, scheduler.waiting_queue)
|
||||||
|
|
||||||
def test_shared_admission_rechecks_after_prefix_lock(self):
|
def test_shared_admission_rechecks_after_prefix_lock(self):
|
||||||
adder = self.create_shared_adder()
|
adder = self.create_shared_adder()
|
||||||
self.assertIsNotNone(adder.token_to_kv_pool_allocator.alloc(24))
|
self.assertIsNotNone(adder.token_to_kv_pool_allocator.alloc(24))
|
||||||
|
|||||||
@@ -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,261 @@
|
|||||||
|
"""Vision inputs under prefill CP merge on the full extend layout before the shard."""
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
from contextlib import contextmanager
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
from sglang.srt.layers.cp.base import init_cp_strategy
|
||||||
|
from sglang.srt.layers.cp.utils import prepare_cp_forward
|
||||||
|
from sglang.srt.managers import mm_schedule
|
||||||
|
from sglang.srt.managers.schedule_batch import (
|
||||||
|
Modality,
|
||||||
|
MultimodalDataItem,
|
||||||
|
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 DeepseekV4ForCausalLM
|
||||||
|
from sglang.srt.runtime_context import get_parallel
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
HIDDEN = 8
|
||||||
|
VOCAB = 64
|
||||||
|
IMAGE_TOKEN_ID = 7
|
||||||
|
CP_SIZE = 4
|
||||||
|
# (prefix_len, extend_len) per request. Request 1 carries one image whose span
|
||||||
|
# [2, 8] starts inside its prefix, so only span rows 1..6 land in this chunk.
|
||||||
|
CHUNKS = [(0, 7), (3, 9), (1, 5)]
|
||||||
|
IMAGE_OFFSET = (2, 8)
|
||||||
|
IMAGE_HASH = 12345
|
||||||
|
NUM_TOKENS = sum(extend_len for _, extend_len in CHUNKS)
|
||||||
|
# 21 tokens over 4 ranks give logical [6, 5, 5, 5], padded to the CP alignment.
|
||||||
|
PHYSICAL_ROWS = 8
|
||||||
|
IMAGE_ROWS = torch.arange(7, 13)
|
||||||
|
POSITIONS = torch.cat([torch.arange(p, p + n) for p, n in CHUNKS])
|
||||||
|
|
||||||
|
|
||||||
|
def _image_span(item: MultimodalDataItem) -> torch.Tensor:
|
||||||
|
start, end = item.offsets[0]
|
||||||
|
rows = end - start + 1
|
||||||
|
return torch.arange(rows * HIDDEN, dtype=torch.float32).view(rows, HIDDEN) + 100.0
|
||||||
|
|
||||||
|
|
||||||
|
def _pad(x: torch.Tensor) -> torch.Tensor:
|
||||||
|
return torch.cat([x, x.new_zeros(PHYSICAL_ROWS - x.shape[0], *x.shape[1:])])
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingBody:
|
||||||
|
def __init__(self, embed: nn.Embedding):
|
||||||
|
self.embed = embed
|
||||||
|
self.calls = []
|
||||||
|
|
||||||
|
def get_input_embeddings(self):
|
||||||
|
return self.embed
|
||||||
|
|
||||||
|
def __call__(self, input_ids, positions, forward_batch, input_embeds=None):
|
||||||
|
self.calls.append(
|
||||||
|
SimpleNamespace(
|
||||||
|
input_ids=input_ids,
|
||||||
|
positions=positions,
|
||||||
|
input_embeds=input_embeds,
|
||||||
|
input_ids_global=forward_batch.input_ids_global,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return input_embeds, input_embeds
|
||||||
|
|
||||||
|
|
||||||
|
class _VisionStub(DeepseekV4ForCausalLM):
|
||||||
|
def __init__(self, embed: nn.Embedding):
|
||||||
|
nn.Module.__init__(self)
|
||||||
|
self.config = SimpleNamespace(image_token_id=IMAGE_TOKEN_ID)
|
||||||
|
self.vision = object()
|
||||||
|
self.tp_size = 1
|
||||||
|
self.mm_owner_group = None
|
||||||
|
self.model = _RecordingBody(embed)
|
||||||
|
self.pp_group = SimpleNamespace(is_last_rank=True)
|
||||||
|
self.lm_head = object()
|
||||||
|
self.capture_aux_hidden_states = False
|
||||||
|
self.logits_calls = []
|
||||||
|
|
||||||
|
def get_image_feature(self, items):
|
||||||
|
return [_image_span(item) for item in items]
|
||||||
|
|
||||||
|
def logits_processor(
|
||||||
|
self,
|
||||||
|
input_ids,
|
||||||
|
hidden_states,
|
||||||
|
lm_head,
|
||||||
|
logits_metadata,
|
||||||
|
aux_hidden_states=None,
|
||||||
|
hidden_states_before_norm=None,
|
||||||
|
):
|
||||||
|
self.logits_calls.append(
|
||||||
|
SimpleNamespace(
|
||||||
|
input_ids=input_ids,
|
||||||
|
hidden_states=hidden_states,
|
||||||
|
logits_metadata=logits_metadata,
|
||||||
|
hidden_states_before_norm=hidden_states_before_norm,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return object()
|
||||||
|
|
||||||
|
|
||||||
|
def _build_batch():
|
||||||
|
item = MultimodalDataItem(
|
||||||
|
modality=Modality.IMAGE, feature=torch.zeros(1), offsets=[IMAGE_OFFSET]
|
||||||
|
)
|
||||||
|
item.set_hash(IMAGE_HASH)
|
||||||
|
ids = list(range(10, 17))
|
||||||
|
ids += [item.pad_value] * len(IMAGE_ROWS) + [20, 21, 22]
|
||||||
|
ids += list(range(30, 35))
|
||||||
|
forward_batch = SimpleNamespace(
|
||||||
|
forward_mode=ForwardMode.EXTEND,
|
||||||
|
mm_inputs=[
|
||||||
|
MultimodalInputs(mm_items=[]),
|
||||||
|
MultimodalInputs(mm_items=[item], im_token_id=IMAGE_TOKEN_ID),
|
||||||
|
None,
|
||||||
|
],
|
||||||
|
extend_prefix_lens_cpu=[prefix for prefix, _ in CHUNKS],
|
||||||
|
extend_seq_lens_cpu=[extend_len for _, extend_len in CHUNKS],
|
||||||
|
seq_lens_cpu=[prefix + extend_len for prefix, extend_len in CHUNKS],
|
||||||
|
input_ids=torch.tensor(ids, dtype=torch.long),
|
||||||
|
positions=POSITIONS.clone(),
|
||||||
|
mm_input_embeds=None,
|
||||||
|
attn_cp_metadata=None,
|
||||||
|
global_num_tokens_cpu=None,
|
||||||
|
out_cache_loc=None,
|
||||||
|
input_ids_global=torch.zeros(1, dtype=torch.long),
|
||||||
|
)
|
||||||
|
return forward_batch, item
|
||||||
|
|
||||||
|
|
||||||
|
def _expected_embeds(embed, scheduler_ids, item):
|
||||||
|
with torch.no_grad():
|
||||||
|
full = embed(scheduler_ids.clamp(max=VOCAB - 1))
|
||||||
|
full[IMAGE_ROWS] = _image_span(item)[1:7]
|
||||||
|
return full
|
||||||
|
|
||||||
|
|
||||||
|
def _canonical(scheduler_ids):
|
||||||
|
canonical = scheduler_ids.clone()
|
||||||
|
canonical[IMAGE_ROWS] = IMAGE_TOKEN_ID
|
||||||
|
return canonical
|
||||||
|
|
||||||
|
|
||||||
|
class TestDeepseekV41VisionPrefillCPInputs(CustomTestCase):
|
||||||
|
def setUp(self):
|
||||||
|
mm_schedule.init_mm_embedding_cache(1 << 20)
|
||||||
|
init_cp_strategy(
|
||||||
|
enable_prefill_cp=True, cp_size=CP_SIZE, cp_strategy="interleave"
|
||||||
|
)
|
||||||
|
torch.manual_seed(0)
|
||||||
|
self.embed = nn.Embedding(VOCAB, HIDDEN)
|
||||||
|
self.model = _VisionStub(self.embed)
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
init_cp_strategy(enable_prefill_cp=False, cp_size=1, cp_strategy="interleave")
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def _cp_collectives(self, full: torch.Tensor, rank: int):
|
||||||
|
def all_gather(output, input_tensor):
|
||||||
|
# Peers contribute their expected shards; this rank's rows come from
|
||||||
|
# what the runner actually handed to the collective.
|
||||||
|
output.zero_()
|
||||||
|
for peer in range(CP_SIZE):
|
||||||
|
rows = full[peer::CP_SIZE]
|
||||||
|
output[peer * PHYSICAL_ROWS : peer * PHYSICAL_ROWS + rows.shape[0]] = (
|
||||||
|
rows
|
||||||
|
)
|
||||||
|
output[rank * PHYSICAL_ROWS : (rank + 1) * PHYSICAL_ROWS] = input_tensor
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("torch.cuda.current_stream", return_value=None),
|
||||||
|
patch(
|
||||||
|
"sglang.srt.layers.cp.interleave.attn_cp_all_gather_into_tensor",
|
||||||
|
side_effect=all_gather,
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
"sglang.srt.layers.cp.interleave.is_allocation_symmetric",
|
||||||
|
return_value=False,
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
"sglang.srt.layers.cp.interleave.use_symmetric_memory",
|
||||||
|
return_value=torch.no_grad(),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
yield
|
||||||
|
|
||||||
|
def _prepare(self, forward_batch, input_embeds=None):
|
||||||
|
with torch.no_grad():
|
||||||
|
return self.model.prepare_model_inputs(
|
||||||
|
input_ids=forward_batch.input_ids,
|
||||||
|
forward_batch=forward_batch,
|
||||||
|
input_embeds=input_embeds,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_cp_runner_merges_before_shard(self):
|
||||||
|
runner = EagerRunner.__new__(EagerRunner)
|
||||||
|
runner.model_runner = SimpleNamespace(model=self.model)
|
||||||
|
padded = torch.zeros(CP_SIZE * PHYSICAL_ROWS, dtype=torch.long)
|
||||||
|
|
||||||
|
for rank in range(CP_SIZE):
|
||||||
|
forward_batch, item = _build_batch()
|
||||||
|
routing_sentinel = forward_batch.input_ids_global
|
||||||
|
scheduler_ids = forward_batch.input_ids.clone()
|
||||||
|
canonical = _canonical(scheduler_ids)
|
||||||
|
full = _expected_embeds(self.embed, scheduler_ids, item)
|
||||||
|
padded[:NUM_TOKENS] = canonical
|
||||||
|
rank_major_ids = padded.view(-1, CP_SIZE).T.flatten()
|
||||||
|
self.model.model.calls.clear()
|
||||||
|
self.model.logits_calls.clear()
|
||||||
|
|
||||||
|
with (
|
||||||
|
get_parallel().override(
|
||||||
|
attn_cp_rank=rank, attn_cp_size=CP_SIZE, attn_cp_group=object()
|
||||||
|
),
|
||||||
|
self._cp_collectives(full, rank),
|
||||||
|
torch.no_grad(),
|
||||||
|
):
|
||||||
|
prepare_cp_forward(forward_batch)
|
||||||
|
runner._execute_extend_cp(forward_batch, {})
|
||||||
|
|
||||||
|
with self.subTest(rank=rank):
|
||||||
|
metadata = forward_batch.attn_cp_metadata
|
||||||
|
self.assertEqual(metadata.per_rank_actual_token, [PHYSICAL_ROWS] * 4)
|
||||||
|
(body,) = self.model.model.calls
|
||||||
|
self.assertTrue(
|
||||||
|
torch.equal(body.input_ids, _pad(canonical[rank::CP_SIZE]))
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
torch.equal(body.positions, _pad(POSITIONS[rank::CP_SIZE]))
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
torch.equal(body.input_embeds, _pad(full[rank::CP_SIZE]))
|
||||||
|
)
|
||||||
|
self.assertTrue(torch.equal(body.input_ids_global, rank_major_ids))
|
||||||
|
|
||||||
|
(logits,) = self.model.logits_calls
|
||||||
|
self.assertTrue(torch.equal(logits.input_ids, canonical))
|
||||||
|
self.assertTrue(torch.equal(logits.hidden_states, full))
|
||||||
|
self.assertTrue(torch.equal(logits.hidden_states_before_norm, full))
|
||||||
|
self.assertIs(logits.logits_metadata, forward_batch)
|
||||||
|
|
||||||
|
self.assertTrue(torch.equal(forward_batch.mm_input_embeds, full))
|
||||||
|
self.assertTrue(torch.equal(forward_batch.input_ids, scheduler_ids))
|
||||||
|
self.assertIs(forward_batch.input_ids_global, routing_sentinel)
|
||||||
|
|
||||||
|
def test_external_embeddings_with_images_are_rejected(self):
|
||||||
|
forward_batch, _ = _build_batch()
|
||||||
|
with self.assertRaisesRegex(ValueError, "Cannot combine"):
|
||||||
|
self._prepare(forward_batch, input_embeds=torch.zeros(NUM_TOKENS, HIDDEN))
|
||||||
|
|
||||||
|
|
||||||
|
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()
|
||||||
@@ -29,6 +29,7 @@ from sglang.srt.arg_groups.cuda_graph_hook import (
|
|||||||
finalize_cuda_graph_prefill_max_context,
|
finalize_cuda_graph_prefill_max_context,
|
||||||
handle_cuda_graph_config,
|
handle_cuda_graph_config,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.arg_groups.deepseek_v4_hook import validate_deepseek_v41_features
|
||||||
from sglang.srt.arg_groups.hicache_hook import (
|
from sglang.srt.arg_groups.hicache_hook import (
|
||||||
handle_hicache,
|
handle_hicache,
|
||||||
handle_hicache_ratio_default,
|
handle_hicache_ratio_default,
|
||||||
@@ -45,6 +46,7 @@ from sglang.srt.arg_groups.kv_cache_hook import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.arg_groups.mamba_hook import handle_mamba_backend
|
from sglang.srt.arg_groups.mamba_hook import handle_mamba_backend
|
||||||
from sglang.srt.arg_groups.memory_hook import handle_gpu_memory_settings
|
from sglang.srt.arg_groups.memory_hook import handle_gpu_memory_settings
|
||||||
|
from sglang.srt.arg_groups.model_hook import handle_model_specific_adjustments
|
||||||
from sglang.srt.arg_groups.model_path_hook import handle_load_format
|
from sglang.srt.arg_groups.model_path_hook import handle_load_format
|
||||||
from sglang.srt.arg_groups.moe_hook import (
|
from sglang.srt.arg_groups.moe_hook import (
|
||||||
handle_a2a_moe,
|
handle_a2a_moe,
|
||||||
@@ -4069,5 +4071,86 @@ class TestLazyReexports(CustomTestCase):
|
|||||||
server_args_module.NotAThing
|
server_args_module.NotAThing
|
||||||
|
|
||||||
|
|
||||||
|
class TestDeepseekV41VisionPrefillCPArgs(CustomTestCase):
|
||||||
|
def _args(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
vision_n_layers=2,
|
||||||
|
prefill_backend=Backend.DISABLED,
|
||||||
|
lock_prefill_backend=False,
|
||||||
|
**overrides,
|
||||||
|
):
|
||||||
|
fields = dict(
|
||||||
|
model_path="dummy",
|
||||||
|
enable_prefill_cp=True,
|
||||||
|
cp_strategy="interleave",
|
||||||
|
tp_size=2,
|
||||||
|
)
|
||||||
|
fields.update(overrides)
|
||||||
|
server_args = ServerArgs(**fields)
|
||||||
|
server_args._model_config = SimpleNamespace(
|
||||||
|
hf_config=SimpleNamespace(
|
||||||
|
architectures=["DeepseekV4ForCausalLM"],
|
||||||
|
model_type="deepseek_v41",
|
||||||
|
vision_n_layers=vision_n_layers,
|
||||||
|
),
|
||||||
|
nvfp4_moe_meta=None,
|
||||||
|
is_fp4_experts=False,
|
||||||
|
)
|
||||||
|
# The dummy path does not initialize phase configs.
|
||||||
|
server_args.cuda_graph_config = CudaGraphConfig(
|
||||||
|
decode=PhaseConfig(backend=Backend.FULL, max_bs=512),
|
||||||
|
prefill=PhaseConfig(backend=prefill_backend, max_bs=512),
|
||||||
|
)
|
||||||
|
server_args._resolved_overrides = []
|
||||||
|
server_args._cuda_graph_config_locked = (
|
||||||
|
{(Phase.PREFILL, "backend")} if lock_prefill_backend else set()
|
||||||
|
)
|
||||||
|
return server_args
|
||||||
|
|
||||||
|
@override_platform(is_cuda=True, is_hip=False)
|
||||||
|
def test_encoder_swa_replay_is_rejected_in_model_hook_order(self):
|
||||||
|
"""The V4.1 validator runs before the CP validator declares attn_cp_size,
|
||||||
|
so encoder SWA replay used to pass resolution with vision prefill CP."""
|
||||||
|
args = self._args(
|
||||||
|
enable_encoder_swa_bounded_replay=True,
|
||||||
|
max_running_requests=4,
|
||||||
|
chunked_prefill_size=128,
|
||||||
|
)
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
ValueError,
|
||||||
|
"encoder-swa-bounded-replay does not support context parallelism",
|
||||||
|
):
|
||||||
|
handle_model_specific_adjustments(args)
|
||||||
|
|
||||||
|
def test_zigzag_is_rejected_only_with_vision(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "requires --cp-strategy interleave"):
|
||||||
|
validate_deepseek_v41_features(self._args(cp_strategy="zigzag"))
|
||||||
|
validate_deepseek_v41_features(
|
||||||
|
self._args(cp_strategy="zigzag", vision_n_layers=0)
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_prefill_graph_explicit_rejects_and_default_resolves_eager(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "runs eager prefill"):
|
||||||
|
validate_deepseek_v41_features(
|
||||||
|
self._args(prefill_backend=Backend.BREAKABLE, lock_prefill_backend=True)
|
||||||
|
)
|
||||||
|
args = self._args(prefill_backend=Backend.BREAKABLE)
|
||||||
|
validate_deepseek_v41_features(args)
|
||||||
|
self.assertEqual(
|
||||||
|
resolution_result(args, "cuda_graph_config").prefill.backend,
|
||||||
|
Backend.DISABLED,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_dspark_with_decoder_swa_bounded_replay_is_rejected(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "DSpark.*decoder-swa-bounded-replay"):
|
||||||
|
validate_deepseek_v41_features(
|
||||||
|
self._args(
|
||||||
|
speculative_algorithm="DSPARK",
|
||||||
|
enable_decoder_swa_bounded_replay=True,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user