Add reproducible DGX Spark deployment for Qwen3.8 Flash Next NVFP4
This commit is contained in:
@@ -0,0 +1,6 @@
|
|||||||
|
**
|
||||||
|
!Dockerfile
|
||||||
|
!patches/
|
||||||
|
!patches/spark_ngram_adapter.py
|
||||||
|
!vendor/
|
||||||
|
!vendor/vllm_ple_mmap.py
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
# Copy to .env on the Spark; use absolute paths. Do not commit .env.
|
||||||
|
HF_CACHE_DIR=/path/to/huggingface
|
||||||
|
API_KEY_FILE=/path/to/qwen.key
|
||||||
|
API_PORT=8000
|
||||||
|
BIND_ADDRESS=0.0.0.0
|
||||||
|
MODEL_HTTP_PROXY=
|
||||||
|
MODEL_HTTPS_PROXY=
|
||||||
|
GPU_MEMORY_UTILIZATION=0.80
|
||||||
|
MAX_NUM_SEQS=4
|
||||||
|
# Optional offline build: local tag that resolves to the pinned base digest.
|
||||||
|
# BASE_IMAGE=local/vllm-base:0bfc7a15
|
||||||
+13
@@ -0,0 +1,13 @@
|
|||||||
|
.env
|
||||||
|
.env.*
|
||||||
|
!.env.example
|
||||||
|
*.key
|
||||||
|
*.pem
|
||||||
|
secrets/
|
||||||
|
runtime-cache/
|
||||||
|
__pycache__/
|
||||||
|
*.pyc
|
||||||
|
*.log
|
||||||
|
*.safetensors
|
||||||
|
*.gguf
|
||||||
|
.DS_Store
|
||||||
+22
@@ -0,0 +1,22 @@
|
|||||||
|
ARG BASE_IMAGE=vllm/vllm-openai@sha256:c4392d76e3eec8983fa152651365158cb062e348fd40398963f499d5867b9e28
|
||||||
|
FROM ${BASE_IMAGE}
|
||||||
|
COPY vendor/vllm_ple_mmap.py patches/spark_ngram_adapter.py /usr/local/lib/python3.12/dist-packages/
|
||||||
|
RUN python3 - <<'PY'
|
||||||
|
from pathlib import Path
|
||||||
|
import ast
|
||||||
|
root = Path('/usr/local/lib/python3.12/dist-packages/vllm')
|
||||||
|
p = root / 'models/qwen4_exp/nvidia/ngram_embedding.py'
|
||||||
|
s = p.read_text()
|
||||||
|
assert 'class Qwen4ExpPLEPinnedHostEmbedding' in s
|
||||||
|
s += '\nfrom spark_ngram_adapter import apply as _spark_disk_apply\n_spark_disk_apply(Qwen4ExpNGramEmbedding)\n'
|
||||||
|
ast.parse(s)
|
||||||
|
p.write_text(s)
|
||||||
|
for file, before, after in [
|
||||||
|
('third_party/flash_linear_attention/ops/utils.py', 'DEFAULT = 102400', 'DEFAULT = 101376'),
|
||||||
|
('third_party/flash_linear_attention/ops/chunk_delta_h.py', 'for num_warps in [2, 4]', 'for num_warps in [2]'),
|
||||||
|
]:
|
||||||
|
p = root / file
|
||||||
|
s = p.read_text()
|
||||||
|
assert before in s, file
|
||||||
|
p.write_text(s.replace(before, after))
|
||||||
|
PY
|
||||||
@@ -0,0 +1,88 @@
|
|||||||
|
# Qwen3.8 Flash Next NVFP4 on DGX Spark
|
||||||
|
|
||||||
|
在单台 NVIDIA DGX Spark（GB10 / ARM64 / 128 GB 统一内存）上部署
|
||||||
|
`nvidia/Qwen3.8-Flash-Next-NVFP4` 的可复现配置。
|
||||||
|
来自 2026-09-17 的实际部署：262144 上下文上限、MTP 2 tokens、OpenAI 兼容 API。
|
||||||
|
PLE 查找表通过磁盘映射按需读取；原始模型权重未重新量化。
|
||||||
|
|
||||||
|
> 这是针对固定 nightly 的社区适配，非 NVIDIA/vLLM 官方支持方案。
|
||||||
|
> 已验证短问答和工具调用；未压测完整 262K 输入。仅适用于此模型、单 GPU / ETP=1。
|
||||||
|
|
||||||
|
## 快速开始
|
||||||
|
|
||||||
|
需要 Linux ARM64、GB10、Docker Compose（支持 `gpus: all`）、NVIDIA Container Toolkit，
|
||||||
|
以及至少约 150 GiB 的模型下载空间，另留 Docker 镜像和缓存空间。
|
||||||
|
实测驱动为 580.173.02。首次启动约需 10–13 分钟。
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cp .env.example .env
|
||||||
|
# 编辑 .env：模型缓存目录、现有 API Key 文件、代理和端口
|
||||||
|
# API_KEY_FILE 指向非空的纯文本密钥文件，不把密钥写入 .env。
|
||||||
|
./scripts/download.sh
|
||||||
|
./scripts/build.sh
|
||||||
|
./scripts/start.sh
|
||||||
|
./scripts/wait.sh
|
||||||
|
./scripts/test.sh
|
||||||
|
```
|
||||||
|
|
||||||
|
已有 Hugging Face 缓存时，`HF_CACHE_DIR` 应指向包含 `hub/` 的目录。
|
||||||
|
下载脚本固定模型 revision，支持断点续传；服务离线读取该快照。
|
||||||
|
已有本项目所固定的基础镜像时，可在 `.env` 设置 `BASE_IMAGE=local/vllm-base:0bfc7a15`。
|
||||||
|
请先核实本地 tag 对应 [版本清单](docs/provenance.md) 中的 digest，不能用任意 nightly 替代。
|
||||||
|
|
||||||
|
API 地址：`http://<Spark-IP>:8000/v1`；模型名：`qwen3.8-flash-next`。
|
||||||
|
使用 `Authorization: Bearer <已有密钥>`。代理仅供模型下载/容器出网使用。
|
||||||
|
|
||||||
|
## 当前配置
|
||||||
|
|
||||||
|
| 参数 | 设置 |
|
||||||
|
|---|---|
|
||||||
|
| 上下文上限 | 262144 tokens |
|
||||||
|
| MTP | 2 speculative tokens |
|
||||||
|
| 同时调度请求数 | 4 |
|
||||||
|
| 批处理 token 上限 | 2048 |
|
||||||
|
| GPU 内存比例 | 0.80 |
|
||||||
|
| KV 精度 | auto（本配置为 BF16） |
|
||||||
|
| 执行 | eager，未启用 CUDA Graph |
|
||||||
|
| Prefix caching | 关闭 |
|
||||||
|
| API | 8000，启用密钥、reasoning parser、工具调用 |
|
||||||
|
| 自动重启 | unless-stopped |
|
||||||
|
|
||||||
|
0.80 留出的内存供系统和 PLE 文件缓存使用，并非浪费。
|
||||||
|
增加 KV 分配会挤占页缓存，不一定提高响应速度，见 [内存说明](docs/architecture.md)。
|
||||||
|
`mem_limit: 112g` 沿用实测配置；统一内存/CUDA 占用不能只用 Docker stats 判断，需结合主机内存。
|
||||||
|
|
||||||
|
## 运维与回退
|
||||||
|
|
||||||
|
```bash
|
||||||
|
docker compose logs -f vllm
|
||||||
|
docker compose ps
|
||||||
|
docker compose stop vllm
|
||||||
|
# 切换为已验证的 32K、不启用 MTP 的基础配置（会重建服务）
|
||||||
|
./scripts/start.sh baseline
|
||||||
|
./scripts/wait.sh
|
||||||
|
./scripts/test.sh
|
||||||
|
# 恢复默认 262K + MTP
|
||||||
|
./scripts/start.sh
|
||||||
|
```
|
||||||
|
|
||||||
|
首次迁移到本项目时，先停止旧部署中占用同一端口/模型内存的服务；不要并行启动两份。
|
||||||
|
项目脚本不会自动停止 MinerU 或其他业务容器。原部署的容器与回退配置可以保留。
|
||||||
|
|
||||||
|
`tests/test_disk_adapter.py` 验证 FP8 字节、边界/重复索引、缩放、原始哈希和 torch.compile。
|
||||||
|
可用 `scripts/test-adapter.sh` 单独执行，需要真实 CUDA GPU 和额外内存。
|
||||||
|
该测试通过不代表完整模型的 CUDA Graph 已验证。
|
||||||
|
|
||||||
|
## 文件说明
|
||||||
|
|
||||||
|
- `Dockerfile`：固定基础镜像，安装 PLE 适配和两项 GB10 FLA 修改。
|
||||||
|
- `patches/`：针对当前 nightly 的适配层。
|
||||||
|
- `vendor/`：保持原样的社区 mmap 代码及许可证。
|
||||||
|
- `configs/baseline-32k.yaml`：基础配置的 Compose override。
|
||||||
|
- `docs/validation.md`：实测结果和未验证边界。
|
||||||
|
- `docs/troubleshooting.md`：故障依据、上游修复与排查步骤。
|
||||||
|
- `docs/provenance.md`：模型/镜像/第三方源码的固定版本和校验值。
|
||||||
|
|
||||||
|
`.env`、密钥、权重、日志、运行缓存均不提交。仓库不包含模型权重；模型使用条件以
|
||||||
|
[NVIDIA 模型页](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-NVFP4) 为准。
|
||||||
|
第三方代码按其原许可证使用，见 [第三方声明](THIRD_PARTY_NOTICES.md)。
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
# Third-party notices
|
||||||
|
|
||||||
|
`vendor/vllm_ple_mmap.py` is from blazux/qwen3.8-Flash-DGX, copyright 2026 blazux,
|
||||||
|
licensed under Apache-2.0. The upstream license notice is preserved in `vendor/LICENSE`.
|
||||||
|
The full Apache-2.0 text is included in `vendor/Apache-2.0.txt`.
|
||||||
|
See docs/provenance.md for the exact vendored file hash.
|
||||||
|
|
||||||
|
vLLM and its dependencies remain in the upstream container with their respective licenses.
|
||||||
|
Model weights are downloaded separately and governed by their upstream terms.
|
||||||
|
No blanket license is assigned here to user-authored deployment files.
|
||||||
@@ -0,0 +1,70 @@
|
|||||||
|
name: qwen38-flash-dgx-spark
|
||||||
|
services:
|
||||||
|
vllm:
|
||||||
|
image: local/qwen38-flash-spark:nightly-0bfc7a15
|
||||||
|
build:
|
||||||
|
context: .
|
||||||
|
args:
|
||||||
|
BASE_IMAGE: ${BASE_IMAGE:-vllm/vllm-openai@sha256:c4392d76e3eec8983fa152651365158cb062e348fd40398963f499d5867b9e28}
|
||||||
|
gpus: all
|
||||||
|
ipc: host
|
||||||
|
restart: unless-stopped
|
||||||
|
mem_limit: 112g
|
||||||
|
memswap_limit: 112g
|
||||||
|
ports:
|
||||||
|
- "${BIND_ADDRESS:-0.0.0.0}:${API_PORT:-8000}:8000"
|
||||||
|
environment:
|
||||||
|
HF_HUB_OFFLINE: "1"
|
||||||
|
VLLM_PLE_CPU_OFFLOAD: "0"
|
||||||
|
VLLM_PLE_MMAP_FAST_ROWS: "0"
|
||||||
|
VLLM_PLE_MMAP_WORKERS: "32"
|
||||||
|
VLLM_PLE_MMAP_PROMETHEUS: "0"
|
||||||
|
VLLM_USE_V2_MODEL_RUNNER: "1"
|
||||||
|
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||||
|
CUTE_DSL_ARCH: sm_121a
|
||||||
|
TORCHINDUCTOR_COMPILE_THREADS: "2"
|
||||||
|
HTTP_PROXY: ${MODEL_HTTP_PROXY:-}
|
||||||
|
HTTPS_PROXY: ${MODEL_HTTPS_PROXY:-}
|
||||||
|
NO_PROXY: localhost,127.0.0.1
|
||||||
|
volumes:
|
||||||
|
- ${HF_CACHE_DIR:?Set HF_CACHE_DIR in .env}:/root/.cache/huggingface:ro
|
||||||
|
- type: bind
|
||||||
|
source: ${API_KEY_FILE:?Set API_KEY_FILE in .env}
|
||||||
|
target: /run/secrets/qwen_api_key
|
||||||
|
read_only: true
|
||||||
|
bind:
|
||||||
|
create_host_path: false
|
||||||
|
- ./runtime-cache:/root/.cache/vllm
|
||||||
|
entrypoint: ["/bin/bash", "-lc"]
|
||||||
|
command:
|
||||||
|
- >-
|
||||||
|
exec vllm serve /root/.cache/huggingface/hub/models--nvidia--Qwen3.8-Flash-Next-NVFP4/snapshots/fc694b54fb0174e0913e6adf86691ef85a4ead47
|
||||||
|
--served-model-name qwen3.8-flash-next
|
||||||
|
--host 0.0.0.0 --port 8000
|
||||||
|
--tensor-parallel-size 1
|
||||||
|
--dtype bfloat16 --kv-cache-dtype auto
|
||||||
|
--gpu-memory-utilization ${GPU_MEMORY_UTILIZATION:-0.80}
|
||||||
|
--max-model-len 262144 --max-num-seqs ${MAX_NUM_SEQS:-4} --max-num-batched-tokens 2048
|
||||||
|
--enable-chunked-prefill --no-enable-prefix-caching
|
||||||
|
--speculative-config '{"method":"mtp","num_speculative_tokens":2}'
|
||||||
|
--enforce-eager --no-enable-flashinfer-autotune
|
||||||
|
--load-format safetensors
|
||||||
|
--reasoning-parser qwen3 --tool-call-parser qwen3_xml --enable-auto-tool-choice
|
||||||
|
--api-key "$$(cat /run/secrets/qwen_api_key)"
|
||||||
|
healthcheck:
|
||||||
|
test: ["CMD", "python3", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=5)"]
|
||||||
|
interval: 30s
|
||||||
|
timeout: 10s
|
||||||
|
retries: 5
|
||||||
|
start_period: 30m
|
||||||
|
download:
|
||||||
|
profiles: [tools]
|
||||||
|
image: ${BASE_IMAGE:-vllm/vllm-openai@sha256:c4392d76e3eec8983fa152651365158cb062e348fd40398963f499d5867b9e28}
|
||||||
|
environment:
|
||||||
|
HTTP_PROXY: ${MODEL_HTTP_PROXY:-}
|
||||||
|
HTTPS_PROXY: ${MODEL_HTTPS_PROXY:-}
|
||||||
|
NO_PROXY: localhost,127.0.0.1
|
||||||
|
volumes:
|
||||||
|
- ${HF_CACHE_DIR:?Set HF_CACHE_DIR in .env}:/root/.cache/huggingface
|
||||||
|
- ./scripts/download.py:/opt/download.py:ro
|
||||||
|
entrypoint: [python3, -u, /opt/download.py]
|
||||||
@@ -0,0 +1,16 @@
|
|||||||
|
services:
|
||||||
|
vllm:
|
||||||
|
command:
|
||||||
|
- >-
|
||||||
|
exec vllm serve /root/.cache/huggingface/hub/models--nvidia--Qwen3.8-Flash-Next-NVFP4/snapshots/fc694b54fb0174e0913e6adf86691ef85a4ead47
|
||||||
|
--served-model-name qwen3.8-flash-next
|
||||||
|
--host 0.0.0.0 --port 8000
|
||||||
|
--tensor-parallel-size 1
|
||||||
|
--dtype bfloat16 --kv-cache-dtype auto
|
||||||
|
--gpu-memory-utilization ${GPU_MEMORY_UTILIZATION:-0.80}
|
||||||
|
--max-model-len 32768 --max-num-seqs ${MAX_NUM_SEQS:-4} --max-num-batched-tokens 2048
|
||||||
|
--enable-chunked-prefill --no-enable-prefix-caching
|
||||||
|
--enforce-eager --no-enable-flashinfer-autotune
|
||||||
|
--load-format safetensors
|
||||||
|
--reasoning-parser qwen3 --tool-call-parser qwen3_xml --enable-auto-tool-choice
|
||||||
|
--api-key "$$(cat /run/secrets/qwen_api_key)"
|
||||||
@@ -0,0 +1,24 @@
|
|||||||
|
# PLE 与统一内存
|
||||||
|
|
||||||
|
GB10 的 CPU/GPU 共用统一内存。把 PLE 固定放到 CPU 内存仍会占用同一内存池。
|
||||||
|
本模型约 47.7 GiB 的 PLE 表不需要每一步完整读取，因此映射原始 safetensors，按索引收集所需行。
|
||||||
|
Linux 页缓存保留访问过的数据，未命中才访问 SSD。此方式不重新量化、不截断表、不修改哈希。
|
||||||
|
|
||||||
|
`patches/spark_ngram_adapter.py` 适配 nightly 新的 ngram_embedding 接口，
|
||||||
|
复用上游哈希/缩放和 vendor 的 mmap/custom-op；只支持此 FP8 PLE 布局和 ETP=1。
|
||||||
|
它修改内部类，因此不能随意升级基础镜像。失败时应修复版本兼容，不能吞掉异常继续运行。
|
||||||
|
|
||||||
|
实测默认配置模型权重占 76.36 GiB，KV 缓存约 14.52 GiB。
|
||||||
|
主机空闲时 used 约 102 GiB、buff/cache 约 19 GiB、available 约 18 GiB、free 约 1.6 GiB。
|
||||||
|
这些是不同口径，available 包含可回收缓存，不能与 buff/cache 相加；也不能认定所有缓存都是 PLE。
|
||||||
|
主机约 227 MiB 的 swap 存量不等同于持续发生 swap I/O。
|
||||||
|
|
||||||
|
增大 `GPU_MEMORY_UTILIZATION` 主要增加 KV 空间，适合更多并发或长输入，并不直接提高单流解码速度。
|
||||||
|
它会减少 PLE 页缓存空间。应比较 TTFT、tokens/s、排队时间、KV 使用率、页缺失和 swap I/O，
|
||||||
|
而不是追求 100% 内存占用。社区默认同样为 0.80，但社区性能数字不是此配置的测试结果。
|
||||||
|
|
||||||
|
完整 PLE 运算包含 GPU 同步，日志中的 op 总时长不能全部归因于磁盘；分析时应区分 gather 与 gpu-wait。
|
||||||
|
当前使用 eager 执行，未开启 prefix caching，也没有采用社区可选的二次 FP8 量化或词表裁剪。
|
||||||
|
|
||||||
|
来源：[社区方案](https://github.com/blazux/qwen3.8-Flash-DGX)、
|
||||||
|
[Linux 内存指标](https://docs.kernel.org/filesystems/proc.html)。
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
# 项目整理验证（2026-09-17）
|
||||||
|
|
||||||
|
- 全部 Python 文件通过语法解析；全部 Shell 脚本通过 bash -n。
|
||||||
|
- vendor/vllm_ple_mmap.py 的 SHA256 与部署时使用的文件一致。
|
||||||
|
- Spark 上 docker compose config 验证默认及 baseline 两套配置成功。
|
||||||
|
- 将两套配置展开后的 command、entrypoint 与原部署逐项比较，完全一致。
|
||||||
|
- 新目录结构在 Spark 上成功 docker build，临时镜像为 local/qwen38-git-check:0bfc7a15。
|
||||||
|
构建使用已经核对的本地基础镜像 tag，没有重新拉取浮动 nightly。
|
||||||
|
- 新版 smoke-test.py 对现有线上容器执行成功：算术 1.76 秒、中文回答 5.18 秒、工具调用 1.92 秒。
|
||||||
|
- 检查 Git 待提交文件：不含 .env、密钥文件、原始日志、模型权重或原主机/代理地址。
|
||||||
|
|
||||||
|
线上服务没有被重建。此轮验证覆盖项目构建、配置等价性与测试脚本；
|
||||||
|
没有从空缓存重新下载 124 GiB 权重，也没有再次完整加载模型。
|
||||||
@@ -0,0 +1,22 @@
|
|||||||
|
# 固定版本与来源
|
||||||
|
|
||||||
|
- 模型：`nvidia/Qwen3.8-Flash-Next-NVFP4`
|
||||||
|
- 模型 revision：`fc694b54fb0174e0913e6adf86691ef85a4ead47`
|
||||||
|
- 基础镜像：`vllm/vllm-openai@sha256:c4392d76e3eec8983fa152651365158cb062e348fd40398963f499d5867b9e28`
|
||||||
|
- 基础镜像源码 commit：`0bfc7a15d095fe83ecc82b50561a93c177fece2d`
|
||||||
|
- 本地构建 tag：`local/qwen38-flash-spark:nightly-0bfc7a15`
|
||||||
|
|
||||||
|
nightly 的包版本曾报告 `0.3.1.dev3+g0bfc7a15d`，因此使用镜像 digest/源码 commit 标识，
|
||||||
|
不把它当作正式版 0.29.0。BASE_IMAGE override 仅用于已核对 digest 的本地镜像别名。
|
||||||
|
|
||||||
|
## 第三方文件
|
||||||
|
|
||||||
|
`vendor/vllm_ple_mmap.py` 原样获取自：
|
||||||
|
https://raw.githubusercontent.com/blazux/qwen3.8-Flash-DGX/main/src/vllm_ple_mmap.py
|
||||||
|
|
||||||
|
获取日期 2026-09-17。原 URL 使用 main，不是永久版本链接；本仓库保存文件本体，以下 hash 固定实际采用内容：
|
||||||
|
|
||||||
|
`SHA256 02adb76b789f9fccf472e922e3be7b4b2a95962ae7d7c3eb99bfeac6052dccad`
|
||||||
|
|
||||||
|
不要重新下载 main 后仍宣称是同一版本。许可证按获取时的文件保存在 vendor/LICENSE。
|
||||||
|
适配层和 Dockerfile 是本次部署的本地改动，第三方 helper 未修改。
|
||||||
@@ -0,0 +1,37 @@
|
|||||||
|
# 故障排查与上游依据
|
||||||
|
|
||||||
|
## 原 Qwen3.6 升级后启动失败
|
||||||
|
|
||||||
|
旧部署使用主模型 marlin、MTP triton，但 V2 runner 曾未正确应用 draft MoE backend。
|
||||||
|
对应修复：[vLLM #54788](https://github.com/vllm-project/vllm/pull/54788)，
|
||||||
|
[GB10 复现评论](https://github.com/vllm-project/vllm/pull/54788#issuecomment-5554777896)。
|
||||||
|
不能把所有 NVFP4/MTP 错误都归为同一问题；本项目 NVIDIA Flash Next 的 MTP 是 block FP8。
|
||||||
|
对应加载支持修复：[vLLM #55513](https://github.com/vllm-project/vllm/pull/55513)。
|
||||||
|
固定 nightly 已包含这些代码；本项目没有把主、草稿模型强制绑定同一 MoE backend。
|
||||||
|
|
||||||
|
## 模型下载超时
|
||||||
|
|
||||||
|
先确认 `.env` 中 MODEL_HTTP_PROXY / MODEL_HTTPS_PROXY 是否沿用了可用代理。
|
||||||
|
下载容器单独使用这些变量，HF 缓存写入主机；服务以 HF_HUB_OFFLINE=1 读取固定 revision。
|
||||||
|
不要为网络故障修改推理参数。已有权重需要位于正确的 HF snapshot 目录并具备完整 blob 链接。
|
||||||
|
|
||||||
|
## 权重加载慢
|
||||||
|
|
||||||
|
本模型约 123.57 GiB，首次主模型加载约 9 分钟，MTP 再约 1 分钟。
|
||||||
|
healthcheck 的 start_period 为 30 分钟。用 `docker compose logs -f vllm` 确认分片进度。
|
||||||
|
如果出现异常退出，先读堆栈；不要只依赖自动重启或盲目提高内存比例。
|
||||||
|
|
||||||
|
## 内存不足或变慢
|
||||||
|
|
||||||
|
先看主机 `free -h`、`vmstat 1`，再看服务日志的 KV 分配和 PLE gather 时间。
|
||||||
|
不要只看 Docker stats，也不要把 available 当作完全空闲。
|
||||||
|
确认没有旧模型或其他 GPU 容器并行运行。默认 0.80 留给页缓存的空间有实际用途。
|
||||||
|
必要时使用 `scripts/start.sh baseline` 回退，再逐项增加配置。
|
||||||
|
|
||||||
|
## 更新 nightly
|
||||||
|
|
||||||
|
不要只改镜像 tag。先核对 ngram_embedding 接口和两处 FLA 源码结构，重跑 adapter GPU 测试、
|
||||||
|
启动测试和实际问答。Dockerfile 对被替换的内容使用断言，接口变化应让构建失败以便审查。
|
||||||
|
|
||||||
|
社区适配来源：[blazux/qwen3.8-Flash-DGX](https://github.com/blazux/qwen3.8-Flash-DGX)。
|
||||||
|
本项目只使用其 mmap helper 和两项 FLA 修改，没有搬用所有社区优化。
|
||||||
@@ -0,0 +1,36 @@
|
|||||||
|
# 验证记录
|
||||||
|
|
||||||
|
日期：2026-09-17。硬件为单台 DGX Spark / GB10，ARM64，121 GiB 可见内存。
|
||||||
|
系统 DGX OS 7.6，内核 7.0.0-1019-nvidia，驱动 580.173.02，驱动支持 CUDA 13.0。
|
||||||
|
版本见 provenance.md；这不是可泛化到所有 nightly 或所有 GPU 的结果。
|
||||||
|
|
||||||
|
## 原部署的实测结果
|
||||||
|
|
||||||
|
| 配置/测试 | 结果 |
|
||||||
|
|---|---|
|
||||||
|
| 32K / 不启用 MTP / eager | 成功启动 |
|
||||||
|
| 32K：17 × 19 | 323，3.66 秒，56 completion tokens |
|
||||||
|
| 32K：中文解释天空颜色 | 正常，5.19 秒，81 completion tokens |
|
||||||
|
| 262K 上限 / MTP 2 / eager | 成功启动，健康检查通过 |
|
||||||
|
| MTP：17 × 19 | 323，2.40 秒，56 completion tokens |
|
||||||
|
| MTP：中文解释天空颜色 | 正常，5.46 秒，150 completion tokens |
|
||||||
|
| 自动工具调用 | get_weather，city=上海，约 2.25 秒 |
|
||||||
|
| 局域网健康检查 | HTTP 200 |
|
||||||
|
| 适配层 CUDA 测试 | PASS，逐字节 FP8 查表及缩放一致 |
|
||||||
|
| 适配层 torch.compile 测试 | COMPILE_PASS |
|
||||||
|
|
||||||
|
completion tokens 包括思考 token。以上是小样本验证，不是性能基准；不同回答长度不能直接比较速度。
|
||||||
|
测试工具调用只检查模型产生的结构化调用，不访问真实天气服务。
|
||||||
|
|
||||||
|
262K 配置日志：模型权重 76.36 GiB，KV 缓存 14.52 GiB，报告 524288 总 token 容量；
|
||||||
|
后者不是“两个满长度请求已压测”的承诺。实际可用空间随运行环境变化。
|
||||||
|
主模型选择 FLASHINFER_CUTLASS；FP8 MTP 选择 DEEPGEMM，实测推理成功。
|
||||||
|
|
||||||
|
## 未验证
|
||||||
|
|
||||||
|
- 完整 262144-token 请求及长时间高并发压力。
|
||||||
|
- 多模态输入、完整模型 CUDA Graph 和 prefix caching。
|
||||||
|
- 其他 GPU、其他模型 revision、其他 nightly、ETP > 1。
|
||||||
|
- 重启机器后的恢复耗时（已设置 unless-stopped，但未为测试重启机器）。
|
||||||
|
|
||||||
|
仓库整理后的构建/配置验证另记录于 packaging-validation.md。
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
"""PLE mmap adapter for pinned nightly 0bfc7a15, single GPU only.
|
||||||
|
|
||||||
|
Uses the reviewed blazux mmap reader, retaining upstream hashing and dequantization.
|
||||||
|
"""
|
||||||
|
import torch
|
||||||
|
from vllm_ple_mmap import _MmapNgramEmbedding, _setup_table_v029, _REGISTRY, _register_op
|
||||||
|
|
||||||
|
|
||||||
|
class DiskEmbedding(_MmapNgramEmbedding):
|
||||||
|
supports_prefetch = False
|
||||||
|
|
||||||
|
def __init__(self, n, d, **kwargs):
|
||||||
|
super().__init__(n, d)
|
||||||
|
self.register_buffer("weight", torch.empty(0, dtype=torch.float8_e4m3fn), persistent=False)
|
||||||
|
|
||||||
|
def dequantize(self, embeddings, output_dtype):
|
||||||
|
if self.table is None:
|
||||||
|
raise RuntimeError("PLE disk table not loaded")
|
||||||
|
return embeddings.to(output_dtype) * self.weight_scale.to(output_dtype)
|
||||||
|
|
||||||
|
def start_prefetch(self, *args):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def apply(cls):
|
||||||
|
import sys
|
||||||
|
from vllm.config import get_current_vllm_config
|
||||||
|
from vllm.distributed import get_etp_group
|
||||||
|
mod = sys.modules[cls.__module__]
|
||||||
|
original_init, original_load = cls.__init__, cls.load_weights
|
||||||
|
|
||||||
|
def init(self, *args, **kwargs):
|
||||||
|
if get_etp_group().world_size != 1:
|
||||||
|
raise RuntimeError("Spark disk adapter supports only ETP=1")
|
||||||
|
device_cls = mod.Qwen4ExpPLEDeviceEmbedding
|
||||||
|
host_cls = mod.Qwen4ExpPLEPinnedHostEmbedding
|
||||||
|
mod.Qwen4ExpPLEDeviceEmbedding = mod.Qwen4ExpPLEPinnedHostEmbedding = DiskEmbedding
|
||||||
|
try:
|
||||||
|
original_init(self, *args, **kwargs)
|
||||||
|
finally:
|
||||||
|
mod.Qwen4ExpPLEDeviceEmbedding, mod.Qwen4ExpPLEPinnedHostEmbedding = device_cls, host_cls
|
||||||
|
self._ple_mmap_prefix = kwargs["prefix"]
|
||||||
|
self._ple_mmap_model_path = get_current_vllm_config().model_config.model
|
||||||
|
_REGISTRY[self._ple_mmap_prefix] = self
|
||||||
|
|
||||||
|
def load(self, weights):
|
||||||
|
loaded = set()
|
||||||
|
def filtered():
|
||||||
|
for name, tensor in weights:
|
||||||
|
if name.startswith("ngram_embedding.shard_") and name.endswith(".weight"):
|
||||||
|
continue
|
||||||
|
if name == "ngram_embedding.weight_scale":
|
||||||
|
self.register_buffer("_offload_weight_scale", tensor.detach().to("cuda"), persistent=False)
|
||||||
|
continue
|
||||||
|
yield name, tensor
|
||||||
|
loaded.update(original_load(self, filtered()))
|
||||||
|
_setup_table_v029(self)
|
||||||
|
self.ngram_embedding.weight_scale = self._offload_weight_scale
|
||||||
|
return loaded
|
||||||
|
|
||||||
|
def forward(self, hidden_states, input_ids, query_start_loc, ngram_context):
|
||||||
|
ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context)
|
||||||
|
output = torch.empty((ids.shape[0], self.embedding_dim),
|
||||||
|
dtype=torch.float8_e4m3fn, device=ids.device)
|
||||||
|
torch.ops.vllm.ple_mmap_lookup_ids(ids, output, self._ple_mmap_prefix)
|
||||||
|
return output
|
||||||
|
|
||||||
|
_register_op()
|
||||||
|
cls.__init__, cls.load_weights, cls.forward = init, load, forward
|
||||||
Executable
+3
@@ -0,0 +1,3 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
|
||||||
|
"${compose[@]}" build vllm
|
||||||
Executable
+9
@@ -0,0 +1,9 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
set -euo pipefail
|
||||||
|
PROJECT_ROOT=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)
|
||||||
|
cd "$PROJECT_ROOT"
|
||||||
|
if [[ ! -f .env ]]; then
|
||||||
|
echo 'Copy .env.example to .env and configure paths/proxy first.' >&2
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
compose=(docker compose --env-file .env -f compose.yaml)
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
from huggingface_hub import snapshot_download
|
||||||
|
|
||||||
|
print(snapshot_download(
|
||||||
|
repo_id="nvidia/Qwen3.8-Flash-Next-NVFP4",
|
||||||
|
revision="fc694b54fb0174e0913e6adf86691ef85a4ead47",
|
||||||
|
max_workers=4,
|
||||||
|
), flush=True)
|
||||||
Executable
+3
@@ -0,0 +1,3 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
|
||||||
|
"${compose[@]}" --profile tools run --rm download
|
||||||
@@ -0,0 +1,42 @@
|
|||||||
|
import json
|
||||||
|
import time
|
||||||
|
import urllib.request
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
key = Path('/run/secrets/qwen_api_key').read_text().strip()
|
||||||
|
base = 'http://127.0.0.1:8000'
|
||||||
|
headers = {'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json'}
|
||||||
|
|
||||||
|
def request(path, payload=None):
|
||||||
|
data = None if payload is None else json.dumps(payload).encode()
|
||||||
|
with urllib.request.urlopen(urllib.request.Request(base + path, data=data, headers=headers), timeout=900) as response:
|
||||||
|
return json.load(response)
|
||||||
|
|
||||||
|
print(json.dumps({'models': request('/v1/models')}, ensure_ascii=False), flush=True)
|
||||||
|
cases = [
|
||||||
|
('arithmetic', '计算 17 × 19，只给出数字。', '323'),
|
||||||
|
('chinese', '请用中文简短说明为什么天空是蓝色的。', None),
|
||||||
|
]
|
||||||
|
for name, prompt, expected in cases:
|
||||||
|
start = time.monotonic()
|
||||||
|
result = request('/v1/chat/completions', {
|
||||||
|
'model': 'qwen3.8-flash-next',
|
||||||
|
'messages': [{'role': 'user', 'content': prompt}],
|
||||||
|
'temperature': 0, 'max_tokens': 1024, 'reasoning_effort': 'low',
|
||||||
|
})
|
||||||
|
choice = result['choices'][0]
|
||||||
|
content = choice['message'].get('content') or ''
|
||||||
|
record = {'test': name, 'seconds': round(time.monotonic()-start, 2),
|
||||||
|
'content': content, 'finish_reason': choice['finish_reason'], 'usage': result.get('usage')}
|
||||||
|
print(json.dumps(record, ensure_ascii=False), flush=True)
|
||||||
|
assert content.strip(), 'No final answer'
|
||||||
|
if expected:
|
||||||
|
assert expected in content, 'Arithmetic mismatch'
|
||||||
|
|
||||||
|
|
||||||
|
p={'model':'qwen3.8-flash-next','messages':[{'role':'user','content':'请调用 get_weather 查询上海的天气，不要自行编造。'}],'tools':[{'type':'function','function':{'name':'get_weather','description':'查询城市天气','parameters':{'type':'object','properties':{'city':{'type':'string'}},'required':['city']}}}],'tool_choice':'auto','temperature':0,'max_tokens':512,'reasoning_effort':'low'}
|
||||||
|
t=time.monotonic();r=request('/v1/chat/completions', p);m=r['choices'][0]['message'];print(json.dumps({'test':'tool_call','seconds':time.monotonic()-t,'message':m},ensure_ascii=False),flush=True)
|
||||||
|
assert m.get('tool_calls') and m['tool_calls'][0]['function']['name']=='get_weather'
|
||||||
|
print('TOOL_TEST_PASS',flush=True)
|
||||||
|
|
||||||
|
print("SMOKE_TEST_PASS", flush=True)
|
||||||
Executable
+10
@@ -0,0 +1,10 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
|
||||||
|
if [[ "${1:-}" == baseline ]]; then
|
||||||
|
compose+=(-f configs/baseline-32k.yaml)
|
||||||
|
elif [[ $# -gt 0 ]]; then
|
||||||
|
echo "Usage: $0 [baseline]" >&2; exit 2
|
||||||
|
fi
|
||||||
|
"${compose[@]}" config --quiet
|
||||||
|
"${compose[@]}" up -d --no-build vllm
|
||||||
|
echo 'Loading takes about 10-13 minutes. Run scripts/wait.sh, then scripts/test.sh.'
|
||||||
Executable
+3
@@ -0,0 +1,3 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
|
||||||
|
"${compose[@]}" run --rm --no-deps -T --entrypoint python3 vllm -u - < tests/test_disk_adapter.py
|
||||||
Executable
+3
@@ -0,0 +1,3 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
|
||||||
|
"${compose[@]}" exec -T vllm python3 -u - < scripts/smoke-test.py
|
||||||
Executable
+10
@@ -0,0 +1,10 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
|
||||||
|
for ((i=0; i<120; i++)); do
|
||||||
|
if "${compose[@]}" exec -T vllm python3 -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/health',timeout=3)" >/dev/null 2>&1; then
|
||||||
|
echo READY; exit 0
|
||||||
|
fi
|
||||||
|
sleep 15
|
||||||
|
done
|
||||||
|
echo 'Startup timed out. Inspect docker compose logs vllm.' >&2
|
||||||
|
exit 1
|
||||||
@@ -0,0 +1,49 @@
|
|||||||
|
import tempfile
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from pathlib import Path
|
||||||
|
import torch
|
||||||
|
import vllm.config
|
||||||
|
import vllm.distributed
|
||||||
|
from safetensors.torch import save_file
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as folder:
|
||||||
|
fake_config = SimpleNamespace(engram_config=None, model_config=SimpleNamespace(model=folder))
|
||||||
|
vllm.config.get_current_vllm_config = lambda: fake_config
|
||||||
|
vllm.distributed.get_etp_group = lambda: SimpleNamespace(world_size=1)
|
||||||
|
from vllm.models.qwen4_exp.nvidia.ngram_embedding import Qwen4ExpNGramEmbedding
|
||||||
|
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||||
|
cfg = SimpleNamespace(ngram_size=3, heads_per_ngram=2, eos_token_id=0,
|
||||||
|
vocab_size=100, split_ngram_parts=2, seed=1234,
|
||||||
|
ngram_vocab_size_base=17, make_ngram_vocab_size_divisible_by=8)
|
||||||
|
prefix = 'model.language_model.layers.0.ple.ple_embedding'
|
||||||
|
with torch.device('cuda'):
|
||||||
|
layer = Qwen4ExpNGramEmbedding(cfg, 8, 0, 32,
|
||||||
|
data_parallel_rank=0, prefix=prefix,
|
||||||
|
quant_config=Fp8Config(is_checkpoint_fp8_serialized=True))
|
||||||
|
count = layer.ngram_embedding.org_vocab_size
|
||||||
|
full = (torch.arange(count * 2).reshape(count, 2) % 13).to(torch.float8_e4m3fn)
|
||||||
|
tensors = {}
|
||||||
|
for i, part in enumerate(full.chunk(2)):
|
||||||
|
tensors[f'{prefix}.ngram_embedding.shard_{i}.weight'] = part.contiguous()
|
||||||
|
tensors[f'{prefix}.ngram_embedding.weight_scale'] = torch.tensor(0.25)
|
||||||
|
save_file(tensors, str(Path(folder) / 'ple.safetensors'))
|
||||||
|
layer.load_weights((name.removeprefix(prefix + '.'), value) for name, value in tensors.items())
|
||||||
|
ids = torch.tensor([[0, 1, 1, count - 1], [5, 2, 9, 7]], device='cuda')
|
||||||
|
out = layer.ngram_embedding(ids)
|
||||||
|
expected = full.view(torch.uint8)[ids.cpu()].view(torch.float8_e4m3fn).cuda()
|
||||||
|
assert torch.equal(out.view(torch.uint8), expected.view(torch.uint8))
|
||||||
|
dequant = layer.ngram_embedding.dequantize(out, torch.bfloat16)
|
||||||
|
assert torch.equal(dequant, expected.to(torch.bfloat16) * 0.25)
|
||||||
|
assert layer.ngram_embedding.weight.numel() == 0
|
||||||
|
tokens = torch.tensor([3, 7, 9], device='cuda')
|
||||||
|
starts = torch.tensor([0, 3], dtype=torch.int32, device='cuda')
|
||||||
|
context = torch.tensor([[0, 0]], device='cuda')
|
||||||
|
hashed = layer.compute_ngram_ids(tokens, starts, context)
|
||||||
|
actual = layer(None, tokens, starts, context)
|
||||||
|
reference = full.view(torch.uint8)[hashed.cpu()].view(torch.float8_e4m3fn).cuda().flatten(-2)
|
||||||
|
assert torch.equal(actual.view(torch.uint8), reference.view(torch.uint8))
|
||||||
|
print('PASS: disk rows, repeated IDs, boundary IDs, FP8 bytes and scaling; no resident PLE table')
|
||||||
|
compiled = torch.compile(lambda t, s, c: layer(None, t, s, c), fullgraph=True)
|
||||||
|
actual_compiled = compiled(tokens, starts, context)
|
||||||
|
assert torch.equal(actual_compiled.view(torch.uint8), reference.view(torch.uint8))
|
||||||
|
print('COMPILE_PASS: full graph matches exact FP8 table lookup')
|
||||||
Vendored
+202
@@ -0,0 +1,202 @@
|
|||||||
|
|
||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright [yyyy] [name of copyright owner]
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
||||||
Vendored
+19
@@ -0,0 +1,19 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
Copyright 2026 blazux
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
||||||
|
|
||||||
|
Full text: https://www.apache.org/licenses/LICENSE-2.0.txt
|
||||||
Vendored
+908
@@ -0,0 +1,908 @@
|
|||||||
|
"""vllm_ple_mmap — serve the Qwen3.8-Flash-Next N-gram (PLE) table from NVMe via mmap.
|
||||||
|
|
||||||
|
Why: the 51B-parameter n-gram table is 47.7 GiB (51.2 GB) in FP8 and vLLM keeps it resident
|
||||||
|
(GPU, or pinned host RAM with VLLM_PLE_CPU_OFFLOAD). On a DGX Spark / GX10 the
|
||||||
|
host and the GPU share one 121 GiB pool, so neither fits next to the 78 GiB main
|
||||||
|
model. But a token only ever touches 16 rows x 160 bytes of that table, so the
|
||||||
|
table can live on disk and be served through the page cache — exactly what
|
||||||
|
llama.cpp does with its GGUF mmap.
|
||||||
|
|
||||||
|
How: with VLLM_PLE_MMAP=1 this module patches ``Qwen3_8FlashNextNGramEmbedding``:
|
||||||
|
* ``__init__`` swaps the 44/95 GiB ``VocabParallelEmbedding`` for a tiny
|
||||||
|
placeholder whose ``forward(ids)`` gathers rows from ``np.memmap`` views of the
|
||||||
|
checkpoint's ``model-plefp8-*.safetensors`` shards (zero-copy, page-cache backed);
|
||||||
|
* ``load_weights`` drops the 128 shard tensors on the floor, keeps the global FP8
|
||||||
|
``weight_scale`` (as ``_offload_weight_scale``, which the untouched
|
||||||
|
``Qwen3_8FlashNextPLELayer._dequantize_embeddings`` already consumes) and opens
|
||||||
|
the memmaps.
|
||||||
|
* ``forward_impl`` (hashing + lookup) is wrapped in a custom op
|
||||||
|
``vllm::ple_mmap_lookup`` so that (a) torch.compile treats it as opaque — the
|
||||||
|
stock version trips an Inductor int64 indexing assert on sm_121 — and (b) it can
|
||||||
|
be listed in ``-cc.splitting_ops`` and run OUTSIDE piecewise CUDA graphs: the
|
||||||
|
gather is CPU work + a pageable H2D copy, which cannot live inside a capture.
|
||||||
|
Use ``-cc.cudagraph_mode=PIECEWISE`` (not FULL*) with the splitting op list in
|
||||||
|
serve-flashnext-vllm.sh, or ``--enforce-eager``.
|
||||||
|
Nothing else in vLLM changes: the n-gram hashing, the short-conv, the dequant path
|
||||||
|
are the stock ones.
|
||||||
|
|
||||||
|
Fast gather hot path (CPU dedup -> persistent pinned staging buffer -> async H2D ->
|
||||||
|
GPU-side inverse expansion, plus a no-threadpool fast path for decode-sized
|
||||||
|
batches), bf16/f16 table support, VLLM_PLE_MMAP_DIR and the periodic stats line
|
||||||
|
were contributed by @Saren-Arterius (github.com/Saren-Arterius/qwen3.8-Flash-DGX-AutoRound).
|
||||||
|
|
||||||
|
Knobs (env):
|
||||||
|
VLLM_PLE_MMAP=1 enable
|
||||||
|
VLLM_PLE_MMAP_WORKERS=32 gather threads (page faults overlap across threads)
|
||||||
|
VLLM_PLE_MMAP_CHUNK=2048 rows per gather task
|
||||||
|
VLLM_PLE_MMAP_MADVISE=random madvise on the shard mmaps: random (default, no readahead) or normal
|
||||||
|
VLLM_PLE_MMAP_PROMETHEUS=1 0 = do not register the vllm:ple_mmap_* counters
|
||||||
|
VLLM_PLE_MMAP_PREWARM=0 1 = stream the whole table once at load to fill the
|
||||||
|
page cache with whatever memory is free (harmless,
|
||||||
|
evictable; ~10 s at 4.7 GB/s)
|
||||||
|
|
||||||
|
Install: the Dockerfile copies this file next to vllm and appends
|
||||||
|
``_ple_mmap_apply(Qwen3_8FlashNextNGramEmbedding)`` to the end of
|
||||||
|
``vllm/models/qwen3_8_flash_next/nvidia/ple_layer.py``. See the repo README.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import glob
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import math
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import struct
|
||||||
|
import sys
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
from typing import Iterable
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
logger = logging.getLogger("vllm.ple_mmap")
|
||||||
|
|
||||||
|
ENV_ENABLE = "VLLM_PLE_MMAP"
|
||||||
|
_FP8_DTYPES = {
|
||||||
|
"F8_E4M3": torch.float8_e4m3fn,
|
||||||
|
"F8_E5M2": torch.float8_e5m2,
|
||||||
|
}
|
||||||
|
# 16-bit tables need no weight_scale: the stock _dequantize_embeddings is a
|
||||||
|
# no-op for non-FP8 rows.
|
||||||
|
_TABLE_DTYPES = {
|
||||||
|
**_FP8_DTYPES,
|
||||||
|
"BF16": torch.bfloat16,
|
||||||
|
"F16": torch.float16,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def enabled() -> bool:
|
||||||
|
return os.environ.get(ENV_ENABLE, "0").lower() in ("1", "true", "yes")
|
||||||
|
|
||||||
|
|
||||||
|
def _madvise(mm: np.memmap, kind: str) -> None:
|
||||||
|
"""Best-effort madvise on the mmap behind a np.memmap (Linux, Python >= 3.8)."""
|
||||||
|
try:
|
||||||
|
import mmap as _mmap
|
||||||
|
|
||||||
|
flag = {"random": _mmap.MADV_RANDOM, "normal": _mmap.MADV_NORMAL}[kind]
|
||||||
|
raw = getattr(mm, "_mmap", None)
|
||||||
|
if raw is not None:
|
||||||
|
raw.madvise(flag)
|
||||||
|
_MADVISED.append(kind)
|
||||||
|
except Exception as exc: # pragma: no cover - platform dependent
|
||||||
|
logger.warning("PLE mmap: madvise(%s) failed: %s", kind, exc)
|
||||||
|
|
||||||
|
|
||||||
|
_MADVISED: list[str] = []
|
||||||
|
|
||||||
|
|
||||||
|
def _env_int(name: str, default: int) -> int:
|
||||||
|
try:
|
||||||
|
return int(os.environ.get(name, default))
|
||||||
|
except ValueError:
|
||||||
|
return default
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# safetensors header parsing (no dependency on the safetensors package: we need
|
||||||
|
# raw file offsets, which its Python API does not expose)
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
def parse_safetensors_header(path: str) -> tuple[dict, int]:
|
||||||
|
"""Return (header_dict, data_start_offset) of a safetensors file."""
|
||||||
|
with open(path, "rb") as f:
|
||||||
|
(header_len,) = struct.unpack("<Q", f.read(8))
|
||||||
|
header = json.loads(f.read(header_len))
|
||||||
|
header.pop("__metadata__", None)
|
||||||
|
return header, 8 + header_len
|
||||||
|
|
||||||
|
|
||||||
|
class MmapPleTable:
|
||||||
|
"""Row gather over a table split into ``split_ngram_parts`` shard files.
|
||||||
|
|
||||||
|
``shards``: {shard_index: (path, absolute_byte_offset, rows)}. Shard ``i``
|
||||||
|
holds global rows ``[i*shard_size, i*shard_size + rows)`` (vLLM's
|
||||||
|
``copy_ple_embedding_shard_`` layout).
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
shards: dict[int, tuple[str, int, int]],
|
||||||
|
shard_size: int,
|
||||||
|
row_bytes: int,
|
||||||
|
torch_dtype: torch.dtype,
|
||||||
|
workers: int = 32,
|
||||||
|
chunk: int = 2048,
|
||||||
|
) -> None:
|
||||||
|
if not shards:
|
||||||
|
raise ValueError("no PLE shards")
|
||||||
|
self.shard_size = int(shard_size)
|
||||||
|
self.row_bytes = int(row_bytes)
|
||||||
|
self.torch_dtype = torch_dtype
|
||||||
|
self.chunk = max(1, int(chunk))
|
||||||
|
self.paths: list[str | None] = [None] * (max(shards) + 1)
|
||||||
|
self.mm: list[np.memmap | None] = [None] * (max(shards) + 1)
|
||||||
|
self.rows_total = 0
|
||||||
|
advise = os.environ.get("VLLM_PLE_MMAP_MADVISE", "random").strip().lower()
|
||||||
|
for idx, (path, offset, rows) in shards.items():
|
||||||
|
self.paths[idx] = path
|
||||||
|
self.mm[idx] = np.memmap(
|
||||||
|
path, dtype=np.uint8, mode="r", offset=offset, shape=(rows, row_bytes)
|
||||||
|
)
|
||||||
|
# Row lookups are 160-byte reads at hashed (random) addresses. Without
|
||||||
|
# MADV_RANDOM the kernel's mmap readahead pulls a window of pages around
|
||||||
|
# every faulting row and fills the page cache with neighbours that are
|
||||||
|
# never used; with it a cold row costs one page. PREWARM reads the file
|
||||||
|
# through a separate descriptor, so it is not affected.
|
||||||
|
if advise in ("random", "1"):
|
||||||
|
_madvise(self.mm[idx], "random")
|
||||||
|
self.rows_total += rows
|
||||||
|
self.pool = ThreadPoolExecutor(max_workers=max(1, int(workers)))
|
||||||
|
self.fast_rows = _env_int("VLLM_PLE_MMAP_FAST_ROWS", 512)
|
||||||
|
|
||||||
|
def gather(self, ids: np.ndarray) -> np.ndarray:
|
||||||
|
"""ids: int64 [N] global row ids -> uint8 [N, row_bytes] (a fresh array)."""
|
||||||
|
import time as _time
|
||||||
|
|
||||||
|
t0 = _time.perf_counter()
|
||||||
|
try:
|
||||||
|
return self._gather(ids)
|
||||||
|
finally:
|
||||||
|
dt = _time.perf_counter() - t0
|
||||||
|
n = int(np.asarray(ids).size)
|
||||||
|
_STATS["gather_ms"] += dt * 1e3
|
||||||
|
_STATS["rows"] += n
|
||||||
|
_STATS["bytes"] += n * self.row_bytes
|
||||||
|
_prom_add(gather_s=dt, rows=n, bytes=n * self.row_bytes)
|
||||||
|
|
||||||
|
def _gather(self, ids: np.ndarray) -> np.ndarray:
|
||||||
|
ids = np.ascontiguousarray(ids, dtype=np.int64).reshape(-1)
|
||||||
|
if ids.size == 0:
|
||||||
|
return np.empty((0, self.row_bytes), dtype=np.uint8)
|
||||||
|
if ids.size <= self.fast_rows:
|
||||||
|
# Decode-sized batches: thread-pool dispatch costs more than the
|
||||||
|
# reads themselves (~50 tasks for ~65 rows). Gather inline instead.
|
||||||
|
if ids.min() < 0 or ids.max() >= self.rows_total:
|
||||||
|
raise IndexError(
|
||||||
|
f"PLE row id out of range: [{ids.min()}, {ids.max()}] "
|
||||||
|
f"for {self.rows_total} rows"
|
||||||
|
)
|
||||||
|
shard = ids // self.shard_size
|
||||||
|
local = ids - shard * self.shard_size
|
||||||
|
out = np.empty((ids.size, self.row_bytes), dtype=np.uint8)
|
||||||
|
for si in np.unique(shard):
|
||||||
|
mask = shard == si
|
||||||
|
out[mask] = self.mm[si][local[mask]]
|
||||||
|
return out
|
||||||
|
# Dedupe + sort: repeated n-grams are common, and sorted rows improve
|
||||||
|
# locality inside a shard.
|
||||||
|
uniq, inverse = np.unique(ids, return_inverse=True)
|
||||||
|
if uniq[0] < 0 or uniq[-1] >= self.rows_total:
|
||||||
|
raise IndexError(
|
||||||
|
f"PLE row id out of range: [{uniq[0]}, {uniq[-1]}] "
|
||||||
|
f"for {self.rows_total} rows"
|
||||||
|
)
|
||||||
|
shard = uniq // self.shard_size
|
||||||
|
local = uniq - shard * self.shard_size
|
||||||
|
out = np.empty((uniq.size, self.row_bytes), dtype=np.uint8)
|
||||||
|
|
||||||
|
bounds = np.flatnonzero(np.diff(shard)) + 1
|
||||||
|
starts = np.concatenate(([0], bounds))
|
||||||
|
ends = np.concatenate((bounds, [uniq.size]))
|
||||||
|
tasks: list[tuple[int, int, int]] = []
|
||||||
|
for s, e in zip(starts.tolist(), ends.tolist()):
|
||||||
|
si = int(shard[s])
|
||||||
|
for c in range(s, e, self.chunk):
|
||||||
|
tasks.append((si, c, min(c + self.chunk, e)))
|
||||||
|
|
||||||
|
def run(task: tuple[int, int, int]) -> None:
|
||||||
|
si, a, b = task
|
||||||
|
mm = self.mm[si]
|
||||||
|
if mm is None:
|
||||||
|
raise IndexError(f"PLE shard {si} missing")
|
||||||
|
# Fancy indexing on a memmap: page faults do the I/O; NumPy releases
|
||||||
|
# the GIL for the copy, so tasks overlap across threads.
|
||||||
|
out[a:b] = mm[local[a:b]]
|
||||||
|
|
||||||
|
if len(tasks) == 1:
|
||||||
|
run(tasks[0])
|
||||||
|
else:
|
||||||
|
for _ in self.pool.map(run, tasks):
|
||||||
|
pass
|
||||||
|
return out[inverse]
|
||||||
|
|
||||||
|
def prewarm(self) -> None:
|
||||||
|
"""Stream every shard once so the page cache holds as much as it can."""
|
||||||
|
block = 64 << 20
|
||||||
|
for path, mm in zip(self.paths, self.mm):
|
||||||
|
if path is None or mm is None:
|
||||||
|
continue
|
||||||
|
start = mm.offset
|
||||||
|
end = start + mm.shape[0] * mm.shape[1]
|
||||||
|
with open(path, "rb", buffering=0) as f:
|
||||||
|
pos = start
|
||||||
|
while pos < end:
|
||||||
|
n = f.readinto(bytearray(min(block, end - pos))) # noqa: F841
|
||||||
|
if not n:
|
||||||
|
break
|
||||||
|
pos += n
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Placeholder that stands in for VocabParallelEmbedding
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
class _MmapNgramEmbedding(nn.Module):
|
||||||
|
"""Duck-types the bits of VocabParallelEmbedding the PLE code reads.
|
||||||
|
|
||||||
|
No ``weight`` attribute on purpose: ``Qwen3_8FlashNextPLELayer`` then falls
|
||||||
|
back to ``ple_embedding._offload_weight_scale`` for the FP8 scale.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, num_embeddings: int, embedding_dim: int) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.num_embeddings = int(num_embeddings)
|
||||||
|
self.org_vocab_size = int(num_embeddings)
|
||||||
|
self.embedding_dim = int(embedding_dim)
|
||||||
|
self.table: MmapPleTable | None = None
|
||||||
|
self._zeros_dtype = torch.bfloat16
|
||||||
|
|
||||||
|
def _pinned_buf(self, rows: int, row_bytes: int) -> torch.Tensor | None:
|
||||||
|
"""Persistent pinned staging buffer for async H2D (grown as needed)."""
|
||||||
|
buf = getattr(self, "_pinned", None)
|
||||||
|
if buf is None or buf.shape[0] < rows or buf.shape[1] != row_bytes:
|
||||||
|
try:
|
||||||
|
cap = max(rows + rows // 2, 4096)
|
||||||
|
buf = torch.empty((cap, row_bytes), dtype=torch.uint8, pin_memory=True)
|
||||||
|
except RuntimeError: # no CUDA (CPU tests) or pinning unavailable
|
||||||
|
buf = None
|
||||||
|
self._pinned = buf
|
||||||
|
return buf
|
||||||
|
|
||||||
|
def forward(self, ids: torch.Tensor) -> torch.Tensor:
|
||||||
|
table = self.table
|
||||||
|
if table is None:
|
||||||
|
# Weights never loaded (e.g. --load-format dummy): keep the plumbing
|
||||||
|
# alive with zeros so kernel tests can run without the 48 GiB table.
|
||||||
|
return torch.zeros(
|
||||||
|
(*ids.shape, self.embedding_dim),
|
||||||
|
dtype=self._zeros_dtype,
|
||||||
|
device=ids.device,
|
||||||
|
)
|
||||||
|
import time as _time
|
||||||
|
|
||||||
|
wait_s = 0.0
|
||||||
|
if ids.device.type == "cuda":
|
||||||
|
# The blocking copy below would wait here anyway, for every kernel queued ahead of it.
|
||||||
|
# Synchronizing first adds no latency and keeps that GPU time out of the lookup's own.
|
||||||
|
ts = _time.perf_counter()
|
||||||
|
torch.cuda.current_stream(ids.device).synchronize()
|
||||||
|
wait_s = _time.perf_counter() - ts
|
||||||
|
t1 = _time.perf_counter()
|
||||||
|
ids_np = ids.detach().to("cpu", non_blocking=False).numpy().reshape(-1)
|
||||||
|
# Dedup on CPU, gather only unique rows, expand on the GPU: fewer disk
|
||||||
|
# reads AND fewer H2D bytes (repeated n-grams are the common case).
|
||||||
|
uniq, inverse = np.unique(ids_np, return_inverse=True)
|
||||||
|
t2 = _time.perf_counter()
|
||||||
|
rows = table.gather(uniq) # uint8 [U, row_bytes], fresh & writable
|
||||||
|
t3 = _time.perf_counter()
|
||||||
|
u = rows.shape[0]
|
||||||
|
buf = self._pinned_buf(u, table.row_bytes) if ids.device.type == "cuda" else None
|
||||||
|
if buf is not None:
|
||||||
|
buf[:u].numpy()[:] = rows
|
||||||
|
dev = buf[:u].to(ids.device, non_blocking=True)
|
||||||
|
else:
|
||||||
|
dev = torch.from_numpy(rows).to(ids.device)
|
||||||
|
inv = torch.from_numpy(inverse.reshape(-1)).to(ids.device, non_blocking=True)
|
||||||
|
out = dev.view(table.torch_dtype)[inv]
|
||||||
|
t4 = _time.perf_counter()
|
||||||
|
_STATS["wait_ms"] += wait_s * 1e3
|
||||||
|
_STATS["dedup_ms"] += (t2 - t1) * 1e3
|
||||||
|
_STATS["stage_ms"] += (t4 - t3) * 1e3
|
||||||
|
_prom_add(gpu_wait_s=wait_s, dedup_s=t2 - t1, stage_s=t4 - t3)
|
||||||
|
return out.reshape(*ids.shape, self.embedding_dim)
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Patch
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
def _find_shards(
|
||||||
|
model_path: str, layer_idx: int
|
||||||
|
) -> tuple[dict[int, tuple[str, int, int]], str | None, tuple[str, int, int, str] | None, int | None]:
|
||||||
|
"""Locate ``layers.<idx>.ple.ple_embedding.ngram_embedding.shard_N.weight``.
|
||||||
|
|
||||||
|
Returns (shards, dtype_str, scale_entry, cols), where scale_entry is
|
||||||
|
(path, abs_offset, nbytes, dtype_str) of ``ngram_embedding.weight_scale`` or
|
||||||
|
None, and cols is the row width shared by all shards.
|
||||||
|
"""
|
||||||
|
shard_re = re.compile(
|
||||||
|
rf"layers\.{layer_idx}\.ple\.ple_embedding\.ngram_embedding\.shard_(\d+)\.weight$"
|
||||||
|
)
|
||||||
|
scale_re = re.compile(
|
||||||
|
rf"layers\.{layer_idx}\.ple\.ple_embedding\.ngram_embedding\.weight_scale$"
|
||||||
|
)
|
||||||
|
index_path = os.path.join(model_path, "model.safetensors.index.json")
|
||||||
|
if os.path.exists(index_path):
|
||||||
|
with open(index_path) as f:
|
||||||
|
weight_map = json.load(f)["weight_map"]
|
||||||
|
files = sorted(
|
||||||
|
{
|
||||||
|
os.path.join(model_path, fn)
|
||||||
|
for name, fn in weight_map.items()
|
||||||
|
if shard_re.search(name) or scale_re.search(name)
|
||||||
|
}
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
files = sorted(glob.glob(os.path.join(model_path, "*.safetensors")))
|
||||||
|
|
||||||
|
shards: dict[int, tuple[str, int, int]] = {}
|
||||||
|
dtype_str: str | None = None
|
||||||
|
scale_entry: tuple[str, int, int, str] | None = None
|
||||||
|
cols: int | None = None
|
||||||
|
for path in files:
|
||||||
|
header, data_start = parse_safetensors_header(path)
|
||||||
|
for name, meta in header.items():
|
||||||
|
m = shard_re.search(name)
|
||||||
|
if m:
|
||||||
|
start, end = meta["data_offsets"]
|
||||||
|
rows, cols = meta["shape"]
|
||||||
|
if dtype_str is None:
|
||||||
|
dtype_str = meta["dtype"]
|
||||||
|
elif meta["dtype"] != dtype_str:
|
||||||
|
raise ValueError("PLE shards have mixed dtypes")
|
||||||
|
if end - start != rows * cols * _itemsize(dtype_str):
|
||||||
|
raise ValueError(f"PLE shard {name}: size/shape mismatch")
|
||||||
|
shards[int(m.group(1))] = (path, data_start + start, rows)
|
||||||
|
elif scale_re.search(name):
|
||||||
|
start, end = meta["data_offsets"]
|
||||||
|
scale_entry = (path, data_start + start, end - start, meta["dtype"])
|
||||||
|
return shards, dtype_str, scale_entry, cols
|
||||||
|
|
||||||
|
|
||||||
|
def _itemsize(dtype_str: str) -> int:
|
||||||
|
return {
|
||||||
|
"F8_E4M3": 1,
|
||||||
|
"F8_E5M2": 1,
|
||||||
|
"U8": 1,
|
||||||
|
"I8": 1,
|
||||||
|
"BF16": 2,
|
||||||
|
"F16": 2,
|
||||||
|
"F32": 4,
|
||||||
|
}[dtype_str]
|
||||||
|
|
||||||
|
|
||||||
|
def _read_scale(entry: tuple) -> torch.Tensor:
|
||||||
|
path, offset, nbytes, dtype_str = entry
|
||||||
|
with open(path, "rb") as f:
|
||||||
|
f.seek(offset)
|
||||||
|
raw = f.read(nbytes)
|
||||||
|
if dtype_str == "F32":
|
||||||
|
return torch.tensor(struct.unpack("<f", raw[:4])[0], dtype=torch.float32)
|
||||||
|
if dtype_str == "BF16":
|
||||||
|
u16 = struct.unpack("<H", raw[:2])[0]
|
||||||
|
return torch.tensor(u16 << 16, dtype=torch.int32).view(torch.float32).squeeze()
|
||||||
|
if dtype_str == "F16":
|
||||||
|
return torch.frombuffer(bytearray(raw[:2]), dtype=torch.float16).clone().squeeze()
|
||||||
|
raise ValueError(f"unsupported weight_scale dtype {dtype_str}")
|
||||||
|
|
||||||
|
|
||||||
|
_REGISTRY: dict[str, nn.Module] = {}
|
||||||
|
_OP_NAME = "ple_mmap_lookup"
|
||||||
|
|
||||||
|
# Aggregate stats, logged every VLLM_PLE_MMAP_STATS_SEC seconds (0 = off). op_ms is wall time in the
|
||||||
|
# lookup op, and the op starts with a blocking device->host copy of the row ids. That copy waits for
|
||||||
|
# every kernel queued ahead of it on the stream: on v0.29 the n-gram hashing op and the layers before
|
||||||
|
# the PLE layer, on the preview image the hashing inside the op. So op_ms mixes that GPU compute with
|
||||||
|
# the lookup's own cost (a prefill window measured 165 ms/op of which 8 ms was the gather). The phases:
|
||||||
|
# wait_ms stream synchronize before the copy: GPU work queued ahead, not PLE cost
|
||||||
|
# dedup_ms copying the ids to the host and np.unique
|
||||||
|
# gather_ms the row reads (page cache or NVMe)
|
||||||
|
# stage_ms pinned-buffer copy, launching the H2D copy and the GPU-side expansion
|
||||||
|
# op_ms - wait_ms is what the lookup itself costs the step.
|
||||||
|
_STATS = {"calls": 0, "op_ms": 0.0, "gather_ms": 0.0, "rows": 0, "bytes": 0,
|
||||||
|
"wait_ms": 0.0, "dedup_ms": 0.0, "stage_ms": 0.0}
|
||||||
|
|
||||||
|
# The same numbers as monotonic Prometheus counters, so the table's behaviour is
|
||||||
|
# visible on a dashboard instead of only in the log line below (which is windowed,
|
||||||
|
# reset every period, and destroyed with the container). Nothing else in this
|
||||||
|
# recipe exposes how the mmapped table is coping: it is the one component whose
|
||||||
|
# cost depends on runtime state (page-cache residency) rather than on config, so
|
||||||
|
# it is exactly the thing worth graphing.
|
||||||
|
#
|
||||||
|
# These live in the EngineCore process, not the API server, so they only reach the
|
||||||
|
# frontend's /metrics when prometheus_client runs in multiprocess mode: the API
|
||||||
|
# server's MultiProcessCollector then aggregates what every process writes under
|
||||||
|
# PROMETHEUS_MULTIPROC_DIR. vLLM only turns that on for api_server_count > 1;
|
||||||
|
# scripts/serve.sh sets it when PROM_MULTIPROC=1 (opt-in). They are created lazily on first use,
|
||||||
|
# because the env var must be set before the first metric is constructed.
|
||||||
|
#
|
||||||
|
# Derived views worth having:
|
||||||
|
# (rate(vllm:ple_mmap_op_seconds_total[5m]) - rate(vllm:ple_mmap_gpu_wait_seconds_total[5m]))
|
||||||
|
# / rate(vllm:ple_mmap_lookup_ops_total[5m]) host seconds per lookup (the PLE's own cost)
|
||||||
|
# rate(vllm:ple_mmap_gpu_wait_seconds_total[5m])
|
||||||
|
# / rate(vllm:ple_mmap_lookup_ops_total[5m]) GPU work queued ahead of the lookup, per lookup
|
||||||
|
# rate(vllm:ple_mmap_gather_seconds_total[5m])
|
||||||
|
# / (rate(vllm:ple_mmap_op_seconds_total[5m]) - rate(vllm:ple_mmap_gpu_wait_seconds_total[5m]))
|
||||||
|
# fraction of the lookup's own time spent on disk
|
||||||
|
# rate(vllm:ple_mmap_bytes_total[5m]) NVMe read bandwidth from the table
|
||||||
|
# The disk fraction is the page-cache health signal: it climbs as the cache is squeezed and falls as
|
||||||
|
# the hot region settles in. Divide by op_seconds alone and it also moves with GPU load, which says
|
||||||
|
# nothing about the cache.
|
||||||
|
_PROM: dict[str, object] | None = None
|
||||||
|
_PROM_TRIED = False
|
||||||
|
|
||||||
|
|
||||||
|
def _prom() -> dict[str, object] | None:
|
||||||
|
"""Prometheus counters, or None if unavailable. Never raises, tried once."""
|
||||||
|
global _PROM, _PROM_TRIED
|
||||||
|
if _PROM_TRIED:
|
||||||
|
return _PROM
|
||||||
|
_PROM_TRIED = True
|
||||||
|
if os.environ.get("VLLM_PLE_MMAP_PROMETHEUS", "1").lower() in ("0", "false", "no"):
|
||||||
|
return None
|
||||||
|
if not os.environ.get("PROMETHEUS_MULTIPROC_DIR"):
|
||||||
|
# The default. vLLM only enables multiprocess metrics for api_server_count > 1,
|
||||||
|
# and exporting these costs vLLM its *_created samples, so it is opt-in: say how
|
||||||
|
# to turn it on rather than warn about the expected configuration.
|
||||||
|
logger.info(
|
||||||
|
"PLE mmap: PROMETHEUS_MULTIPROC_DIR is unset, so the vllm:ple_mmap_* counters "
|
||||||
|
"stay in this process and will not appear on /metrics while the engine runs "
|
||||||
|
"in its own process. scripts/serve.sh exports them with PROM_MULTIPROC=1."
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
from prometheus_client import Counter
|
||||||
|
|
||||||
|
_PROM = {
|
||||||
|
"ops": Counter(
|
||||||
|
"vllm:ple_mmap_lookup_ops_total",
|
||||||
|
"PLE mmap lookups (hash + gather + H2D) executed.",
|
||||||
|
),
|
||||||
|
"op_s": Counter(
|
||||||
|
"vllm:ple_mmap_op_seconds_total",
|
||||||
|
"Cumulative seconds in the PLE mmap lookup op, including the GPU wait "
|
||||||
|
"(vllm:ple_mmap_gpu_wait_seconds_total).",
|
||||||
|
),
|
||||||
|
"gather_s": Counter(
|
||||||
|
"vllm:ple_mmap_gather_seconds_total",
|
||||||
|
"Cumulative seconds in the PLE mmap row gather (the disk reads).",
|
||||||
|
),
|
||||||
|
"rows": Counter(
|
||||||
|
"vllm:ple_mmap_rows_total",
|
||||||
|
"Rows gathered from the mmapped PLE table.",
|
||||||
|
),
|
||||||
|
"bytes": Counter(
|
||||||
|
"vllm:ple_mmap_bytes_total",
|
||||||
|
"Bytes read from the mmapped PLE table (page cache or NVMe).",
|
||||||
|
),
|
||||||
|
"gpu_wait_s": Counter(
|
||||||
|
"vllm:ple_mmap_gpu_wait_seconds_total",
|
||||||
|
"Cumulative seconds the lookup waited for GPU work queued ahead of it "
|
||||||
|
"(not PLE cost; subtract from op_seconds).",
|
||||||
|
),
|
||||||
|
"dedup_s": Counter(
|
||||||
|
"vllm:ple_mmap_dedup_seconds_total",
|
||||||
|
"Cumulative seconds copying row ids to the host and deduplicating them.",
|
||||||
|
),
|
||||||
|
"stage_s": Counter(
|
||||||
|
"vllm:ple_mmap_stage_seconds_total",
|
||||||
|
"Cumulative seconds staging gathered rows for the GPU (pinned copy, H2D launch).",
|
||||||
|
),
|
||||||
|
}
|
||||||
|
logger.info("PLE mmap: Prometheus counters registered")
|
||||||
|
except Exception as exc: # pragma: no cover - metrics must never break serving
|
||||||
|
logger.warning("PLE mmap: Prometheus counters unavailable: %s", exc)
|
||||||
|
_PROM = None
|
||||||
|
return _PROM
|
||||||
|
|
||||||
|
|
||||||
|
def _prom_add(**kw: float) -> None:
|
||||||
|
"""Best-effort counter increment; a metrics failure must not fail a request."""
|
||||||
|
p = _prom()
|
||||||
|
if not p:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
for key, value in kw.items():
|
||||||
|
p[key].inc(value) # type: ignore[attr-defined]
|
||||||
|
except Exception: # pragma: no cover
|
||||||
|
pass
|
||||||
|
_STATS_LAST = [0.0]
|
||||||
|
_STATS_SEC = _env_int("VLLM_PLE_MMAP_STATS_SEC", 30)
|
||||||
|
|
||||||
|
|
||||||
|
def _stats_log() -> None:
|
||||||
|
import time as _time
|
||||||
|
|
||||||
|
now = _time.monotonic()
|
||||||
|
if _STATS_SEC <= 0 or now - _STATS_LAST[0] < _STATS_SEC:
|
||||||
|
return
|
||||||
|
elapsed = now - _STATS_LAST[0] if _STATS_LAST[0] else float(_STATS_SEC)
|
||||||
|
_STATS_LAST[0] = now
|
||||||
|
s = _STATS
|
||||||
|
if not s["calls"]:
|
||||||
|
return
|
||||||
|
# The prefix up to "MiB read" is unchanged, for anything that already parses it.
|
||||||
|
n = s["calls"]
|
||||||
|
logger.info(
|
||||||
|
"PLE mmap stats (last %.0fs): %d ops, op %.0f ms total (%.2f ms/op), "
|
||||||
|
"gather %.0f ms total (%.2f ms/op), %d rows, %.1f MiB read, "
|
||||||
|
"gpu-wait %.2f ms/op, host %.2f ms/op (dedup %.2f, gather %.2f, stage %.2f)",
|
||||||
|
elapsed, n, s["op_ms"], s["op_ms"] / n,
|
||||||
|
s["gather_ms"], s["gather_ms"] / n,
|
||||||
|
s["rows"], s["bytes"] / 2**20,
|
||||||
|
s["wait_ms"] / n, max(0.0, s["op_ms"] - s["wait_ms"]) / n,
|
||||||
|
s["dedup_ms"] / n, s["gather_ms"] / n, s["stage_ms"] / n,
|
||||||
|
)
|
||||||
|
s.update(calls=0, op_ms=0.0, gather_ms=0.0, rows=0, bytes=0, wait_ms=0.0, dedup_ms=0.0, stage_ms=0.0)
|
||||||
|
|
||||||
|
|
||||||
|
def _lookup_impl(
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
query_start_loc: torch.Tensor,
|
||||||
|
ngram_context: torch.Tensor,
|
||||||
|
output: torch.Tensor,
|
||||||
|
layer_name: str,
|
||||||
|
) -> None:
|
||||||
|
import time as _time
|
||||||
|
|
||||||
|
t0 = _time.perf_counter()
|
||||||
|
layer = _REGISTRY[layer_name]
|
||||||
|
result = layer._ple_mmap_orig_forward_impl(
|
||||||
|
None, input_ids, query_start_loc, ngram_context
|
||||||
|
)
|
||||||
|
output[: result.shape[0]].copy_(result.to(output.dtype))
|
||||||
|
dt = _time.perf_counter() - t0
|
||||||
|
_STATS["calls"] += 1
|
||||||
|
_STATS["op_ms"] += dt * 1e3
|
||||||
|
_prom_add(ops=1, op_s=dt)
|
||||||
|
_stats_log()
|
||||||
|
|
||||||
|
|
||||||
|
def _lookup_fake(
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
query_start_loc: torch.Tensor,
|
||||||
|
ngram_context: torch.Tensor,
|
||||||
|
output: torch.Tensor,
|
||||||
|
layer_name: str,
|
||||||
|
) -> None:
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
|
_OP_NAME_IDS = "ple_mmap_lookup_ids"
|
||||||
|
|
||||||
|
|
||||||
|
def _lookup_ids_impl(ngram_ids: torch.Tensor, output: torch.Tensor, layer_name: str) -> None:
|
||||||
|
"""v0.29 layout: gather rows for already-hashed ids; output is (N, ngram_heads * head_dim)."""
|
||||||
|
import time as _time
|
||||||
|
|
||||||
|
t0 = _time.perf_counter()
|
||||||
|
layer = _REGISTRY[layer_name]
|
||||||
|
rows = layer.ngram_embedding(ngram_ids) # (N, heads, head_dim), table dtype (or zeros)
|
||||||
|
output.copy_(rows.reshape(rows.shape[0], -1).to(output.dtype))
|
||||||
|
dt = _time.perf_counter() - t0
|
||||||
|
_STATS["calls"] += 1
|
||||||
|
_STATS["op_ms"] += dt * 1e3
|
||||||
|
_prom_add(ops=1, op_s=dt)
|
||||||
|
_stats_log()
|
||||||
|
|
||||||
|
|
||||||
|
def _lookup_ids_fake(ngram_ids: torch.Tensor, output: torch.Tensor, layer_name: str) -> None:
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
|
def _register_op() -> None:
|
||||||
|
if hasattr(torch.ops.vllm, _OP_NAME):
|
||||||
|
return
|
||||||
|
from vllm.utils.torch_utils import direct_register_custom_op
|
||||||
|
|
||||||
|
direct_register_custom_op(
|
||||||
|
op_name=_OP_NAME,
|
||||||
|
op_func=_lookup_impl,
|
||||||
|
mutates_args=["output"],
|
||||||
|
fake_impl=_lookup_fake,
|
||||||
|
)
|
||||||
|
direct_register_custom_op(
|
||||||
|
op_name=_OP_NAME_IDS,
|
||||||
|
op_func=_lookup_ids_impl,
|
||||||
|
mutates_args=["output"],
|
||||||
|
fake_impl=_lookup_ids_fake,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def _setup_table_v029(self) -> None:
|
||||||
|
if self.ngram_embedding.table is not None:
|
||||||
|
return
|
||||||
|
# VLLM_PLE_MMAP_DIR: serve the table from a different directory than the
|
||||||
|
# checkpoint (e.g. an FP8 copy of the table on local NVMe).
|
||||||
|
model_path = os.environ.get("VLLM_PLE_MMAP_DIR") or self._ple_mmap_model_path
|
||||||
|
if not model_path or not os.path.isdir(model_path):
|
||||||
|
raise RuntimeError(
|
||||||
|
f"PLE mmap: table path {model_path!r} is not a local directory; "
|
||||||
|
"point --model at the downloaded snapshot or set VLLM_PLE_MMAP_DIR"
|
||||||
|
)
|
||||||
|
m = re.search(r"layers\.(\d+)\.", self._ple_mmap_prefix)
|
||||||
|
if not m:
|
||||||
|
raise RuntimeError(f"PLE mmap: cannot find layer index in {self._ple_mmap_prefix!r}")
|
||||||
|
layer_idx = int(m.group(1))
|
||||||
|
shards, dtype_str, scale_entry, cols = _find_shards(model_path, layer_idx)
|
||||||
|
if not shards:
|
||||||
|
raise RuntimeError(f"PLE mmap: no shard tensors for layer {layer_idx} under {model_path}")
|
||||||
|
if cols != self.head_dim:
|
||||||
|
raise RuntimeError(f"PLE mmap: shard width {cols} != head_dim {self.head_dim}")
|
||||||
|
if dtype_str not in _TABLE_DTYPES:
|
||||||
|
raise RuntimeError(f"PLE mmap: unsupported shard dtype {dtype_str}")
|
||||||
|
if dtype_str in _FP8_DTYPES and not hasattr(self, "_offload_weight_scale"):
|
||||||
|
if scale_entry is None:
|
||||||
|
raise RuntimeError("PLE mmap: FP8 shards without ngram_embedding.weight_scale")
|
||||||
|
self.register_buffer(
|
||||||
|
"_offload_weight_scale",
|
||||||
|
_read_scale(scale_entry).to(torch.accelerator.current_accelerator()),
|
||||||
|
persistent=False,
|
||||||
|
)
|
||||||
|
parts = int(self.split_ngram_parts)
|
||||||
|
vocab = int(self.ngram_embedding.org_vocab_size)
|
||||||
|
shard_size = math.ceil(vocab / parts)
|
||||||
|
for idx, (_p, _o, rows) in shards.items():
|
||||||
|
expected = max(0, min(shard_size, vocab - idx * shard_size))
|
||||||
|
if rows != expected:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"PLE mmap: shard {idx} has {rows} rows, expected {expected}"
|
||||||
|
)
|
||||||
|
table = MmapPleTable(
|
||||||
|
shards, shard_size, cols * _itemsize(dtype_str), _TABLE_DTYPES[dtype_str],
|
||||||
|
workers=_env_int("VLLM_PLE_MMAP_WORKERS", 32),
|
||||||
|
chunk=_env_int("VLLM_PLE_MMAP_CHUNK", 2048),
|
||||||
|
)
|
||||||
|
if _env_int("VLLM_PLE_MMAP_PREWARM", 0):
|
||||||
|
logger.info("PLE mmap: prewarming page cache (%.1f GiB)...", table.rows_total * table.row_bytes / 2**30)
|
||||||
|
table.prewarm()
|
||||||
|
self.ngram_embedding.table = table
|
||||||
|
logger.info(
|
||||||
|
"PLE mmap: layer %d, %d shards, %d rows x %d B (%.1f GiB on disk), dtype %s, %d workers",
|
||||||
|
layer_idx, len(shards), table.rows_total, table.row_bytes,
|
||||||
|
table.rows_total * table.row_bytes / 2**30, dtype_str, table.pool._max_workers,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def apply(cls: type) -> None:
|
||||||
|
"""Patch the n-gram embedding class (pass the class) when enabled.
|
||||||
|
|
||||||
|
Two layouts are supported: the preview image's ``Qwen3_8FlashNextNGramEmbedding``
|
||||||
|
(hashing + lookup in ``forward_impl``) and vLLM >= 0.29's ``Qwen4ExpNGramEmbedding``
|
||||||
|
(hashing in the ``qwen4_exp_compute_ple_ngram_ids`` op, lookup through a
|
||||||
|
``PLEVocabParallelEmbedding`` whose ``weight_scale`` the PLE layer reads).
|
||||||
|
"""
|
||||||
|
if not enabled():
|
||||||
|
return
|
||||||
|
if getattr(cls, "_ple_mmap_patched", False):
|
||||||
|
return
|
||||||
|
if not hasattr(cls, "forward_impl"):
|
||||||
|
_apply_v029(cls)
|
||||||
|
return
|
||||||
|
mod = sys.modules[cls.__module__]
|
||||||
|
orig_init = cls.__init__
|
||||||
|
orig_load_weights = cls.load_weights
|
||||||
|
|
||||||
|
def __init__(self, config, embedding_dim, ple_dense_layer_id, max_total_tokens,
|
||||||
|
max_num_reqs, prefix, quant_config=None, params_dtype=None):
|
||||||
|
# Run the stock constructor (hash buffers, workspaces, ...) with the
|
||||||
|
# embedding class swapped for our placeholder so nothing large is
|
||||||
|
# allocated. quant_config=None keeps the stock code from selecting an
|
||||||
|
# FP8 quant method that would create an FP8 weight parameter.
|
||||||
|
real_embedding_cls = mod.VocabParallelEmbedding
|
||||||
|
mod.VocabParallelEmbedding = lambda n, d, **_kw: _MmapNgramEmbedding(n, d)
|
||||||
|
try:
|
||||||
|
orig_init(self, config, embedding_dim, ple_dense_layer_id,
|
||||||
|
max_total_tokens, max_num_reqs, prefix,
|
||||||
|
quant_config=None, params_dtype=params_dtype)
|
||||||
|
finally:
|
||||||
|
mod.VocabParallelEmbedding = real_embedding_cls
|
||||||
|
self._ple_mmap_prefix = prefix
|
||||||
|
_REGISTRY[prefix] = self
|
||||||
|
self._ple_mmap_model_path = None
|
||||||
|
try:
|
||||||
|
from vllm.config import get_current_vllm_config
|
||||||
|
self._ple_mmap_model_path = get_current_vllm_config().model_config.model
|
||||||
|
except Exception as exc: # pragma: no cover - defensive
|
||||||
|
logger.warning("PLE mmap: cannot read model path from vllm config: %s", exc)
|
||||||
|
if params_dtype is not None:
|
||||||
|
self.ngram_embedding._zeros_dtype = params_dtype
|
||||||
|
logger.info(
|
||||||
|
"PLE mmap: %s -> placeholder embedding (%d rows x %d), table will be mmapped",
|
||||||
|
prefix, self.ngram_embedding.org_vocab_size, self.head_dim,
|
||||||
|
)
|
||||||
|
|
||||||
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||||
|
loaded: set[str] = set()
|
||||||
|
rest: list[tuple[str, torch.Tensor]] = []
|
||||||
|
for name, w in weights:
|
||||||
|
if name.startswith("ngram_embedding.shard_") and name.endswith(".weight"):
|
||||||
|
loaded.add(name) # served from disk, never materialised
|
||||||
|
continue
|
||||||
|
if name == "ngram_embedding.weight_scale":
|
||||||
|
self.register_buffer(
|
||||||
|
"_offload_weight_scale",
|
||||||
|
w.detach().to(device=torch.accelerator.current_accelerator()),
|
||||||
|
persistent=False,
|
||||||
|
)
|
||||||
|
loaded.add(name)
|
||||||
|
continue
|
||||||
|
rest.append((name, w))
|
||||||
|
loaded.update(orig_load_weights(self, rest))
|
||||||
|
_setup_table(self)
|
||||||
|
return loaded
|
||||||
|
|
||||||
|
def _setup_table(self) -> None:
|
||||||
|
if self.ngram_embedding.table is not None:
|
||||||
|
return
|
||||||
|
# VLLM_PLE_MMAP_DIR: serve the table from a different directory than the
|
||||||
|
# checkpoint (e.g. an FP8 copy of the table on local NVMe).
|
||||||
|
model_path = os.environ.get("VLLM_PLE_MMAP_DIR") or self._ple_mmap_model_path
|
||||||
|
if not model_path or not os.path.isdir(model_path):
|
||||||
|
raise RuntimeError(
|
||||||
|
f"PLE mmap: table path {model_path!r} is not a local directory; "
|
||||||
|
"point --model at the downloaded snapshot or set VLLM_PLE_MMAP_DIR"
|
||||||
|
)
|
||||||
|
m = re.search(r"layers\.(\d+)\.", self._ple_mmap_prefix)
|
||||||
|
if not m:
|
||||||
|
raise RuntimeError(f"PLE mmap: cannot find layer index in {self._ple_mmap_prefix!r}")
|
||||||
|
layer_idx = int(m.group(1))
|
||||||
|
shards, dtype_str, scale_entry, cols = _find_shards(model_path, layer_idx)
|
||||||
|
if not shards:
|
||||||
|
raise RuntimeError(f"PLE mmap: no shard tensors for layer {layer_idx} under {model_path}")
|
||||||
|
if cols != self.head_dim:
|
||||||
|
raise RuntimeError(f"PLE mmap: shard width {cols} != head_dim {self.head_dim}")
|
||||||
|
if dtype_str not in _TABLE_DTYPES:
|
||||||
|
raise RuntimeError(f"PLE mmap: unsupported shard dtype {dtype_str}")
|
||||||
|
if dtype_str in _FP8_DTYPES and not hasattr(self, "_offload_weight_scale"):
|
||||||
|
if scale_entry is None:
|
||||||
|
raise RuntimeError("PLE mmap: FP8 shards without ngram_embedding.weight_scale")
|
||||||
|
self.register_buffer(
|
||||||
|
"_offload_weight_scale",
|
||||||
|
_read_scale(scale_entry).to(torch.accelerator.current_accelerator()),
|
||||||
|
persistent=False,
|
||||||
|
)
|
||||||
|
parts = int(self.split_ngram_parts)
|
||||||
|
vocab = int(self.ngram_embedding.org_vocab_size)
|
||||||
|
shard_size = math.ceil(vocab / parts)
|
||||||
|
for idx, (_p, _o, rows) in shards.items():
|
||||||
|
expected = max(0, min(shard_size, vocab - idx * shard_size))
|
||||||
|
if rows != expected:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"PLE mmap: shard {idx} has {rows} rows, expected {expected}"
|
||||||
|
)
|
||||||
|
table = MmapPleTable(
|
||||||
|
shards, shard_size, cols * _itemsize(dtype_str), _TABLE_DTYPES[dtype_str],
|
||||||
|
workers=_env_int("VLLM_PLE_MMAP_WORKERS", 32),
|
||||||
|
chunk=_env_int("VLLM_PLE_MMAP_CHUNK", 2048),
|
||||||
|
)
|
||||||
|
if _env_int("VLLM_PLE_MMAP_PREWARM", 0):
|
||||||
|
logger.info("PLE mmap: prewarming page cache (%.1f GiB)...", table.rows_total * table.row_bytes / 2**30)
|
||||||
|
table.prewarm()
|
||||||
|
self.ngram_embedding.table = table
|
||||||
|
logger.info(
|
||||||
|
"PLE mmap: layer %d, %d shards, %d rows x %d B (%.1f GiB on disk), dtype %s, %d workers",
|
||||||
|
layer_idx, len(shards), table.rows_total, table.row_bytes,
|
||||||
|
table.rows_total * table.row_bytes / 2**30, dtype_str, table.pool._max_workers,
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward_impl(self, hidden_states, input_ids, query_start_loc, ngram_context,
|
||||||
|
output_buffer=None):
|
||||||
|
del hidden_states, output_buffer
|
||||||
|
num_tokens = input_ids.reshape(-1).shape[0]
|
||||||
|
table = self.ngram_embedding.table
|
||||||
|
dtype = table.torch_dtype if table is not None else self.ngram_embedding._zeros_dtype
|
||||||
|
output = torch.empty(
|
||||||
|
(num_tokens, self.embedding_dim), dtype=dtype, device=input_ids.device
|
||||||
|
)
|
||||||
|
getattr(torch.ops.vllm, _OP_NAME)(
|
||||||
|
input_ids, query_start_loc, ngram_context, output, self._ple_mmap_prefix
|
||||||
|
)
|
||||||
|
return output
|
||||||
|
|
||||||
|
_register_op()
|
||||||
|
cls._ple_mmap_orig_forward_impl = cls.forward_impl
|
||||||
|
cls.forward_impl = forward_impl
|
||||||
|
cls.__init__ = __init__
|
||||||
|
cls.load_weights = load_weights
|
||||||
|
cls._setup_table = _setup_table
|
||||||
|
cls._ple_mmap_patched = True
|
||||||
|
logger.info("PLE mmap patch applied to %s.%s", cls.__module__, cls.__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _apply_v029(cls: type) -> None:
|
||||||
|
"""vLLM >= 0.29 layout (``vllm/models/qwen4_exp``)."""
|
||||||
|
mod = sys.modules[cls.__module__]
|
||||||
|
orig_init = cls.__init__
|
||||||
|
orig_load_weights = cls.load_weights
|
||||||
|
embed_attr = "PLEVocabParallelEmbedding"
|
||||||
|
if not hasattr(mod, embed_attr):
|
||||||
|
raise RuntimeError(f"PLE mmap: {mod.__name__} has no {embed_attr}; layout not recognised")
|
||||||
|
|
||||||
|
def __init__(self, config, embedding_dim, ple_dense_layer_id, max_total_tokens,
|
||||||
|
max_num_reqs, prefix, layer_name, quant_config=None, params_dtype=None):
|
||||||
|
real_cls = getattr(mod, embed_attr)
|
||||||
|
setattr(mod, embed_attr, lambda n, d, **_kw: _MmapNgramEmbedding(n, d))
|
||||||
|
try:
|
||||||
|
orig_init(self, config, embedding_dim, ple_dense_layer_id, max_total_tokens,
|
||||||
|
max_num_reqs, prefix, layer_name, quant_config=None,
|
||||||
|
params_dtype=params_dtype)
|
||||||
|
finally:
|
||||||
|
setattr(mod, embed_attr, real_cls)
|
||||||
|
self._ple_mmap_prefix = prefix
|
||||||
|
_REGISTRY[prefix] = self
|
||||||
|
self._ple_mmap_model_path = None
|
||||||
|
try:
|
||||||
|
from vllm.config import get_current_vllm_config
|
||||||
|
self._ple_mmap_model_path = get_current_vllm_config().model_config.model
|
||||||
|
except Exception as exc: # pragma: no cover - defensive
|
||||||
|
logger.warning("PLE mmap: cannot read model path from vllm config: %s", exc)
|
||||||
|
if params_dtype is not None:
|
||||||
|
self.ngram_embedding._zeros_dtype = params_dtype
|
||||||
|
logger.info(
|
||||||
|
"PLE mmap (v0.29 layout): %s -> placeholder embedding (%d rows x %d), table will be mmapped",
|
||||||
|
prefix, self.ngram_embedding.org_vocab_size, self.head_dim,
|
||||||
|
)
|
||||||
|
|
||||||
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||||
|
loaded: set[str] = set()
|
||||||
|
rest: list[tuple[str, torch.Tensor]] = []
|
||||||
|
dev = torch.accelerator.current_accelerator()
|
||||||
|
for name, w in weights:
|
||||||
|
if name.startswith("ngram_embedding.shard_") and name.endswith(".weight"):
|
||||||
|
loaded.add(name) # served from disk, never materialised
|
||||||
|
continue
|
||||||
|
if name == "ngram_embedding.weight_scale":
|
||||||
|
scale = w.detach().to(device=dev)
|
||||||
|
self.register_buffer("_offload_weight_scale", scale, persistent=False)
|
||||||
|
# Qwen4ExpPLELayer._get_embedding_weight_scale reads ngram_embedding.weight_scale
|
||||||
|
self.ngram_embedding.weight_scale = scale
|
||||||
|
loaded.add(name)
|
||||||
|
continue
|
||||||
|
rest.append((name, w))
|
||||||
|
loaded.update(orig_load_weights(self, rest))
|
||||||
|
self._setup_table()
|
||||||
|
if getattr(self.ngram_embedding, "weight_scale", None) is None and hasattr(self, "_offload_weight_scale"):
|
||||||
|
self.ngram_embedding.weight_scale = self._offload_weight_scale
|
||||||
|
return loaded
|
||||||
|
|
||||||
|
def forward(self, input_ids, query_start_loc, ngram_context):
|
||||||
|
ngram_ids = input_ids.new_empty((input_ids.shape[0], self.ngram_heads), dtype=torch.long)
|
||||||
|
torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids(
|
||||||
|
input_ids, query_start_loc, ngram_context, ngram_ids, self.layer_name
|
||||||
|
)
|
||||||
|
table = self.ngram_embedding.table
|
||||||
|
dtype = table.torch_dtype if table is not None else self.ngram_embedding._zeros_dtype
|
||||||
|
output = torch.empty((ngram_ids.shape[0], self.embedding_dim), dtype=dtype, device=input_ids.device)
|
||||||
|
getattr(torch.ops.vllm, _OP_NAME_IDS)(ngram_ids, output, self._ple_mmap_prefix)
|
||||||
|
return output
|
||||||
|
|
||||||
|
_register_op()
|
||||||
|
cls.__init__ = __init__
|
||||||
|
cls.load_weights = load_weights
|
||||||
|
cls.forward = forward
|
||||||
|
cls._setup_table = _setup_table_v029
|
||||||
|
cls._ple_mmap_patched = True
|
||||||
|
logger.info("PLE mmap patch (v0.29 layout) applied to %s.%s", cls.__module__, cls.__name__)
|
||||||
Reference in New Issue
Block a user