From 5f3030260ee8932a02aecedfb7b7a8ff8c008890 Mon Sep 17 00:00:00 2001 From: ruihan Date: Thu, 17 Sep 2026 16:45:50 +0800 Subject: [PATCH] Add reproducible DGX Spark deployment for Qwen3.8 Flash Next NVFP4 --- .dockerignore | 6 + .env.example | 11 + .gitignore | 13 + Dockerfile | 22 + README.md | 88 ++++ THIRD_PARTY_NOTICES.md | 10 + compose.yaml | 70 +++ configs/baseline-32k.yaml | 16 + docs/architecture.md | 24 + docs/packaging-validation.md | 13 + docs/provenance.md | 22 + docs/troubleshooting.md | 37 ++ docs/validation.md | 36 ++ patches/spark_ngram_adapter.py | 69 +++ scripts/build.sh | 3 + scripts/common.sh | 9 + scripts/download.py | 7 + scripts/download.sh | 3 + scripts/smoke-test.py | 42 ++ scripts/start.sh | 10 + scripts/test-adapter.sh | 3 + scripts/test.sh | 3 + scripts/wait.sh | 10 + tests/test_disk_adapter.py | 49 ++ vendor/Apache-2.0.txt | 202 ++++++++ vendor/LICENSE | 19 + vendor/vllm_ple_mmap.py | 908 +++++++++++++++++++++++++++++++++ 27 files changed, 1705 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 THIRD_PARTY_NOTICES.md create mode 100644 compose.yaml create mode 100644 configs/baseline-32k.yaml create mode 100644 docs/architecture.md create mode 100644 docs/packaging-validation.md create mode 100644 docs/provenance.md create mode 100644 docs/troubleshooting.md create mode 100644 docs/validation.md create mode 100644 patches/spark_ngram_adapter.py create mode 100755 scripts/build.sh create mode 100755 scripts/common.sh create mode 100644 scripts/download.py create mode 100755 scripts/download.sh create mode 100644 scripts/smoke-test.py create mode 100755 scripts/start.sh create mode 100755 scripts/test-adapter.sh create mode 100755 scripts/test.sh create mode 100755 scripts/wait.sh create mode 100644 tests/test_disk_adapter.py create mode 100644 vendor/Apache-2.0.txt create mode 100644 vendor/LICENSE create mode 100644 vendor/vllm_ple_mmap.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..bf9e0ee --- /dev/null +++ b/.dockerignore @@ -0,0 +1,6 @@ +** +!Dockerfile +!patches/ +!patches/spark_ngram_adapter.py +!vendor/ +!vendor/vllm_ple_mmap.py diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..7257afd --- /dev/null +++ b/.env.example @@ -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 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..cac2e53 --- /dev/null +++ b/.gitignore @@ -0,0 +1,13 @@ +.env +.env.* +!.env.example +*.key +*.pem +secrets/ +runtime-cache/ +__pycache__/ +*.pyc +*.log +*.safetensors +*.gguf +.DS_Store diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..2afb7eb --- /dev/null +++ b/Dockerfile @@ -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 diff --git a/README.md b/README.md new file mode 100644 index 0000000..f34bc2b --- /dev/null +++ b/README.md @@ -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://: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)。 diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..2952600 --- /dev/null +++ b/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. diff --git a/compose.yaml b/compose.yaml new file mode 100644 index 0000000..e27e8d9 --- /dev/null +++ b/compose.yaml @@ -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] diff --git a/configs/baseline-32k.yaml b/configs/baseline-32k.yaml new file mode 100644 index 0000000..fe8e46a --- /dev/null +++ b/configs/baseline-32k.yaml @@ -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)" diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 0000000..b61ef17 --- /dev/null +++ b/docs/architecture.md @@ -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)。 diff --git a/docs/packaging-validation.md b/docs/packaging-validation.md new file mode 100644 index 0000000..6ff03bf --- /dev/null +++ b/docs/packaging-validation.md @@ -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 权重，也没有再次完整加载模型。 diff --git a/docs/provenance.md b/docs/provenance.md new file mode 100644 index 0000000..3971a00 --- /dev/null +++ b/docs/provenance.md @@ -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 未修改。 diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md new file mode 100644 index 0000000..014867d --- /dev/null +++ b/docs/troubleshooting.md @@ -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 修改，没有搬用所有社区优化。 diff --git a/docs/validation.md b/docs/validation.md new file mode 100644 index 0000000..cb4c9b9 --- /dev/null +++ b/docs/validation.md @@ -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。 diff --git a/patches/spark_ngram_adapter.py b/patches/spark_ngram_adapter.py new file mode 100644 index 0000000..00b0979 --- /dev/null +++ b/patches/spark_ngram_adapter.py @@ -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 diff --git a/scripts/build.sh b/scripts/build.sh new file mode 100755 index 0000000..73da053 --- /dev/null +++ b/scripts/build.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash +source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh" +"${compose[@]}" build vllm diff --git a/scripts/common.sh b/scripts/common.sh new file mode 100755 index 0000000..6b5b45f --- /dev/null +++ b/scripts/common.sh @@ -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) diff --git a/scripts/download.py b/scripts/download.py new file mode 100644 index 0000000..960eacb --- /dev/null +++ b/scripts/download.py @@ -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) diff --git a/scripts/download.sh b/scripts/download.sh new file mode 100755 index 0000000..6d717f4 --- /dev/null +++ b/scripts/download.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash +source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh" +"${compose[@]}" --profile tools run --rm download diff --git a/scripts/smoke-test.py b/scripts/smoke-test.py new file mode 100644 index 0000000..4f66906 --- /dev/null +++ b/scripts/smoke-test.py @@ -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) diff --git a/scripts/start.sh b/scripts/start.sh new file mode 100755 index 0000000..4ebfeb0 --- /dev/null +++ b/scripts/start.sh @@ -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.' diff --git a/scripts/test-adapter.sh b/scripts/test-adapter.sh new file mode 100755 index 0000000..a33f3c7 --- /dev/null +++ b/scripts/test-adapter.sh @@ -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 diff --git a/scripts/test.sh b/scripts/test.sh new file mode 100755 index 0000000..d800ec8 --- /dev/null +++ b/scripts/test.sh @@ -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 diff --git a/scripts/wait.sh b/scripts/wait.sh new file mode 100755 index 0000000..481f061 --- /dev/null +++ b/scripts/wait.sh @@ -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 diff --git a/tests/test_disk_adapter.py b/tests/test_disk_adapter.py new file mode 100644 index 0000000..f06aca2 --- /dev/null +++ b/tests/test_disk_adapter.py @@ -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') diff --git a/vendor/Apache-2.0.txt b/vendor/Apache-2.0.txt new file mode 100644 index 0000000..d645695 --- /dev/null +++ b/vendor/Apache-2.0.txt @@ -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. diff --git a/vendor/LICENSE b/vendor/LICENSE new file mode 100644 index 0000000..5c9f8c7 --- /dev/null +++ b/vendor/LICENSE @@ -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 diff --git a/vendor/vllm_ple_mmap.py b/vendor/vllm_ple_mmap.py new file mode 100644 index 0000000..a8ba8d8 --- /dev/null +++ b/vendor/vllm_ple_mmap.py @@ -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(" 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..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("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__)