Compare commits

Author SHA1 Message Date
minke.yu 4f0ca95477 deploy: b300 build scripts with pip cache mount + no-build-isolation fast path
Base image already has torch 2.13, but pyproject build-system.requires
pulls it into pip's isolated build env on every build (~6GB via throttled
mirror, 47min). --no-build-isolation reuses the base env; build.sh picks
Dockerfile.fast (--no-deps) when python/pyproject.toml is unchanged.
2026-09-24 13:11:41 +08:00
minke.yu b91137ab98 scheduler: CP-symmetric idle check for health-check admission
build-sglang-image / build (push) Successful in 27m52s
Health-check admit/skip used is_fully_idle(), which includes rank-local
hicache drain queues; ranks diverge right after activity, so one rank
dispatched the health-check generate while others piggyback-skipped,
deadlocking CP (hicache drain all_reduce vs CP request broadcast).
Seen on cp2/cp4 + hicache L3 after router health checks.

Recovered from b300 /data/ymk/sglang working copy (uncommitted WIP).
2026-09-24 12:01:57 +08:00
minke.yu 67d8368a84 deploy: archive b300 ds41 compose files; skip CI for deploy/ and .gitea/ changes 2026-09-24 11:49:37 +08:00
minke.yu db7d2cb7db fix: vision check uses get_parallel().pp_group (get_pp_group undefined in this tree)
build-sglang-image / build (push) Successful in 27m58s
2026-09-23 18:05:57 +08:00
injet f23179ce99 更新 .gitea/workflows/build-image.yaml
build-sglang-image / build (push) Successful in 26m22s
2026-09-23 17:54:53 +08:00
injet 8266769b2d 更新 .gitea/workflows/build-image.yaml
build-sglang-image / build (push) Failing after 1m12s
2026-09-23 17:49:30 +08:00
injet 3b671d6086 更新 .gitea/workflows/build-image.yaml 2026-09-23 17:48:46 +08:00
minke.yu 7925735a3e ci: also trigger image build on dsv41-pd-visioncp pushes
build-sglang-image / build (push) Successful in 29m50s
2026-09-23 16:46:06 +08:00
Xinyuan Tong 6833498646 model: prune comments and redundant tests in dsv41 vision CP 2026-09-23 14:36:01 +08:00
Xinyuan Tong b48e2cb1eb model: TP-wide single-owner image encoding for DeepSeek V4.1
ViT and Aligner are replicated per TP rank, encoding each image eight
times with TP8 on both CP1 and CP8. Elect one owner per image and use
ordered full-span broadcasts with a six-phase agreement protocol.
Both CP1 and CP8 benefit while local cache hits preserve collective order.
2026-09-23 14:36:01 +08:00
Xinyuan Tong bfeb7cd9b2 model: support DeepSeek V4.1 vision with interleave prefill CP
The CP runner bypassed the vision merge and used bare text embeddings.
Merge image features before sharding so request-global offsets stay valid.
Canonicalize model IDs separately to preserve scheduler hash IDs.
Keep unsupported combinations guarded and isolate embedding overrides
from multimodal prefills without starving queued FCFS requests.
2026-09-23 14:36:00 +08:00
Yuwei An ddf5207630 [Fix] Handle chunked paged MQA metadata in DSV4.1 eager forwards (#40637)
build-sglang-image / build (push) Successful in 32m15s
2026-09-23 13:35:55 +08:00
minke.yu 104218d9ed ci: tag builds with -ci-<UTC datetime> suffix (naming rule: <branch>-<sha9>-<ci|local>-<UTC>)
build-sglang-image / build (push) Successful in 30m44s
2026-09-22 20:37:57 +08:00
minke.yu 92632a60ba ci: drop gha cache backend (runner cache server 404s on this instance)
build-sglang-image / build (push) Successful in 50m51s
2026-09-22 18:49:20 +08:00
minke.yu 4b3b367b63 ci: hardcode Gitea registry (vars context falls back to docker.io on this instance)
build-sglang-image / build (push) Failing after 24m13s
2026-09-22 18:23:35 +08:00
minke.yu a78da9b524 ci: Gitea Actions workflow to build sglang image on push to dsv41-pd
build-sglang-image / build (push) Failing after 27m21s
Overseas node: original docker.io base (lmsysorg/sglang:dev-dsv41) and
default pypi. Tags follow the ymkymx convention <branch>-<sha9>-<datetime>
plus a moving <branch>-latest. Registry defaults to the Gitea instance's
container registry, override with the REGISTRY variable.
2026-09-22 16:59:13 +08:00
minke.yu b081dd3d23 Merge branch 'main' into dsv41-pd 2026-09-22 14:56:35 +08:00
Kevin Mi 8ac19cc19f [AMD][Kimi-K3] Fix deferred KDA gate projection and update DCP cookbook (#39066)
PR Test (Arm64) / check-changes (push) Successful in 10s
PR Test (NPU) / set-image-config (push) Successful in 1s
PR Test (NPU) / Recommend tests from coverage (push) Skipped
PR Test (NPU) / check-changes (push) Successful in 12s
PR Test (sgl-router) / gate (push) Successful in 8s
PR Test (Xeon) / check-changes (push) Successful in 9s
PR Test (XPU) / check-changes (push) Successful in 16s
pr-test-arm64.yml / pr-gate (push) Successful in 3s
PR Test (Arm64) / pr-gate (push) Successful in 3s
pr-test-npu.yml / pr-gate (push) Successful in 2s
PR Test (NPU) / pr-gate (push) Successful in 2s
PR Test (sgl-router) / tier-1 — lint (push) Failing after 33s
PR Test (sgl-router) / tier-2 — build + test (push) Skipped
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Skipped
PR Test (sgl-router) / tier-3 — k8s integration (push) Skipped
PR Test (sgl-router) / tier-3 — e2e (push) Skipped
pr-test-xpu.yml / pr-gate (push) Successful in 2s
PR Test (XPU) / pr-gate (push) Successful in 2s
pr-test-xeon.yml / pr-gate (push) Successful in 3s
PR Test (Xeon) / pr-gate (push) Successful in 3s
PR Test (sgl-router) / finish (push) Successful in 1s
Lint / lint (push) Failing after 2m51s
PR Test (NPU) / base-a-test-1-npu-a2 (push) Canceled after 0s
PR Test (NPU) / base-b-test-1-npu-a3 (push) Canceled after 0s
PR Test (NPU) / base-b-test-2-npu-a3 (push) Canceled after 0s
PR Test (NPU) / base-b-test-4-npu-a3 (push) Canceled after 0s
PR Test (NPU) / base-b-test-8-npu-a3 (push) Canceled after 0s
PR Test (NPU) / base-b-test-16-npu-a3 (push) Canceled after 0s
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (0) (push) Canceled after 0s
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (1) (push) Canceled after 0s
PR Test (NPU) / multimodal-gen-test-4-npu-a3 (0) (push) Canceled after 0s
PR Test (NPU) / multimodal-gen-test-4-npu-a3 (1) (push) Canceled after 0s
PR Test (NPU) / base-c-test-acc-2-npu-a3 (push) Canceled after 0s
PR Test (NPU) / base-c-test-acc-16-npu-a3 (push) Canceled after 0s
PR Test (NPU) / base-c-test-perf-2-npu-a3 (push) Canceled after 0s
PR Test (NPU) / base-c-test-perf-16-npu-a3 (push) Canceled after 0s
PR Test (NPU) / Analyze failure report (push) Canceled after 0s
PR Test (NPU) / setup-covstub (push) Canceled after 0s
PR Test (NPU) / pr-test-npu-finish (push) Canceled after 0s
pr-test-npu.yml / run (${{ fromJson(inputs.partitions).arr }}) (push) Canceled after 0s
PR Test (XPU) / finish (push) Canceled after 0s
PR Test (Arm64) / build-test (push) Canceled after 0s
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Canceled after 0s
PR Test (XPU) / multimodal-gen-test-1-gpu-xpu (push) Canceled after 0s
PR Test (Xeon) / build-test (gnr, gnr, xeon-gnr, stage-a-tp-test-cpu-intel) (push) Canceled after 0s
PR Test (Xeon) / build-test (spr1, 0, 3, spr, xeon-spr, stage-a-test-cpu-intel,stage-b-test-cpu-intel) (push) Canceled after 0s
PR Test (Xeon) / build-test (spr2, 1, 3, spr, xeon-spr, stage-a-test-cpu-intel,stage-b-test-cpu-intel) (push) Canceled after 0s
PR Test (Xeon) / build-test (spr3, 2, 3, spr, xeon-spr, stage-a-test-cpu-intel,stage-b-test-cpu-intel) (push) Canceled after 0s
2026-09-22 06:35:12 +00:00
Mohammad Miadh AngkadandMohammad Angkad 4c81cd1b09 [KDA] Fix missing beta sigmoid in PTX prefill (#40685)
Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>
2026-09-21 23:24:38 -07:00
Khoa Pham bc22e1de9e [DSpark] Fix draft CUDA graph stream explosion (#40658) 2026-09-21 23:05:45 -07:00
Brayden ZhongandXinyuan Tong a0781f2714 [Docs] GLM-5.3/5.3-Flash cookbooks: enable reasoning/tool-call parsers by default via auto (#40497)
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
2026-09-22 13:54:08 +08:00
Guangda LiuandGuangda Liu 04c0913434 [HiSparse] Add MHA hisparse support for MiniMax M3 (#31446)
Co-authored-by: Guangda Liu <bingps@users.noreply.github.com>
2026-09-22 13:28:03 +08:00
095e45100b [AMD] [GLM-5.3-Flash Day 0] Route mHC through AITER on gfx950 (#38545)
Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: Kevin Mi <45493463+kevin-mii@users.noreply.github.com>
Co-authored-by: Kevin Mi <mikevin920@yahoo.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-21 22:24:29 -07:00
kangwangamd 264da63319 [AMD] Update ROCm AITER pin to acf8fdf9 (#39965) 2026-09-21 22:01:17 -07:00
Piotr Mazurek b01961e295 [LFM2-VL] Add DSpark speculative decoding (#40651) 2026-09-21 21:49:51 -07:00
YAMY 9b59fc5db5 [ModelOpt][PP] Keep BF16 shared experts out of the NVFP4 fusion so TP1 pipeline stages can load (#40628) 2026-09-21 21:45:58 -07:00
2032f3a071 [Router] Abort the engine when a client disconnects mid-request (#39461)
Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: Kan Wu <wukanustc@gmail.com>
2026-09-22 12:35:56 +08:00
minke.yu c74a4037fb mooncake: probe pointers via cuda.bindings instead of ctypes libcudart
The ctypes-based probing loaded a second libcudart copy into the process,
which corrupted CUDA/runtime state and segfaulted the scheduler processes
shortly after registration (reproducible on prefill ranks, intermittent on
decode). cuda.bindings.runtime.cudaPointerGetAttributes is a properly
typed binding and boots cleanly on both roles.
2026-09-21 12:46:25 +08:00
minke.yu b963295489 mooncake: support host-memory payloads with intra-node NVLink transport
IntraNodeNvlinkTransport can only register/reach device memory. sglang's
register list includes host regions (aux buffers, some state components);
one host region failed the whole registerLocalMemoryBatch and the engine
rolled back every region, leaving the segment descriptor empty so all KV
transfers failed with 'Requested address ... not found'.

When the intra-node NVLink transport is active (MC_INTRANODE_NVLINK etc.):
- register only device-memory regions (probe via cudaPointerGetAttributes)
- route transfer blocks whose local source is host memory over the ordered
  per-endpoint zmq channel (new STATE_DATA message) instead of the engine;
  the receiver validates and writes the block into its buffer, relying on
  the same ordering guarantee as the aux TCP path
- force aux over the TCP socket path automatically
2026-09-21 11:53:26 +08:00
abing fa826e08b1 Reuse shared compressed KV dequantization in DeepSeek V4.1 CP prefill
(cherry picked from commit 44eb378e94610e4ccc6b990bf09c105f8c8f9ee9)
2026-09-20 22:12:47 +08:00
SYChen123 3810f531a8 Support V4.1 decode vision MegaMoE
(cherry picked from commit 98aa79972d6d31f30a3c5bc83b9d7e3009ae638a)
2026-09-20 22:12:43 +08:00
SYChen123 2580c24d1b Support V4.1 DP attention in DP-only DSpark PD
(cherry picked from commit 6ef0ebba730b3fffde4b59b953c5c9b58a35c4ed)
2026-09-20 22:11:48 +08:00
abing 4f22146e51 Avoid host synchronization in DeepSeek V4.1 CP prefill
(cherry picked from commit 510596019fd5782d494c9b5227bf784d46661f6f)
2026-09-20 22:07:30 +08:00
abing 12e3b82e52 run pass cp+megamoe+bcg
(cherry picked from commit 85b8a150ef69c35cbd8c3fe940e41b414b903757)
2026-09-20 22:07:29 +08:00
abing 8305f66fc8 fix bug
(cherry picked from commit 60d31cb0c268c4b733ab290e133e3763b65f090e)
2026-09-20 22:07:24 +08:00
abing 21a4a16b4b update code
(cherry picked from commit 1453150a0bdf2cce425e9e2da83ef16b8e8f2859)
2026-09-20 22:07:24 +08:00
abing fc954b7e08 add test
(cherry picked from commit 094abd5e90af77507acd7c91513f69b5b29b685b)
2026-09-20 22:07:24 +08:00
abing c2059c4fb2 run pass llm cp
(cherry picked from commit 1d85394563d96cf22d5d84f15c0e9043582419e4)
2026-09-20 22:07:23 +08:00
134 changed files with 8637 additions and 471 deletions
+57
View File
@@ -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
+11
View File
@@ -0,0 +1,11 @@
# Gitea Actions 用(海外节点):原版 docker.io 基底 + 原版 pypi
# 与 b300 离线版(/data/ymk/build/Dockerfile)的区别:
# - 基底直用 docker.io 的 lmsysorg/sglang:dev-dsv41(不走 umirror/DaoCloud)
# - 不设 PIP_INDEX_URL,用默认 pypi.org
# - 源码由 CI checkout 后经 COPY 进镜像(不用 in-image clone)
FROM lmsysorg/sglang:dev-dsv41
# checkout 含 .git,setuptools-scm 可打戳(main 系分支 version 显示 dev 属正常)
COPY . /sgl-workspace/sglang
RUN pip install --no-cache-dir -e /sgl-workspace/sglang/python
+21
View File
@@ -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
+15
View File
@@ -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
+33
View File
@@ -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)"
+33
View File
@@ -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` 引用
+171
View File
@@ -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
+43
View File
@@ -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}'
+43
View File
@@ -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}'
+42
View File
@@ -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}'
+160
View File
@@ -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
+44
View File
@@ -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}'
+44
View File
@@ -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}'
+62
View File
@@ -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
+42
View File
@@ -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}'
+44
View File
@@ -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
+47
View File
@@ -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
+50
View File
@@ -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}'
+55
View File
@@ -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
+52
View File
@@ -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}'
+57
View File
@@ -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
+54
View File
@@ -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
+67
View File
@@ -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
+65
View File
@@ -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
+59
View File
@@ -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}'
+92
View File
@@ -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
+43
View File
@@ -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}'
+6
View File
@@ -0,0 +1,6 @@
{
"master_server_address": "127.0.0.1:50051",
"metadata_server": "P2PHANDSHAKE",
"global_segment_size": "300gb",
"protocol": "tcp"
}
+8 -8
View File
@@ -61,7 +61,7 @@ ENV BUILD_TRITON="0"
ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT_DEFAULT="4ad99832823dde2315b361cbd3b54b1c5c12acd5"
ENV AITER_COMMIT_DEFAULT="acf8fdf9307431ece8ee275971c41cb3d1a7020b"
# ===============================
# Base image 942 with rocm720 and args
@@ -71,7 +71,7 @@ ENV BUILD_TRITON="1"
ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT_DEFAULT="4ad99832823dde2315b361cbd3b54b1c5c12acd5"
ENV AITER_COMMIT_DEFAULT="acf8fdf9307431ece8ee275971c41cb3d1a7020b"
ENV TRITON_COMMIT_DEFAULT="42270451990532c67e69d753fbd026f28fcc4840"
# ===============================
@@ -82,7 +82,7 @@ ENV BUILD_TRITON="1"
ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT_DEFAULT="4ad99832823dde2315b361cbd3b54b1c5c12acd5"
ENV AITER_COMMIT_DEFAULT="acf8fdf9307431ece8ee275971c41cb3d1a7020b"
# Pin the ROCm torch stack for every pip invocation in this flavor. The file is
# filled in after the torch 2.11 upgrade below; it must already exist (empty is
# valid) because pip reads PIP_CONSTRAINT from the first pip call onwards.
@@ -106,7 +106,7 @@ ENV BUILD_TRITON="0"
ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT_DEFAULT="4ad99832823dde2315b361cbd3b54b1c5c12acd5"
ENV AITER_COMMIT_DEFAULT="acf8fdf9307431ece8ee275971c41cb3d1a7020b"
# ===============================
# Base image 950 with rocm720 and args
@@ -116,7 +116,7 @@ ENV BUILD_TRITON="1"
ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT_DEFAULT="4ad99832823dde2315b361cbd3b54b1c5c12acd5"
ENV AITER_COMMIT_DEFAULT="acf8fdf9307431ece8ee275971c41cb3d1a7020b"
ENV TRITON_COMMIT_DEFAULT="42270451990532c67e69d753fbd026f28fcc4840"
# ===============================
@@ -127,7 +127,7 @@ ENV BUILD_TRITON="1"
ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT_DEFAULT="4ad99832823dde2315b361cbd3b54b1c5c12acd5"
ENV AITER_COMMIT_DEFAULT="acf8fdf9307431ece8ee275971c41cb3d1a7020b"
# Pin the ROCm torch stack for every pip invocation in this flavor. The file is
# filled in after the torch 2.11 upgrade below; it must already exist (empty is
# valid) because pip reads PIP_CONSTRAINT from the first pip call onwards.
@@ -286,7 +286,7 @@ ENV BUILD_TRITON="0"
ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT_DEFAULT="4ad99832823dde2315b361cbd3b54b1c5c12acd5"
ENV AITER_COMMIT_DEFAULT="acf8fdf9307431ece8ee275971c41cb3d1a7020b"
# Same reasoning as the rocm724 stages: keep pip from resolving the image's
# ROCm torch away to a PyPI CUDA build. Populated after the stack is in place.
ENV PIP_CONSTRAINT="/etc/sglang/constraints/torch-rocm.txt"
@@ -300,7 +300,7 @@ ENV BUILD_TRITON="0"
ENV BUILD_LLVM="0"
ENV BUILD_AITER_ALL="1"
ENV BUILD_MOONCAKE="1"
ENV AITER_COMMIT_DEFAULT="4ad99832823dde2315b361cbd3b54b1c5c12acd5"
ENV AITER_COMMIT_DEFAULT="acf8fdf9307431ece8ee275971c41cb3d1a7020b"
ENV PIP_CONSTRAINT="/etc/sglang/constraints/torch-rocm.txt"
RUN mkdir -p /etc/sglang/constraints && : > /etc/sglang/constraints/torch-rocm.txt
@@ -10,10 +10,10 @@ tag: NEW
<Accordion title="Install SGLang">
Use an SGLang build that includes GLM-5.3-Flash support.
Use an SGLang build that includes GLM-5.3-Flash support (v0.5.20 or later).
```bash Command
docker pull lmsysorg/sglang:glm-5.3-flash
docker pull lmsysorg/sglang:latest
```
The deployment panel can render a complete `docker run` command for the selected hardware and options. See [Install SGLang with Docker](/docs/get-started/install#method-3-using-docker) for host setup.
@@ -104,7 +104,7 @@ The **Speculative** card in the Playground changes the algorithm without leaving
- **EAGLE / MTP 5-1-6** is exactly what Low Latency serves, so a Low Latency base starts on this chip. Pick it from a High Throughput base to keep that recipe's other settings and add the MTP head.
- **Off (greedy)** strips the whole `--speculative-*` family, which is what High Throughput already starts from.
- **DFlash2** swaps the in-checkpoint MTP head for the trained block-diffusion draft in [`incoai/GLM-5.3-Flash-DFlash2`](https://huggingface.co/incoai/GLM-5.3-Flash-DFlash2). The draft proposes a whole block per step and the target verifies it in one forward pass, so output quality stays the target's. Its block size comes from the draft checkpoint, and the draft runs on `fa4` rather than the target's DSA backends. It needs a build that carries the GLM-5.3-Flash hidden-state capture from [PR #36708](https://github.com/sgl-project/sglang/pull/36708), which is merged into the [PR #36507](https://github.com/sgl-project/sglang/pull/36507) support branch (`xinyuan/glm-5.3-flash-support`) rather than into `main`, so the image pinned above is not enough on its own — pull that branch at its current head, or add #36708's commit on top of an older checkout. The draft repository is also access-gated: request access on its model page, then download it alongside the target before serving. This combination is not yet measured on the cookbook hardware, so treat it as a starting point.
- **DFlash2** swaps the in-checkpoint MTP head for the trained block-diffusion draft in [`incoai/GLM-5.3-Flash-DFlash2`](https://huggingface.co/incoai/GLM-5.3-Flash-DFlash2). The draft proposes a whole block per step and the target verifies it in one forward pass, so output quality stays the target's. Its block size comes from the draft checkpoint, and the draft runs on `fa4` rather than the target's DSA backends. The hidden-state capture it needs ([PR #36708](https://github.com/sgl-project/sglang/pull/36708)) shipped with the GLM-5.3-Flash support in v0.5.20, so the image pinned above is enough. The draft repository is access-gated: request access on its model page, then download it alongside the target before serving. This combination is not yet measured on the cookbook hardware, so treat it as a starting point.
Neither algorithm runs with DP-Attention; the card disables the affected chips and names the reason.
@@ -134,13 +134,13 @@ The default multimodal feature transport is automatic, and on a single CUDA node
### 3.1 Reasoning
Thinking is enabled by the checkpoint's generation configuration, and generated commands enable `--reasoning-parser glm45` by default. The OpenAI-compatible API then places thinking in `message.reasoning_content` and the final answer in `message.content`. You can disable **Reasoning Parser** in the Playground when an integration needs the raw response format.
Thinking is enabled by the checkpoint's generation configuration, and generated commands enable `--reasoning-parser auto` (which resolves to `glm45` for GLM-5.3-Flash) by default. The OpenAI-compatible API then places thinking in `message.reasoning_content` and the final answer in `message.content`. You can disable **Reasoning Parser** in the Playground when an integration needs the raw response format.
To disable thinking for a request, pass `chat_template_kwargs: {"thinking": false}` in the request body.
### 3.2 Tool calling
Generated commands enable `--tool-call-parser glm47` by default, so structured calls are returned in `message.tool_calls`. You can disable **Tool Call Parser** in the Playground when tool calling is not needed. On follow-up turns, read both `reasoning_content` and `content` because a thinking model can use either field around tool execution.
Generated commands enable `--tool-call-parser auto` (which resolves to `glm47` for GLM-5.3-Flash) by default, so structured calls are returned in `message.tool_calls`. You can disable **Tool Call Parser** in the Playground when tool calling is not needed. On follow-up turns, read both `reasoning_content` and `content` because a thinking model can use either field around tool execution.
### 3.3 Multimodal serving
+4 -4
View File
@@ -103,7 +103,7 @@ import { Playground } from "/src/snippets/_playground.jsx";
- **DeepSeek Sparse Attention (DSA).** GLM-5.3 uses the `glm_moe_dsa` architecture; SGLang auto-selects the DSA attention backends (`flashmla_sparse` prefill, `fa3` decode, `sgl-kernel` indexer topk). No attention-backend flag is needed on the supported hardware. SGLang also auto-selects the KV-cache dtype for DSA models — `fp8_e4m3` on Blackwell (B200/GB300/B300, which then routes DSA through the TensorRT-LLM backend) and `bf16` on Hopper (H200) — so no `--kv-cache-dtype` flag is required. On Hopper, pairing `--kv-cache-dtype fp8_e4m3` with `--dsa-prefill-backend flashmla_sparse_q8 --dsa-decode-backend flashmla_kv` selects the native FP8 sparse prefill kernel (computes directly on the fp8 KV cache with no fp8→bf16 dequantization round-trip; GLM-5.3's 64 query heads match the kernel's native tile) — see the [DeepSeek-V3.2 page](../DeepSeek/DeepSeek-V3_2) for kernel details; the optional `SGLANG_ENABLE_DSA_Q8KV8_*` performance env vars are documented in `python/sglang/srt/environ.py`.
- **MTP / speculative decoding.** The checkpoint ships one nextn layer. Enable EAGLE MTP for lower latency (`--speculative-algorithm EAGLE --speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6` for low-latency; `1-1-2` for balanced). The config's `index_share_for_mtp_iteration` reuses the DSA indexer's topk across draft steps (effective only at `--speculative-eagle-topk 1`). Watch the server's reported **accept length** and adjust `--speculative-num-steps` / `--speculative-num-draft-tokens`: lower the draft length when rejected draft tokens create excess verification work.
- **DFlash2 (block-diffusion draft).** The **Speculative** card in the [Playground above](#playground) also offers **DFlash2**, which replaces the in-checkpoint MTP layer with the separately trained block-diffusion drafter [`incoai/GLM-5.3-DFlash2`](https://huggingface.co/incoai/GLM-5.3-DFlash2). It proposes a whole block per step and the target verifies the block in one forward pass, so output quality stays the target's. The block size — 8, i.e. 7 draft tokens per verification step — comes from the draft checkpoint's own `dflash_config`, so no `--speculative-num-draft-tokens` is passed; the draft is a small dense model and runs on `fa4` instead of the target's DSA backends. Two prerequisites: the DFlash2 drafter ([PR #35371](https://github.com/sgl-project/sglang/pull/35371)) merged **after v0.5.18**, so install SGLang from `main` (or use a nightly image) rather than the release this page pins; and DFLASH runs on **CUDA/NPU only** and rejects **DP-Attention**, so turn DP-Attention off in the **Attention** card before selecting it on a high-throughput base. The draft repository is public but licensed CC BY-NC-ND 4.0 for research and evaluation.
- **DFlash2 (block-diffusion draft).** The **Speculative** card in the [Playground above](#playground) also offers **DFlash2**, which replaces the in-checkpoint MTP layer with the separately trained block-diffusion drafter [`incoai/GLM-5.3-DFlash2`](https://huggingface.co/incoai/GLM-5.3-DFlash2). It proposes a whole block per step and the target verifies the block in one forward pass, so output quality stays the target's. The block size — 8, i.e. 7 draft tokens per verification step — comes from the draft checkpoint's own `dflash_config`, so no `--speculative-num-draft-tokens` is passed; the draft is a small dense model and runs on `fa4` instead of the target's DSA backends. Note that DFLASH runs on **CUDA/NPU only** and rejects **DP-Attention**, so turn DP-Attention off in the **Attention** card before selecting it on a high-throughput base. The draft repository is public but licensed CC BY-NC-ND 4.0 for research and evaluation.
- **Memory.** The FP8 weights are large (MoE total, not active params). Start around `--mem-fraction-static 0.8` on H200 (TP8) and tune up; raise it for the 4-GPU GB300 single-node layout (TP4).
- **DP-Attention + DeepEP** for the balanced/high-throughput strategies spreads attention across data-parallel ranks and routes MoE through DeepEP.
- **BF16 weights need more GPUs.** The full-precision build (`zai-org/GLM-5.3-BF16`, ~1.5 TB) does not fit a single 8×H200 / 8×B200 / 4×GB300 node. It fits single-node on **8×B300** (TP8, ~2.1 TB HBM); on the smaller GPUs it needs a **multi-node** layout (e.g. 2×8×H200 or 2×8×B200 at TP16, 2×4×GB300 at TP8). FP8 is the recommended deployment. Use the same DSA / MTP / chunked-prefill guidance as FP8.
@@ -117,7 +117,7 @@ import { Playground } from "/src/snippets/_playground.jsx";
### 3.1 Reasoning
GLM-5.3 is a reasoning model. Enable the `glm45` reasoning parser (toggle **Reasoning Parser** in the **Parsers** card of the [Playground above](#playground)) to separate thinking from the final answer — thinking lands in `message.reasoning_content`, the answer in `message.content`. The chat template defaults `clear_thinking` to `false`; for multi-turn chat, pass `chat_template_kwargs: {"clear_thinking": True}` so previous reasoning is cleared before the next response.
GLM-5.3 is a reasoning model, and generated commands enable `--reasoning-parser auto` (which resolves to `glm45` for GLM-5.3) by default so thinking is separated from the final answer — thinking lands in `message.reasoning_content`, the answer in `message.content`. Without the parser the server returns the thinking and the answer as one `content` string with a stray `</think>` between them, because the chat template opens `<think>` in the generation prompt. You can disable **Reasoning Parser** in the **Parsers** card of the [Playground above](#playground) when an integration needs that raw format. The chat template defaults `clear_thinking` to `false`; for multi-turn chat, pass `chat_template_kwargs: {"clear_thinking": True}` so previous reasoning is cleared before the next response.
**Reasoning effort.** Pass `chat_template_kwargs: {"reasoning_effort": ...}` to select `low`, `high`, or `max`. If you omit it or pass another value, the template uses `max`.
@@ -164,7 +164,7 @@ Here is how you can calculate it:
### 3.2 Tool Calling
Enable the `glm47` tool-call parser (toggle **Tool Call Parser** in the **Parsers** card of the [Playground above](#playground)) to surface structured tool calls via `message.tool_calls`. GLM-5.3 emits the newer `<tool_call>…<arg_key>…<arg_value>…` format, so it needs the **`glm47`** parser — the older `glm45` parser does not parse it (the call would be left as raw text in `content`). On thinking mode the turn also fills `reasoning_content`, so print both fields.
Generated commands enable `--tool-call-parser auto` by default, so structured calls are returned in `message.tool_calls` with `finish_reason: "tool_calls"`. `auto` resolves to **`glm47`** for GLM-5.3: the model emits the newer `<tool_call>…<arg_key>…<arg_value>…` format, which the older `glm45` parser does not parse (the call would be left as raw text in `content`). Running with no tool-call parser fails the same way, and `finish_reason` stays `"stop"`, so an agent loop never sees the call. You can disable **Tool Call Parser** in the **Parsers** card of the [Playground above](#playground) when tool calling is not needed. On thinking mode the turn also fills `reasoning_content`, so print both fields.
<Accordion title="Tool Calling Example (Python)">
@@ -218,7 +218,7 @@ For long-context, prefix-heavy workloads, enable hierarchical KV caching to spil
### 3.4 Claude Code Integration
GLM-5.3's strong reasoning + tool-calling makes it a good backend for [Claude Code](https://code.claude.com/docs/en/overview), Anthropic's agentic CLI. SGLang exposes the Anthropic-compatible `/v1/messages` endpoint on every server, so Claude Code can talk to a GLM-5.3 server with only environment variables — no code change. Launch the server with `--reasoning-parser glm45 --tool-call-parser glm47` (any recipe from the Deployment panel above works), then:
GLM-5.3's strong reasoning + tool-calling makes it a good backend for [Claude Code](https://code.claude.com/docs/en/overview), Anthropic's agentic CLI. SGLang exposes the Anthropic-compatible `/v1/messages` endpoint on every server, so Claude Code can talk to a GLM-5.3 server with only environment variables — no code change. Launch the server with `--reasoning-parser auto --tool-call-parser auto` (any recipe from the Deployment panel above works), then:
```bash Command
export ANTHROPIC_BASE_URL="http://127.0.0.1:30000"
@@ -129,6 +129,7 @@ The NVIDIA Blackwell recipes are validated single-node: **B200 at `--tp 8`** and
- **Memory**: `--mem-fraction-static` reserves GPU memory for weights + KV pool; the rest is prefill **activation headroom**. The value scales with *free* memory per GPU (card capacity minus per-GPU weight), so it tracks the card more than the TP degree: **`0.65` on B200** (180 GB — less headroom once weights are resident) and **`0.75` on the larger-memory B300 / GB300** (`0.80` on AMD). Lower TP packs more weight per GPU, so a tighter config needs a *lower* value — B200 needs `0.65` even at `--tp 4`. Raising it past the validated value is fine only for low-concurrency single-stream serving; it OOMs under high concurrency or long context.
- **Long context (32K+)**: keep `--mem-fraction-static` at the platform default and raise `--chunked-prefill-size` to `16384`. Decode TPOT stays roughly flat in context length thanks to sparse attention; 1K–128K prompts are validated.
- **HiSparse for decode capacity**: on NVIDIA CUDA, HiSparse keeps the three dense layers on GPU, moves the 57 sparse-layer K/V caches to pinned host memory, and feeds selected block IDs directly to the swap-in kernel. For the released four-KV-head model, use `--tp 4` or greater, `--disable-radix-cache`, and `device_buffer_size >= 2048`. Enable it with `--enable-hisparse --hisparse-config='{"device_buffer_size":4096,"host_to_device_ratio":2}'` on the Triton launch command.
- **Scaling TP**: B200 is documented at `--tp 8`; B300 / GB200 / GB300 at `--tp 4` (the single-node cross-family common denominator). On an 8-GPU B300 host you can also raise to `--tp 8` for more throughput / KV headroom.
- **Expert parallelism**: to trade latency for throughput add `--ep` (see [Expert Parallelism Deployment](../../../docs/advanced_features/expert_parallelism)). On AMD, set `--ep` equal to `--tp`. Shared-experts fusion is automatically disabled when EP > 1; on AMD standard EP the server also disables `--enable-aiter-allreduce-fusion` automatically to preserve accuracy.
- `--trust-remote-code` is required to load the MiniMax config / processor classes.
@@ -30,7 +30,7 @@ Then run the **Python** output of the command panel below in that environment.
```bash Command
docker pull lmsysorg/sglang:latest # NVIDIA (CUDA)
docker pull lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260910 # AMD MI350X / MI355X (ROCm)
docker pull lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260916 # AMD MI350X / MI355X (ROCm)
```
For how to launch the image, see [Install → Method 3: Using Docker](../../../docs/get-started/install#method-3-using-docker). Substitute the inner `sglang serve ...` with what the command generator below produces.
@@ -65,7 +65,7 @@ Pick your hardware, then the deployment shape and operating point. Node count fo
**Strategy** — the operating point within that shape:
- **Low-Latency** — no DCP, so the MLA KV stays TP-replicated. For chat. B200 splits its two nodes into PP2 × TP8; every other platform is flat TP.
- **Balanced** — the accuracy-preserving default: PP2 × DCPEP8 on B200 (the two pipeline stages and DCP8 split KV and KDA state), TP16/DCP16 on GB200, TP8/DCP8 on B300/GB300, TP8 ROCm/AITER on MI35x.
- **Balanced** — the accuracy-preserving default: PP2 × DCPEP8 on B200 (the two pipeline stages and DCP8 split KV and KDA state), TP16/DCP16 on GB200, TP8/DCP8 on B300/GB300, TP8/DCP8 ROCm/AITER on MI35x.
- **High-Throughput** — the large-scale lane: pick a **Cluster Size** and **Large-Scale Preset** in the Playground ([details](#large-scale-presets)). The cell itself is Balanced, except on H100 (plus `extra_buffer_lazy`) and H200 (widens to 4×8 TP32/EP32 at `--mem-fraction-static 0.90`).
`Long-Context` appears only under the `Prefill` PD mode; for long-context unified serving on B200, start from High-Throughput and raise `--context-length`.
@@ -93,6 +93,22 @@ import { KimiK3MambaRatioCalculator } from "/src/snippets/_kimi_k3_mamba_ratio_c
NVFP4 NOSPEC / NVFP4 DSPARK), which is why no point past concurrency 64 is published for Balanced.
</Note>
### AMD AITER with DCP8
The MI350X/MI355X unified Balanced recipe uses TP8/DCP8 with AITER prefill and
decode attention. DCP shards the target MLA KV cache; RadixArk DSPARK's draft KV
remains replicated. The pinned `v0.5.19-rocm720-mi35x-20260916` image records
SGLang revision `e7f7447333`, which includes
[AITER DCP support (#34432)](https://github.com/sgl-project/sglang/pull/34432) and
the [DCP KV-free fix (#38941)](https://github.com/sgl-project/sglang/pull/38941).
No source overlay is required for DCP.
Keep `SGLANG_K3_KDA_FUSED_BACKEND` unset with this image. The separate fused-KDA
opt-in requires the [deferred-gate fix (#39066)](https://github.com/sgl-project/sglang/pull/39066),
which is not included in this image. This updated recipe remains **Final
Verification In Progress**; the recorded speed numbers use their original
configurations and do not validate the new image or DCP8 recipe.
### Mamba ratio calculator
<KimiK3MambaRatioCalculator />
@@ -180,11 +196,11 @@ Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request
| GB200 4×4 | TP16/DCP16 | MNNVL auto-detected |
| H200 2×8 (4×8 on Unified High-Throughput) | TP16/EP16 + symm-mem, Marlin + FlashMLA; High-Throughput widens to TP32/EP32 over 4 nodes at mem-frac 0.90 with `extra_buffer_lazy` | same block on every node; export the cross-node NIC (`GLOO_SOCKET_IFNAME` / `NCCL_SOCKET_IFNAME`, `SGLANG_HOST_IP`); keep `NCCL_MNNVL_ENABLE=1 NCCL_CUMEM_ENABLE=1` |
| H100 4×8 | TP32/EP32, Marlin + FlashMLA | SM90a build of the K3 image; pin NCCL/Gloo to the same NIC on all nodes; least post-weight headroom (80 GB) |
| MI350X/MI355X 1×8 | TP8 ROCm/AITER | AITER A8W4 FlyDSL MoE, Triton attention (`SGLANG_MLA_DECODE_TUNE=1` for gfx950 MLA decode geometry), graph bs up to 256, fp8 kvcache; DSPARK supported. Activation-quant and fused-KDA-decode knobs: [AMD ROCm/AITER environment](#amd-env) |
| MI350X/MI355X 1×8 | TP8/DCP8 ROCm/AITER (Unified Balanced) | AITER A8W4 FlyDSL MoE, AITER prefill/decode attention with sharded target MLA KV, graph bs up to 256, fp8 kvcache; DSPARK supported. Activation-quant and fused-KDA-decode knobs: [AMD ROCm/AITER environment](#amd-env) |
| Ascend A3 Series 4×8 (32 cards / 64 dies) | TP64/DP4 + DeepEP | PD-mixed `Unified` only; DSPARK baked in; pin `GLOO`/`HCCL_SOCKET_IFNAME` on every node |
| Ascend 950PR/DT Series 4×8 | TP32/dp1 + DeepEP | PD-mixed `Unified` only; DSPARK baked in; shared experts / dense MLP shard over attention-TP (`--shared-experts-tp-size 4`); radix cache off; pin `GLOO`/`HCCL_SOCKET_IFNAME` on every node |
**DCP notes** — the DCP cells are Balanced and High-Throughput on every Blackwell platform, in both the `Unified` and `Decode` roles:
**Blackwell DCP notes** — the DCP cells are Balanced and High-Throughput on every Blackwell platform, in both the `Unified` and `Decode` roles:
- DCP is the only axis that shards the TP-replicated MLA KV; Low-Latency skips it.
- Leave `--dcp-comm-backend` unset (fabric-resolved: `fi_a2a` on GB200/GB300, `a2a` on B200/B300).
+10 -1
View File
@@ -6,7 +6,7 @@ metatags:
HiSparse reduces per-request GPU memory consumption during the decode phase by maintaining only a small "hot" KV buffer on GPU while keeping complete KV data in CPU pinned memory. Combined with PD disaggregation, it enables significantly higher decode concurrency.
> **Prerequisites**: HiSparse works with models that use **DeepSeek Sparse Attention (DSA)** architectures (e.g., DeepSeek-V3.2, GLM-5.1) and **DeepSeek V4**. These models natively select a subset of tokens for attention, making it possible to keep only the top-k KV on GPU while storing the full KV in host memory — without accuracy loss. Additionally, HiSparse currently requires **PD disaggregation mode** and is enabled on the **decode instance** only.
> **Prerequisites**: HiSparse works with models that use **DeepSeek Sparse Attention (DSA)** architectures (e.g., DeepSeek-V3.2, GLM-5.1), **DeepSeek V4**, and **MiniMax M3**. These models natively select a subset of tokens for attention, making it possible to keep only the top-k KV on GPU while storing the full KV in host memory — without accuracy loss. Additionally, HiSparse currently requires **PD disaggregation mode** and is enabled on the **decode instance** only.
## Why HiSparse?
@@ -165,6 +165,15 @@ python3 -m sglang.launch_server \
> **Note**: For DSA models, `--kv-cache-dtype` defaults to `auto`, which resolves to `fp8_e4m3` on SM100+ (Blackwell) and `bfloat16` on older architectures. The DSA decode backend is automatically selected based on KV dtype (`bfloat16` → `flashmla_sparse`, `fp8_e4m3` → `flashmla_kv`), except for GLM DSA models on SM120/SM121 with `fp8_e4m3`, which use `flashinfer_sparse_mla`. DSA backend flags apply only to DSA models; DeepSeek V4 uses its own `dsv4` attention backend.
### MiniMax M3
Dense-layer K/V and index K stay on GPU; sparse-layer K/V use host memory plus a GPU working set.
- Use TP 4 or greater, with the same TP size and PP 1 on both PD instances.
- Use `--attention-backend triton`, `--mm-attention-backend triton_attn`, `--disable-prefill-cuda-graph`, and `--disable-radix-cache`.
- Set `device_buffer_size` to at least 2048 in `--hisparse-config`; `top_k` does not override the model's selection width.
- PD retraction backup is not supported. Use `--num-reserved-decode-tokens` to reserve capacity for the expected output length.
### Benchmark
```bash Command
@@ -478,8 +478,8 @@ export const config = {
gb200: "lmsysorg/sglang:kimi-k3",
// 20260903 or newer: the AITER SiTU A4W4/A8W4 layout fix (sgl-project/sglang#33838,
// merged Sep 3) and the fused gfx950 KDA decode boundary (#34198) first ship here.
mi350x: "lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260910",
mi355x: "lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260910",
mi350x: "lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260916",
mi355x: "lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260916",
// NVFP4 needs a build with sgl-project/sglang#35077; the purpose-built dev
// image is cut from that PR's head (CUDA 13).
"b300|nvfp4": "lmsysorg/sglang:dev-dev-kimi-k3-nvfp4",
@@ -1217,7 +1217,7 @@ export const config = {
],
},
{
// MI350X and MI355X use the same single-node TP8 ROCm/AITER profile.
// MI350X and MI355X use the same single-node TP8/DCP8 ROCm/AITER profile.
match: { hw: "mi350x", pdMode: "unified", strategy: "balanced" },
nnodes: 1,
verified: false,
@@ -1232,7 +1232,10 @@ export const config = {
"--model-path {{MODEL_NAME}}",
"--trust-remote-code",
"--tp-size 8",
"--attention-backend triton",
"--dcp-size 8",
"--dcp-comm-backend a2a",
"--prefill-attention-backend aiter",
"--decode-attention-backend aiter",
"--kv-cache-dtype fp8_e4m3",
"--dtype bfloat16",
"--mem-fraction-static 0.85",
@@ -1259,7 +1262,10 @@ export const config = {
"--model-path {{MODEL_NAME}}",
"--trust-remote-code",
"--tp-size 8",
"--attention-backend triton",
"--dcp-size 8",
"--dcp-comm-backend a2a",
"--prefill-attention-backend aiter",
"--decode-attention-backend aiter",
"--kv-cache-dtype fp8_e4m3",
"--dtype bfloat16",
"--mem-fraction-static 0.85",
@@ -119,6 +119,32 @@ export const config = {
},
],
},
// Parser flags live in one overlay dim so every generated command gets
// them without per-cell duplication; the Parsers card toggles derive
// on/off from the composed flags. `auto` needs the GLM-5.3 template
// detection (v0.5.20+).
{
id: "parsers",
title: "Parsers",
default: "auto",
options: [
{
id: "auto",
label: "Auto (glm45 + glm47)",
stripPrefixes: ["--reasoning-parser", "--tool-call-parser"],
flags: [
"--reasoning-parser auto",
"--tool-call-parser auto",
],
},
{
id: "off",
label: "Off",
stripPrefixes: ["--reasoning-parser", "--tool-call-parser"],
flags: [],
},
],
},
],
modelNames: {
@@ -174,15 +200,16 @@ sgl-eval run gsm8k \\
["aime2026_pct", "AIME 2026", "%"],
],
// Support is not in a public sglang release yet, so the nightly images do
// not work; every NVIDIA lane uses the purpose-built CUDA 13 image.
// v0.5.20 (= latest) carries GLM-5.3-Flash support (#36507) and the GLM-5.3
// template parser detection (#38297) that `--*-parser auto` needs; the old
// glm-5.3-flash dev image (2026-09-03) predates #38297 and misdetects.
dockerImages: {
gb300: "lmsysorg/sglang:glm-5.3-flash",
h100: "lmsysorg/sglang:glm-5.3-flash",
h200: "lmsysorg/sglang:glm-5.3-flash",
b200: "lmsysorg/sglang:glm-5.3-flash",
b300: "lmsysorg/sglang:glm-5.3-flash",
gb200: "lmsysorg/sglang:glm-5.3-flash",
gb300: "lmsysorg/sglang:latest",
h100: "lmsysorg/sglang:latest",
h200: "lmsysorg/sglang:latest",
b200: "lmsysorg/sglang:latest",
b300: "lmsysorg/sglang:latest",
gb200: "lmsysorg/sglang:latest",
},
github: {
@@ -259,8 +286,8 @@ sgl-eval run gsm8k \\
parsers: {
items: [
{ id: "reasoning", label: "Reasoning Parser", flag: "--reasoning-parser glm45" },
{ id: "toolCall", label: "Tool Call Parser", flag: "--tool-call-parser glm47" },
{ id: "reasoning", label: "Reasoning Parser", flag: "--reasoning-parser auto" },
{ id: "toolCall", label: "Tool Call Parser", flag: "--tool-call-parser auto" },
],
},
@@ -305,12 +332,7 @@ sgl-eval run gsm8k \\
"--speculative-draft-model-path incoai/GLM-5.3-Flash-DFlash2",
"--speculative-draft-attention-backend fa4",
],
// DFLASH needs this model's hidden-state capture, which landed on the
// GLM-5.3-Flash support branch (PR #36708 into #36507's
// xinyuan/glm-5.3-flash-support), not on main — so it postdates the
// image the Install accordion pins. Drop this note once #36507 merges
// and a published image carries it.
note: "⚠️ Needs the GLM-5.3-Flash hidden-state capture from PR #36708. It is merged into the PR #36507 support branch (xinyuan/glm-5.3-flash-support), not into main, so pull that branch at its current head — or add #36708's commit on top of an older checkout — before serving. The lmsysorg/sglang:glm-5.3-flash image alone is not enough.",
note: "⚠️ The draft checkpoint incoai/GLM-5.3-Flash-DFlash2 is access-gated: request access on its Hugging Face page, then download it alongside the target before serving.",
disable: [
{
when: { dpAttnOn: [true] },
@@ -346,8 +368,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -371,8 +391,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend trtllm",
"--kv-cache-dtype fp8_e4m3",
"--moe-runner-backend flashinfer_trtllm",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -406,8 +424,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--mem-fraction-static 0.85",
"--cuda-graph-max-bs-decode 32",
"--host {{HOST_IP}}",
@@ -434,8 +450,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend tilelang",
"--kv-cache-dtype bfloat16",
"--moe-runner-backend flashinfer_cutlass",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--mem-fraction-static 0.85",
"--host {{HOST_IP}}",
"--port {{PORT}}",
@@ -461,8 +475,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--mem-fraction-static 0.85",
"--cuda-graph-max-bs-decode 32",
"--host {{HOST_IP}}",
@@ -482,8 +494,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend tilelang",
"--kv-cache-dtype bfloat16",
"--moe-runner-backend flashinfer_cutlass",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--mem-fraction-static 0.85",
"--host {{HOST_IP}}",
"--port {{PORT}}",
@@ -506,8 +516,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--mem-fraction-static 0.85",
"--cuda-graph-max-bs-decode 32",
"--host {{HOST_IP}}",
@@ -527,8 +535,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend tilelang",
"--kv-cache-dtype bfloat16",
"--moe-runner-backend flashinfer_cutlass",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--mem-fraction-static 0.85",
"--host {{HOST_IP}}",
"--port {{PORT}}",
@@ -551,8 +557,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--mem-fraction-static 0.85",
"--cuda-graph-max-bs-decode 32",
"--host {{HOST_IP}}",
@@ -572,8 +576,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend tilelang",
"--kv-cache-dtype bfloat16",
"--moe-runner-backend flashinfer_cutlass",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--mem-fraction-static 0.85",
"--host {{HOST_IP}}",
"--port {{PORT}}",
@@ -603,8 +605,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -626,8 +626,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend tilelang",
"--kv-cache-dtype bfloat16",
"--moe-runner-backend deep_gemm",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -655,8 +653,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -677,8 +673,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend tilelang",
"--kv-cache-dtype bfloat16",
"--moe-runner-backend deep_gemm",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -701,8 +695,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -722,8 +714,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend trtllm",
"--kv-cache-dtype fp8_e4m3",
"--moe-runner-backend flashinfer_trtllm",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -746,8 +736,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -767,8 +755,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend trtllm",
"--kv-cache-dtype fp8_e4m3",
"--moe-runner-backend flashinfer_trtllm",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -789,8 +775,6 @@ sgl-eval run gsm8k \\
"--speculative-num-steps 5",
"--speculative-eagle-topk 1",
"--speculative-num-draft-tokens 6",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
@@ -807,8 +791,6 @@ sgl-eval run gsm8k \\
"--dsa-decode-backend trtllm",
"--kv-cache-dtype fp8_e4m3",
"--moe-runner-backend flashinfer_trtllm",
"--reasoning-parser glm45",
"--tool-call-parser glm47",
"--host {{HOST_IP}}",
"--port {{PORT}}",
],
+36 -9
View File
@@ -96,15 +96,45 @@ sgl-eval run aime25 \\
b200: "lmsysorg/sglang:latest",
gb300: "lmsysorg/sglang:latest",
b300: "lmsysorg/sglang:latest",
mi355x: "lmsysorg/sglang-rocm:v0.5.13.post1-rocm720-mi35x-20260618",
mi325x: "lmsysorg/sglang-rocm:v0.5.13.post1-rocm700-mi30x-20260616",
mi300x: "lmsysorg/sglang-rocm:v0.5.13.post1-rocm700-mi30x-20260616",
// >= v0.5.20 so `--*-parser auto` detects GLM-5.3 (#38297); the rocm700
// line stopped at v0.5.19, so mi30x moves to the rocm720 build.
mi355x: "lmsysorg/sglang-rocm:v0.5.20-rocm720-mi35x-20260920",
mi325x: "lmsysorg/sglang-rocm:v0.5.20-rocm720-mi30x-20260920",
mi300x: "lmsysorg/sglang-rocm:v0.5.20-rocm720-mi30x-20260920",
},
github: {
cookbookModel: "zai-org/glm-5.3",
},
// Parser flags live in one overlay dim so every generated command gets them
// without per-cell duplication; the Parsers card toggles derive on/off from
// the composed flags. `auto` needs the GLM-5.3 template detection (v0.5.20+).
overlayDims: [
{
id: "parsers",
title: "Parsers",
default: "auto",
options: [
{
id: "auto",
label: "Auto (glm45 + glm47)",
stripPrefixes: ["--reasoning-parser", "--tool-call-parser"],
flags: [
"--reasoning-parser auto",
"--tool-call-parser auto",
],
},
{
id: "off",
label: "Off",
stripPrefixes: ["--reasoning-parser", "--tool-call-parser"],
flags: [],
},
],
},
],
playgroundFeatures: {
// ----- Card 1: "Attention Parallelism" -----
@@ -162,8 +192,8 @@ sgl-eval run aime25 \\
// ----- Card 3: "Parsers" -----
parsers: {
items: [
{ id: "reasoning", label: "Reasoning Parser", flag: "--reasoning-parser glm45" },
{ id: "toolCall", label: "Tool Call Parser", flag: "--tool-call-parser glm47" },
{ id: "reasoning", label: "Reasoning Parser", flag: "--reasoning-parser auto" },
{ id: "toolCall", label: "Tool Call Parser", flag: "--tool-call-parser auto" },
],
},
@@ -194,10 +224,7 @@ sgl-eval run aime25 \\
flags: ["--speculative-algorithm DFLASH",
"--speculative-draft-model-path incoai/GLM-5.3-DFlash2",
"--speculative-draft-attention-backend fa4"],
// The DFlash2 drafter (PR #35371) merged after v0.5.18, so neither the
// release wheel nor the lmsysorg/sglang:latest image this page pins
// carries it. Drop this note once a release ships it.
note: "⚠️ Needs a nightly image: the DFlash2 drafter (PR #35371) is not in the release wheel nor the lmsysorg/sglang:latest image this page pins — install SGLang from main or use a lmsysorg/sglang:dev image. The draft is a separate checkpoint, so fetch incoai/GLM-5.3-DFlash2 alongside the target; it is public but licensed CC BY-NC-ND 4.0 for research and evaluation.",
note: "⚠️ The draft is a separate checkpoint: fetch incoai/GLM-5.3-DFlash2 alongside the target. It is public but licensed CC BY-NC-ND 4.0 for research and evaluation.",
disable: [
{ when: { dpAttnOn: [true] },
reason: "DFLASH speculative decoding does not support DP-Attention — the server rejects the combination at startup. Turn DP-Attention off in the Attention card above (the high-throughput recipes enable it)." },
@@ -0,0 +1,47 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use axum::http::{header::AUTHORIZATION, HeaderMap};
use reqwest::{Client, RequestBuilder, Url};
use std::time::Duration;
/// Cancels unfinished engine work without delaying request cleanup.
pub(super) struct AbortOnDrop(Option<RequestBuilder>);
impl AbortOnDrop {
// Only router-minted IDs are safe: the engine aborts by prefix.
pub(super) fn new(
client: &Client,
worker: &Url,
headers: &HeaderMap,
rid: Option<&str>,
) -> Self {
Self(rid.filter(|rid| !rid.is_empty()).map(|rid| {
let mut request = client
.post(worker.join("/abort_request").expect("validated worker URL"))
.json(&serde_json::json!({"rid": rid, "abort_all": false}))
.timeout(Duration::from_secs(5));
if let Some(auth) = headers.get(AUTHORIZATION) {
request = request.header(AUTHORIZATION, auth);
}
request
}))
}
pub(super) fn disarm(&mut self) {
self.0 = None;
}
}
impl Drop for AbortOnDrop {
fn drop(&mut self) {
let Some(request) = self.0.take() else { return };
if let Ok(runtime) = tokio::runtime::Handle::try_current() {
runtime.spawn(async move {
if let Err(error) = request.send().await.and_then(|r| r.error_for_status()) {
tracing::warn!(%error, "engine abort failed");
}
});
}
}
}
+24
View File
@@ -3,8 +3,11 @@
//! HTTP proxy — forwards requests to the upstream SGLang worker.
mod abort;
pub mod sse;
use abort::AbortOnDrop;
use crate::health::circuit_breaker::CircuitBreaker;
use crate::server::error::ApiError;
use crate::server::header_utils::should_forward_request_header;
@@ -203,6 +206,7 @@ impl Proxy {
/// path concatenation (no double-slash) and pass a typed URL to the
/// split error variants (`UpstreamUnreachable` / `UpstreamTimeout` /
/// `UpstreamStatus`).
#[allow(clippy::too_many_arguments)]
pub async fn forward_json_to(
&self,
worker_url: &str,
@@ -211,6 +215,7 @@ impl Proxy {
path: &str,
headers: &HeaderMap,
body: Bytes,
abort_rid: Option<&str>,
) -> Result<Response<Body>, ApiError> {
let permit = breaker.acquire().ok_or_else(|| ApiError::BreakerOpen {
worker: worker_url.to_string(),
@@ -228,6 +233,8 @@ impl Proxy {
req = req
.header("content-type", "application/json")
.timeout(self.request_timeout);
let mut abort =
AbortOnDrop::new(self.client_for(protocol), &worker_url, headers, abort_rid);
let resp = req.send().await.map_err(|e| {
breaker.record_failure();
Self::classify_reqwest_error_for(worker_url.clone(), e, path)
@@ -257,6 +264,7 @@ impl Proxy {
return Err(ApiError::UpstreamStatus { status });
}
};
abort.disarm();
match breaker_outcome(status) {
BreakerOutcome::Failure => breaker.record_failure(),
BreakerOutcome::Success => breaker.record_success(),
@@ -301,6 +309,7 @@ impl Proxy {
path: &str,
headers: &HeaderMap,
body: Bytes,
abort_rid: Option<&str>,
stream_guards: Option<Box<dyn Send + 'static>>,
on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>,
on_stream_end: Option<Box<dyn FnOnce(sse::StreamEnd) + Send + 'static>>,
@@ -322,11 +331,16 @@ impl Proxy {
req = req
.header("content-type", "application/json")
.header("accept", "text/event-stream");
let mut abort =
AbortOnDrop::new(self.client_for(protocol), &worker_url, headers, abort_rid);
let resp = req.send().await.map_err(|e| {
breaker.record_failure();
Self::classify_reqwest_error_for(worker_url.clone(), e, path)
})?;
let status = resp.status();
if !status.is_success() {
abort.disarm();
}
let upstream_ct = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
@@ -366,6 +380,9 @@ impl Proxy {
BreakerOutcome::Success => {
let breaker_for_hook = Arc::clone(breaker);
Some(Box::new(move |end| {
if end.reason == sse::StreamEndReason::Completed {
abort.disarm();
}
match stream_breaker_outcome(end) {
BreakerOutcome::Success => breaker_for_hook.record_success(),
BreakerOutcome::Failure => breaker_for_hook.record_failure(),
@@ -465,6 +482,7 @@ mod tests {
"/chat",
&headers,
Bytes::new(),
None,
)
.now_or_never()
.is_none());
@@ -481,6 +499,7 @@ mod tests {
None,
None,
None,
None,
)
.now_or_never()
.is_none());
@@ -583,6 +602,7 @@ mod tests {
None,
None,
None,
None,
expiration,
)
.await
@@ -694,6 +714,7 @@ mod tests {
"/v1/chat/completions",
&headers,
Bytes::from_static(b"{}"),
None,
)
.await
.expect("dispatch should reach the worker (breaker must stay closed)");
@@ -737,6 +758,7 @@ mod tests {
"/v1/chat/completions",
&headers,
Bytes::from_static(b"{}"),
None,
)
.await;
}
@@ -780,6 +802,7 @@ mod tests {
"/v1/chat/completions",
&headers,
Bytes::from_static(b"{}"),
None,
)
.await
.expect("the half-open probe must be admitted and reach the worker");
@@ -819,6 +842,7 @@ mod tests {
None,
None,
None,
None,
)
.await
.expect("streaming dispatch should reach the worker");
@@ -150,6 +150,9 @@ async fn access_log_and_record(
.unwrap_or_else(|| outcome_from_status(status.as_u16()))
.as_str(),
worker = log_ctx.map(|c| c.worker_url.as_str()).unwrap_or(""),
engine_rid = log_ctx
.and_then(|c| c.engine_rid.as_deref())
.unwrap_or(""),
model = log_ctx.map(|c| c.model_id.as_str()).unwrap_or(""),
stream = log_ctx.is_some_and(|c| c.streaming),
latency_ms,
@@ -446,6 +449,7 @@ mod tests {
model_id: "tiny".into(),
streaming: false,
outcome: RequestOutcome::Cancelled,
engine_rid: Some("1f0c2b7a4e9d4f3ab6c5d8e7f0a1b2c3".into()),
});
resp
}),
@@ -468,6 +472,11 @@ mod tests {
logs.contains("worker=\"http://worker-a:30000\"") && logs.contains("model=\"tiny\""),
"a routed request must be logged with its worker and model; captured:\n{logs}",
);
assert!(
logs.contains("engine_rid=\"1f0c2b7a4e9d4f3ab6c5d8e7f0a1b2c3\"")
&& logs.contains("request_id="),
"the minted rid must be logged beside the caller's request id; captured:\n{logs}",
);
// The handler's outcome must win over the status-derived fallback —
// otherwise the log and `worker_requests_total` can disagree about a
// request the handler classified itself (here, a cancellation served
@@ -223,6 +223,8 @@ pub struct RequestLogContext {
/// line and `worker_requests_total` cannot disagree — the middleware can
/// only see the status, which cannot express a router-side cancellation.
pub outcome: RequestOutcome,
/// Router-minted engine ID, logged beside the caller's correlation ID.
pub engine_rid: Option<String>,
}
/// Final outcome of a 2xx SSE stream.
@@ -81,7 +81,12 @@ pub(super) async fn forward_chat_request(
};
(decode, bootstrap)
});
let body = request.into_outgoing_body(ctx, pd.as_ref().map(|(_, bootstrap)| bootstrap))?;
let engine_rid = request.engine_rid(pd.is_some());
let body = request.into_outgoing_body(
ctx,
pd.as_ref().map(|(_, bootstrap)| bootstrap),
engine_rid.as_deref(),
)?;
let prefill_load_guards = (worker_load_guard, active_request_guard);
// In PD mode, prefill runs independently and decode supplies the client response.
@@ -112,6 +117,7 @@ pub(super) async fn forward_chat_request(
&response_worker,
&headers,
body,
engine_rid.as_deref(),
response_load_guards,
&metrics,
expiration_token.clone(),
@@ -124,7 +130,7 @@ pub(super) async fn forward_chat_request(
model: metrics.model.clone(),
}),
};
let log_context = metrics.record_dispatch_result(&result);
let log_context = metrics.record_dispatch_result(&result, engine_rid);
// Materialize dispatch errors here so the access log retains the selected worker.
let mut response = match result {
Ok(mut response) => {
@@ -171,6 +177,7 @@ fn spawn_prefill_request(
CHAT_PATH,
&headers,
body,
None,
)
.await
{
@@ -188,11 +195,13 @@ fn spawn_prefill_request(
});
}
#[allow(clippy::too_many_arguments)]
async fn forward_to_response_worker(
ctx: &AppContext,
worker: &Worker,
headers: &HeaderMap,
body: Bytes,
engine_rid: Option<&str>,
load_guards: LoadGuards,
metrics: &DispatchMetrics,
expiration: CancellationToken,
@@ -209,6 +218,7 @@ async fn forward_to_response_worker(
CHAT_PATH,
headers,
body,
engine_rid,
Some(stream_guards),
Some(metrics.first_byte_callback()),
Some(metrics.stream_end_callback(worker.url.clone())),
@@ -226,6 +236,7 @@ async fn forward_to_response_worker(
CHAT_PATH,
headers,
body,
engine_rid,
)
.await
}
@@ -295,6 +306,7 @@ impl DispatchMetrics {
fn record_dispatch_result(
&self,
result: &Result<Response<Body>, ApiError>,
engine_rid: Option<String>,
) -> RequestLogContext {
let http_status = match result {
Ok(response) => response.status().as_u16(),
@@ -326,6 +338,7 @@ impl DispatchMetrics {
model_id: self.model.clone(),
streaming: self.streaming,
outcome,
engine_rid,
}
}
}
@@ -26,6 +26,8 @@ pub(super) struct PreparedChatRequest {
pub(super) tokens: Option<RequestTokens>,
/// Token count for routing/load accounting; estimated from body size when unavailable.
pub(super) input_token_count: usize,
caller_set_rid: bool,
fans_out: bool,
can_forward_input_ids: bool,
parsed_body: Option<Value>,
sampling_defaults: Vec<(SamplingField, Number)>,
@@ -69,16 +71,27 @@ impl PreparedChatRequest {
body,
tokens,
input_token_count,
caller_set_rid: fields.caller_set_rid,
fans_out: requests_multiple_samples(&fields, &sampling_defaults),
can_forward_input_ids,
parsed_body,
sampling_defaults,
})
}
pub(super) fn engine_rid(&self, pd_mode: bool) -> Option<String> {
// Caller IDs are unsafe for prefix aborts; fan-out regenerates IDs; PD must finish KV transfer.
if self.caller_set_rid || self.fans_out || pd_mode {
return None;
}
Some(uuid::Uuid::new_v4().simple().to_string())
}
pub(super) fn into_outgoing_body(
self,
ctx: &AppContext,
bootstrap: Option<&BootstrapFields>,
engine_rid: Option<&str>,
) -> Result<Bytes, ApiError> {
// Routing tokens can replace engine tokenization only for supported chat templates.
let input_ids = match (self.tokens.as_ref(), self.parsed_body.as_ref()) {
@@ -104,6 +117,7 @@ impl PreparedChatRequest {
input_ids,
bootstrap,
&self.sampling_defaults,
engine_rid,
)
}
}
@@ -116,6 +130,8 @@ pub(super) struct RoutingFields {
max_tokens: Option<u64>,
max_completion_tokens: Option<u64>,
sampling: [SamplingValue; SamplingField::ALL.len()],
// Preserve both string and list IDs without retaining their contents.
caller_set_rid: bool,
}
/// Null is absent; unrepresentable values are rejected only under a sampling contract.
@@ -248,6 +264,7 @@ impl RoutingKey {
enum RequestKey {
Routing(RoutingKey),
Sampling(SamplingField),
Rid,
Other,
}
@@ -267,6 +284,7 @@ impl<'de> Deserialize<'de> for RequestKey {
"model" => RequestKey::Routing(RoutingKey::Model),
"max_tokens" => RequestKey::Routing(RoutingKey::MaxTokens),
"max_completion_tokens" => RequestKey::Routing(RoutingKey::MaxCompletionTokens),
"rid" => RequestKey::Rid,
other => match SamplingField::from_wire_name(other) {
Some(field) => RequestKey::Sampling(field),
None => RequestKey::Other,
@@ -324,6 +342,9 @@ impl<'de> serde::de::Visitor<'de> for RoutingFieldsVisitor {
SamplingValue::Unusable
};
}
RequestKey::Rid => {
fields.caller_set_rid = map.next_value::<Option<IgnoredAny>>()?.is_some();
}
RequestKey::Other => {
// Validate unrelated JSON without retaining its contents.
map.next_value::<IgnoredAny>()?;
@@ -374,6 +395,22 @@ fn should_tokenize_request(
can_forward_input_ids || policy_needs_request_tokens || bucket_routing_enabled
}
// Use the effective n, including injected defaults; unreadable values opt out.
fn requests_multiple_samples(
fields: &RoutingFields,
sampling_defaults: &[(SamplingField, Number)],
) -> bool {
match fields.sampling_field(SamplingField::N) {
SamplingValue::Number(n) => n > 1.0,
SamplingValue::Unusable => true,
SamplingValue::Absent => sampling_defaults
.iter()
.find(|(field, _)| *field == SamplingField::N)
.and_then(|(_, value)| value.as_f64())
.is_some_and(|n| n > 1.0),
}
}
fn estimate_prefill_tokens(body: &Bytes) -> usize {
// Never 0: a zero-load entry is invisible to the cache-aware imbalance fast path.
(body.len() / BYTES_PER_TOKEN_ESTIMATE).max(1)
@@ -391,9 +428,10 @@ pub(super) struct BootstrapFields {
}
/// Append before the closing brace so injected values win over explicit nulls.
fn append_sampling_defaults(
fn append_top_level_fields(
body: &Bytes,
sampling_defaults: &[(SamplingField, Number)],
rid: Option<&str>,
) -> Option<Bytes> {
use std::io::Write as _;
@@ -405,13 +443,23 @@ fn append_sampling_defaults(
let has_members = body[open + 1..close]
.iter()
.any(|b| !b.is_ascii_whitespace());
let mut output = Vec::with_capacity(body.len() + 24 * sampling_defaults.len() + 1);
let rid_budget = rid.map_or(0, |rid| rid.len() + ",\"rid\":\"\"".len());
let mut output = Vec::with_capacity(body.len() + 24 * sampling_defaults.len() + rid_budget + 1);
output.extend_from_slice(&body[..close]);
for (i, (field, value)) in sampling_defaults.iter().enumerate() {
if has_members || i > 0 {
let mut wrote_any = has_members;
for (field, value) in sampling_defaults {
if wrote_any {
output.push(b',');
}
write!(output, "\"{}\":{}", field.wire_name(), value).ok()?;
wrote_any = true;
}
if let Some(rid) = rid {
if wrote_any {
output.push(b',');
}
output.extend_from_slice(b"\"rid\":");
serde_json::to_writer(&mut output, rid).ok()?;
}
output.extend_from_slice(&body[close..]);
Some(Bytes::from(output))
@@ -424,14 +472,15 @@ fn build_outgoing_body(
input_ids: Option<&[u32]>,
bootstrap: Option<&BootstrapFields>,
sampling_defaults: &[(SamplingField, Number)],
rid: Option<&str>,
) -> Result<Bytes, ApiError> {
let sampling_only = input_ids.is_none() && bootstrap.is_none();
if sampling_only && sampling_defaults.is_empty() {
let needs_parse = input_ids.is_some() || bootstrap.is_some();
if !needs_parse && sampling_defaults.is_empty() && rid.is_none() {
// Cloning Bytes shares the original allocation when no injection is needed.
return Ok(body.clone());
}
if sampling_only {
if let Some(spliced) = append_sampling_defaults(body, sampling_defaults) {
if !needs_parse {
if let Some(spliced) = append_top_level_fields(body, sampling_defaults, rid) {
return Ok(spliced);
}
}
@@ -445,6 +494,9 @@ fn build_outgoing_body(
return Err(invalid_request());
}
};
if let Some(rid) = rid {
body_fields.insert("rid".into(), Value::String(rid.to_owned()));
}
for (field, default) in sampling_defaults {
body_fields.insert(field.wire_name().into(), default.clone().into());
}
@@ -739,7 +791,8 @@ mod tests {
.unwrap()
.extend(fields.as_object().unwrap().clone());
for value in [None, Some(original.clone())] {
let out = build_outgoing_body(&body, value, ids, bootstrap.as_ref(), &[]).unwrap();
let out =
build_outgoing_body(&body, value, ids, bootstrap.as_ref(), &[], None).unwrap();
assert_eq!(serde_json::from_slice::<Value>(&out).unwrap(), expected);
}
}
@@ -750,7 +803,7 @@ mod tests {
for raw in [r#"{"model":"x"}"#, r#"{"model":"x","messages":[]}"#] {
let body = Bytes::copy_from_slice(raw.as_bytes());
for value in [None, Some(serde_json::from_slice(&body).unwrap())] {
let out = build_outgoing_body(&body, value, None, None, &[]).unwrap();
let out = build_outgoing_body(&body, value, None, None, &[], None).unwrap();
assert_eq!(out, body);
assert_eq!(out.as_ptr(), body.as_ptr());
}
@@ -1108,7 +1161,7 @@ mod tests {
&metrics(),
)
.unwrap();
let out = build_outgoing_body(&body, None, Some(&[1, 2, 3]), None, &inject).unwrap();
let out = build_outgoing_body(&body, None, Some(&[1, 2, 3]), None, &inject, None).unwrap();
assert_eq!(
serde_json::from_slice::<Value>(&out).unwrap(),
json!({
@@ -1240,7 +1293,7 @@ mod tests {
let inject = resolve_sampling_defaults(&config, &fields_of(raw), &metrics()).unwrap();
assert_eq!(inject.len(), 2);
for value in [None, Some(serde_json::from_slice(&body).unwrap())] {
let out = build_outgoing_body(&body, value, None, None, &inject).unwrap();
let out = build_outgoing_body(&body, value, None, None, &inject, None).unwrap();
assert_eq!(std::str::from_utf8(&out).unwrap(), expected, "{raw}");
let parsed: Value = serde_json::from_slice(&out).unwrap();
assert_eq!(parsed["temperature"], json!(1.0));
@@ -1464,4 +1517,35 @@ mod tests {
"a value at the cap is still read"
);
}
#[test]
fn abort_opt_outs_follow_caller_rid_and_effective_sample_count() {
for raw in [r#"{"rid":"abc"}"#, r#"{"rid":["a","b"]}"#] {
assert!(fields_of(raw).caller_set_rid);
}
assert!(!fields_of(r#"{"rid":null}"#).caller_set_rid);
for (raw, fan_out) in [
(r#"{}"#, false),
(r#"{"n":1}"#, false),
(r#"{"n":2}"#, true),
(r#"{"n":"3"}"#, true),
(r#"{"n":[2]}"#, true),
] {
assert_eq!(
requests_multiple_samples(&fields_of(raw), &[]),
fan_out,
"{raw}"
);
}
for (config, fan_out) in [(r#"{"n":1}"#, false), (r#"{"n":4}"#, true)] {
let fields = fields_of("{}");
let defaults = resolve_sampling_defaults(
&overrides_of(ConflictPolicy::Reject, config),
&fields,
&metrics(),
)
.unwrap();
assert_eq!(requests_multiple_samples(&fields, &defaults), fan_out);
}
}
}
@@ -198,7 +198,18 @@ async fn caller_input_ids_are_used_for_routing_and_preserved() {
send(Arc::clone(&ctx), request.clone()).await,
StatusCode::OK
);
assert_eq!(captured(&mock), request, "body must be forwarded untouched");
let mut forwarded = captured(&mock);
let rid = forwarded
.as_object_mut()
.expect("a forwarded chat body is an object")
.remove("rid");
assert!(
rid.as_ref()
.and_then(Value::as_str)
.is_some_and(crate::common::is_engine_shaped_rid),
"plain mode must mint an abort rid; got {rid:?}",
);
assert_eq!(forwarded, request, "body must be forwarded untouched");
}
// Bypasses are not rendering failures.
assert!(!ctx
@@ -17,10 +17,14 @@ use sgl_router::workers::{WireProtocol, Worker, WorkerRegistry};
use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use sgl_router::state::load_monitor::router_inflight_load::{
spawn_janitor, JanitorHandle, RouterInflightLoadRegistry,
};
use std::sync::Arc;
use std::time::Duration;
use tower::ServiceExt;
mod cancellation;
mod reorg;
const TEST_TIMEOUT: Duration = Duration::from_secs(5);
@@ -74,6 +78,36 @@ fn build_ctx_with_worker(url: &str) -> Arc<AppContext> {
Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies))
}
/// Expire requests after 50ms; keep the janitor handle alive during the test.
fn build_ctx_with_janitor(url: &str) -> (Arc<AppContext>, JanitorHandle) {
let cfg = config_for(url);
let registry = Arc::new(WorkerRegistry::default());
let _ = registry.add(WorkerSpec {
id: WorkerId("w1".into()),
url: url.to_string(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
});
let policies = Arc::new(build_policy_registry(&cfg).unwrap());
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap());
let router_inflight_load = RouterInflightLoadRegistry::new(
Arc::new(sgl_router::state::load_monitor::router_inflight_load::SystemTimeClock),
Duration::from_millis(50),
);
let janitor = spawn_janitor(Arc::clone(&router_inflight_load), Duration::from_millis(20));
let ctx = Arc::new(AppContext::with_router_inflight_load(
cfg,
tokenizers,
proxy,
registry,
policies,
router_inflight_load,
));
(ctx, janitor)
}
#[tokio::test]
async fn non_streaming_returns_200() {
let worker = crate::common::mock_worker::MockWorker::start(vec![]).await;
@@ -1033,6 +1067,7 @@ async fn forward_json_to_records_failure_on_body_drop() {
"/v1/chat/completions",
&headers,
body,
None,
)
.await;
assert!(res.is_err(), "body drop should surface as ApiError");
@@ -1090,6 +1125,7 @@ async fn forward_json_to_records_success_only_after_body_completes() {
"/v1/chat/completions",
&headers,
bytes::Bytes::from_static(b"{}"),
None,
)
.await;
assert!(res.is_ok(), "clean OK call must succeed: {res:?}");
@@ -1145,6 +1181,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() {
None,
None,
None,
None,
)
.await;
@@ -1248,6 +1285,7 @@ async fn forward_json_to_records_failure_on_5xx() {
"/v1/chat/completions",
&headers,
body,
None,
)
.await;
@@ -1282,6 +1320,7 @@ async fn forward_json_to_rejects_when_breaker_open() {
"/v1/chat/completions",
&headers,
body,
None,
)
.await;
@@ -1320,6 +1359,7 @@ async fn forward_json_to_malformed_url_returns_worker_misconfigured_and_trips_br
"/v1/chat/completions",
&headers,
body,
None,
)
.await;
@@ -1607,58 +1647,17 @@ async fn streaming_active_load_drops_on_client_disconnect() {
/// returns; cancellation fires; handler returns 504.
#[tokio::test]
async fn janitor_expiry_returns_504_stale_request_expired() {
use sgl_router::state::load_monitor::router_inflight_load::{
spawn_janitor, RouterInflightLoadRegistry,
};
// Upstream that takes 2s to respond — longer than our 50ms
// stale_request_timeout.
// Upstream that takes 2s to respond — longer than the helper's 50ms
// stale_request_timeout, so the janitor sweeps before it answers.
let worker =
crate::common::mock_worker::MockWorker::start_hanging(Duration::from_secs(2)).await;
let cfg = config_for(&worker.url);
let registry = Arc::new(WorkerRegistry::default());
let _ = registry.add(WorkerSpec {
id: WorkerId("w1".into()),
url: worker.url.clone(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
});
let policies = Arc::new(build_policy_registry(&cfg).unwrap());
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap());
// Aggressive 50ms timeout: the janitor will sweep on the next
// tick (every 20ms) and fire the cancellation token before the
// upstream returns.
let router_inflight_load = RouterInflightLoadRegistry::new(
Arc::new(sgl_router::state::load_monitor::router_inflight_load::SystemTimeClock),
Duration::from_millis(50),
);
let _janitor = spawn_janitor(Arc::clone(&router_inflight_load), Duration::from_millis(20));
let ctx = Arc::new(AppContext::with_router_inflight_load(
cfg,
tokenizers,
proxy,
registry,
policies,
router_inflight_load,
));
let (ctx, _janitor) = build_ctx_with_janitor(&worker.url);
let app = build_router(ctx);
let req = Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(
serde_json::to_vec(&serde_json::json!({
"model": "tiny",
"messages": [{"role": "user", "content": "hi"}],
"stream": false,
}))
.unwrap(),
))
let res = app
.oneshot(cancellation::request(serde_json::json!({})))
.await
.unwrap();
let res = app.oneshot(req).await.unwrap();
assert_eq!(
res.status(),
StatusCode::GATEWAY_TIMEOUT,
@@ -1677,6 +1676,53 @@ async fn janitor_expiry_returns_504_stale_request_expired() {
body_str.contains("\"code\":\"stale_request_expired\""),
"504 body must encode the same code in the JSON envelope: {body_str}",
);
assert_engine_abort(&worker).await;
}
#[tokio::test]
async fn janitor_expiry_aborts_before_headers_and_mid_stream() {
use crate::common::mock_worker::MockWorker;
for before_headers in [true, false] {
let worker = if before_headers {
MockWorker::start_hanging(Duration::from_secs(2)).await
} else {
MockWorker::start_slow_stream(vec!["data: a\n\n"], Duration::from_secs(2)).await
};
let (ctx, _janitor) = build_ctx_with_janitor(&worker.url);
let response = build_router(ctx)
.oneshot(cancellation::request(serde_json::json!({"stream":true})))
.await
.unwrap();
assert_eq!(
response.status(),
if before_headers {
StatusCode::GATEWAY_TIMEOUT
} else {
StatusCode::OK
}
);
let result = response.into_body().collect().await;
assert_eq!(result.is_ok(), before_headers);
assert_engine_abort(&worker).await;
}
}
async fn assert_engine_abort(worker: &crate::common::mock_worker::MockWorker) {
tokio::time::timeout(TEST_TIMEOUT, async {
while worker.abort_log.lock().unwrap().is_empty() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap();
let forwarded: serde_json::Value =
serde_json::from_slice(worker.captured.lock().unwrap().last_body.as_ref().unwrap())
.unwrap();
assert_eq!(
*worker.abort_log.lock().unwrap(),
vec![serde_json::json!({"rid":forwarded["rid"], "abort_all":false})]
);
}
/// Task A: a non-streaming request that errors out (upstream
@@ -0,0 +1,205 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use super::*;
use axum::{extract::State, http::HeaderMap, routing::post, Json, Router};
use serde_json::{json, Value};
use tokio::sync::mpsc;
type Event = (&'static str, Value);
type Events = mpsc::UnboundedSender<Event>;
async fn chat(State(events): State<Events>, Json(body): Json<Value>) -> (StatusCode, Body) {
events.send(("chat", body.clone())).unwrap();
if body["before_headers"] == true || (body["hold"] == true && body["stream"] != true) {
std::future::pending::<()>().await;
}
let response = if body["hold"] == true {
Body::from_stream(futures::stream::pending::<
Result<bytes::Bytes, std::io::Error>,
>())
} else if body["stream"] == true {
Body::from("data: [DONE]\n\n")
} else {
Body::from("{}")
};
(
StatusCode::from_u16(body["status"].as_u64().unwrap_or(200) as u16).unwrap(),
response,
)
}
async fn abort(
State(events): State<Events>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> StatusCode {
assert_eq!(headers["authorization"], "Bearer test");
events.send(("abort", body)).unwrap();
StatusCode::INTERNAL_SERVER_ERROR // Abort failures must not affect the worker's breaker.
}
struct Harness {
ctx: Arc<AppContext>,
events: mpsc::UnboundedReceiver<Event>,
server: tokio::task::JoinHandle<()>,
}
impl Harness {
async fn new() -> Self {
let (events, rx) = mpsc::unbounded_channel();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let ctx = build_ctx_with_worker(&format!("http://{}", listener.local_addr().unwrap()));
let server = tokio::spawn(async move {
axum::serve(
listener,
Router::new()
.route("/v1/chat/completions", post(chat))
.route("/abort_request", post(abort))
.with_state(events),
)
.await
.unwrap();
});
Self {
ctx,
events: rx,
server,
}
}
async fn event(&mut self, expected: &str) -> Value {
let (kind, body) = tokio::time::timeout(TEST_TIMEOUT, self.events.recv())
.await
.unwrap()
.unwrap();
assert_eq!(kind, expected);
body
}
async fn quiet(&mut self) {
assert!(
tokio::time::timeout(Duration::from_millis(50), self.events.recv())
.await
.is_err()
);
}
}
impl Drop for Harness {
fn drop(&mut self) {
self.server.abort();
}
}
pub(super) fn request(mut body: Value) -> Request<Body> {
body["model"] = json!("tiny");
Request::post("/v1/chat/completions")
.header("content-type", "application/json")
.header("authorization", "Bearer test")
.header("x-request-id", "reused-gateway-id")
.body(Body::from(body.to_string()))
.unwrap()
}
#[tokio::test]
async fn only_unfinished_requests_abort() {
let mut h = Harness::new().await;
for (stream, hold, before_headers) in [
(false, false, false),
(true, false, false),
(false, true, false),
(true, true, true),
(true, true, false),
] {
let task = tokio::spawn(build_router(h.ctx.clone()).oneshot(request(json!({
"stream": stream, "hold": hold, "before_headers": before_headers,
}))));
let forwarded = h.event("chat").await;
assert!(crate::common::is_engine_shaped_rid(
forwarded["rid"].as_str().unwrap()
));
let worker = h.ctx.registry.get(&WorkerId("w1".into())).unwrap();
if hold && (!stream || before_headers) {
worker.breaker.record_failure();
worker.breaker.record_failure();
task.abort();
assert!(task.await.unwrap_err().is_cancelled());
} else {
let response = tokio::time::timeout(TEST_TIMEOUT, task)
.await
.unwrap()
.unwrap()
.unwrap();
if hold {
drop(response); // Silent upstream: cancellation must not wait for a token.
} else {
response.into_body().collect().await.unwrap();
}
}
if hold {
assert_eq!(
h.event("abort").await,
json!({"rid": forwarded["rid"], "abort_all": false})
);
}
h.quiet().await;
assert_eq!(worker.breaker.snapshot().state_code, 0);
worker.breaker.record_success();
}
}
#[tokio::test]
async fn caller_ids_fan_out_and_rejected_streams_do_not_abort() {
let mut h = Harness::new().await;
for fields in [
json!({"rid":"a"}),
json!({"rid":["a","b"]}),
json!({"n":2}),
json!({"status":400}),
json!({"status":429}),
json!({"status":500}),
json!({"status":503}),
] {
let mut body = json!({"stream":true, "hold":true});
body.as_object_mut()
.unwrap()
.extend(fields.as_object().unwrap().clone());
let response = build_router(h.ctx.clone())
.oneshot(request(body))
.await
.unwrap();
let forwarded = h.event("chat").await;
if fields.get("status").is_none() {
assert_eq!(forwarded.get("rid"), fields.get("rid"));
}
drop(response);
h.quiet().await;
}
}
#[tokio::test]
async fn concurrent_requests_with_the_same_header_get_distinct_abort_ids() {
let mut h = Harness::new().await;
let mut tasks = Vec::new();
for _ in 0..2 {
tasks.push(tokio::spawn(
build_router(h.ctx.clone()).oneshot(request(json!({"hold":true}))),
));
}
let mut rids = Vec::new();
for _ in 0..2 {
rids.push(h.event("chat").await["rid"].clone());
}
assert_ne!(rids[0], rids[1]);
for task in tasks {
task.abort();
assert!(task.await.unwrap_err().is_cancelled());
}
for _ in 0..2 {
let aborted = h.event("abort").await;
let index = rids.iter().position(|rid| rid == &aborted["rid"]).unwrap();
rids.remove(index);
}
h.quiet().await;
}
@@ -230,7 +230,19 @@ async fn length_selects_plain_bucket_before_engine_selection() {
let response = app.clone().oneshot(request(body("hi"))).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let _ = response.into_body().collect().await.unwrap();
assert!(short_worker.captured.lock().unwrap().last_body.is_some());
let forwarded: serde_json::Value = serde_json::from_slice(
short_worker
.captured
.lock()
.unwrap()
.last_body
.as_ref()
.unwrap(),
)
.unwrap();
assert!(crate::common::is_engine_shaped_rid(
forwarded["rid"].as_str().unwrap()
));
assert!(long_worker.captured.lock().unwrap().last_body.is_none());
let response = app
@@ -293,6 +305,7 @@ async fn pd_picks_both_groups_from_selected_bucket_and_shares_bootstrap() {
let d: serde_json::Value =
serde_json::from_slice(decode.captured.lock().unwrap().last_body.as_ref().unwrap())
.unwrap();
assert!(p.get("rid").is_none() && d.get("rid").is_none());
assert!(p["bootstrap_room"].is_number());
assert_eq!(p["bootstrap_room"], d["bootstrap_room"]);
let calls = policy.calls.lock().unwrap();
@@ -39,9 +39,22 @@ pub struct MockWorker {
// Used in header_forwarding_test; not every test file reads captured headers.
#[allow(dead_code)]
pub captured: Arc<Mutex<CapturedHeaders>>,
#[allow(dead_code)]
pub abort_log: Arc<Mutex<Vec<Value>>>,
_shutdown: oneshot::Sender<()>,
}
#[allow(dead_code)] // shared across all axum variants
fn abort_request_route<S>(log: Arc<Mutex<Vec<Value>>>) -> axum::routing::MethodRouter<S>
where
S: Clone + Send + Sync + 'static,
{
post(move |Json(body): Json<Value>| async move {
log.lock().unwrap().push(body);
StatusCode::OK
})
}
impl MockWorker {
/// Bind to a random port on 127.0.0.1 and start serving.
///
@@ -50,6 +63,7 @@ impl MockWorker {
#[allow(dead_code)] // Only used by some test files.
pub async fn start(stream_chunks: Vec<&'static str>) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let state = MockWorkerState {
captured: captured.clone(),
stream_chunks: Arc::new(stream_chunks),
@@ -60,6 +74,7 @@ impl MockWorker {
let app = axum::Router::new()
.route("/v1/chat/completions", post(chat))
.route("/server_info", get(serve_tiny_server_info))
.route("/abort_request", abort_request_route(abort_log.clone()))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -77,6 +92,7 @@ impl MockWorker {
Self {
url,
captured,
abort_log,
_shutdown: tx,
}
}
@@ -88,6 +104,7 @@ impl MockWorker {
#[allow(dead_code)]
pub async fn start_hanging(delay: Duration) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
#[derive(Clone)]
struct HangState {
@@ -127,6 +144,7 @@ impl MockWorker {
let app = axum::Router::new()
.route("/v1/chat/completions", post(hang_handler))
.route("/server_info", get(serve_tiny_server_info))
.route("/abort_request", abort_request_route(abort_log.clone()))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -144,6 +162,7 @@ impl MockWorker {
Self {
url,
captured,
abort_log,
_shutdown: tx,
}
}
@@ -154,6 +173,7 @@ impl MockWorker {
#[allow(dead_code)]
pub async fn start_slow_stream(chunks: Vec<&'static str>, delay: Duration) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
#[derive(Clone)]
struct SlowState {
@@ -207,6 +227,7 @@ impl MockWorker {
let app = axum::Router::new()
.route("/v1/chat/completions", post(slow_chat))
.route("/server_info", get(serve_tiny_server_info))
.route("/abort_request", abort_request_route(abort_log.clone()))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -224,6 +245,7 @@ impl MockWorker {
Self {
url,
captured,
abort_log,
_shutdown: tx,
}
}
@@ -248,6 +270,7 @@ impl MockWorker {
partial_body_bytes: &'static [u8],
) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap();
let url = format!("http://{addr}");
@@ -309,6 +332,7 @@ impl MockWorker {
Self {
url,
captured,
abort_log,
_shutdown: tx,
}
}
@@ -319,6 +343,7 @@ impl MockWorker {
#[allow(dead_code)]
pub async fn start_returning_error(status: StatusCode, body: Value) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let body_arc = Arc::new(body.to_string());
#[derive(Clone)]
@@ -360,6 +385,7 @@ impl MockWorker {
let app = axum::Router::new()
.route("/v1/chat/completions", post(error_handler))
.route("/server_info", get(serve_tiny_server_info))
.route("/abort_request", abort_request_route(abort_log.clone()))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -377,6 +403,7 @@ impl MockWorker {
Self {
url,
captured,
abort_log,
_shutdown: tx,
}
}
@@ -6,3 +6,11 @@
pub mod cache_aware_fixture;
pub mod mock_worker;
pub mod streaming;
#[allow(dead_code)] // not every test file inspects forwarded rids
pub fn is_engine_shaped_rid(rid: &str) -> bool {
rid.len() == 32
&& rid
.bytes()
.all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase())
}
@@ -72,6 +72,7 @@ async fn h2c_client_reaches_http2_only_worker() {
"/v1/chat/completions",
&axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"),
None,
)
.await
.expect("h2c client must reach an HTTP/2-only worker");
@@ -96,6 +97,7 @@ async fn http1_client_cannot_reach_http2_only_worker() {
"/v1/chat/completions",
&axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"),
None,
)
.await;
assert!(
@@ -164,6 +166,7 @@ async fn h2c_client_streams_sse_from_http2_only_worker() {
&axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"),
None,
None,
Some(Box::new(move || {
flag.store(true, std::sync::atomic::Ordering::SeqCst);
})),
@@ -376,3 +376,61 @@ async fn pd_mode_prefill_5xx_does_not_poison_decode_response() {
let pv = parse_body(&prefill_body);
assert_eq!(bootstrap_port(&pv), Some(8997));
}
fn streaming_chat_request() -> Request<Body> {
Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(
serde_json::to_vec(&json!({
"model": "tiny",
"messages": [{"role": "user", "content": "hi"}],
"stream": true,
}))
.unwrap(),
))
.unwrap()
}
#[tokio::test]
async fn pd_mode_disconnect_does_not_abort_either_worker() {
let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start_slow_stream(
vec!["data: a\n\n", "data: b\n\n", "data: c\n\n"],
Duration::from_millis(50),
)
.await;
let ctx = build_ctx(vec![
WorkerSpec {
id: WorkerId("p1".into()),
url: prefill.url.clone(),
mode: WorkerMode::Prefill,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: Some(8997),
},
WorkerSpec {
id: WorkerId("d1".into()),
url: decode.url.clone(),
mode: WorkerMode::Decode,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
},
]);
let app = build_router(ctx);
let res = app.oneshot(streaming_chat_request()).await.unwrap();
assert_eq!(res.status(), StatusCode::OK);
use futures::StreamExt;
let mut data_stream = res.into_body().into_data_stream();
assert!(data_stream.next().await.is_some());
drop(data_stream);
tokio::time::sleep(Duration::from_millis(100)).await;
for worker in [&prefill, &decode] {
assert!(worker.abort_log.lock().unwrap().is_empty());
let body = await_captured_body(worker, Duration::from_secs(2), "PD worker").await;
assert!(parse_body(&body).get("rid").is_none());
}
}
@@ -106,9 +106,23 @@ fn without_forwarding(mut cfg: Config, policy: PolicyKind) -> Config {
cfg
}
fn without_minted_rid(mut body: Value) -> Value {
let rid = body
.as_object_mut()
.expect("a forwarded chat body is an object")
.remove("rid");
assert!(
rid.as_ref()
.and_then(Value::as_str)
.is_some_and(crate::common::is_engine_shaped_rid),
"plain mode must mint an abort rid; got {rid:?}",
);
body
}
async fn assert_forwarded_unchanged(ctx: &Arc<AppContext>, mock: &MockWorker, request: &Value) {
assert_eq!(send(Arc::clone(ctx), request.clone()).await, StatusCode::OK);
assert_eq!(captured(mock), *request);
assert_eq!(without_minted_rid(captured(mock)), *request);
assert!(!ctx
.metrics
.render()
@@ -428,6 +442,6 @@ async fn kimi_ids_forward_with_engine_rendering_fallback() {
} else {
assert!(ids.is_none());
}
assert_eq!(captured(&mock), request);
assert_eq!(without_minted_rid(captured(&mock)), request);
}
}
@@ -303,13 +303,26 @@ template <int NUM_TOP_K, int HOT_BUFFER_SIZE>
struct SmemLayout {
static constexpr int HASH_SIZE = NUM_TOP_K * 2;
static constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + WARP_SIZE - 1) / WARP_SIZE;
// int32_t region: top_k_tokens + chunk_offset + evict_chunk_offset + hash_keys + total_hits + newest_hit
// int32_t region: top_k_tokens + chunk offsets + hash keys + hit counters
static constexpr int TOTAL_INT32 = NUM_TOP_K + (NUM_BUFFER_CHUNKS + 1) + (NUM_BUFFER_CHUNKS + 1) + HASH_SIZE + 2;
// int16_t region: lru_slots_out + hash_vals
static constexpr int TOTAL_INT16 = HOT_BUFFER_SIZE + HASH_SIZE;
static constexpr size_t BYTES = TOTAL_INT32 * sizeof(int32_t) + TOTAL_INT16 * sizeof(int16_t);
};
template <int SPARSE_BLOCK_SIZE, bool TopKIsBlocks>
__device__ __forceinline__ int32_t resolve_selected_token(const int32_t* top_k, int32_t token_index) {
if constexpr (TopKIsBlocks) {
const int32_t block_index = top_k[token_index / SPARSE_BLOCK_SIZE];
if (block_index < 0) {
return -1;
}
return block_index * SPARSE_BLOCK_SIZE + token_index % SPARSE_BLOCK_SIZE;
} else {
return top_k[token_index];
}
}
// Each block processes one request
// req_pool_indices and seq_lens can each be int32_t or int64_t
// Layout: [HOT_BUFFER_SIZE slots for LRU] + [page_size slots for newest token]
@@ -319,23 +332,28 @@ struct SmemLayout {
// false -> generic byte-stride: device + host both linear, stride = item_size_bytes
// true -> DSv4 page-padded device + page-padded host (kvcacheio.cuh constants)
//
// TopKIsBlocks makes the kernel consume block ids directly. It resolves token
// positions in registers and writes the flattened token-slot table expected by
// sparse attention without materializing an intermediate token-index tensor.
// RecordMissPlan records this step's miss plan (miss_src/dst = host/device loc
// per miss, miss_count per request) for shared-index skip layers to replay via
// copy_cache_planned_kernel. SkipIO elides only the KV byte movement (timing
// probe; output is garbage). Both are compile-time flags so the production
// (false, false) instantiation stays byte-identical.
// probe; output is garbage). These are compile-time flags, so inactive paths
// are removed from each specialization.
template <
int BLOCK_SIZE,
int NUM_TOP_K,
int HOT_BUFFER_SIZE,
bool IsMLA,
bool IsDsv4Layout,
int SPARSE_BLOCK_SIZE,
bool TopKIsBlocks,
bool RecordMissPlan,
bool SkipIO,
typename SeqLensT,
typename ReqPoolIndicesT>
__global__ void load_cache_to_device_buffer_kernel(
const int32_t* __restrict__ top_k_tokens,
const int32_t* __restrict__ top_k,
int32_t* __restrict__ device_buffer_tokens,
const int64_t* __restrict__ host_cache_locs,
const int32_t* __restrict__ device_buffer_locs,
@@ -351,7 +369,7 @@ __global__ void load_cache_to_device_buffer_kernel(
int64_t buffer_stride_0,
int64_t host_stride,
int64_t lru_slot_stride_0,
int64_t top_k_tokens_stride,
int64_t top_k_stride,
int64_t top_k_device_locs_stride,
int64_t page_size,
int64_t item_size_bytes,
@@ -360,9 +378,12 @@ __global__ void load_cache_to_device_buffer_kernel(
int32_t* __restrict__ miss_count_out,
int64_t plan_stride) {
static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA).");
// todo hisparse: support page wise sparsity
static_assert(SPARSE_BLOCK_SIZE > 0, "SPARSE_BLOCK_SIZE must be positive.");
// Cache residency and LRU replacement remain token-granular even when the
// sparse-attention selection arrives as block ids.
constexpr int NUM_TOP_K_TOKENS = NUM_TOP_K * (TopKIsBlocks ? SPARSE_BLOCK_SIZE : 1);
constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE;
constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K + WARP_SIZE - 1) / WARP_SIZE;
constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K_TOKENS + WARP_SIZE - 1) / WARP_SIZE;
constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + WARP_SIZE - 1) / WARP_SIZE;
const int bid = blockIdx.x;
@@ -372,7 +393,7 @@ __global__ void load_cache_to_device_buffer_kernel(
// CUDA graph pads the batch to a captured size. Keep padded output rows
// invalid without a separate fill kernel.
if (bid >= num_real_reqs[0]) {
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
req_top_k_device_locs[i] = -1;
}
return;
@@ -386,7 +407,7 @@ __global__ void load_cache_to_device_buffer_kernel(
const int64_t seq_len = seq_lens[bid];
// Calculate offsets for this request
const int32_t* req_top_k_tokens = top_k_tokens + bid * top_k_tokens_stride;
const int32_t* req_top_k = top_k + bid * top_k_stride;
const int64_t buffer_offset = rid * buffer_stride_0;
int32_t* req_device_buffer_tokens = device_buffer_tokens + buffer_offset;
@@ -396,14 +417,16 @@ __global__ void load_cache_to_device_buffer_kernel(
// Fast path: short sequences have all tokens in the device buffer in order.
if (seq_len <= HOT_BUFFER_SIZE) {
const int count = (seq_len < NUM_TOP_K) ? static_cast<int>(seq_len) : NUM_TOP_K;
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
const int count = (seq_len < NUM_TOP_K_TOKENS) ? static_cast<int>(seq_len) : NUM_TOP_K_TOKENS;
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
int32_t device_loc = -1;
if (i < count) {
int32_t token_pos = req_top_k_tokens[i];
if (token_pos >= 0) {
const int32_t token_pos = resolve_selected_token<SPARSE_BLOCK_SIZE, TopKIsBlocks>(req_top_k, i);
if constexpr (TopKIsBlocks) {
if (token_pos >= 0 && token_pos < seq_len) {
device_loc = req_device_buffer_locs[token_pos];
}
} else if (i < count && token_pos >= 0) {
device_loc = req_device_buffer_locs[token_pos];
}
req_top_k_device_locs[i] = device_loc;
}
@@ -418,21 +441,21 @@ __global__ void load_cache_to_device_buffer_kernel(
// Dynamic shared memory layout: int32_t arrays first, then int16_t arrays.
extern __shared__ char smem_raw[];
using Layout = SmemLayout<NUM_TOP_K, HOT_BUFFER_SIZE>;
using Layout = SmemLayout<NUM_TOP_K_TOKENS, HOT_BUFFER_SIZE>;
constexpr int HASH_SIZE = Layout::HASH_SIZE;
int32_t* smem_i32 = reinterpret_cast<int32_t*>(smem_raw);
// Top-k token positions; reused as miss-token scratch in the copy phase
int32_t* s_top_k_tokens = smem_i32;
// Prefix-sum offsets for hit counting and miss counting
int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K;
int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K_TOKENS;
// Prefix-sum offsets for evictable counting
int32_t* s_evict_chunk_offset = s_chunk_offset + (NUM_BUFFER_CHUNKS + 1);
// Open-addressing hash table: top-k token_id -> top-k index (keys)
int32_t* s_hash_keys = s_evict_chunk_offset + (NUM_BUFFER_CHUNKS + 1);
// Scalar counters
int32_t& s_total_hits = s_hash_keys[HASH_SIZE];
int32_t& s_newest_hit = s_hash_keys[HASH_SIZE + 1];
int32_t& s_total_misses = s_hash_keys[HASH_SIZE + 1];
int16_t* smem_i16 = reinterpret_cast<int16_t*>(smem_i32 + Layout::TOTAL_INT32);
// Compacted slot ordering: [hits fwd-> ... <-evictables bwd]
@@ -443,7 +466,7 @@ __global__ void load_cache_to_device_buffer_kernel(
// Initialize shared memory: counters, hash table, prefix-sum offsets.
if (tid == 0) {
s_total_hits = 0;
s_newest_hit = 0;
s_total_misses = 0;
}
for (int i = tid; i < HASH_SIZE; i += BLOCK_SIZE) {
s_hash_keys[i] = HASH_EMPTY;
@@ -458,14 +481,20 @@ __global__ void load_cache_to_device_buffer_kernel(
const int32_t newest_token = seq_len - 1;
// Insert top-k tokens into shared-memory hash table.
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
int32_t token_idx = req_top_k_tokens[i];
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
const int32_t token_idx = resolve_selected_token<SPARSE_BLOCK_SIZE, TopKIsBlocks>(req_top_k, i);
if constexpr (TopKIsBlocks) {
if (token_idx < 0 || token_idx >= seq_len) {
s_top_k_tokens[i] = TOKEN_HIT;
req_top_k_device_locs[i] = -1;
continue;
}
}
if (token_idx == newest_token) {
// If topk includes the latest token, bind its canonical occurrence to newest_slot (at HOT_BUFFER_SIZE) and mark
// it as a hit. newest_slot is at the first position of the extra page, excluded from LRU tracking.
s_top_k_tokens[i] = TOKEN_HIT;
req_top_k_device_locs[i] = req_device_buffer_locs[newest_slot];
s_newest_hit = 1;
} else {
int slot = hash_slot(token_idx, HASH_SIZE);
while (true) {
@@ -580,7 +609,7 @@ __global__ void load_cache_to_device_buffer_kernel(
const int chunk_token_start = chunk_idx * WARP_SIZE;
const int my_token_idx = chunk_token_start + lane_id;
const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K);
const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K_TOKENS);
int32_t my_token = 0;
bool is_miss = false;
@@ -611,6 +640,9 @@ __global__ void load_cache_to_device_buffer_kernel(
#else
total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses);
#endif
if (tid == 0) {
s_total_misses = total_misses;
}
}
__syncthreads();
@@ -632,7 +664,7 @@ __global__ void load_cache_to_device_buffer_kernel(
}
__syncthreads();
total_misses = NUM_TOP_K - s_total_hits - s_newest_hit;
total_misses = s_total_misses;
if constexpr (RecordMissPlan) {
if (tid == 0) {
miss_count_out[bid] = total_misses;
@@ -695,10 +727,12 @@ template <
int HOT_BUFFER_SIZE,
bool IsMLA,
bool IsDsv4Layout,
int SPARSE_BLOCK_SIZE,
bool TopKIsBlocks,
bool RecordMissPlan,
bool SkipIO>
void load_cache_to_device_buffer(
tvm::ffi::TensorView top_k_tokens,
tvm::ffi::TensorView top_k,
tvm::ffi::TensorView device_buffer_tokens,
tvm::ffi::TensorView host_cache_locs,
tvm::ffi::TensorView device_buffer_locs,
@@ -718,7 +752,8 @@ void load_cache_to_device_buffer(
tvm::ffi::TensorView miss_count_out) {
using namespace host;
const int64_t bs = top_k_tokens.shape()[0];
constexpr int NUM_TOP_K_TOKENS = NUM_TOP_K * (TopKIsBlocks ? SPARSE_BLOCK_SIZE : 1);
const int64_t bs = top_k.shape()[0];
const int64_t host_stride = host_cache_locs.shape()[1];
// Miss-plan side outputs; 0-dim sentinels when RecordMissPlan is false.
int64_t* const miss_src_ptr = RecordMissPlan ? static_cast<int64_t*>(miss_src_out.data_ptr()) : nullptr;
@@ -730,9 +765,9 @@ void load_cache_to_device_buffer(
}
const int64_t buffer_stride_0 = device_buffer_tokens.strides()[0];
const int64_t lru_slot_stride_0 = lru_slots.strides()[0];
const int64_t top_k_tokens_stride = top_k_tokens.strides()[0];
const int64_t top_k_stride = top_k.strides()[0];
const int64_t top_k_device_locs_stride = top_k_device_locs.strides()[0];
const auto kernel_device = top_k_tokens.device();
const auto kernel_device = top_k.device();
const auto device = LaunchKernel::resolve_device(kernel_device);
const void* const host_cache_k_ptr = runtime::get_device_accessible_ptr(host_cache_k);
const void* const host_cache_v_ptr =
@@ -741,7 +776,7 @@ void load_cache_to_device_buffer(
// Generic lambda: int32/int64 kernel variants are compiled for both
// seq_lens and req_pool_indices; the correct combo is selected at runtime.
auto launch = [&](auto kernel_fn, const auto* seq_lens_ptr, const auto* req_pool_indices_ptr) {
constexpr size_t smem_bytes = SmemLayout<NUM_TOP_K, HOT_BUFFER_SIZE>::BYTES;
constexpr size_t smem_bytes = SmemLayout<NUM_TOP_K_TOKENS, HOT_BUFFER_SIZE>::BYTES;
#ifndef USE_ROCM
if constexpr (smem_bytes > 48u * 1024u) {
cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
@@ -749,7 +784,7 @@ void load_cache_to_device_buffer(
#endif
LaunchKernel(bs, BLOCK_SIZE, device, smem_bytes)(
kernel_fn,
static_cast<const int32_t*>(top_k_tokens.data_ptr()),
static_cast<const int32_t*>(top_k.data_ptr()),
static_cast<int32_t*>(device_buffer_tokens.data_ptr()),
static_cast<const int64_t*>(host_cache_locs.data_ptr()),
static_cast<const int32_t*>(device_buffer_locs.data_ptr()),
@@ -765,7 +800,7 @@ void load_cache_to_device_buffer(
buffer_stride_0,
host_stride,
lru_slot_stride_0,
top_k_tokens_stride,
top_k_stride,
top_k_device_locs_stride,
page_size,
item_size_bytes,
@@ -788,6 +823,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int64_t,
@@ -802,6 +839,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int64_t,
@@ -816,6 +855,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int32_t,
@@ -830,6 +871,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int32_t,
@@ -23,6 +23,7 @@ from ..common.utils import (
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
"HAS_SINK": lambda args: args["sink_ptr"] is not None,
"BATCH_SIZE_BUCKET": lambda args: triton.next_power_of_2(args["batch_size"]),
"HAS_HISPARSE_SLOTS": lambda args: args["hisparse_slots_ptr"] is not None,
}
)
@triton.autotune(
@@ -43,6 +44,7 @@ def _gqa_share_sparse_decode_kernel(
idx_ptr, # topk index: qh x b x topk
o_ptr, # O partial: c x b x qh x d
lse_ptr, # lse partial: c x b x qh
hisparse_slots_ptr, # pre-resolved device slots: kh x b x (topk * block)
seq_lens,
slot_ids,
# shape
@@ -52,6 +54,8 @@ def _gqa_share_sparse_decode_kernel(
head_dim,
max_topk,
max_kv_len,
hisparse_slots_stride_h,
hisparse_slots_stride_b,
# sm_scale
sm_scale,
# per-tensor KV dequant scales (1.0 when the cache is unit-scaled)
@@ -89,6 +93,7 @@ def _gqa_share_sparse_decode_kernel(
NUM_TOPK_CHUNKS: tl.constexpr,
HAS_SINK: tl.constexpr,
IS_FP8: tl.constexpr,
HAS_HISPARSE_SLOTS: tl.constexpr,
):
# decode program ids: split-K over the topk dimension to give every SM
# something to do at small batch. pid(0) folds (batch, chunk) together so
@@ -161,18 +166,30 @@ def _gqa_share_sparse_decode_kernel(
# only iterate over this chunk's topk slice. the load must respect the
# per-chunk start offset.
cur_idx_ptr = idx_base + chunk_start_topk * stride_ti_t
hisparse_topk_counter = chunk_start_topk
for _ in tl.range(chunk_start_topk, chunk_end_topk):
# load index
c = tl.load(cur_idx_ptr).to(tl.int32) * BLOCK_SIZE_N
cur_idx_ptr = cur_idx_ptr + stride_ti_t
# resolve slots for this block via req_to_token
pos = c + off_n
pos_mask = pos < seq_len
slots = tl.load(
req_to_token_ptr + sid * stride_r2t_b + pos,
mask=pos_mask,
other=0,
).to(tl.int64)
if HAS_HISPARSE_SLOTS:
slots = tl.load(
hisparse_slots_ptr
+ pid_kh * hisparse_slots_stride_h
+ pid_b * hisparse_slots_stride_b
+ hisparse_topk_counter * BLOCK_SIZE_N
+ off_n,
mask=off_n < BLOCK_SIZE_N,
other=0,
).to(tl.int64)
hisparse_topk_counter = hisparse_topk_counter + 1
else:
slots = tl.load(
req_to_token_ptr + sid * stride_r2t_b + pos,
mask=pos_mask,
other=0,
).to(tl.int64)
slots = (slots + max_slots) % max_slots # safety against negative
# load K as (head_dim, BLOCK_SIZE_N) via indirect addressing
k_off = (
@@ -321,6 +338,7 @@ def flash_decode_with_gqa_share_sparse(
q_scale: Optional[float] = None,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
hisparse_slots: Optional[torch.Tensor] = None,
) -> torch.Tensor:
triton.set_allocator(robust_allocator)
is_fp8 = check_sparse_kv_fp8(q, k_cache, v_cache, label="decode")
@@ -384,6 +402,7 @@ def flash_decode_with_gqa_share_sparse(
topk_idx,
o_partial,
lse_partial,
hisparse_slots,
seq_lens,
slot_ids,
max_slots,
@@ -392,6 +411,8 @@ def flash_decode_with_gqa_share_sparse(
head_dim,
max_topk,
max_kv_len,
hisparse_slots.stride(0) if hisparse_slots is not None else 0,
hisparse_slots.stride(1) if hisparse_slots is not None else 0,
sm_scale,
k_scale,
v_scale,
@@ -28,6 +28,7 @@ from ..common.utils import (
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
"BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"] * args["BLOCK_SIZE_H"],
"HAS_SINK": lambda args: args["sink_ptr"] is not None,
"HAS_LOC_MAPPING": lambda args: args["loc_mapping_ptr"] is not None,
}
)
@triton.autotune(
@@ -55,6 +56,7 @@ def _gqa_share_sparse_fwd_kernel(
t_ptr, # topk_idx: kh x n x k
o_ptr, # O: n x h x d
req_to_token_ptr, # req_to_token: max_reqs x max_kv_len
loc_mapping_ptr, # logical slot to HiSparse device slot
# seqlens
cu_seqlens_q,
cu_seqblocks_q,
@@ -106,6 +108,7 @@ def _gqa_share_sparse_fwd_kernel(
HAS_SINK: tl.constexpr,
USE_TMA: tl.constexpr,
IS_FP8: tl.constexpr,
HAS_LOC_MAPPING: tl.constexpr,
):
sm_scale_log2e = sm_scale * 1.4426950409
# get batch id and head id
@@ -199,6 +202,12 @@ def _gqa_share_sparse_fwd_kernel(
mask=pos_mask,
other=0,
).to(tl.int64)
if HAS_LOC_MAPPING:
slots = tl.load(
loc_mapping_ptr + slots,
mask=pos_mask,
other=0,
).to(tl.int64)
slots = (slots + max_slots) % max_slots # safety against negative
# k shape: [BLOCK_SIZE_KD, BLOCK_SIZE_K] (transposed for tl.dot)
k = tl.load(
@@ -289,6 +298,7 @@ def flash_prefill_with_gqa_share_sparse(
q_scale: Optional[float] = None,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
loc_mapping: Optional[torch.Tensor] = None,
) -> torch.Tensor:
triton.set_allocator(robust_allocator)
is_fp8 = check_sparse_kv_fp8(q, k_cache, v_cache, label="prefill")
@@ -340,6 +350,7 @@ def flash_prefill_with_gqa_share_sparse(
topk_idx,
o,
req_to_token,
loc_mapping,
cu_seqlens,
cu_seqblocks_q,
seq_lens,
+84 -1
View File
@@ -174,6 +174,8 @@ def _jit_sparse_module(
hot_buffer_size: int,
is_mla: bool = False,
is_dsv4_layout: bool = False,
top_k_block_size: int = 1,
top_k_is_blocks: bool = False,
record_miss_plan: bool = False,
skip_io: bool = False,
) -> Module:
@@ -185,6 +187,8 @@ def _jit_sparse_module(
hot_buffer_size,
is_mla,
is_dsv4_layout,
top_k_block_size,
top_k_is_blocks,
record_miss_plan,
skip_io,
)
@@ -195,6 +199,8 @@ def _jit_sparse_module(
hot_buffer_size,
is_mla,
is_dsv4_layout,
top_k_block_size,
top_k_is_blocks,
record_miss_plan,
skip_io,
)
@@ -308,7 +314,7 @@ def _load_cache_to_device_buffer_mla(
skip_io=skip_io,
)
empty = torch.empty(0)
empty = torch.empty(0, device=top_k_tokens.device)
if num_real_reqs is None:
num_real_reqs = torch.tensor(
@@ -399,6 +405,83 @@ def load_cache_to_device_buffer_mla(
)
def load_blocks_to_device_buffer_mha(
top_k_blocks: torch.Tensor,
device_buffer_tokens: torch.Tensor,
host_cache_locs: torch.Tensor,
device_buffer_locs: torch.Tensor,
host_cache_k: torch.Tensor,
host_cache_v: torch.Tensor,
device_buffer_k: torch.Tensor,
device_buffer_v: torch.Tensor,
top_k_device_locs: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
lru_slots: torch.Tensor,
item_size_bytes: int,
hot_buffer_size: int,
sparse_block_size: int,
page_size: int = 1,
block_size: int = 256,
num_real_reqs: torch.Tensor | None = None,
skip_io: bool = False,
) -> None:
"""Swap block-selected MHA K/V into the HiSparse device pool."""
num_top_k_blocks = top_k_blocks.size(1)
num_top_k_tokens = num_top_k_blocks * sparse_block_size
assert hot_buffer_size >= num_top_k_tokens, (
f"hot_buffer_size ({hot_buffer_size}) must be >= selected tokens "
f"({num_top_k_tokens})"
)
assert top_k_device_locs.size(1) >= num_top_k_tokens
k_stride = host_cache_k.stride(0) * host_cache_k.element_size()
v_stride = host_cache_v.stride(0) * host_cache_v.element_size()
assert k_stride == v_stride == item_size_bytes, (
"K/V token strides must equal item_size_bytes: "
f"k_stride={k_stride}, v_stride={v_stride}, "
f"item_size_bytes={item_size_bytes}"
)
module = _jit_sparse_module(
item_size_bytes,
block_size,
num_top_k_blocks,
hot_buffer_size,
is_mla=False,
is_dsv4_layout=False,
top_k_block_size=sparse_block_size,
top_k_is_blocks=True,
record_miss_plan=False,
skip_io=skip_io,
)
empty = torch.empty(0, device=top_k_blocks.device)
if num_real_reqs is None:
num_real_reqs = torch.tensor(
[top_k_blocks.size(0)], dtype=torch.int32, device=top_k_blocks.device
)
module.load_cache_to_device_buffer(
top_k_blocks,
device_buffer_tokens,
host_cache_locs,
device_buffer_locs,
host_cache_k,
host_cache_v,
device_buffer_k,
device_buffer_v,
top_k_device_locs,
req_pool_indices,
seq_lens,
lru_slots,
num_real_reqs,
page_size,
item_size_bytes,
empty,
empty,
empty,
)
def copy_cache_planned_mla(
*,
miss_src: torch.Tensor,
+139 -4
View File
@@ -18,10 +18,99 @@ from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_interleave
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.utils.common import strict_contiguous
from sglang.srt.runtime_context import get_parallel, get_platform
from sglang.srt.utils import is_gfx95_supported, is_hip
from sglang.srt.utils.common import is_gfx1250_supported
logger = logging.getLogger(__name__)
_AITER_MHC_RUNTIME_DISABLED = False
_AITER_MHC_ACTIVE_LOGGED = False
def _use_aiter_mhc() -> bool:
return (
not _AITER_MHC_RUNTIME_DISABLED
and is_gfx95_supported()
and envs.SGLANG_USE_AITER.get()
)
def _try_aiter_mhc_pre(
residual: torch.Tensor,
fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
rms_eps: float,
hc_pre_eps: float,
hc_sinkhorn_eps: float,
hc_post_mult_value: float,
sinkhorn_repeat: int,
norm_weight: torch.Tensor | None,
norm_eps: float | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
global _AITER_MHC_RUNTIME_DISABLED, _AITER_MHC_ACTIVE_LOGGED
try:
from aiter.ops.mhc import mhc_pre as aiter_mhc_pre
except Exception as err:
logger.warning("AITER mHC pre is unavailable, falling back: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
kwargs = {}
if norm_weight is not None:
kwargs["norm_weight"] = norm_weight
kwargs["norm_eps"] = norm_eps if norm_eps is not None else rms_eps
try:
result = aiter_mhc_pre(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
**kwargs,
)
except Exception as err:
logger.warning("AITER mHC pre failed, disabling fast path: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
if not _AITER_MHC_ACTIVE_LOGGED:
logger.info("Using AITER gfx950 mHC pre/post kernels")
_AITER_MHC_ACTIVE_LOGGED = True
return result
def _try_aiter_mhc_post(
x: torch.Tensor,
residual: torch.Tensor,
post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor,
) -> torch.Tensor | None:
global _AITER_MHC_RUNTIME_DISABLED
try:
from aiter.ops.mhc import mhc_post as aiter_mhc_post
except Exception as err:
logger.warning("AITER mHC post is unavailable, falling back: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
out = torch.empty_like(residual)
try:
aiter_mhc_post(out, x, residual, post_layer_mix, comb_res_mix)
except Exception as err:
logger.warning("AITER mHC post failed, disabling fast path: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
return out
# This module is imported during model-registry discovery. Do not import the real
# TileLang package here: it loads native CUDA stubs. The proxy below lets
# module-level @tilelang.jit declarations parse, then imports and applies real
@@ -119,6 +208,24 @@ pass_configs = {
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
}
def _use_deep_gemm_hc_prenorm() -> bool:
if is_hip() or not envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
return False
from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
return ENABLE_JIT_DEEPGEMM
def _use_tilelang_mhc_pre() -> bool:
return envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() and not is_hip()
def _use_tilelang_mhc_post() -> bool:
return envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get() and not is_hip()
FP8 = "float8_e4m3"
BF16 = "bfloat16"
FP32 = "float32"
@@ -1041,7 +1148,7 @@ def mhc_pre(
num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
)
if envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
if _use_deep_gemm_hc_prenorm():
n_splits = _compute_num_split_for_mhc_pre(num_tokens, hc_hidden_size)
gemm_out_mul = torch.empty(
@@ -1653,7 +1760,7 @@ def mhc_fused_post_pre(
hidden_size,
)
if envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
if _use_deep_gemm_hc_prenorm():
import deep_gemm
deep_gemm.tf32_hc_prenorm_gemm(
@@ -1847,7 +1954,25 @@ def _mhc_pre_dispatch(
norm_eps: float | None = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:
assert residual.dim() == 3, f"residual must be (s, n, h); got {residual.shape}"
if not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
if _use_aiter_mhc():
result = _try_aiter_mhc_pre(
residual=residual,
fn=fn,
hc_scale=hc_scale,
hc_base=hc_base,
rms_eps=rms_eps,
hc_pre_eps=hc_pre_eps,
hc_sinkhorn_eps=hc_sinkhorn_eps,
hc_post_mult_value=hc_post_mult_value,
sinkhorn_repeat=sinkhorn_repeat,
norm_weight=norm_weight,
norm_eps=norm_eps,
)
if result is not None:
post_mix, comb_mix, layer_input = result
return post_mix, comb_mix, layer_input, norm_weight is not None
if not _use_tilelang_mhc_pre():
post_mix, comb_mix, layer_input = _mhc_pre_torch(
residual=residual,
fn=fn,
@@ -1886,7 +2011,17 @@ def _mhc_post_dispatch(
) -> torch.Tensor:
assert x.dim() == 2 and residual.dim() == 3
assert post_layer_mix.dim() == 3 and comb_res_mix.dim() == 3
if not envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get():
if _use_aiter_mhc():
result = _try_aiter_mhc_post(
x=x,
residual=residual,
post_layer_mix=post_layer_mix,
comb_res_mix=comb_res_mix,
)
if result is not None:
return result
if not _use_tilelang_mhc_post():
return _mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
return mhc_post(x, residual, post_layer_mix, comb_res_mix)
@@ -181,12 +181,15 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
)
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:
raise ValueError(
"--enable-encoder-swa-bounded-replay requires DeepSeek-V4.1"
)
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:
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,
),
("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),
("unified memory", cfg.enable_unified_memory),
("PD disaggregation", cfg.disaggregation_mode != "null"),
@@ -249,23 +253,33 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
if (
read_ragged_verify_mode() is not RaggedVerifyMode.STATIC
or cfg.disaggregation_transfer_backend != "mooncake"
or cfg.dp_size != 1
or cfg.enable_dp_attention
or cfg.attn_cp_size != 1
or cfg.dcp_size != 1
):
raise ValueError(
"DeepSeek-V4.1 DSpark PD requires static verify, Mooncake, "
"DP=1 and CP=1. Both servers must enable DSpark with the same "
"block size and TP size."
"and CP=1 on both servers. DP attention is supported when "
"both servers use the same block size and target/draft KV layout."
)
from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase
prefill_graph = cfg.cuda_graph_config.prefill
if prefill_graph.backend != Backend.DISABLED and prefill_graph.max_seq_len is None:
# The captured low-ratio indexer scores a static context width; 16k
# keeps it inside the candidate window at under 1 ms per layer.
cp_breakable_prefill = (
cfg.enable_prefill_cp
and cfg.cp_strategy == "interleave"
and cfg.tp_size > 1
and prefill_graph.backend == Backend.BREAKABLE
)
if (
prefill_graph.backend != Backend.DISABLED
and prefill_graph.max_seq_len is None
and not cp_breakable_prefill
):
# The non-CP captured low-ratio indexer scores a static context width.
# CP BCG runs these sources eagerly with live prefix metadata, so this
# default would only force long-prefix CP batches back to eager.
# Explicit max_seq_len values still constrain both paths.
declare_resolution(
server_args,
"validate_deepseek_v41_features",
@@ -296,3 +310,40 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
"--enable-decoder-swa-bounded-replay cannot be combined with "
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."
)
@@ -86,14 +86,16 @@ def validate_hisparse(server_args: ServerArgs) -> None:
from sglang.srt.configs.model_config import (
is_deepseek_dsa,
is_deepseek_v4,
is_minimax_sparse,
)
hf_config = model_config_of(server_args).hf_config
is_v4_hisparse = is_deepseek_v4(hf_config)
is_m3_hisparse = is_minimax_sparse(hf_config)
is_hip = get_platform().is_hip
assert is_deepseek_dsa(hf_config) or is_v4_hisparse, (
assert is_deepseek_dsa(hf_config) or is_v4_hisparse or is_m3_hisparse, (
"--enable-hisparse is only supported for DSA (DeepSeek Sparse Attention) "
"models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. "
"models (e.g., DeepSeek V3.2, GLM-5), DeepSeek V4, and MiniMax M3 now. "
)
assert cfg.disable_radix_cache, (
@@ -121,6 +123,10 @@ def validate_hisparse(server_args: ServerArgs) -> None:
)
return
# MiniMax M3 uses its own Triton sparse kernels.
if is_m3_hisparse:
return
if resolved_view(server_args).kv_cache_dtype not in (
"bfloat16",
"auto",
+4 -2
View File
@@ -966,12 +966,14 @@ def handle_language_model_only(server_args: Any):
):
if flag:
raise ValueError(f"--language-model-only cannot be combined with {name}")
if cfg.disaggregation_mode != "null":
hf_config = model_config_of(server_args).hf_config
# V4.1 text-only workers use the standard PD KV transfer path.
if cfg.disaggregation_mode != "null" and hf_config.model_type != "deepseek_v41":
raise ValueError(
"--language-model-only is incompatible with --disaggregation-mode "
"prefill/decode"
)
architectures = model_config_of(server_args).hf_config.architectures
architectures = hf_config.architectures
if not any(
a in server_args.LANGUAGE_MODEL_ONLY_ARCHITECTURES for a in architectures
):
@@ -24,6 +24,7 @@ class StateType(str, enum.Enum):
# only the live subrange of that row for the current open pool.
DSA_TAIL = "dsa_tail"
MINIMAX_INDEX_K = "minimax_index_k"
MINIMAX_DENSE_KV = "minimax_dense_kv"
# DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot
# (req_pool_idx * ring_stride + pos % ring_stride), needs its own component.
SWA_RING = "swa_ring"
@@ -945,9 +945,33 @@ class CommonKVManager(BaseKVManager):
"enable DSpark with the same block size and target/draft KV "
"layout. Upgrade both servers together."
)
if info.attn_tp_size != self.attn_tp_size:
same_tp_with_prefill_cp = (
info.attn_cp_size > 1
and (self.is_mla_backend or self.is_hybrid_mla_backend)
and self.attn_cp_size == 1
and info.attn_tp_size * info.attn_cp_size == self.attn_tp_size
)
# Combined branch (40323-series + 40177): prefill CP can also pair
# with a DP-attention decode server. MLA KV is replicated across
# prefill CP ranks, so per-rank layouts match when attn_tp matches.
dp_decode_with_prefill_cp = (
info.attn_cp_size > 1
and self.attn_cp_size == 1
and (self.is_mla_backend or self.is_hybrid_mla_backend)
and info.attn_tp_size == self.attn_tp_size
)
non_cp_mla_layout = info.attn_cp_size == self.attn_cp_size == 1 and (
self.is_mla_backend or self.is_hybrid_mla_backend
)
if info.attn_tp_size != self.attn_tp_size and not (
same_tp_with_prefill_cp
or dp_decode_with_prefill_cp
or non_cp_mla_layout
):
raise RuntimeError(
"DeepSeek-V4.1 DSpark PD requires the same TP size on both servers"
"DeepSeek-V4.1 DSpark PD requires matching attention TP "
"unless both servers use CP=1 with an MLA KV layout, "
"or prefill runs CP with an MLA KV layout"
)
if self.dcp_size > 1:
+4 -2
View File
@@ -62,6 +62,7 @@ from sglang.srt.disaggregation.utils import (
build_staging_slot_metadata,
get_dsa_tail_state_indices,
get_kv_class,
get_kv_transfer_buf_infos,
get_qsa_pending_state_indices,
is_mla_backend,
is_unadmitted_reject,
@@ -575,8 +576,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
if self.scheduler.enable_hisparse
else self.token_to_kv_pool
)
kv_data_ptrs, kv_data_lens, kv_item_lens = (
transfer_kv_pool.get_contiguous_buf_infos()
kv_data_ptrs, kv_data_lens, kv_item_lens = get_kv_transfer_buf_infos(
transfer_kv_pool
)
kv_data_mem_kinds = (
["DRAM"] * len(kv_data_ptrs)
@@ -1579,6 +1580,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
StateType.DSA: _full_kv_pages_payload,
StateType.DSA_TAIL: _dsa_tail_payload,
StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
StateType.MINIMAX_DENSE_KV: _full_kv_pages_payload,
StateType.SWA_RING: _swa_ring_payload,
StateType.DSV4_REQUEST_STATE: _request_state_payload,
StateType.BLOCK_SCALE: _full_kv_pages_payload,
+289 -13
View File
@@ -82,6 +82,115 @@ FAILED_SESSION_RECOVERIES = Counter(
)
# ---------------------------------------------------------------------------
# Intra-node NVLink transport helpers.
#
# Mooncake's IntraNodeNvlinkTransport can only register and reach *device*
# memory (it IPC-opens the remote cudaMalloc segments). Host-resident regions
# (aux buffers, some state components) cannot be registered: one host region
# makes the whole registerLocalMemoryBatch fail, and the engine then rolls
# back *every* region, leaving the segment descriptor empty and all KV
# transfers failing with "Requested address ... not found". When the
# intra-node NVLink transport is active we therefore
# 1. register only device-memory regions, and
# 2. route blocks whose source is host memory over the ordered zmq channel
# (same ordering guarantee the aux TCP path relies on) instead of the
# transfer engine.
# ---------------------------------------------------------------------------
import ctypes as _ctypes
_CUDA_MEMORY_TYPE_DEVICE = 2
try:
from cuda.bindings import runtime as _cudart
except ImportError: # pragma: no cover - cuda-python is always present in images
_cudart = None
def _is_device_pointer(ptr: int) -> bool:
"""Probe a *local* pointer with cudaPointerGetAttributes.
Only valid for pointers owned by this process (never probe remote
segment addresses). Returns False on any error so the caller falls back
to the safe host path.
"""
if _cudart is None:
# Cannot tell; assume device so behavior stays unchanged.
return True
err, attr = _cudart.cudaPointerGetAttributes(int(ptr))
if int(err) != 0:
# Clear the error so subsequent CUDA calls are not poisoned.
_cudart.cudaGetLastError()
return False
return int(attr.type) == _CUDA_MEMORY_TYPE_DEVICE
def _read_bytes_from_address(addr: int, length: int) -> Optional[bytes]:
if length <= 0:
return b""
if _is_device_pointer(addr):
buf = bytearray(length)
# cudaMemcpyDeviceToHost = 2; synchronous default-stream copy.
err, = _cudart.cudaMemcpy(
_ctypes.addressof((_ctypes.c_char * length).from_buffer(buf)),
int(addr),
length,
2,
)
if int(err) != 0:
logger.error(
f"cudaMemcpy D2H failed (err={err}) for addr {hex(addr)} len {length}"
)
return None
return bytes(buf)
return _ctypes.string_at(int(addr), length)
def _write_bytes_to_address(addr: int, data: bytes) -> bool:
if not data:
return True
if _is_device_pointer(addr):
buf = _ctypes.create_string_buffer(data, len(data))
# cudaMemcpyHostToDevice = 1; synchronous default-stream copy.
err, = _cudart.cudaMemcpy(
int(addr), _ctypes.addressof(buf), len(data), 1
)
if int(err) != 0:
logger.error(
f"cudaMemcpy H2D failed (err={err}) for addr {hex(addr)} "
f"len {len(data)}"
)
return False
return True
_ctypes.memmove(int(addr), data, len(data))
return True
_NVLINK_INTRA_ACTIVE = None
def _nvlink_intra_transport_active() -> bool:
"""Whether mooncake installed the intra-node NVLink transport.
Mirrors the env probing in mooncake's transfer_engine_impl.cpp: the
transport is installed iff MC_INTRANODE_NVLINK is set (any value), or an
equivalent protocol selection was made.
"""
global _NVLINK_INTRA_ACTIVE
if _NVLINK_INTRA_ACTIVE is None:
active = bool(
os.environ.get("MC_INTRANODE_NVLINK")
or os.environ.get("MC_INTRA_NVLINK")
)
if not active:
proto = (os.environ.get("MOONCAKE_PROTOCOL") or "").strip().lower()
active = proto in ("nvlink_intra", "nvlink-intra", "intra_nvlink")
_NVLINK_INTRA_ACTIVE = active
return _NVLINK_INTRA_ACTIVE
# decode
@dataclasses.dataclass
class TransferInfo:
@@ -214,6 +323,7 @@ class KVArgsRegisterInfo:
class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
AUX_DATA_HEADER = b"AUX_DATA"
STATE_DATA_HEADER = b"STATE_DATA"
# Implements teardown() below, so runtime PD role switching is supported.
supports_role_switch = True
@@ -227,6 +337,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
self.init_engine()
self.register_buffer_to_engine()
# session_id -> (endpoint, dst_port, room), used to route host-memory
# transfer blocks over zmq when the intra-node NVLink transport is
# active (it cannot reach host memory). Populated on bootstrap.
self._session_endpoint_map = {}
self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
self.max_transfer_batch_indices = (
envs.SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES.get()
@@ -322,6 +436,13 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
Deduped because the unified memory pool reports one raw buffer as both
its KV and its mamba state component, and double registration fails in
the engine.
When the intra-node NVLink transport is active, host-memory regions
(aux buffers, some state components) are skipped: the transport only
accepts device memory, and a single host region fails the whole batch
and triggers a full engine-side rollback that would unregister the KV
pools too. Host-resident payloads are instead exchanged over the
ordered zmq channel (see _transfer_data / send_aux).
"""
regions: List[Tuple[int, int]] = []
seen: Set[Tuple[int, int]] = set()
@@ -338,6 +459,24 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
):
add(ptrs, lens)
if _nvlink_intra_transport_active():
device_regions = []
skipped = []
for ptr, length in regions:
if _is_device_pointer(ptr):
device_regions.append((ptr, length))
else:
skipped.append((ptr, length))
if skipped:
logger.info(
"Intra-node NVLink transport: skipping %d host-memory "
"regions from engine registration (they will be exchanged "
"over the zmq channel instead): %s",
len(skipped),
[(hex(p), l) for p, l in skipped[:8]],
)
regions = device_regions
return regions
def register_buffer_to_engine(self):
@@ -748,10 +887,63 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
if not transfer_blocks:
return 0
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
return self.engine.batch_transfer_sync(
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
if not _nvlink_intra_transport_active():
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
return self.engine.batch_transfer_sync(
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
# Intra-node NVLink transport can only move device memory. Partition
# blocks by the *local source* pointer (probing a local pointer is
# safe; the remote dst is never probed): device-sourced blocks go
# through the engine as usual, host-sourced blocks are shipped over
# the ordered zmq channel and written into the peer's buffer by the
# receiver (see _handle_state_data). This mirrors the aux TCP path.
device_blocks = []
host_blocks = []
for src, dst, length in transfer_blocks:
if _is_device_pointer(src):
device_blocks.append((src, dst, length))
else:
host_blocks.append((src, dst, length))
rc = 0
if device_blocks:
src_addrs, dst_addrs, lengths = zip(*device_blocks)
rc = self.engine.batch_transfer_sync(
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
if rc == 0 and host_blocks:
rc = self._send_host_blocks_tcp(mooncake_session_id, host_blocks)
return rc
def _send_host_blocks_tcp(self, mooncake_session_id, host_blocks):
target = self._session_endpoint_map.get(mooncake_session_id)
if target is None:
logger.error(
f"No zmq endpoint known for mooncake session "
f"{mooncake_session_id}; cannot deliver {len(host_blocks)} "
"host-memory transfer blocks"
)
return -1
endpoint, dst_port, room = target
na = NetworkAddress(endpoint, dst_port)
for src, dst, length in host_blocks:
data = _read_bytes_from_address(src, length)
if data is None:
return -1
self._send_multipart_locked(
na.to_tcp(),
[
MooncakeKVManager.STATE_DATA_HEADER,
str(room).encode("ascii"),
str(int(dst)).encode("ascii"),
struct.pack(">I", len(data)),
data,
],
is_ipv6=na.is_ipv6,
)
return 0
def _send_kvcache_generic(
self,
@@ -1398,7 +1590,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
if src_layer_ids or dst_layer_ids:
# Draft buffers break the flat [K block, V block] layout, so pair by
# layer ID instead of the half-split used by get_mha_kv_ptrs_with_pp.
if any(l != src_kv_item_len for l in self.kv_args.kv_item_lens):
if any(
item_len != src_kv_item_len for item_len in self.kv_args.kv_item_lens
):
logger.error(
f"[{mooncake_session_id}] head-sliced transfer assumes one item "
f"length for every KV entry, got {set(self.kv_args.kv_item_lens)}"
@@ -1480,8 +1674,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
):
# TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free
if (
self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK"
) or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get():
(self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK")
or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get()
or _nvlink_intra_transport_active()
):
return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs)
transfer_blocks = []
@@ -1564,6 +1760,71 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}"
)
def _host_transfer_regions(self):
"""Address ranges this process published as transfer targets.
Used to validate STATE_DATA writes. Built lazily because kv_args is
fully populated only after registration.
"""
regions = getattr(self, "_host_transfer_regions_cache", None)
if regions is None:
regions = []
for ptr, length in zip(
self.kv_args.kv_data_ptrs or [], self.kv_args.kv_data_lens or []
):
regions.append((int(ptr), int(ptr) + int(length)))
for ptr, length in zip(
self.kv_args.aux_data_ptrs or [], self.kv_args.aux_data_lens or []
):
regions.append((int(ptr), int(ptr) + int(length)))
for ptrs, lens in zip(
self.kv_args.state_data_ptrs or [], self.kv_args.state_data_lens or []
):
for ptr, length in zip(ptrs or [], lens or []):
regions.append((int(ptr), int(ptr) + int(length)))
self._host_transfer_regions_cache = regions
return regions
def _handle_state_data(self, msg: List[bytes]):
"""Handle STATE_DATA messages received by the decode thread.
Carries one host-memory transfer block that could not go through the
intra-node NVLink transport. Written directly into the local buffer at
the destination address; ordering against the final status message is
guaranteed by the shared per-endpoint zmq socket.
"""
room = int(msg[1].decode("ascii"))
dst_addr = int(msg[2].decode("ascii"))
data_length = struct.unpack(">I", msg[3])[0]
data = msg[4]
if len(data) != data_length:
logger.error(f"STATE_DATA length mismatch for bootstrap_room {room}")
return
in_region = any(
start <= dst_addr and dst_addr + len(data) <= end
for start, end in self._host_transfer_regions()
)
if not in_region:
logger.error(
f"STATE_DATA for bootstrap_room {room} targets unknown region "
f"{hex(dst_addr)}..{hex(dst_addr + len(data))}; dropping"
)
return
if not _write_bytes_to_address(dst_addr, data):
logger.error(
f"STATE_DATA write failed for bootstrap_room {room} at "
f"{hex(dst_addr)} len {len(data)}"
)
return
logger.debug(
f"Received STATE_DATA for bootstrap_room {room} at {hex(dst_addr)} "
f"with length:{len(data)}"
)
def _get_dsa_cache_transfer_skip_flags(
self, info: Optional[KVArgsRegisterInfo]
) -> Tuple[bool, bool]:
@@ -1852,12 +2113,11 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
)
or rc
)
elif st == StateType.MINIMAX_INDEX_K:
# Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
# lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
elif st in (StateType.MINIMAX_INDEX_K, StateType.MINIMAX_DENSE_KV):
# Compacted layer lists require equal TP and PP=1 on both peers.
if self.pp_size is not None and self.pp_size > 1:
raise RuntimeError(
"PD disagg: PP>1 not supported for MiniMax sparse index yet."
"PD disagg: PP>1 not supported for MiniMax state yet."
)
if (
target_rank_registration_info is not None
@@ -1866,11 +2126,17 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
):
raise RuntimeError(
"PD disagg: heterogeneous TP not supported for MiniMax "
"sparse index yet."
"state yet."
)
src_indices = list(indices)
dst_indices_local = list(dst_indices)
if len(src_indices) > len(dst_indices_local):
if st == StateType.MINIMAX_DENSE_KV:
if len(src_indices) != len(dst_indices_local):
raise RuntimeError(
f"{st.value} state index length mismatch: "
f"prefill={len(src_indices)}, dst={len(dst_indices_local)}"
)
elif len(src_indices) > len(dst_indices_local):
src_indices = src_indices[: len(dst_indices_local)]
elif len(src_indices) < len(dst_indices_local):
dst_indices_local = dst_indices_local[: len(src_indices)]
@@ -2405,6 +2671,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
):
self._staging_outstanding.pop(kv_chunk.room, None)
if kv_chunk.room in self.transfer_infos:
for sid in self.transfer_infos[kv_chunk.room]:
self._session_endpoint_map.pop(sid, None)
self.transfer_infos.pop(kv_chunk.room)
self.req_to_decode_prefix_len.pop(kv_chunk.room, None)
if self.enable_staging:
@@ -2550,6 +2818,11 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
self.transfer_infos[room][mooncake_session_id] = (
TransferInfo.from_zmq(waiting_req_bytes)
)
self._session_endpoint_map[mooncake_session_id] = (
self.transfer_infos[room][mooncake_session_id].endpoint,
self.transfer_infos[room][mooncake_session_id].dst_port,
room,
)
# NOTE: after bootstrapping we can mark the req as waiting for input
if len(self.transfer_infos[room]) == required_dst_info_num:
self.resolve_kv_replica_factor(self.transfer_infos[room])
@@ -2578,6 +2851,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
self._handle_aux_data(msg)
continue
if msg[0] == MooncakeKVManager.STATE_DATA_HEADER:
self._handle_state_data(msg)
continue
# Staging: prefill notifies a chunk written to staging buffer
if msg[0] == b"CHUNK_READY":
@@ -1282,6 +1282,7 @@ class MoriKVManager(CommonKVManager):
"swa_ring",
"c128_state",
"minimax_index_k",
"minimax_dense_kv",
):
statuses.extend(
self._send_swa_dsa_state(
@@ -1409,7 +1410,12 @@ class MoriKVManager(CommonKVManager):
f"PD state transfer does not support TP-mismatched non-MLA SWA models "
f"(prefill_tp_size={self.attn_tp_size}, decode_tp_size={peer_info.decode_tp_size})"
)
if state_type in ("qsa_pending", "qsa_compressed", "minimax_index_k"):
if state_type in (
"qsa_pending",
"qsa_compressed",
"minimax_index_k",
"minimax_dense_kv",
):
if self.pp_size is not None and self.pp_size > 1:
# MORI registration does not exchange state_layer_ids. Compact
# sparse-state lists therefore cannot be paired safely across
@@ -1445,6 +1451,7 @@ class MoriKVManager(CommonKVManager):
"qsa_compressed",
"swa_ring",
"c128_state",
"minimax_dense_kv",
):
raise RuntimeError(
f"{state_type.upper()} state index length mismatch: "
@@ -2696,17 +2696,16 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
dst_layer_ids=dst_lids,
dst_item_lens=dst_lens,
)
elif st == StateType.MINIMAX_INDEX_K:
# Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
# lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
elif st in (StateType.MINIMAX_INDEX_K, StateType.MINIMAX_DENSE_KV):
# Compacted layer lists require equal TP and PP=1 on both peers.
if self.pp_size is not None and self.pp_size > 1:
raise RuntimeError(
"PD disagg: PP>1 not supported for MiniMax sparse index yet."
"PD disagg: PP>1 not supported for MiniMax state yet."
)
if self.attn_tp_size != decode_tp_size:
raise RuntimeError(
"PD disagg: heterogeneous TP not supported for MiniMax "
"sparse index yet."
"state yet."
)
if len(src_indices) != len(dst_indices):
raise RuntimeError(
+4 -2
View File
@@ -53,6 +53,7 @@ from sglang.srt.disaggregation.utils import (
build_staging_slot_metadata,
get_dsa_tail_state_indices,
get_kv_class,
get_kv_transfer_buf_infos,
get_qsa_pending_state_indices,
is_aborted,
is_mla_backend,
@@ -256,8 +257,8 @@ class PrefillBootstrapQueue:
hf_text_config=self.scheduler.model_config.hf_text_config,
)
)
kv_data_ptrs, kv_data_lens, kv_item_lens = (
self.token_to_kv_pool.get_contiguous_buf_infos()
kv_data_ptrs, kv_data_lens, kv_item_lens = get_kv_transfer_buf_infos(
self.token_to_kv_pool
)
kv_args.prefill_end_layer = (
kv_args.prefill_start_layer + len(kv_data_ptrs)
@@ -1424,6 +1425,7 @@ class SchedulerDisaggregationPrefillMixin:
StateType.DSA: _full_kv_pages_payload,
StateType.DSA_TAIL: _dsa_tail_payload,
StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
StateType.MINIMAX_DENSE_KV: _full_kv_pages_payload,
StateType.SWA_RING: _swa_ring_payload,
StateType.DSV4_REQUEST_STATE: _request_state_payload,
StateType.BLOCK_SCALE: _full_kv_pages_payload,
+13
View File
@@ -1310,6 +1310,14 @@ def build_dsa_tail_transfer_blocks(
return transfer_blocks
def get_kv_transfer_buf_infos(pool):
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
if isinstance(pool, MiniMaxSparseKVPool):
return pool.get_sparse_kv_buf_infos()
return pool.get_contiguous_buf_infos()
def setup_state_kv_args(
kv_args: KVArgs,
token_to_kv_pool,
@@ -1375,6 +1383,11 @@ def setup_state_kv_args(
if token_to_kv_pool.index_k_pool is not None:
dp, dl, il = token_to_kv_pool.get_index_k_state_buf_infos()
append_state_component(kv_args, StateType.MINIMAX_INDEX_K, dp, dl, il)
append_state_component(
kv_args,
StateType.MINIMAX_DENSE_KV,
*token_to_kv_pool.get_dense_kv_state_buf_infos(),
)
elif hasattr(token_to_kv_pool, "get_state_buf_infos"):
data_ptrs, data_lens, item_lens = token_to_kv_pool.get_state_buf_infos()
@@ -1518,9 +1518,7 @@ class GroupCoordinator:
if self.world_size == 1:
return input_
# Always use pynccl to avoid capturing hip graph failure on torch
# version smaller than or equal to 2.11
if is_hip() and self.pynccl_comm is not None and not self.pynccl_comm.disabled:
if self.pynccl_comm is not None and not self.pynccl_comm.disabled:
self.pynccl_comm.broadcast(input_, src=src)
else:
# Broadcast.
@@ -736,7 +736,14 @@ class DSV4AttnMetadata:
if src_val is None and dst_val is None:
continue
assert dst_val is not None, f"{field_name=} {src_val=} {dst_val=}"
dst_val.copy_(src_val)
shape_mismatch = dst_val.shape != src_val.shape
assert not shape_mismatch or field_name in self._CP_GLOBAL_FIELDS, (
f"Only CP-global replay metadata may use a shorter live prefix, "
f"got {field_name=} {src_val.shape=} {dst_val.shape=}"
)
_copy_tensor_allowing_storage_alias(
dst_val, src_val, pad_value=0 if shape_mismatch else None
)
# These fields are safe to replace because captured kernels only need
# the current per-replay objects, or the field is produced inside the
@@ -988,6 +995,27 @@ def _prefill_graph_max_seq_len() -> Optional[int]:
return get_exec().graph.cuda_graph_config.prefill.max_seq_len
def _copy_tensor_allowing_storage_alias(
dst: torch.Tensor, src: torch.Tensor, *, pad_value: Optional[int] = None
) -> None:
"""Copy replay metadata while preserving capture-stable destination addresses."""
if dst is src:
return
if dst.untyped_storage().data_ptr() == src.untyped_storage().data_ptr():
src = src.clone()
if dst.shape == src.shape:
dst.copy_(src)
return
assert (
pad_value is not None
and dst.ndim == src.ndim
and dst.shape[0] >= src.shape[0]
and dst.shape[1:] == src.shape[1:]
), f"Cannot copy replay metadata from {src.shape=} to {dst.shape=}"
dst.fill_(pad_value)
dst[: src.shape[0]].copy_(src)
@dataclass
class DSV4Metadata:
core_attn_metadata: DSV4AttnMetadata
@@ -1254,6 +1282,11 @@ class DeepseekV4AttnBackend(
] = None
self.online_c128_mtp = OnlineC128MTPController(self)
self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device)
# CP V4.1 consumers share compressed KV across layers. Separate ratio
# workspaces keep those prefixes intact while each layer refreshes SWA.
self.shared_compressed_prefill_workspaces = {
ratio: SparsePrefillWorkspace(self.device) for ratio in (1, 2)
}
spec_alg = model_runner.spec_algorithm
self.needs_cpu_seq_lens = not spec_alg.is_dspark() and (
not _is_cuda or self.online_c128_mtp.enabled()
@@ -1566,8 +1599,12 @@ class DeepseekV4AttnBackend(
@property
def low_ratio_prefill_graph(self) -> bool:
"""Whether ratio-1/2 sources use captured projections and indexer metadata."""
return (
bool(self.low_ratios) and _has_dense_fp4_indexer() and _is_sm100_or_newer()
bool(self.low_ratios)
and _has_dense_fp4_indexer()
and _is_sm100_or_newer()
and get_parallel().attn_cp_size == 1
)
def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool:
@@ -2870,7 +2907,7 @@ class DeepseekV4AttnBackend(
q_lora[:num_local],
positions[:num_local].to(torch.int64),
forward_batch,
torch.tensor(q_lens_cpu, dtype=torch.int32, device=x.device),
self._move_to_device(q_lens_cpu),
q_lens_cpu,
)
@@ -3236,7 +3273,9 @@ class DeepseekV4AttnBackend(
continue
j = torch.arange(lc, device=device)
slot_chunks.append(
self.req_to_token[req_pool_indices[r], j * ratio].to(torch.int64)
self.req_to_token[req_pool_indices[r : r + 1], j * ratio].to(
torch.int64
)
// ratio
)
start += lc
@@ -3261,7 +3300,7 @@ class DeepseekV4AttnBackend(
weights = indexer.head_weights(x).float()
compress_lens = ((pos + 1) // ratio).to(torch.int32)
ks = torch.repeat_interleave(
torch.tensor(starts, dtype=torch.int32, device=device),
self._move_to_device(starts),
q_lens.to(torch.int64),
output_size=num_tokens,
)
@@ -3447,17 +3486,44 @@ class DeepseekV4AttnBackend(
self.candidate_indexer.publish_decode(inputs, page_indices, raw_indices)
)
return
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)
if isinstance(metadata.deep_gemm_metadata, list):
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(
(q_fp4[rows], q_sf[rows]),
k_cache,
weights[rows],
metadata.compressed_seq_lens[rows],
metadata.page_table[rows],
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
# 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
)
n_compressed = flat_token_ids.shape[0]
workspace = self.sparse_prefill_workspace.get(
n_compressed + cache.swa_token_ids.shape[0]
reuse_compressed = compress_ratio in (1, 2) and is_cp_active(forward_batch)
workspace_pool = (
self.shared_compressed_prefill_workspaces[compress_ratio]
if reuse_compressed
else self.sparse_prefill_workspace
)
workspace = workspace_pool.get(n_compressed + cache.swa_token_ids.shape[0])
compressed_slice = workspace[:n_compressed]
swa_slice = workspace[n_compressed:]
if compressed_slice is not None:
dequantize_k_cache_paged(
extra_k_cache,
flat_token_ids,
page_size=extra_page_size,
out=compressed_slice,
layout=token_to_kv_pool.get_extra_key_layout(layer_id),
)
source_key = None
if reuse_compressed:
source_layer = token_to_kv_pool.source_layer_of(layer_id)
source_key = (source_layer, workspace.data_ptr())
gather = cache.compressed[compress_ratio]
# A source layer may have just updated its cache in place. Consumer
# layers only reuse the compressed prefix; their top-k and SWA stay live.
if (
source_key is None
or layer_id == source_key[0]
or gather.dequantized_source != source_key
):
dequantize_k_cache_paged(
extra_k_cache,
flat_token_ids,
page_size=extra_page_size,
out=compressed_slice,
layout=token_to_kv_pool.get_extra_key_layout(layer_id),
)
if source_key is not None:
gather.dequantized_source = source_key
dequantize_k_cache_paged(
token_to_kv_pool.get_swa_key_buffer_radix(layer_id),
cache.swa_token_ids,
@@ -21,6 +21,7 @@ from sglang.srt.layers.attention.dsv4.candidate_indexer import (
)
from sglang.srt.layers.attention.dsv4.indexer import (
deep_gemm_fp4_paged_mqa_logits,
topk_transform_paged_from_metadata,
)
CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16
@@ -176,6 +177,10 @@ class DeepGemmCandidateIndexer:
metadata."""
metadata = inputs.metadata
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(
(inputs.q_fp4, inputs.q_sf),
inputs.k_cache,
@@ -231,6 +236,83 @@ class DeepGemmCandidateIndexer:
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:
return sparse_logits(
inputs.q_fp4,
@@ -475,27 +475,40 @@ def topk_transform_paged_from_metadata(
metadata,
page_indices: torch.Tensor,
raw_indices: Optional[torch.Tensor] = None,
*,
rows: Optional[slice] = None,
topk_metadata: Optional[torch.Tensor] = None,
) -> None:
"""Pool slots into ``page_indices`` (``-1`` past the valid count) and, when given,
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:
topk_transform_paged_v2(
logits,
metadata.compressed_seq_lens,
metadata.page_table,
page_indices,
seq_lens,
page_table,
out_page_indices,
metadata.compressed_page_size,
metadata.topk_metadata,
raw_indices,
metadata.topk_metadata if topk_metadata is None else topk_metadata,
out_raw_indices,
)
else:
topk_transform_paged(
logits,
metadata.compressed_seq_lens,
metadata.page_table,
page_indices,
seq_lens,
page_table,
out_page_indices,
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 (
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
logger = logging.getLogger(__name__)
@@ -306,7 +307,8 @@ class PagedIndexerMetadata:
):
return None
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_tc_piecewise_cuda_graph()
):
@@ -325,14 +327,27 @@ class PagedIndexerMetadata:
def row_chunks(self):
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(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(
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):
# 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:
"""Backend-owned scratch storage for sparse prefill KV dequantization.
The workspace contents are fully overwritten before every attention call,
so token buckets and compression ratios can safely share one buffer. Sparse
prefill executes eagerly and serially on the supported paths, which makes it
safe to replace the scratch allocation when a larger extent is needed.
Callers normally overwrite the entire workspace. Shared compressed-KV
callers keep separate workspaces per ratio and track prefix validity in the
per-forward gather cache, including the allocation address. Sparse prefill
executes eagerly and serially on the supported paths, so the allocation can
be replaced when a larger extent is needed.
"""
def __init__(self, device: torch.device):
@@ -275,14 +276,18 @@ class CompressedGather:
# chunk-invariant per request; subsequent layers only overwrite that prefix.
combined_indices: Optional[torch.Tensor] = None
combined_lens: Optional[torch.Tensor] = None
# Valid only for this forward's gather layout. Each ratio has its own
# workspace; its compressed prefix survives consumer layers' SWA writes.
dequantized_source: Optional[tuple[int, int]] = None
@dataclass
class SparsePrefillChunkCache:
"""Cache prefill-chunk metadata shared across layers.
Fields depend on request/token mappings and compressed page tables, not
per-layer k_cache; per-layer top-k combinations are recomputed into reused
Gather layouts depend on request/token mappings and compressed page tables.
Shared-source dequantization keys live only for this forward; per-layer
top-k combinations are recomputed into reused
buffers.
"""
@@ -19,7 +19,7 @@ Correctness-sensitive cases stay on Triton:
Single-sequence token counts that are not a multiple of the kernel's 64-token
chunk are padded up to a bucket (1k/2k/4k/8k/16k/32k) in a persistent staging
buffer, which bounds the resident workspace set. Pad rows are state-neutral:
k/v/beta zero => no rank-1 update; raw gate -1000 => transformed decay of
k/v zero => no rank-1 update, even with beta sigmoid; raw gate -1000 => decay of
exactly 1. Multi-sequence batches go through the kernel's own varlen grid
(real cu_seqlens, no padding), so their shapes are whatever the scheduler
produces and each distinct shape can retain another workspace.
@@ -327,6 +327,7 @@ class PtxKDAKernel(LinearAttnKernelBase):
dt_bias=self._flat_param(dt_bias),
return_intermediate_states=return_intermediate_states,
use_qk_l2norm_in_kernel=True,
use_beta_sigmoid_in_kernel=kwargs.get("beta_is_raw", False),
)
out, final_state, h = result[0], result[1], result[10]
ssm_states.index_copy_(0, slot, final_state.to(ssm_states.dtype))
@@ -116,6 +116,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
assert isinstance(runner.token_to_kv_pool, MiniMaxSparseKVPool)
self.is_npu = is_npu()
self.kv_pool = runner.token_to_kv_pool
self.hisparse_coordinator = runner.hisparse_coordinator
self.token_to_kv_pool = runner.token_to_kv_pool # alias for TboAttnBackend
self.req_to_token_pool = runner.req_to_token_pool # pool obj for TboAttnBackend
self.req_to_token = runner.req_to_token_pool.req_to_token
@@ -176,6 +177,18 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
local_tokens + self.block_size_k - 1
) // self.block_size_k + 1
self.topk_blocks = sparse_cfg["sparse_topk_blocks"]
if self.hisparse_coordinator is not None:
selected_tokens = self.topk_blocks * self.block_size_k
assert selected_tokens <= self.hisparse_coordinator.device_buffer_size, (
f"MiniMax M3 selects {selected_tokens} sparse-attention tokens, "
"but the HiSparse device buffer holds only "
f"{self.hisparse_coordinator.device_buffer_size}."
)
self._loc_mapping = (
self.kv_pool.main_pool.full_to_hisparse_device_index_mapping
)
else:
self._loc_mapping = None
# MSA (fmha_sm100) is SM100-only; fall back to the Triton sparse path when
# the kernel is unavailable or its constraints don't hold.
@@ -209,6 +222,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
)
self.use_msa = (
not envs.SGLANG_DISABLE_MSA.get()
and self.hisparse_coordinator is None
and msa_available()
and self.block_size_k == 128
and self.kv_pool.page_size == self.block_size_k
@@ -245,6 +259,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
self.page_size = self.kv_pool.page_size
self.use_dense_sparse_decode = (
(not self.is_npu)
and self.hisparse_coordinator is None
and envs.SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE.get()
and self.block_size_k % self.page_size == 0
# _dense_sparse_main_decode calls trtllm decode with a bf16 q and
@@ -326,6 +341,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
f"msa_owns_decode={self._msa_owns_decode}, "
f"decode_cuda_graph={_decode_cuda_graph}, "
f"fp8_attn_gemm={self.fp8_attn_gemm}, "
f"hisparse={'enabled' if self._loc_mapping is not None else 'disabled'}, "
f"npu_native_attn={'on' if (self._native_sparse_ok and _native_attn_enabled()) else 'off'}, "
f"disable_value_layers={sorted(self.disable_value_layer_ids)})"
)
@@ -336,6 +352,22 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
"take minutes; compiles serialize across TP ranks)."
)
def _hisparse_swap_in_blocks(
self,
forward_batch: ForwardBatch,
topk_idx: torch.Tensor,
layer_id: int,
) -> torch.Tensor:
assert topk_idx.size(0) == 1
top_k_device_locs = self.hisparse_coordinator.swap_in_selected_blocks(
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
top_k_blocks=topk_idx[0],
layer_id=layer_id,
sparse_block_size=self.block_size_k,
)
return top_k_device_locs.unsqueeze(0)
@staticmethod
def _choose_decode_score_max_chunks(batch_size: int) -> int:
"""Score chunk count per graph bucket.
@@ -1549,6 +1581,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
idx_v_scale=layer.idx_v_scale_float,
cached_topk_idx=cached_topk_idx,
return_topk_idx=want_topk,
loc_mapping=self._loc_mapping,
)
if want_topk:
idx_o, o, reduced_topk_idx = result
@@ -1702,6 +1735,16 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
else:
_cached_topk = _topk_buf
hisparse_swap_in_fn = None
if self.hisparse_coordinator is not None:
def hisparse_swap_in_fn(topk_idx):
return self._hisparse_swap_in_blocks(
forward_batch=forward_batch,
topk_idx=topk_idx,
layer_id=layer.layer_id,
)
idx_o, o = minimax_sparse_decode(
q,
None,
@@ -1735,6 +1778,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
idx_v_scale=layer.idx_v_scale_float,
cached_topk_idx=_cached_topk,
topk_out=_topk_buf if _want_topk else None,
hisparse_swap_in_fn=hisparse_swap_in_fn,
)
return (
None if idx_o is None else idx_o.reshape(q.shape[0], -1).contiguous(),
@@ -75,6 +75,7 @@ def minimax_sparse_prefill(
idx_v_scale: Optional[float] = None,
cached_topk_idx: Optional[torch.Tensor] = None,
return_topk_idx: bool = False,
loc_mapping: Optional[torch.Tensor] = None,
):
"""Run MiniMax-M3 sparse prefill.
@@ -146,7 +147,7 @@ def minimax_sparse_prefill(
# Step 3: Sparse attention using topk index (main head). The MSA path only
# replaces this step; the indexer above is unchanged. MSA has no attn-sink
# input, so keep the Triton path when sink is present.
if use_msa and sink is None:
if use_msa and sink is None and loc_mapping is None:
from .msa import MSAUnavailableError, msa_sparse_prefill_main
try:
@@ -188,6 +189,7 @@ def minimax_sparse_prefill(
q_scale=q_scale,
k_scale=k_scale,
v_scale=v_scale,
loc_mapping=loc_mapping,
)
else:
o = flash_prefill_with_gqa_share_sparse(
@@ -210,6 +212,7 @@ def minimax_sparse_prefill(
q_scale=q_scale,
k_scale=k_scale,
v_scale=v_scale,
loc_mapping=loc_mapping,
)
if return_topk_idx:
return idx_o, o, reduced_topk_idx
@@ -255,6 +258,7 @@ def minimax_sparse_decode(
idx_v_scale: Optional[float] = None,
cached_topk_idx: Optional[torch.Tensor] = None,
topk_out: Optional[torch.Tensor] = None,
hisparse_swap_in_fn: Optional[Callable] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
# Index top-k sharing for DECODE. A group's source layer passes ``topk_out``
# (a persistent buffer) and publishes its reduced top-k there; the group's
@@ -319,9 +323,12 @@ def minimax_sparse_decode(
f"reduced top-k shape {tuple(topk_idx.shape)}"
)
topk_out.copy_(topk_idx)
hisparse_slots = (
hisparse_swap_in_fn(topk_idx) if hisparse_swap_in_fn is not None else None
)
# Step 3: Sparse attention using topk index (main head). The MSA path
# only replaces this step; keep the Triton path when sink is present.
if use_msa and sink is None:
if use_msa and sink is None and hisparse_slots is None:
from .msa import MSAUnavailableError, msa_sparse_decode_main
try:
@@ -357,6 +364,7 @@ def minimax_sparse_decode(
q_scale=q_scale,
k_scale=k_scale,
v_scale=v_scale,
hisparse_slots=hisparse_slots,
)
else:
o = flash_decode_with_gqa_share_sparse(
@@ -373,5 +381,6 @@ def minimax_sparse_decode(
q_scale=q_scale,
k_scale=k_scale,
v_scale=v_scale,
hisparse_slots=hisparse_slots,
)
return idx_o, o
+97 -10
View File
@@ -23,17 +23,22 @@ import torch
from sglang.srt.arg_groups.overrides import (
attention_backends_of,
model_config_of,
resolved_view,
resolving_view,
)
from sglang.srt.configs.model_config import is_deepseek_v4
from sglang.srt.layers.cp.base import get_cp_strategy
from sglang.srt.layers.cp.interleave import InterleaveCPStrategy
from sglang.srt.layers.cp.padding import get_cp_padding_align_size
from sglang.srt.layers.cp.utils import (
cp_gather_after_forward,
cp_shard_hidden_states,
cp_split_before_forward,
prepare_cp_forward,
)
from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
from sglang.srt.layers.logits_processor import LogitsMetadata
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
if TYPE_CHECKING:
@@ -50,12 +55,18 @@ def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool:
cfg = resolving_view(server_args)
resolved = resolved_view(server_args)
prefill_attention_backend, _ = attention_backends_of(resolved_view(server_args))
supports_layout = (
cfg.cp_strategy == "zigzag" and prefill_attention_backend == "trtllm_mha"
) or (
cfg.cp_strategy == "interleave"
and prefill_attention_backend == "dsv4"
and is_deepseek_v4(model_config_of(server_args).hf_config)
)
return (
cfg.enable_prefill_cp
and cfg.pp_size == 1
and resolved.attn_cp_size == cfg.tp_size
and cfg.cp_strategy == "zigzag"
and prefill_attention_backend == "trtllm_mha"
and supports_layout
)
@@ -67,8 +78,12 @@ def enable_cp_bcg_capture(server_args: ServerArgs) -> bool:
def filter_prefill_cp_bcg_capture_num_tokens(
capture_num_tokens: list[int], server_args: ServerArgs
) -> list[int]:
"""Keep only token buckets where the zigzag CP strategy can run."""
min_num_tokens = resolved_view(server_args).attn_cp_size * 2
"""Keep only token buckets where the configured CP strategy can run."""
cfg = resolving_view(server_args)
cp_segments_per_token_block = 2 if cfg.cp_strategy == "zigzag" else 1
min_num_tokens = (
resolved_view(server_args).attn_cp_size * cp_segments_per_token_block
)
filtered = [size for size in capture_num_tokens if size >= min_num_tokens]
if not filtered:
raise ValueError(
@@ -96,6 +111,8 @@ class PrefillCPBCGInput:
input_embeds: torch.Tensor
positions: torch.Tensor
input_ids: Optional[torch.Tensor] = None
num_token_non_padded: Optional[torch.Tensor] = None
bucket_local_tokens: Dict[int, int] = field(default_factory=dict)
live_local_tokens: int = 0
@@ -114,12 +131,22 @@ class PrefillCPBCGInput:
(runner.max_num_tokens,),
dtype=torch.int64,
),
input_ids=torch.zeros((runner.max_num_tokens,), dtype=torch.int64),
num_token_non_padded=torch.zeros((), dtype=torch.int32),
)
def required_local_tokens(self, extend_seq_lens: Any) -> Optional[int]:
"""Return the aligned CP-local rows required by a live zigzag layout."""
"""Return the aligned CP-local rows required by the active layout."""
strategy = get_cp_strategy()
if not isinstance(strategy, ZigzagCPStrategy) or extend_seq_lens is None:
if extend_seq_lens is None:
return None
if isinstance(strategy, InterleaveCPStrategy):
logical_tokens = (
sum(int(length) for length in extend_seq_lens) + strategy.cp_size - 1
) // strategy.cp_size
align_size = get_cp_padding_align_size()
return (logical_tokens + align_size - 1) // align_size * align_size
if not isinstance(strategy, ZigzagCPStrategy):
return None
cp_segment_num = strategy.cp_size * 2
@@ -219,6 +246,7 @@ class PrefillCPBCGInput:
raw_tokens = int(forward_batch.extend_num_tokens)
global_input_ids = forward_batch.input_ids[:raw_tokens]
global_positions = forward_batch.positions[:raw_tokens]
local_input_ids = cp_shard_hidden_states(global_input_ids, forward_batch)
global_input_embeds = runner.model_runner.model.get_input_embeddings()(
global_input_ids
)
@@ -249,12 +277,31 @@ class PrefillCPBCGInput:
input_embeds = self.input_embeds[:captured_local_tokens]
positions = self.positions[:captured_local_tokens]
assert self.input_ids is not None
input_ids = self.input_ids[:captured_local_tokens]
input_embeds.zero_()
positions.zero_()
input_ids.zero_()
input_embeds[:live_local_tokens].copy_(local_input_embeds)
positions[:live_local_tokens].copy_(local_positions)
input_ids[:live_local_tokens].copy_(local_input_ids)
forward_batch.input_embeds = input_embeds
forward_batch.positions = positions
forward_batch._cp_positions = positions
# Keep the global input_ids field intact: the runner uses its length to
# select the global capture bucket. The DSV4 body consumes this fixed,
# rank-local view for hash routing and MegaMoE.
forward_batch._cp_input_ids = input_ids
forward_batch.input_ids_global = input_ids
if forward_batch.num_token_non_padded is not None:
assert self.num_token_non_padded is not None
metadata = forward_batch.attn_cp_metadata
logical_tokens = (
metadata.per_rank_logical_token or metadata.per_rank_actual_token
)
strategy = get_cp_strategy()
assert strategy is not None
self.num_token_non_padded.fill_(logical_tokens[strategy.cp_rank])
forward_batch.num_token_non_padded = self.num_token_non_padded
self.live_local_tokens = live_local_tokens
@@ -307,10 +354,50 @@ def execute_prefill_cp_bcg(
static_forward_batch,
torch.cuda.current_stream(),
)
return model.logits_processor(
forward_batch.input_ids,
if aux_hidden_states is not None:
if torch.is_tensor(aux_hidden_states):
aux_hidden_states = cp_gather_after_forward(
aux_hidden_states, static_forward_batch, torch.cuda.current_stream()
)
else:
aux_hidden_states = [
cp_gather_after_forward(
aux, static_forward_batch, torch.cuda.current_stream()
)
for aux in aux_hidden_states
]
hidden_states_before_norm = None
if isinstance(hidden_states, tuple):
assert len(hidden_states) == 2
hidden_states, hidden_states_before_norm = hidden_states
input_ids = forward_batch.input_ids
logits_metadata = forward_batch
tail = None
language_model = getattr(model, "model", None)
if (
capture_aux_hidden_states
and getattr(language_model, "late_layer_start", None) is not None
and forward_batch.forward_mode.is_extend_without_speculative()
):
tail_metadata = runner.model_runner.attn_backend.tail_forward_metadata
tail = tail_metadata.late_layer_tail
input_ids = tail.rows(input_ids)
logits_metadata = LogitsMetadata.from_forward_batch(forward_batch)
logits_metadata.extend_seq_lens = tail.extend_seq_lens
logits_metadata.extend_seq_lens_cpu = tail.extend_seq_lens_cpu
logits_metadata.extend_logprob_start_lens_cpu = tail.extend_seq_lens_cpu
output = model.logits_processor(
input_ids,
hidden_states,
model.lm_head,
forward_batch,
logits_metadata,
aux_hidden_states,
hidden_states_before_norm=(
None if aux_hidden_states is not None else hidden_states_before_norm
),
)
if tail is not None:
output.hidden_states_token_indices = tail.token_indices
return output
+26 -13
View File
@@ -216,20 +216,33 @@ def _run_mega_routed(
if num_tokens > 0:
router_logits = moe.gate(hidden_states, forward_batch=forward_batch)
topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {}
topk_output = moe.topk(
hidden_states,
router_logits,
num_token_non_padded=(
forward_batch.num_token_non_padded
if forward_batch is not None
else None
),
expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
layer_id=moe.layer_id,
),
**topk_kwargs,
num_token_non_padded = (
forward_batch.num_token_non_padded if forward_batch is not None else None
)
if isinstance(
getattr(moe.gate, "e_score_correction_bias_vl", None), torch.Tensor
):
# V4.1 uses a different correction bias for image-token rows. The
# MegaMoE transport consumes the same routed ids/weights as TopK.
from sglang.srt.multimodal.dsv41.vl_routing import vision_topk
topk_output = vision_topk(
moe,
router_logits,
input_ids_global,
num_token_non_padded=num_token_non_padded,
)
else:
topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {}
topk_output = moe.topk(
hidden_states,
router_logits,
num_token_non_padded=num_token_non_padded,
expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
layer_id=moe.layer_id,
),
**topk_kwargs,
)
topk_ids = topk_output.topk_ids
topk_weights = topk_output.topk_weights
else:
@@ -1497,6 +1497,15 @@ class ModelOptFp4Config(ModelOptQuantConfig):
def get_min_capability(cls) -> int:
return 80
def can_fuse_shared_expert(self) -> bool:
# A shared-expert body kept BF16 via exclude_modules cannot share the packed
# FP4 FusedMoE buffers. The shared_expert_gate is a separate linear (kept
# BF16 by e.g. Qwen3-Next NVFP4 checkpoints) and must not veto fusion.
return not any(
"shared_expert" in name and "shared_expert_gate" not in name
for name in self.exclude_modules
)
@staticmethod
def common_group_size(cfg: dict) -> int:
"""Return the unique group_size across the config; raise if missing/mismatched."""
@@ -19,9 +19,16 @@ if is_xpu():
"copy_cache_planned_mla has no AOT sgl_kernel implementation."
)
def load_blocks_to_device_buffer_mha(*args, **kwargs):
raise RuntimeError(
"MiniMax M3 HiSparse block swap-in is unsupported on XPU: "
"load_blocks_to_device_buffer_mha has no AOT sgl_kernel implementation."
)
else:
from sglang.kernels.ops.kvcache.hisparse import (
copy_cache_planned_mla,
load_blocks_to_device_buffer_mha,
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
)
@@ -36,8 +43,9 @@ from sglang.srt.mem_cache.allocator.hisparse import (
from sglang.srt.mem_cache.hisparse_memory_pool import (
HiSparseDSATokenToKVPool,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool, ReqToTokenPool
from sglang.srt.mem_cache.memory_pool_host import DeepSeekV4PagedHostPool
from sglang.srt.mem_cache.pool_host.mha import HiSparseMHATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
device_module = get_device_module()
@@ -157,9 +165,11 @@ class HiSparseCoordinator:
)
self.compress_ratio = self.token_to_kv_pool_allocator.compress_ratio
kvcache = self.token_to_kv_pool_allocator.get_kvcache()
self.is_dsv4_hisparse = isinstance(
self.token_to_kv_pool_allocator, DeepSeekV4HiSparseTokenToKVPoolAllocator
)
self.is_m3_hisparse = isinstance(kvcache, MiniMaxSparseKVPool)
if self.is_dsv4_hisparse:
self.mem_pool_device = self.token_to_kv_pool_allocator.hisparse_kvcache
page_size = self.mem_pool_device.page_size
@@ -184,18 +194,30 @@ class HiSparseCoordinator:
assert isinstance(
self.token_to_kv_pool_allocator, HiSparseTokenToKVPoolAllocator
)
self.mem_pool_device: HiSparseDSATokenToKVPool = (
self.token_to_kv_pool_allocator.get_kvcache()
)
self.mem_pool_host = MLATokenToKVPoolHost(
device_pool=self.mem_pool_device,
host_to_device_ratio=host_to_device_ratio,
host_size=0,
page_size=self.mem_pool_device.page_size,
layout="layer_first",
override_kv_cache_dim=self.mem_pool_device.kv_cache_dim,
)
self.item_size_bytes = self.mem_pool_host.token_stride_size
if self.is_m3_hisparse:
self.mem_pool_device = kvcache.main_pool
assert self.mem_pool_device.head_num == 1, (
"MiniMax M3 HiSparse requires one KV head per TP rank, "
f"got {self.mem_pool_device.head_num}. Increase the "
"tensor-parallel size."
)
self.mem_pool_host = HiSparseMHATokenToKVPoolHost(
device_pool=self.mem_pool_device,
host_to_device_ratio=host_to_device_ratio,
page_size=self.mem_pool_device.page_size,
)
self.item_size_bytes = self.mem_pool_device.bytes_per_token_k
else:
self.mem_pool_device: HiSparseDSATokenToKVPool = kvcache
self.mem_pool_host = MLATokenToKVPoolHost(
device_pool=self.mem_pool_device,
host_to_device_ratio=host_to_device_ratio,
host_size=0,
page_size=self.mem_pool_device.page_size,
layout="layer_first",
override_kv_cache_dim=self.mem_pool_device.kv_cache_dim,
)
self.item_size_bytes = self.mem_pool_host.token_stride_size
self.page_size = self.mem_pool_device.page_size
max_num_req_slots = req_to_token_pool.req_to_token.shape[0]
@@ -263,9 +285,17 @@ class HiSparseCoordinator:
self.device_buffer_size, dtype=torch.int32, device=device
)
# Pre-allocated output buffer for swap_in_selected_pages (CUDA-graph safe)
# Pre-allocated output buffer for swap-in (CUDA-graph safe). MiniMax
# selects blocks, so its flattened token-slot output can occupy any
# prefix up to the full device working-set size.
swap_output_width = (
self.device_buffer_size if self.is_m3_hisparse else self.top_k
)
self.top_k_device_locs_buffer = torch.full(
(max_num_req_slots, self.top_k), -1, dtype=torch.int32, device=device
(max_num_req_slots, swap_output_width),
-1,
dtype=torch.int32,
device=device,
)
self.raw_indices_buffer = torch.full(
(max_num_req_slots, self.top_k), -1, dtype=torch.int32, device=device
@@ -457,7 +487,12 @@ class HiSparseCoordinator:
host_indices = self.req_to_host_pool[req.kv.req_pool_idx, :n]
device_locs = self.req_to_device_buffer[req.kv.req_pool_idx, :n]
for layer_id in range(self.mem_pool_device.layer_num):
layer_ids = (
range(self.mem_pool_device.start_layer, self.mem_pool_device.end_layer)
if self.is_m3_hisparse
else range(self.mem_pool_device.layer_num)
)
for layer_id in layer_ids:
self.mem_pool_host.load_to_device_per_layer(
self.mem_pool_device,
host_indices,
@@ -642,13 +677,9 @@ class HiSparseCoordinator:
compressed_locs = self.token_to_kv_pool_allocator.get_last_loc_compressed(
out_cache_loc
)
# ROCm: the decode remap creates a temporary hisparse device slot per
# new token (via the page_size==1 allocator path). Free the stale
# slot before pointing the mapping at the reserved device-buffer slot,
# otherwise the temporary slots leak and corrupt later swap-in lookups.
# CUDA keeps the original behavior: the swap-in kernel consumes only
# top_k_device_locs, so stale mapping entries are harmless there.
if _is_hip:
# Page-size-one allocation creates a temporary slot before remapping
# the new token into the request's reserved device-buffer slot.
if _is_hip or self.mem_pool_device.page_size == 1:
previous_locs = self.mem_pool_device._translate_loc_to_hisparse_device(
compressed_locs
)
@@ -976,8 +1007,7 @@ class HiSparseCoordinator:
miss plan into self._miss_{src,dst,count} for the skip layers to replay.
"""
num_reqs = req_pool_indices.size(0)
top_k_indices = self.top_k_device_locs_buffer[:num_reqs]
top_k_indices = self.top_k_device_locs_buffer[:num_reqs, : self.top_k]
swap_in_fn = (
load_cache_to_device_buffer_dsv4_mla
if self.is_dsv4_hisparse
@@ -1015,6 +1045,46 @@ class HiSparseCoordinator:
)
return top_k_indices
def swap_in_selected_blocks(
self,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
top_k_blocks: torch.Tensor,
layer_id: int,
sparse_block_size: int,
) -> torch.Tensor:
assert self.is_m3_hisparse
num_reqs = req_pool_indices.size(0)
num_selected_tokens = top_k_blocks.size(1) * sparse_block_size
assert num_selected_tokens <= self.device_buffer_size, (
f"MiniMax M3 selected {num_selected_tokens} tokens, but the "
f"HiSparse device buffer holds only {self.device_buffer_size}."
)
top_k_indices = self.top_k_device_locs_buffer[:num_reqs, :num_selected_tokens]
host_layer = layer_id - self.mem_pool_device.start_layer
load_blocks_to_device_buffer_mha(
top_k_blocks=top_k_blocks,
device_buffer_tokens=self.req_device_buffer_tokens[host_layer],
host_cache_locs=self.req_to_host_pool,
device_buffer_locs=self.req_device_buffer_token_locs[host_layer],
host_cache_k=self.mem_pool_host.k_buffer[host_layer],
host_cache_v=self.mem_pool_host.v_buffer[host_layer],
device_buffer_k=self.mem_pool_device.get_key_buffer(layer_id),
device_buffer_v=self.mem_pool_device.get_value_buffer(layer_id),
top_k_device_locs=top_k_indices,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
lru_slots=self.lru_slots[host_layer],
item_size_bytes=self.item_size_bytes,
hot_buffer_size=self.device_buffer_size,
sparse_block_size=sparse_block_size,
page_size=1,
block_size=self.swap_in_block_size,
num_real_reqs=self.num_real_reqs,
skip_io=self.skip_io,
)
return top_k_indices
def _run_copy_only_kernel(self, num_reqs: int, skip_layer: int) -> None:
"""Replay the anchor's recorded miss plan into a skip layer's buffers
(IO-only; the anchor's slot table stays valid -- lockstep layout)."""
@@ -1045,7 +1115,10 @@ class HiSparseCoordinator:
"""
if not self.enable_prefetch:
return self._run_swap_in_kernel(
req_pool_indices, compressed_seq_lens, top_k_result, layer_id
req_pool_indices,
compressed_seq_lens,
top_k_result,
layer_id,
)
num_reqs = req_pool_indices.size(0)
@@ -1054,7 +1127,7 @@ class HiSparseCoordinator:
# applies (shared index + lockstep buffers).
slot = self._prefetch_slot[layer_id]
self._prefetch_events[slot].wait(device_module.current_stream())
return self.top_k_device_locs_buffer[:num_reqs]
return self.top_k_device_locs_buffer[:num_reqs, : self.top_k]
# Anchor: swap in synchronously (recording the plan), then prefetch the
# skip layers' copies on the side stream.
@@ -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
+83 -37
View File
@@ -5,6 +5,7 @@ from typing import Callable, Dict, List, Optional, Tuple
import torch
from sglang.srt.managers.mm_owner_embedding import ImageSpanRequest, MmOwnerSession
from sglang.srt.managers.schedule_batch import MultimodalDataItem
from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
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]] = {}
hash_to_embedding: Dict[Tuple[Optional[int], int], torch.Tensor] = {}
# Phase 1a: find overlapping items per request and collect cache misses
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:
expected_token_count = end - start + 1
cache_key = (item.hash, expected_token_count)
if cache_key in hash_to_embedding:
# Phase 1a: collect cache misses over the unique overlapping spans
for span in _collect_image_span_requests(per_image_requests):
cache_key = (span.hash, span.span_len)
cached = embedding_cache.get_single(span.hash)
if cached is not None:
cached_embedding = cached.embedding
cached_token_count = _embedding_token_count(cached_embedding)
if cached_token_count == span.span_len:
hash_to_embedding[cache_key] = cached_embedding
continue
cached = embedding_cache.get_single(item.hash)
if cached is not None:
cached_embedding = cached.embedding
cached_token_count = _embedding_token_count(cached_embedding)
if cached_token_count == expected_token_count:
hash_to_embedding[cache_key] = cached_embedding
else:
_discard_mismatched_cached_embedding(
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)
_discard_mismatched_cached_embedding(
span.hash, span.span_len, cached_token_count
)
elif (
span.inside_chunk and span.item.can_defer_cuda_ipc_feature_reconstruction()
):
span.item.model_specific_data[BORROW_CUDA_IPC_FEATURE_KEY] = True
unique_misses[cache_key] = (span.item, span.span_len)
# Phase 1b: single ViT call for all unique cache misses
if unique_misses:
@@ -412,6 +394,52 @@ def _batch_encode_per_image_misses(
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(
data_embedding_func: DataEmbeddingFunc,
embedding_items_per_req: List[MultimodalDataItem],
@@ -537,6 +565,7 @@ def _get_chunked_prefill_embedding(
extend_length: List[int],
items_offset_list: List[List[Tuple[int, int]]],
input_ids: torch.Tensor,
mm_owner: Optional[MmOwnerSession] = None,
) -> tuple[torch.Tensor | None, torch.Tensor]:
"""
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
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(
data_embedding_func, per_image_requests, device
)
@@ -701,6 +745,7 @@ def get_embedding_and_mask(
prefix_length: List[int],
extend_length: List[int],
items_offset_list: List[List[Tuple[int, int]]],
mm_owner: Optional[MmOwnerSession] = None,
) -> Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]:
"""
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,
items_offset_list,
input_ids,
mm_owner=mm_owner,
)
if embedding is None:
return None, None, input_ids
+13 -2
View File
@@ -10,6 +10,7 @@ import pickle
import sys
from abc import abstractmethod
from collections import defaultdict
from contextlib import nullcontext
from multiprocessing import shared_memory
from typing import Any, Dict, List, Optional, Tuple
@@ -24,6 +25,7 @@ from sglang.srt.managers.io_struct import (
TokenizedEmbeddingReqInput,
TokenizedGenerateReqInput,
)
from sglang.srt.managers.mm_owner_embedding import MmOwnerSession
# Preserve the existing initialization import for downstream callers.
from sglang.srt.managers.mm_schedule import (
@@ -397,6 +399,7 @@ def embed_mm_inputs(
data_embedding_func_mapping: Dict[Modality, DataEmbeddingFunc] = None,
placeholder_tokens: dict[Modality, List[int]] = None,
use_deepstack: Dict[Modality, bool] = {},
mm_owner: Optional[MmOwnerSession] = None,
) -> Optional[torch.Tensor]:
"""
Embed multimodal inputs and integrate them with text token embeddings.
@@ -478,6 +481,7 @@ def embed_mm_inputs(
prefix_length=extend_prefix_lens,
extend_length=extend_seq_lens,
items_offset_list=items_offsets,
mm_owner=mm_owner,
)
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.
# There values are useless because their embeddings will be replaced by vision embeddings anyway.
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
if use_deepstack:
@@ -525,7 +534,9 @@ def embed_mm_inputs(
_scatter_mm_embedding(dest=input_embeds, mask=mask, src=embedding)
if use_deepstack.get(modality, None):
_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
@@ -1019,6 +1019,15 @@ class PrefillAdder:
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):
if self.dllm_config is not None:
_rem_tokens = self._get_dllm_remain_tokens()
+34 -3
View File
@@ -2097,9 +2097,13 @@ class Scheduler(
vmm_errors = self._materialize_cuda_vmm_inputs(recv_req)
# 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(
for_health_check=True
):
# NOTE: the admit/skip decision must be identical on every CP/TP rank.
# 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(
getattr(recv_req, "http_worker_ipc", None)
)
@@ -3940,6 +3944,8 @@ class Scheduler(
for req in self.waiting_queue:
if self.enable_lora and not self.can_schedule_lora_req(req, running_loras):
continue
if not adder.can_share_extend_batch(req):
break
running_bs = len(running_batch.reqs)
candidate_beam_width = (
@@ -4980,6 +4986,31 @@ class Scheduler(
else:
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:
# Health check piggybacks on running requests in process_output.
# Only running_batch + waiting_queue guarantee active GPU processing;
@@ -1279,6 +1279,16 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
raise ValueError(
"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
input_token_num = len(input_ids) if input_ids is not None else 0
input_token_num += self.num_reserved_tokens
@@ -1,4 +1,7 @@
from __future__ import annotations
import weakref
from typing import TYPE_CHECKING
import torch
@@ -11,6 +14,9 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
from sglang.srt.utils.common import get_num_new_pages
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def __init__(
@@ -19,7 +25,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
page_size: int,
dtype: torch.dtype,
device: torch.device,
kvcache: HiSparseDSATokenToKVPool,
kvcache: HiSparseDSATokenToKVPool | MiniMaxSparseKVPool,
need_sort: bool,
host_to_device_ratio: int = 2,
):
+4
View File
@@ -155,6 +155,10 @@ def free_kv_row_segments(
def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
if getattr(req, "skip_radix_cache_insert", False):
kv_indices = tree_cache.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, : len(req.get_fill_ids())
]
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
return
tree_cache.cache_unfinished_req(req, **kwargs)
@@ -9,7 +9,7 @@ from sglang.kernels.ops.kvcache.hisparse_slot_mapping import (
translate_padded_hisparse_locations,
)
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, MHATokenToKVPool
from sglang.srt.utils import is_cuda, is_hip, is_xpu
logger = logging.getLogger(__name__)
@@ -147,3 +147,111 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool):
self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
):
raise NotImplementedError("HiSparseDevicePool does not support load_cpu_copy")
class HiSparseMHAMainPool(MHATokenToKVPool):
"""MHA KV pool with HiSparse logical-to-device mapping.
Used by MiniMax M3 HiSparse. The index pools (index_kv_pool, index_k_pool)
stay fully resident on the device and do not use this mapping.
"""
def __init__(
self,
size: int,
page_size: int,
dtype: torch.dtype,
head_num: int,
head_dim: int,
layer_num: int,
device: str,
enable_memory_saver: bool,
start_layer: Optional[int] = None,
end_layer: Optional[int] = None,
):
super().__init__(
size=size,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
layer_num=layer_num,
device=device,
enable_memory_saver=enable_memory_saver,
start_layer=start_layer,
end_layer=end_layer,
)
self.full_to_hisparse_device_index_mapping: Optional[torch.Tensor] = None
self.bytes_per_token_k = head_num * head_dim * self.store_dtype.itemsize
self.bytes_per_token_v = head_num * self.v_head_dim * self.store_dtype.itemsize
def register_mapping(
self, full_to_hisparse_device_index_mapping: torch.Tensor
) -> None:
self.full_to_hisparse_device_index_mapping = (
full_to_hisparse_device_index_mapping
)
def translate_loc_to_hisparse_device(self, indices: torch.Tensor) -> torch.Tensor:
assert self.full_to_hisparse_device_index_mapping is not None
return self.full_to_hisparse_device_index_mapping[indices]
def _translate_loc_to_hisparse_device(self, indices: torch.Tensor) -> torch.Tensor:
assert self.full_to_hisparse_device_index_mapping is not None
return self.full_to_hisparse_device_index_mapping[indices]
def translate_loc_from_full_to_hisparse_device(
self, full_indices: torch.Tensor
) -> torch.Tensor:
assert self.full_to_hisparse_device_index_mapping is not None
return self.full_to_hisparse_device_index_mapping[full_indices]
def translate_loc_from_full_to_compressed(
self, full_indices: torch.Tensor
) -> torch.Tensor:
return full_indices
def set_kv_buffer(
self,
layer: RadixAttention,
loc,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
*args,
**kwargs,
):
from sglang.srt.mem_cache.memory_pool import unwrap_write_loc
raw_loc, _, _ = unwrap_write_loc(loc)
translated = self.translate_loc_to_hisparse_device(raw_loc)
super().set_kv_buffer(layer, translated, cache_k, cache_v, *args, **kwargs)
def transfer_values_on_device(
self,
dst_indices: torch.Tensor,
src_indices: torch.Tensor,
) -> None:
transfer_kv_all_layer_mla(
src_layers=self.k_data_ptrs,
dst_layers=self.k_data_ptrs,
src_indices=src_indices,
dst_indices=dst_indices,
item_size=self.bytes_per_token_k,
num_layers=self.layer_num,
)
transfer_kv_all_layer_mla(
src_layers=self.v_data_ptrs,
dst_layers=self.v_data_ptrs,
src_indices=src_indices,
dst_indices=dst_indices,
item_size=self.bytes_per_token_v,
num_layers=self.layer_num,
)
def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
raise NotImplementedError("HiSparseMHAMainPool does not support get_cpu_copy")
def load_cpu_copy(
self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
):
raise NotImplementedError("HiSparseMHAMainPool does not support load_cpu_copy")

Some files were not shown because too many files have changed in this diff Show More