qwen3.5-122B-A10B on DGX Spark: vLLM + DFlash + dense-bandwidth stack, one-shot installer
This commit is contained in:
ent
+20
@@ -0,0 +1,20 @@
|
|||||||
|
.DS_Store
|
||||||
|
*.swp
|
||||||
|
*.swo
|
||||||
|
*~
|
||||||
|
.idea/
|
||||||
|
.vscode/
|
||||||
|
__pycache__/
|
||||||
|
*.pyc
|
||||||
|
|
||||||
|
# Logs / local results
|
||||||
|
*.log
|
||||||
|
/local/
|
||||||
|
/results/
|
||||||
|
|
||||||
|
# Never commit weights or checkpoints
|
||||||
|
*.safetensors
|
||||||
|
*.gguf
|
||||||
|
*.bin
|
||||||
|
/hybrid-ckpt/
|
||||||
|
*-hybrid-int4-fp8/
|
||||||
@@ -0,0 +1,21 @@
|
|||||||
|
MIT License
|
||||||
|
|
||||||
|
Copyright (c) 2026 qwen3.5-122B-A10B-on-spark contributors
|
||||||
|
|
||||||
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
in the Software without restriction, including without limitation the rights
|
||||||
|
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
|
copies of the Software, and to permit persons to whom the Software is
|
||||||
|
furnished to do so, subject to the following conditions:
|
||||||
|
|
||||||
|
The above copyright notice and this permission notice shall be included in all
|
||||||
|
copies or substantial portions of the Software.
|
||||||
|
|
||||||
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||||
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||||
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||||
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||||
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||||
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||||
|
SOFTWARE.
|
||||||
@@ -0,0 +1,219 @@
|
|||||||
|
# qwen3.5-122B-A10B-on-spark
|
||||||
|
|
||||||
|
[`Qwen3.5-122B-A10B`](https://huggingface.co/Intel/Qwen3.5-122B-A10B-int4-AutoRound)
|
||||||
|
(hybrid GDN + mamba + 128-expert MoE, ~10B active) running on a single
|
||||||
|
**NVIDIA DGX Spark** (GB10 / SM121, 128 GiB unified) under **vLLM**, with
|
||||||
|
**[DFlash](https://modal.com/blog/spec-is-all-u-need) block-diffusion speculative
|
||||||
|
decode** and an optional **dense-bandwidth patch stack** — measured end-to-end,
|
||||||
|
with a per-token bandwidth model that explains every number.
|
||||||
|
|
||||||
|
**Status:** Working end-to-end, one-shot install. On real Hermes-agent
|
||||||
|
tool-call turns, **DFlash decode reaches a median ~81 tok/s on GB10** —
|
||||||
|
**~2× the native MTP-2 head (~40 tok/s)** on the same workload, and above the
|
||||||
|
best previously published number for this model on Spark (albond's fully-patched
|
||||||
|
MTP stack, 51.58 tok/s end-to-end). DFlash's acceptance is task-dependent (it
|
||||||
|
block-drafts 12 tokens in one parallel forward), so the win is largest on
|
||||||
|
structured/tool-call/code traffic and collapses to parity on open-ended prose.
|
||||||
|
|
||||||
|
A separate **dense-bandwidth stack** (hybrid INT4+FP8 shared experts + int8
|
||||||
|
lm-head) adds **+28 % to no-spec / base decode** (28.2 → 36.0 tok/s) but, by the
|
||||||
|
amortization law below, washes out to ~null on high-acceptance agent traffic —
|
||||||
|
so it's a lever for *base / low-acceptance* serving, not for the agent path.
|
||||||
|
|
||||||
|
- **Engine:** [`vLLM`](https://github.com/vllm-project/vllm) 0.23, sm121 build with the DFlash PRs, via the prebuilt image `ghcr.io/aeon-7/aeon-vllm-ultimate:2026-06-18-v0.23.0-dflashfix`. No host build — the four runtime patches in [`runtime/`](runtime/) are applied at serve time.
|
||||||
|
- **Target:** [`Intel/Qwen3.5-122B-A10B-int4-AutoRound`](https://huggingface.co/Intel/Qwen3.5-122B-A10B-int4-AutoRound) — INT4 (AutoRound/GPTQ) routed experts + attention, BF16 shared experts/embeddings/head, ~62 GiB. (Safetensors, *not* GGUF — vLLM serves HF checkpoints directly.)
|
||||||
|
- **Drafter:** [`z-lab/Qwen3.5-122B-A10B-DFlash`](https://huggingface.co/z-lab/Qwen3.5-122B-A10B-DFlash) — 0.8B / 6-layer non-causal block-diffusion drafter (block 16), shares the target's `embed_tokens` + `lm_head`, ~1.6 GiB.
|
||||||
|
- **Hardware:** NVIDIA DGX Spark, GB10, SM121, 128 GiB LPDDR5X unified, ~273 GB/s.
|
||||||
|
|
||||||
|
## Quick start
|
||||||
|
|
||||||
|
On a DGX Spark with Docker + the NVIDIA container runtime:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
curl -sSL https://raw.githubusercontent.com/Entrpi/qwen3.5-122B-A10B-on-spark/main/install.sh | bash -s -- --start
|
||||||
|
```
|
||||||
|
|
||||||
|
That one command:
|
||||||
|
|
||||||
|
1. Verifies the host (aarch64, GB10/SM121, Docker GPU access, free disk).
|
||||||
|
2. Pulls the sm121 vLLM image (~40 GiB, one-time).
|
||||||
|
3. Downloads the INT4 target (~62 GiB) + DFlash drafter (~1.6 GiB) into the HF cache.
|
||||||
|
4. Starts the `dflash` profile on `:8000`, waits until READY, and runs the
|
||||||
|
"capital of France" smoke test (asserts "Paris").
|
||||||
|
|
||||||
|
**Already have the model?** Skip the 62 GiB download:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# point at a checkpoint dir you already have (mounted read-only at /model):
|
||||||
|
./install.sh --start --model-dir /path/to/Qwen3.5-122B-A10B-int4-AutoRound
|
||||||
|
# or reuse an existing HF cache (download becomes a no-op if already present):
|
||||||
|
./install.sh --start --hf-home /mnt/big/hf
|
||||||
|
```
|
||||||
|
|
||||||
|
Preview without running: `... | bash -s -- --help`.
|
||||||
|
|
||||||
|
## Hardware requirements
|
||||||
|
|
||||||
|
| | |
|
||||||
|
|---|---|
|
||||||
|
| Validated on | NVIDIA DGX Spark (GB10, SM121, 128 GiB unified) |
|
||||||
|
| Likely to work | other Blackwell with `--force` (untested) |
|
||||||
|
| Runtime | Docker + NVIDIA container runtime (`docker run --gpus all`) |
|
||||||
|
| Disk | ≥ 75 GiB free (image + weights); ≥ 150 GiB if `--build-hybrid` |
|
||||||
|
| OS | aarch64 Linux (Grace) |
|
||||||
|
| Memory | 128 GiB unified is enough for the model + DFlash drafter + KV @ 16k |
|
||||||
|
|
||||||
|
GB10 is detected via `nvidia-smi --query-gpu=compute_cap` returning `12.1`;
|
||||||
|
anything else needs `--force`.
|
||||||
|
|
||||||
|
## What you get — profiles
|
||||||
|
|
||||||
|
Pick with `--profile`:
|
||||||
|
|
||||||
|
| Profile | Stack | Best for | Measured |
|
||||||
|
|---|---|---|---|
|
||||||
|
| **`dflash`** *(default)* | INT4 + DFlash n=12 | agents / tool-calls / code | **~81 tok/s** Hermes · 53.7 albond-bench |
|
||||||
|
| `dense` | hybrid INT4+FP8 + int8 lm-head + DFlash n=12 | base / low-accept serving | 36.0 base (+28%) · 59.0 albond-bench |
|
||||||
|
| `base` | plain INT4, no spec | airtight baseline | 28.2 tok/s c=1 |
|
||||||
|
| `mtp` | INT4 + native MTP-2 head | comparison | ~40 tok/s Hermes |
|
||||||
|
|
||||||
|
The server is OpenAI-compatible (`/v1/chat/completions` with tool calls + SSE,
|
||||||
|
`/v1/completions`, `/v1/models`) and serves under the model name `qwen`.
|
||||||
|
|
||||||
|
## Benchmarks
|
||||||
|
|
||||||
|
All single-stream (c=1), temperature 0, GB10. "Hermes" = regenerating the next
|
||||||
|
assistant turn over 10 real conversations from a live agent's `state.db` (73 %
|
||||||
|
tool-calls); "albond-bench" = albond's own end-to-end harness (completion_tokens
|
||||||
|
/ total wallclock incl. prefill, 5 prompts, run-1 discarded — directly
|
||||||
|
comparable to his published 51.58).
|
||||||
|
|
||||||
|
### DFlash vs MTP, same harness, unpatched
|
||||||
|
|
||||||
|
| Workload (accept len) | base no-spec | MTP-2 | **DFlash n=12** |
|
||||||
|
|---|---|---|---|
|
||||||
|
| Prose (~2.3) | 28.2 | 33.7 | 33.2 *(use n=4)* |
|
||||||
|
| Code (~5.4) | 28.2 | 40.5 | **54.5** |
|
||||||
|
| Counting (~11) | 28.2 | 43.7 *(MTP caps at acc 3)* | **124.5** |
|
||||||
|
| **Hermes, real turns (8.3)** | — | **39.9** | **~81** |
|
||||||
|
| albond-bench e2e (6.5) | — | — | **53.7** |
|
||||||
|
|
||||||
|
MTP-2 drafts 2 tokens *sequentially* (acceptance caps at ~3); DFlash block-drafts
|
||||||
|
12 in **one parallel forward**, so on predictable/agent traffic it accepts 5–11
|
||||||
|
and pulls ~2× ahead. They tie only on low-acceptance prose. **53.7 unpatched
|
||||||
|
already clears albond's fully-patched MTP (51.58)** under his own method.
|
||||||
|
|
||||||
|
### The dense-bandwidth stack (`dense` profile)
|
||||||
|
|
||||||
|
Two independent always-on levers, ported to vLLM 0.23 as runtime patches:
|
||||||
|
hybrid INT4+FP8 (BF16 shared experts → calibrated FP8) and int8 lm-head (the
|
||||||
|
248 320-row vocab projection → int8 w8a16 GEMV, ~2× the bf16 read).
|
||||||
|
|
||||||
|
| Config | base (acc 1) | DFlash spec, albond-bench (acc 6.4) | Hermes (acc 8.3) |
|
||||||
|
|---|---|---|---|
|
||||||
|
| INT4 baseline | 28.2 | 53.7 | ~81 |
|
||||||
|
| + hybrid-FP8 | 30.4 (+7.8%) | 57.0 (+6.1%) | ~80 |
|
||||||
|
| + int8 lm-head | 32.7 (+16%) | — | — |
|
||||||
|
| **+ both** | **36.0 (+28%)** | **59.0 (+10%)** | ~80–87 *(noise)* |
|
||||||
|
|
||||||
|
## The amortization law
|
||||||
|
|
||||||
|
The dense levers cut **always-on** weight reads (shared experts + lm-head, read
|
||||||
|
every token). Under speculative decode the verify forward reads those weights
|
||||||
|
**once and amortizes them across the accepted block**, so the gain shrinks as
|
||||||
|
acceptance rises — monotonically, across the whole curve:
|
||||||
|
|
||||||
|
```
|
||||||
|
dense stack uplift: +28% (base, accept 1)
|
||||||
|
→ +10% (albond-bench, accept ~6.4)
|
||||||
|
→ ~0% (Hermes, accept ~8.3)
|
||||||
|
```
|
||||||
|
|
||||||
|
**Consequence:** for the **agent path (`dflash`)**, DFlash's own high acceptance
|
||||||
|
already saturates the dense levers — its remaining bottleneck is *routed-expert*
|
||||||
|
verify-batch reads, which no dense-weight quant touches. For **base / low-accept
|
||||||
|
serving (`dense`)**, the stack is a real +28 %. See [`docs/FINDINGS.md`](docs/FINDINGS.md).
|
||||||
|
|
||||||
|
## Under the hood: the four runtime patches
|
||||||
|
|
||||||
|
vLLM is unmodified on disk; [`runtime/serve.sh`](runtime/serve.sh) edits the
|
||||||
|
installed package in-place before `vllm serve` (idempotent, sentinel-guarded):
|
||||||
|
|
||||||
|
| Patch | What it does | Needed by |
|
||||||
|
|---|---|---|
|
||||||
|
| [`patch_unify2.py`](runtime/patch_unify2.py) | scale-block KV-cache **unify** so the hybrid GDN+mamba target absorbs the drafter's attention spec (the original assert can't); + `--no-enable-prefix-caching` routes to the no-hash-assert coordinator | **DFlash** (any spec profile) |
|
||||||
|
| [`patch_inc_hybrid.py`](runtime/patch_inc_hybrid.py) | adds an `INCConfig.maybe_update_config` override that detects FP8 dense layers in the hybrid checkpoint and dispatches `Fp8LinearMethod` for `shared_expert` | `dense` |
|
||||||
|
| [`patch_int8_lmhead_v3.py`](runtime/patch_int8_lmhead_v3.py) | replaces the lm-head matmul in `_get_logits` with a batched int8 w8a16 Triton GEMV (keeps bf16 weight for the shared drafter) | `dense` |
|
||||||
|
| [`patch_fla_shmem.py`](runtime/patch_fla_shmem.py) | lets the FLA GDN chunk kernels use big tiles on sm121's 99 KiB shmem (prefill/TTFT only; harmless) | always (free) |
|
||||||
|
|
||||||
|
Why DFlash needs the unify patch at all, why the drafter must run `FLASH_ATTN`
|
||||||
|
(non-causal), and the full vLLM-vs-SGLANG dead-end history are in
|
||||||
|
[`docs/FINDINGS.md`](docs/FINDINGS.md).
|
||||||
|
|
||||||
|
## Repo layout
|
||||||
|
|
||||||
|
```
|
||||||
|
install.sh One-shot installer (curl | bash | --help)
|
||||||
|
runtime/ Mounted read-only at /host inside the container:
|
||||||
|
serve.sh vLLM serve wrapper (applies the patches, then serves)
|
||||||
|
patch_unify2.py DFlash KV-unify fix
|
||||||
|
patch_inc_hybrid.py hybrid INT4+FP8 dispatch
|
||||||
|
patch_int8_lmhead_v3.py int8 lm-head GEMV
|
||||||
|
patch_fla_shmem.py FLA sm121 big-tile (prefill)
|
||||||
|
mtp_serve.sh MTP-2 comparison serve
|
||||||
|
scripts/ Host-side helpers:
|
||||||
|
monitor.sh Container-startup monitor with OOM auto-kill guard
|
||||||
|
bench_decode.py Decode-only tok/s (excludes TTFT)
|
||||||
|
bench_albond.py albond's e2e method (comparable to his 51.58)
|
||||||
|
hermes_bench.py Real agent turns from ~/.hermes/state.db
|
||||||
|
run_bank.sh prose/code/counting/hermes bank on any server
|
||||||
|
tools/
|
||||||
|
build-hybrid-checkpoint.py Build the hybrid INT4+FP8 ckpt (for --build-hybrid)
|
||||||
|
inspect_ckpt.py Which layers are INT4 vs BF16 vs FP8
|
||||||
|
validate_*.py Standalone correctness checks for the patches
|
||||||
|
docs/
|
||||||
|
FINDINGS.md The full investigation, methodology, and the
|
||||||
|
amortization-law derivation
|
||||||
|
```
|
||||||
|
|
||||||
|
## Reproducing
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# default agent path (DFlash) + smoke test:
|
||||||
|
./install.sh --start
|
||||||
|
|
||||||
|
# the dense stack (build the hybrid ckpt once, ~20 min, then serve):
|
||||||
|
./install.sh --build-hybrid
|
||||||
|
./install.sh --start --profile dense
|
||||||
|
|
||||||
|
# benches (run on the host against the server; need: pip install requests):
|
||||||
|
python3 scripts/bench_decode.py --base-url http://127.0.0.1:8000 --model qwen \
|
||||||
|
--prompt "Write a detailed essay about the history of tea."
|
||||||
|
python3 scripts/bench_albond.py http://127.0.0.1:8000 "dflash" # e2e, vs 51.58
|
||||||
|
python3 scripts/hermes_bench.py --base-url http://127.0.0.1:8000 # real agent turns
|
||||||
|
|
||||||
|
# MTP comparison:
|
||||||
|
./install.sh --start --profile mtp
|
||||||
|
```
|
||||||
|
|
||||||
|
## How this fits with related work
|
||||||
|
|
||||||
|
| Piece | Role |
|
||||||
|
|---|---|
|
||||||
|
| [`vLLM`](https://github.com/vllm-project/vllm) | the inference engine; this repo serves Qwen3.5 + DFlash on it, unmodified-on-disk |
|
||||||
|
| [`Intel/...int4-AutoRound`](https://huggingface.co/Intel/Qwen3.5-122B-A10B-int4-AutoRound) · [`z-lab/...DFlash`](https://huggingface.co/z-lab/Qwen3.5-122B-A10B-DFlash) | the target + drafter weights |
|
||||||
|
| [`albond/DGX_Spark_Qwen3.5-122B-A10B-AR-INT4`](https://github.com/albond/DGX_Spark_Qwen3.5-122B-A10B-AR-INT4) | the MTP + hybrid-FP8 + int8-lmhead recipe we benchmarked against and ported the dense levers from |
|
||||||
|
| [`Entrpi/ds4-on-spark`](https://github.com/Entrpi/ds4-on-spark) | sibling repo, same hardware, different model (DeepSeek-V4-Flash via ds4) |
|
||||||
|
| [Modal: *Speculative decoding is all you need*](https://modal.com/blog/spec-is-all-u-need) | the DFlash block-diffusion drafter and the task-dependent-acceptance framing |
|
||||||
|
|
||||||
|
## Acknowledgements
|
||||||
|
|
||||||
|
- [`z-lab`](https://huggingface.co/z-lab) / [Modal](https://modal.com/blog/spec-is-all-u-need) — the DFlash drafter and block-diffusion speculative decode.
|
||||||
|
- [`Intel/AutoRound`](https://huggingface.co/Intel) — the INT4 target quantization.
|
||||||
|
- [`vLLM`](https://github.com/vllm-project/vllm) and the AEON sm121 image maintainers — the engine and the DFlash-enabled GB10 build.
|
||||||
|
- [`albond`](https://github.com/albond/DGX_Spark_Qwen3.5-122B-A10B-AR-INT4) — the MTP/hybrid-FP8/int8-lmhead recipe and the end-to-end benchmark methodology.
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
MIT — see [LICENSE](LICENSE). The patches are original; vendored third-party
|
||||||
|
files (`tools/build-hybrid-checkpoint.py`) retain their upstream attribution.
|
||||||
@@ -0,0 +1,133 @@
|
|||||||
|
# FINDINGS — DFlash + dense levers for Qwen3.5-122B-A10B on DGX Spark
|
||||||
|
|
||||||
|
Single-stream (c=1) decode of `Qwen3.5-122B-A10B` (hybrid GDN + mamba + 128-expert
|
||||||
|
MoE, ~10B active) on GB10 / SM121, 128 GiB unified, ~273 GB/s. The agent this
|
||||||
|
backs (Hermes) is ~73 % tool-calls. All numbers temperature 0.
|
||||||
|
|
||||||
|
## 1. Getting DFlash to run on the hybrid 122B in vLLM
|
||||||
|
|
||||||
|
The DFlash drafter is **non-causal** (it block-drafts 16 tokens in one parallel
|
||||||
|
forward) — only the `FLASH_ATTN` (FA2) backend supports non-causal attention.
|
||||||
|
But the hybrid GDN+mamba+MoE target's KV-cache page geometry won't *unify* with
|
||||||
|
the drafter's attention spec:
|
||||||
|
|
||||||
|
- vLLM auto-aligns the hybrid (attention block 2240, mamba page padded +0.54 % to
|
||||||
|
match), so `max_page_size` is a *padded* value. `unify_kv_cache_spec_page_size`
|
||||||
|
scales the drafter's attention block by `ratio` and then asserts
|
||||||
|
`page == max` — which fails, because `page_size_bytes` ignores `block_size`
|
||||||
|
once `page_size_padded` is set.
|
||||||
|
- **Fix** ([`patch_unify2.py`](../runtime/patch_unify2.py)): keep the *scaled*
|
||||||
|
`block_size` **and** pad the <1 % remainder (mirrors vLLM's own
|
||||||
|
`HiddenStateCacheSpec` handling). The earlier "pad but keep block_size=16"
|
||||||
|
patch mis-strided the drafter KV → acceptance collapsed to 1.47 (a real bug,
|
||||||
|
not a quant mismatch).
|
||||||
|
- `--mamba-block-size 256` (from other Spark recipes) **breaks** the 122B: it
|
||||||
|
makes the mamba group block ≠ cache block, tripping the coordinator hash
|
||||||
|
assert. Omit it.
|
||||||
|
- `--no-enable-prefix-caching` routes to `KVCacheCoordinatorNoPrefixCache`, which
|
||||||
|
has no line-504 hash assert. Prefix caching is irrelevant at c=1 anyway.
|
||||||
|
|
||||||
|
Working stack: `patch_unify2` + prefix-off + INT4 (bf16 KV) target + drafter
|
||||||
|
pinned to `FLASH_ATTN`. **No FA4 shim needed** — vLLM gates FA4 to cap families
|
||||||
|
90/100/110 (excludes 120), so the drafter runs FA2. (The whole fa4-sm120 saga is
|
||||||
|
SGLang-only; SGLang's DFlash works too but its sm121 base decode is ~2× slower
|
||||||
|
than vLLM's, so it loses on absolute throughput.)
|
||||||
|
|
||||||
|
## 2. DFlash vs MTP — acceptance is task-dependent
|
||||||
|
|
||||||
|
MTP-2 (the native head) drafts 2 tokens **sequentially**, so acceptance caps at
|
||||||
|
~3. DFlash block-drafts 12 in **one parallel forward**, so its acceptance fills
|
||||||
|
the block on predictable traffic. Same harness, unpatched, flash_attn:
|
||||||
|
|
||||||
|
| Workload | accept (MTP-2 / DFlash) | tok/s (MTP-2 / DFlash) |
|
||||||
|
|---|---|---|
|
||||||
|
| Prose | 2.24 / 2.3 | 33.7 / 33.2 *(tie; use DFlash n=4)* |
|
||||||
|
| Code | 2.77 / 5.4 | 40.5 / 54.5 |
|
||||||
|
| Counting | **3.00 (maxed)** / 11 | 43.7 / 124.5 |
|
||||||
|
| Hermes (real) | 2.88 / **8.66** | 39.9 / **~81** |
|
||||||
|
|
||||||
|
The "DFlash caps at 2.3 / 33 tok/s" story was a **prose-benchmark artifact**.
|
||||||
|
On agent/code traffic DFlash pulls ~2× ahead because MTP is acceptance-saturated.
|
||||||
|
`n` (num_speculative_tokens) is task-dependent: prose → 4, agent/code → 12+.
|
||||||
|
|
||||||
|
## 3. Methodology — two non-comparable harnesses
|
||||||
|
|
||||||
|
- `bench_decode.py` = **decode-only** tok/s (excludes TTFT). Good for c=1 kernel
|
||||||
|
comparisons. This is the ~81 Hermes number.
|
||||||
|
- `bench_albond.py` = **end-to-end** (completion_tokens / total wallclock incl.
|
||||||
|
prefill, non-streaming, 5 prompts, run-1 discarded). This reproduces albond's
|
||||||
|
own method and is directly comparable to his published **51.58**.
|
||||||
|
|
||||||
|
Apples-to-apples (albond's method): **DFlash n=12 unpatched = 53.7 tok/s
|
||||||
|
cross-prompt mean — already above his fully-patched MTP stack (51.58).**
|
||||||
|
|
||||||
|
## 4. The dense-bandwidth levers and the amortization law
|
||||||
|
|
||||||
|
albond's non-MTP wins are *always-on* bandwidth cuts. We ported the two that
|
||||||
|
transfer to vLLM 0.23 + DFlash:
|
||||||
|
|
||||||
|
- **hybrid INT4+FP8** ([`patch_inc_hybrid.py`](../runtime/patch_inc_hybrid.py)):
|
||||||
|
the Intel base already stores **attention as INT4** (0.5 B/param, *better* than
|
||||||
|
albond's FP8 attention), so the only thing to gain is the BF16 **shared
|
||||||
|
experts** → calibrated FP8 (144 layers, 0.48 GB saved). The dispatch patch adds
|
||||||
|
an `INCConfig.maybe_update_config` override (AEON 0.23's hook signature takes
|
||||||
|
`hf_config=`, unlike albond's 0.19) that detects FP8 dense layers and
|
||||||
|
dispatches `Fp8LinearMethod` for them.
|
||||||
|
- **int8 lm-head** ([`patch_int8_lmhead_v3.py`](../runtime/patch_int8_lmhead_v3.py)):
|
||||||
|
the 248 320-row vocab projection is the single largest dense read (1.5 GB BF16,
|
||||||
|
*every token*). A batched int8 w8a16 Triton GEMV reads it at ~227 GB/s (vs bf16
|
||||||
|
~6.5–8.8 ms) — **~2× faster, argmax-exact**. Prior ports failed not on the
|
||||||
|
kernel but on **integration**: zeroing the lm-head weight corrupted the
|
||||||
|
*drafter-shared* head (garbage), and a per-row loop for B>4 was slower under
|
||||||
|
spec. v3 uses one batched kernel and **keeps** the bf16 weight.
|
||||||
|
|
||||||
|
Why the denominator matters: the 0.48 GB shared-expert saving is **0.7 % of the
|
||||||
|
71 GB on disk** but **~8 % of the ~6 GB *active per-token* footprint** (the disk
|
||||||
|
is mostly sparse routed experts). Shared experts and the lm-head are **dense —
|
||||||
|
read every token** — so at base decode the savings land in full:
|
||||||
|
|
||||||
|
| Config | base (acc 1) | albond-bench (acc 6.4) | Hermes (acc 8.3) |
|
||||||
|
|---|---|---|---|
|
||||||
|
| INT4 baseline | 28.2 | 53.7 | ~81 |
|
||||||
|
| + hybrid-FP8 | 30.4 (+7.8%) | 57.0 (+6.1%) | ~80 |
|
||||||
|
| + int8 lm-head | 32.7 (+16%) | — | — |
|
||||||
|
| **+ both** | **36.0 (+28%)** | **59.0 (+10%)** | ~80–87 (noise) |
|
||||||
|
|
||||||
|
The levers compose additively (step savings 2.6 + 4.9 ≈ 7.7 ms). But the uplift
|
||||||
|
**decays monotonically with acceptance**:
|
||||||
|
|
||||||
|
> Under speculative decode the verify forward processes ~`accept` positions and
|
||||||
|
> reads each dense weight **once**, amortized across them. So a dense-weight cut
|
||||||
|
> that is +X % at base is ~+X/accept % under spec.
|
||||||
|
|
||||||
|
```
|
||||||
|
+28% base (accept 1) → +10% albond-bench (accept 6.4) → ~0% Hermes (accept 8.3)
|
||||||
|
```
|
||||||
|
|
||||||
|
**Consequences**
|
||||||
|
|
||||||
|
- For the **agent path** (`dflash`), DFlash's own high acceptance already
|
||||||
|
amortizes the dense levers to ~null. Its remaining bottleneck is **routed-expert
|
||||||
|
verify-batch reads** (each of ~13 verify positions routes to different experts)
|
||||||
|
— untouched by any dense-weight quant. To push Hermes further you must attack
|
||||||
|
*that*: a smaller/faster drafter, lower `n` at equal acceptance, or sub-INT4
|
||||||
|
routed experts.
|
||||||
|
- For **base / low-acceptance** serving (`dense`), the stack is a real **+28 %**
|
||||||
|
(36 tok/s) and is the recommended config there.
|
||||||
|
|
||||||
|
## 5. Things that did NOT help c=1 decode
|
||||||
|
|
||||||
|
- **FLASHINFER target backend** — null both short-context and Hermes (attention
|
||||||
|
isn't the bottleneck on this GDN/mamba-heavy MoE; most layers are linear
|
||||||
|
attention). albond's "+16 %" was on his dense-attention MTP path.
|
||||||
|
- **b12x / native FP4 MoE** — null at c=1 (a throughput/concurrency lever, not a
|
||||||
|
latency one; at batch 1 the active-expert GEMM is tiny).
|
||||||
|
- **FLA sm121 big-tile shmem fix** — real bug, but prefill/TTFT only; c=1 decode
|
||||||
|
uses the GDN *recurrent* path, a different kernel. Kept (free TTFT win).
|
||||||
|
- **PR#38325 swapAB FP8 GEMM** — marginal (+0.76 %), only with the FP8 checkpoint.
|
||||||
|
|
||||||
|
## Production recommendation
|
||||||
|
|
||||||
|
Ship **`dflash`** for the agent (DFlash unpatched, ~81 tok/s, ~2× MTP, > albond's
|
||||||
|
patched 51.58). Reserve **`dense`** for base / low-acceptance serving (+28 %).
|
||||||
|
The dense patches are upside there, not a requirement for the agent to win.
|
||||||
Executable
+319
@@ -0,0 +1,319 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
# install.sh — Qwen3.5-122B-A10B + DFlash speculative decode on NVIDIA DGX Spark
|
||||||
|
# (GB10 / SM121, 128 GiB unified), via vLLM in Docker.
|
||||||
|
#
|
||||||
|
# curl -sSL https://raw.githubusercontent.com/Entrpi/qwen3.5-122B-A10B-on-spark/main/install.sh | bash
|
||||||
|
# curl -sSL https://raw.githubusercontent.com/Entrpi/qwen3.5-122B-A10B-on-spark/main/install.sh | bash -s -- --help
|
||||||
|
#
|
||||||
|
# What this does (every step idempotent — safe to re-run):
|
||||||
|
#
|
||||||
|
# 1. Verifies the host is a DGX Spark / GB10 (SM121) with Docker + the NVIDIA
|
||||||
|
# container runtime, and enough free disk for the chosen profile.
|
||||||
|
# 2. Pulls the prebuilt sm121 vLLM image (DFlash-enabled, vLLM 0.23).
|
||||||
|
# 3. Downloads the INT4 target + DFlash drafter from Hugging Face into the HF
|
||||||
|
# cache — OR reuses a checkpoint you already have (--model-dir / --hf-home).
|
||||||
|
# 4. (optional) Builds the hybrid INT4+FP8 checkpoint for the "dense" profile.
|
||||||
|
# 5. Starts the vLLM server for the chosen --profile, waits until READY, and
|
||||||
|
# runs the "capital of France" smoke test (expects "Paris").
|
||||||
|
#
|
||||||
|
# The script makes NO changes outside:
|
||||||
|
# - the Docker image cache (the pulled image)
|
||||||
|
# - $HF_HOME (default ~/.cache/huggingface)
|
||||||
|
# - $HYBRID_DIR (only with --build-hybrid)
|
||||||
|
# - the running container named $NAME (only with --start / smoke)
|
||||||
|
#
|
||||||
|
# This repo provides the install + serve + patch + benchmark layer ON TOP of:
|
||||||
|
# - Intel/Qwen3.5-122B-A10B-int4-AutoRound (target weights)
|
||||||
|
# - z-lab/Qwen3.5-122B-A10B-DFlash (block-diffusion drafter)
|
||||||
|
# - ghcr.io/aeon-7/aeon-vllm-ultimate (sm121 DFlash-enabled vLLM)
|
||||||
|
#
|
||||||
|
# License: MIT. Source: https://github.com/Entrpi/qwen3.5-122B-A10B-on-spark
|
||||||
|
|
||||||
|
set -euo pipefail
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# 0. defaults + flag parsing
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
# Prebuilt sm121 vLLM image with the DFlash PRs + the .pth that auto-applies our
|
||||||
|
# KV-unify patch is NOT baked in — we apply patches at serve time from runtime/.
|
||||||
|
IMAGE="${QWEN_IMAGE:-ghcr.io/aeon-7/aeon-vllm-ultimate:2026-06-18-v0.23.0-dflashfix}"
|
||||||
|
|
||||||
|
TARGET_REPO="${TARGET_REPO:-Intel/Qwen3.5-122B-A10B-int4-AutoRound}" # INT4 target (~62 GiB)
|
||||||
|
DRAFT_REPO="${DRAFT_REPO:-z-lab/Qwen3.5-122B-A10B-DFlash}" # 0.8B drafter (~1.6 GiB)
|
||||||
|
FP8_REPO="${FP8_REPO:-Qwen/Qwen3.5-122B-A10B-FP8}" # FP8 donor for --build-hybrid
|
||||||
|
|
||||||
|
HF_HOME="${HF_HOME:-$HOME/.cache/huggingface}"
|
||||||
|
HYBRID_DIR="${HYBRID_DIR:-$HOME/qwen3.5-122b-hybrid-int4-fp8}"
|
||||||
|
MODEL_DIR="" # --model-dir: a pre-downloaded INT4 checkpoint dir
|
||||||
|
|
||||||
|
# This repo's own dir (works for `curl | bash` too: falls back to a clone).
|
||||||
|
REPO_DIR="${REPO_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" 2>/dev/null && pwd || echo "$HOME/code/qwen3.5-122B-A10B-on-spark")}"
|
||||||
|
REPO_URL="${REPO_URL:-https://github.com/Entrpi/qwen3.5-122B-A10B-on-spark.git}"
|
||||||
|
|
||||||
|
NAME="${NAME:-qwen-spark}"
|
||||||
|
PROFILE="dflash" # dflash | dense | base | mtp
|
||||||
|
NSPEC="" # override num_speculative_tokens (default per profile)
|
||||||
|
PORT="${PORT:-8000}"
|
||||||
|
CTX="${CTX:-16384}"
|
||||||
|
GPU_MEM="${GPU_MEM:-0.8}"
|
||||||
|
BACKEND="${BACKEND:-flash_attn}"
|
||||||
|
|
||||||
|
FORCE_HW=0
|
||||||
|
SKIP_PULL=0
|
||||||
|
SKIP_DOWNLOAD=0
|
||||||
|
BUILD_HYBRID=0
|
||||||
|
START_SERVER=0
|
||||||
|
SKIP_SMOKE=0
|
||||||
|
|
||||||
|
usage() {
|
||||||
|
cat <<EOF
|
||||||
|
Usage: $0 [flags]
|
||||||
|
|
||||||
|
Profiles (--profile):
|
||||||
|
dflash INT4 target + DFlash drafter, n=12 (DEFAULT — best for agents/Hermes;
|
||||||
|
~81 tok/s on real tool-call turns)
|
||||||
|
dense hybrid INT4+FP8 + int8 lm-head + DFlash (the dense-bandwidth stack;
|
||||||
|
+28% at base, +10% low-accept spec.
|
||||||
|
Needs --build-hybrid first.)
|
||||||
|
base plain INT4, no speculative decode (~28 tok/s c=1 baseline)
|
||||||
|
mtp INT4 + native MTP-2 head (the albond comparison path)
|
||||||
|
|
||||||
|
Flags:
|
||||||
|
--help Show this help.
|
||||||
|
--profile NAME One of dflash|dense|base|mtp (default: dflash).
|
||||||
|
--start Start the vLLM server + smoke test after setup.
|
||||||
|
--build-hybrid Build the hybrid INT4+FP8 checkpoint (~20 min, needs FP8 donor).
|
||||||
|
--no-pull Skip docker pull (use the local image).
|
||||||
|
--no-download Skip HF download (assume target+drafter already cached).
|
||||||
|
--model-dir DIR Use a pre-downloaded INT4 target checkpoint dir (skip its
|
||||||
|
download; mounted read-only at /model).
|
||||||
|
--hf-home DIR Use/populate this HF cache dir (default: $HF_HOME).
|
||||||
|
--nspec N num_speculative_tokens (default 12 dflash/dense, 2 mtp, 0 base).
|
||||||
|
--port N Server port (default: $PORT).
|
||||||
|
--ctx N max-model-len (default: $CTX).
|
||||||
|
--gpu-mem F gpu-memory-utilization, keep <=0.84 on 128GiB (default: $GPU_MEM).
|
||||||
|
--force Skip the GB10/SM121 host check.
|
||||||
|
--no-smoke Start the server but skip the Paris smoke test.
|
||||||
|
|
||||||
|
Environment equivalents:
|
||||||
|
QWEN_IMAGE TARGET_REPO DRAFT_REPO FP8_REPO HF_HOME HYBRID_DIR
|
||||||
|
NAME PORT CTX GPU_MEM BACKEND REPO_DIR REPO_URL
|
||||||
|
EOF
|
||||||
|
}
|
||||||
|
|
||||||
|
while [[ $# -gt 0 ]]; do
|
||||||
|
case "$1" in
|
||||||
|
--help|-h) usage; exit 0 ;;
|
||||||
|
--profile) PROFILE="$2"; shift 2 ;;
|
||||||
|
--start) START_SERVER=1; shift ;;
|
||||||
|
--build-hybrid) BUILD_HYBRID=1; shift ;;
|
||||||
|
--no-pull) SKIP_PULL=1; shift ;;
|
||||||
|
--no-download) SKIP_DOWNLOAD=1; shift ;;
|
||||||
|
--model-dir) MODEL_DIR="$2"; shift 2 ;;
|
||||||
|
--hf-home) HF_HOME="$2"; shift 2 ;;
|
||||||
|
--nspec) NSPEC="$2"; shift 2 ;;
|
||||||
|
--port) PORT="$2"; shift 2 ;;
|
||||||
|
--ctx) CTX="$2"; shift 2 ;;
|
||||||
|
--gpu-mem) GPU_MEM="$2"; shift 2 ;;
|
||||||
|
--backend) BACKEND="$2"; shift 2 ;;
|
||||||
|
--force) FORCE_HW=1; shift ;;
|
||||||
|
--no-smoke) SKIP_SMOKE=1; shift ;;
|
||||||
|
*) echo "Unknown flag: $1" >&2; usage; exit 2 ;;
|
||||||
|
esac
|
||||||
|
done
|
||||||
|
|
||||||
|
case "$PROFILE" in dflash|dense|base|mtp) ;; *) echo "Bad --profile: $PROFILE" >&2; exit 2 ;; esac
|
||||||
|
|
||||||
|
c_red() { printf '\033[31m%s\033[0m' "$*"; }
|
||||||
|
c_green() { printf '\033[32m%s\033[0m' "$*"; }
|
||||||
|
c_yellow(){ printf '\033[33m%s\033[0m' "$*"; }
|
||||||
|
log() { printf '%s %s\n' "[$(date +%H:%M:%S)]" "$*"; }
|
||||||
|
die() { printf '\n%s %s\n' "$(c_red FATAL:)" "$*" >&2; exit 1; }
|
||||||
|
warn(){ printf '%s %s\n' "$(c_yellow WARN:)" "$*" >&2; }
|
||||||
|
ok() { printf '%s %s\n' "$(c_green OK:)" "$*"; }
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# 1. host verification
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
verify_host() {
|
||||||
|
log "Verifying host..."
|
||||||
|
local m; m=$(uname -m)
|
||||||
|
if [[ "$m" != "aarch64" ]] && [[ "$FORCE_HW" -eq 0 ]]; then
|
||||||
|
die "Expected aarch64 (Grace+Blackwell); got $m. Pass --force to skip."
|
||||||
|
fi
|
||||||
|
command -v docker >/dev/null 2>&1 || die "docker not found. Install Docker + the NVIDIA container runtime."
|
||||||
|
command -v nvidia-smi >/dev/null 2>&1 || die "nvidia-smi not found. Need the NVIDIA driver."
|
||||||
|
local gpu; gpu=$(nvidia-smi --query-gpu=name,compute_cap --format=csv,noheader 2>/dev/null || true)
|
||||||
|
[[ -n "$gpu" ]] || die "nvidia-smi failed to enumerate GPUs."
|
||||||
|
log "GPU: $gpu"
|
||||||
|
if ! echo "$gpu" | grep -qE '12\.1|GB10|Spark'; then
|
||||||
|
[[ "$FORCE_HW" -eq 1 ]] || die "Not detecting GB10/SM12.1. Pass --force (and maybe --backend) to proceed."
|
||||||
|
warn "Host is not GB10/SM121; proceeding under --force (untested)."
|
||||||
|
fi
|
||||||
|
# Docker can see the GPU?
|
||||||
|
if ! docker info 2>/dev/null | grep -qiE 'nvidia|Default Runtime: nvidia' \
|
||||||
|
&& ! docker run --rm --gpus all "$IMAGE" true 2>/dev/null; then
|
||||||
|
warn "Could not confirm Docker GPU access (nvidia-container-toolkit). 'docker run --gpus all' must work."
|
||||||
|
fi
|
||||||
|
# Disk
|
||||||
|
local need=75; [[ "$BUILD_HYBRID" -eq 1 ]] && need=150
|
||||||
|
local free; free=$(df -BG "$HOME" | awk 'NR==2{gsub("G","",$4);print $4}')
|
||||||
|
if (( free < need )) && [[ "$SKIP_DOWNLOAD" -eq 0 ]] && [[ -z "$MODEL_DIR" ]]; then
|
||||||
|
die "Need >= ${need} GiB free under $HOME for profile '$PROFILE'; have ${free} GiB. Use --model-dir / --no-download, or free space."
|
||||||
|
fi
|
||||||
|
ok "Host checks passed (free ${free} GiB)."
|
||||||
|
}
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# 2. pull image
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
pull_image() {
|
||||||
|
if [[ "$SKIP_PULL" -eq 1 ]]; then log "Skipping docker pull (--no-pull)."; return; fi
|
||||||
|
if docker image inspect "$IMAGE" >/dev/null 2>&1; then ok "Image present: $IMAGE"; return; fi
|
||||||
|
log "Pulling $IMAGE (~40 GiB, one-time) ..."
|
||||||
|
docker pull "$IMAGE"
|
||||||
|
ok "Image pulled."
|
||||||
|
}
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# 3. download models (idempotent — snapshot_download no-ops if cached)
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
hf_get() { # repo -> populate HF cache via the image's huggingface_hub
|
||||||
|
local repo="$1"
|
||||||
|
docker run --rm --net=host -e HF_HOME=/hf ${HF_TOKEN:+-e HF_TOKEN="$HF_TOKEN"} \
|
||||||
|
-v "$HF_HOME:/hf" --entrypoint python3 "$IMAGE" \
|
||||||
|
-c "from huggingface_hub import snapshot_download as s; s('$repo')"
|
||||||
|
}
|
||||||
|
|
||||||
|
download_models() {
|
||||||
|
if [[ "$SKIP_DOWNLOAD" -eq 1 ]]; then log "Skipping HF download (--no-download)."; return; fi
|
||||||
|
mkdir -p "$HF_HOME"
|
||||||
|
if [[ -n "$MODEL_DIR" ]]; then
|
||||||
|
[[ -f "$MODEL_DIR/config.json" ]] || die "--model-dir $MODEL_DIR has no config.json"
|
||||||
|
log "Using pre-downloaded target at $MODEL_DIR (skipping target download)."
|
||||||
|
else
|
||||||
|
log "Fetching target $TARGET_REPO into $HF_HOME ..."
|
||||||
|
hf_get "$TARGET_REPO"
|
||||||
|
fi
|
||||||
|
log "Fetching drafter $DRAFT_REPO ..."
|
||||||
|
hf_get "$DRAFT_REPO"
|
||||||
|
ok "Models ready."
|
||||||
|
}
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# 4. optional: build the hybrid INT4+FP8 checkpoint
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
build_hybrid() {
|
||||||
|
[[ "$BUILD_HYBRID" -eq 1 ]] || return
|
||||||
|
if [[ -f "$HYBRID_DIR/model.safetensors.index.json" ]]; then ok "Hybrid ckpt present: $HYBRID_DIR"; return; fi
|
||||||
|
local gptq="$MODEL_DIR"
|
||||||
|
if [[ -z "$gptq" ]]; then
|
||||||
|
gptq=$(docker run --rm -v "$HF_HOME:/hf" -e HF_HOME=/hf --entrypoint python3 "$IMAGE" \
|
||||||
|
-c "from huggingface_hub import snapshot_download as s; print(s('$TARGET_REPO'))" | tail -1)
|
||||||
|
gptq="/hf-snap" # mount the cache; resolve inside the container below
|
||||||
|
fi
|
||||||
|
mkdir -p "$HYBRID_DIR"
|
||||||
|
log "Building hybrid INT4+FP8 checkpoint -> $HYBRID_DIR (~20 min) ..."
|
||||||
|
docker run --rm --net=host -e HF_HOME=/hf ${HF_TOKEN:+-e HF_TOKEN="$HF_TOKEN"} \
|
||||||
|
-v "$HF_HOME:/hf" -v "$HYBRID_DIR:/out" -v "$REPO_DIR/tools:/tools:ro" \
|
||||||
|
${MODEL_DIR:+-v "$MODEL_DIR:/gptq:ro"} \
|
||||||
|
--entrypoint bash "$IMAGE" -c '
|
||||||
|
set -e
|
||||||
|
GPTQ="'"${MODEL_DIR:+/gptq}"'"
|
||||||
|
if [ -z "$GPTQ" ]; then
|
||||||
|
GPTQ=$(python3 -c "from huggingface_hub import snapshot_download as s; print(s(\"'"$TARGET_REPO"'\"))")
|
||||||
|
fi
|
||||||
|
python3 /tools/build-hybrid-checkpoint.py --gptq-dir "$GPTQ" \
|
||||||
|
--fp8-repo "'"$FP8_REPO"'" --output /out --force
|
||||||
|
rm -rf /out/.fp8_cache'
|
||||||
|
ok "Hybrid checkpoint built: $HYBRID_DIR"
|
||||||
|
}
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# 5. start server (+ smoke test)
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
ensure_runtime() { # make sure runtime/ (serve wrapper + patches) is on disk
|
||||||
|
if [[ -f "$REPO_DIR/runtime/serve.sh" ]]; then return; fi
|
||||||
|
log "runtime/ not found next to install.sh — cloning repo to $HOME/code/qwen3.5-122B-A10B-on-spark"
|
||||||
|
REPO_DIR="$HOME/code/qwen3.5-122B-A10B-on-spark"
|
||||||
|
[[ -d "$REPO_DIR/.git" ]] || git clone --depth 1 "$REPO_URL" "$REPO_DIR"
|
||||||
|
[[ -f "$REPO_DIR/runtime/serve.sh" ]] || die "runtime/serve.sh still missing after clone."
|
||||||
|
}
|
||||||
|
|
||||||
|
start_server() {
|
||||||
|
[[ "$START_SERVER" -eq 1 ]] || { log "Setup complete. Re-run with --start to launch the server."; return; }
|
||||||
|
ensure_runtime
|
||||||
|
|
||||||
|
# profile -> serve args + env + mounts
|
||||||
|
local nspec model_env=() mounts=() serve_args
|
||||||
|
case "$PROFILE" in
|
||||||
|
dflash) nspec="${NSPEC:-12}"; serve_args="$nspec $BACKEND" ;;
|
||||||
|
dense) nspec="${NSPEC:-12}"; serve_args="$nspec $BACKEND"
|
||||||
|
[[ -f "$HYBRID_DIR/model.safetensors.index.json" ]] || die "dense profile needs the hybrid ckpt — run with --build-hybrid first."
|
||||||
|
model_env=(-e MODEL=/model -e INC_HYBRID=1 -e INT8_LMHEAD_V3=1)
|
||||||
|
mounts=(-v "$HYBRID_DIR:/model:ro") ;;
|
||||||
|
base) nspec="${NSPEC:-0}"; serve_args="$nspec $BACKEND" ;;
|
||||||
|
mtp) nspec="${NSPEC:-2}"; serve_args="$nspec $BACKEND" ;;
|
||||||
|
esac
|
||||||
|
if [[ -n "$MODEL_DIR" && "$PROFILE" != "dense" ]]; then
|
||||||
|
model_env=(-e MODEL=/model); mounts=(-v "$MODEL_DIR:/model:ro")
|
||||||
|
fi
|
||||||
|
local wrapper="/host/serve.sh"; [[ "$PROFILE" == "mtp" ]] && wrapper="/host/mtp_serve.sh"
|
||||||
|
|
||||||
|
log "Starting profile=$PROFILE (nspec=$nspec, ctx=$CTX, gpu-mem=$GPU_MEM) as container '$NAME' ..."
|
||||||
|
docker rm -f "$NAME" >/dev/null 2>&1 || true
|
||||||
|
# shellcheck disable=SC2086
|
||||||
|
docker run -d --name "$NAME" --gpus all --net=host --ipc=host --ulimit memlock=-1:-1 \
|
||||||
|
-e HF_HOME=/hf -e MAX_MODEL_LEN="$CTX" -e GPU_MEM="$GPU_MEM" ${HF_TOKEN:+-e HF_TOKEN="$HF_TOKEN"} \
|
||||||
|
"${model_env[@]}" \
|
||||||
|
-v "$HF_HOME:/hf" -v "$REPO_DIR/runtime:/host:ro" "${mounts[@]}" \
|
||||||
|
--entrypoint bash "$IMAGE" "$wrapper" $serve_args >/dev/null
|
||||||
|
log "Container started. Model load + compile is ~8-12 min. Tail: docker logs -f $NAME"
|
||||||
|
|
||||||
|
log "Waiting for http://127.0.0.1:$PORT/health ..."
|
||||||
|
local i
|
||||||
|
for i in $(seq 1 180); do
|
||||||
|
if ! docker ps --format '{{.Names}}' | grep -q "^$NAME$"; then
|
||||||
|
docker logs "$NAME" 2>&1 | tail -30; die "Container exited during load. See log above."
|
||||||
|
fi
|
||||||
|
if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
|
||||||
|
ok "Server READY on http://127.0.0.1:$PORT"
|
||||||
|
break
|
||||||
|
fi
|
||||||
|
sleep 5
|
||||||
|
done
|
||||||
|
curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 || die "Server not ready within ~15 min. docker logs $NAME"
|
||||||
|
|
||||||
|
[[ "$SKIP_SMOKE" -eq 1 ]] && { log "Skipping smoke test (--no-smoke)."; return; }
|
||||||
|
log "Smoke test: 'capital of France' ..."
|
||||||
|
local out
|
||||||
|
out=$(curl -s "http://127.0.0.1:$PORT/v1/completions" -H 'Content-Type: application/json' \
|
||||||
|
-d "{\"model\":\"qwen\",\"prompt\":\"What is the capital of France? Answer in one word.\",\"max_tokens\":8,\"temperature\":0}" \
|
||||||
|
| python3 -c "import sys,json;print(json.load(sys.stdin)['choices'][0]['text'])" 2>/dev/null || true)
|
||||||
|
echo " -> $out"
|
||||||
|
echo "$out" | grep -qi paris && ok "Smoke test PASSED — 'Paris'." || die "Smoke test FAILED — 'Paris' not in output."
|
||||||
|
}
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# main
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
verify_host
|
||||||
|
pull_image
|
||||||
|
download_models
|
||||||
|
build_hybrid
|
||||||
|
start_server
|
||||||
|
|
||||||
|
echo
|
||||||
|
ok "Done (profile=$PROFILE)."
|
||||||
|
echo " Server: http://127.0.0.1:$PORT/v1 (model name: qwen)"
|
||||||
|
echo " Logs: docker logs -f $NAME"
|
||||||
|
echo " Stop: docker rm -f $NAME"
|
||||||
|
echo " Bench: python3 scripts/bench_decode.py --base-url http://127.0.0.1:$PORT --model qwen --prompt 'Write an essay about tea.'"
|
||||||
|
echo " Agent: python3 scripts/hermes_bench.py --base-url http://127.0.0.1:$PORT # real tool-call turns"
|
||||||
Executable
+25
@@ -0,0 +1,25 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
# mtp_serve.sh — native qwen3_5 MTP-N head for the comparison (`--profile mtp`).
|
||||||
|
# The MTP head (1 layer, reuses target KV/embed/lm_head) is in the Intel
|
||||||
|
# checkpoint (mtp.layers.0) — no separate drafter, no unify patch needed.
|
||||||
|
# $1 = num_speculative_tokens (default 2 = the "MTP-2" recipe); $2 = backend.
|
||||||
|
set -euo pipefail
|
||||||
|
NSPEC="${1:-2}"
|
||||||
|
BACKEND="${2:-flash_attn}"
|
||||||
|
MODEL="${MODEL:-Intel/Qwen3.5-122B-A10B-int4-AutoRound}"
|
||||||
|
MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
|
||||||
|
GPU_MEM="${GPU_MEM:-0.8}"
|
||||||
|
PORT="${PORT:-8000}"
|
||||||
|
echo "[mtp] qwen3_5_mtp — backend=$BACKEND, num_speculative_tokens=$NSPEC, model=$MODEL"
|
||||||
|
exec vllm serve "$MODEL" \
|
||||||
|
--served-model-name qwen \
|
||||||
|
--host 0.0.0.0 --port "$PORT" \
|
||||||
|
--max-model-len "$MAX_MODEL_LEN" \
|
||||||
|
--max-num-seqs 16 \
|
||||||
|
--max-num-batched-tokens "$MAX_MODEL_LEN" \
|
||||||
|
--gpu-memory-utilization "$GPU_MEM" \
|
||||||
|
--no-enable-prefix-caching \
|
||||||
|
--enable-chunked-prefill \
|
||||||
|
--trust-remote-code \
|
||||||
|
--attention-backend "$BACKEND" \
|
||||||
|
--speculative-config "{\"method\":\"qwen3_5_mtp\",\"num_speculative_tokens\":$NSPEC,\"model\":\"$MODEL\"}"
|
||||||
@@ -0,0 +1,32 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""spark-fla-shmem: let the FLA (flash-linear-attention / GDN) Triton kernels use their
|
||||||
|
BIG tiles on sm121 (GB10 / DGX Spark).
|
||||||
|
|
||||||
|
The FLA Backend gate (vllm/model_executor/layers/fla/ops/utils.py, identical in SGLang)
|
||||||
|
picks big tiles only if device_max_shared_mem >= Backend.DEFAULT (102400 = 100 KiB):
|
||||||
|
cumsum.py : BS_LIST = [32,64] if check_shared_mem() else [16,32]
|
||||||
|
chunk_o.py: BKV_LIST = [64,128] if check_shared_mem() else [32,64]
|
||||||
|
sm121 reports max_shared_mem = 101376 (99 KiB), JUST below 102400 -> check returns False
|
||||||
|
-> small tiles -> slower GDN/linear-attention (the per-token hot path). But 101376 is
|
||||||
|
EXACTLY Backend.ADA, and RTX 4090 (ADA, same 99 KiB) runs the big tiles fine, so they
|
||||||
|
provably fit in 99 KiB. Lower DEFAULT to 101376 so sm121 (and ADA) pass the gate.
|
||||||
|
"""
|
||||||
|
import pathlib, sys
|
||||||
|
|
||||||
|
P = pathlib.Path(
|
||||||
|
"/usr/local/lib/python3.12/site-packages/vllm/model_executor/layers/fla/ops/utils.py"
|
||||||
|
)
|
||||||
|
src = P.read_text()
|
||||||
|
if "spark-fla-shmem" in src:
|
||||||
|
print("[patch_fla_shmem] already patched", flush=True)
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
OLD = " DEFAULT = 102400 # Default"
|
||||||
|
NEW = " DEFAULT = 101376 # spark-fla-shmem: 102400->101376 so sm121 GB10 (99 KiB) uses big GDN tiles (already fit on ADA's identical 99 KiB)"
|
||||||
|
|
||||||
|
if OLD not in src:
|
||||||
|
print("[patch_fla_shmem] ERROR: target line not found — FLA source differs", flush=True)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
P.write_text(src.replace(OLD, NEW))
|
||||||
|
print("[patch_fla_shmem] patched OK (DEFAULT 102400 -> 101376; big GDN tiles on sm121)", flush=True)
|
||||||
@@ -0,0 +1,163 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""spark-dflash-hybrid-fp8: port of albond's hybrid INT4+FP8 dispatch
|
||||||
|
(patches/01-hybrid-int4-fp8/inc.py.patch) onto AEON 0.23's vllm INCConfig.
|
||||||
|
|
||||||
|
Adds, to vllm/model_executor/layers/quantization/inc.py:
|
||||||
|
* INCConfig.fp8_config / fp8_layers fields
|
||||||
|
* maybe_update_config OVERRIDE — AEON 0.23 signature (model_name, hf_config=None,
|
||||||
|
revision=None); the base hook is already CALLED from config/vllm.py:634, so we
|
||||||
|
only supply the override. Scans the checkpoint's safetensors metadata for
|
||||||
|
float8_e4m3fn weights that have a .weight_scale_inv, builds an Fp8Config, and
|
||||||
|
records those layer prefixes.
|
||||||
|
* _is_layer_fp8 — exact + fused + substring match against fp8_layers
|
||||||
|
* FP8 dispatch at BOTH dense short-circuits: get_quant_method's extra_config
|
||||||
|
override AND the apply_*_quant_layer not-quantized blocks.
|
||||||
|
|
||||||
|
Idempotent; sentinel 'spark-dflash-hybrid-fp8'. Mirrors patch_unify2.py's style.
|
||||||
|
"""
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import vllm.model_executor.layers.quantization.inc as inc_mod
|
||||||
|
|
||||||
|
path = inc_mod.__file__
|
||||||
|
src = open(path).read()
|
||||||
|
SENT = "spark-dflash-hybrid-fp8"
|
||||||
|
if SENT in src:
|
||||||
|
print(f"[patch_inc_hybrid] already applied: {path}")
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
# 1. __init__ fields (anchor unique: only INCConfig has pack_factor = Fraction(32,..))
|
||||||
|
a1 = " self.pack_factor = Fraction(32, weight_bits)\n"
|
||||||
|
b1 = a1 + (
|
||||||
|
" # spark-dflash-hybrid-fp8: populated by maybe_update_config\n"
|
||||||
|
" self.fp8_config = None\n"
|
||||||
|
" self.fp8_layers = set()\n"
|
||||||
|
)
|
||||||
|
assert src.count(a1) == 1, f"anchor1 count={src.count(a1)}"
|
||||||
|
src = src.replace(a1, b1)
|
||||||
|
|
||||||
|
# 2. apply_vllm_mapper: remap fp8_layers (HF->vLLM names) after extra_config remap
|
||||||
|
a2 = (
|
||||||
|
" if self.extra_config is not None:\n"
|
||||||
|
" self.extra_config = hf_to_vllm_mapper.apply_dict(self.extra_config)\n"
|
||||||
|
)
|
||||||
|
b2 = a2 + (
|
||||||
|
" if self.fp8_layers: # spark-dflash-hybrid-fp8\n"
|
||||||
|
" self.fp8_layers = set(\n"
|
||||||
|
" hf_to_vllm_mapper.apply_list(list(self.fp8_layers))\n"
|
||||||
|
" )\n"
|
||||||
|
)
|
||||||
|
assert src.count(a2) == 1, f"anchor2 count={src.count(a2)}"
|
||||||
|
src = src.replace(a2, b2)
|
||||||
|
|
||||||
|
# 3. insert maybe_update_config + _is_layer_fp8 before apply_awq_quant_layer
|
||||||
|
a3 = ' def apply_awq_quant_layer(self, layer, prefix: str, backend: str = "auto"):\n'
|
||||||
|
methods = ''' def maybe_update_config( # spark-dflash-hybrid-fp8
|
||||||
|
self,
|
||||||
|
model_name: str,
|
||||||
|
hf_config=None,
|
||||||
|
revision: str | None = None,
|
||||||
|
):
|
||||||
|
"""Detect FP8 dense layers in a hybrid INT4+FP8 checkpoint."""
|
||||||
|
import torch as _torch
|
||||||
|
from safetensors.torch import _TYPES as _SF
|
||||||
|
from vllm.transformers_utils.config import get_safetensors_params_metadata
|
||||||
|
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||||
|
metadata = get_safetensors_params_metadata(model_name, revision=revision)
|
||||||
|
fp8_weights = {}
|
||||||
|
for pn, info in metadata.items():
|
||||||
|
ds = info.get("dtype", None)
|
||||||
|
if ds is None:
|
||||||
|
continue
|
||||||
|
if _SF.get(ds) == _torch.float8_e4m3fn and pn.endswith(".weight"):
|
||||||
|
sn = pn.replace(".weight", ".weight_scale_inv")
|
||||||
|
if sn in metadata:
|
||||||
|
fp8_weights[pn] = info
|
||||||
|
if not fp8_weights:
|
||||||
|
logger.info("spark-dflash-hybrid-fp8: no FP8 dense layers detected")
|
||||||
|
return
|
||||||
|
block_size = None
|
||||||
|
for pn, info in fp8_weights.items():
|
||||||
|
sn = pn.replace(".weight", ".weight_scale_inv")
|
||||||
|
ws = info.get("shape", [])
|
||||||
|
ss = metadata[sn].get("shape", [])
|
||||||
|
if len(ws) == 2 and len(ss) == 2:
|
||||||
|
block_size = [ws[0] // ss[0], ws[1] // ss[1]]
|
||||||
|
break
|
||||||
|
if block_size is None:
|
||||||
|
return
|
||||||
|
self.fp8_config = Fp8Config(
|
||||||
|
is_checkpoint_fp8_serialized=True,
|
||||||
|
activation_scheme="dynamic",
|
||||||
|
weight_block_size=block_size,
|
||||||
|
)
|
||||||
|
self.fp8_layers = {n.rsplit(".weight", 1)[0] for n in fp8_weights}
|
||||||
|
_sample = sorted(self.fp8_layers)[:3]
|
||||||
|
logger.info(
|
||||||
|
"spark-dflash-hybrid-fp8: detected %d FP8 dense layers "
|
||||||
|
"(block_size=%s) e.g. %s",
|
||||||
|
len(self.fp8_layers), block_size, _sample,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _is_layer_fp8(self, prefix: str) -> bool: # spark-dflash-hybrid-fp8
|
||||||
|
if not self.fp8_layers:
|
||||||
|
return False
|
||||||
|
if prefix in self.fp8_layers:
|
||||||
|
return True
|
||||||
|
fused = getattr(self, "packed_modules_mapping", {})
|
||||||
|
proj = prefix.split(".")[-1]
|
||||||
|
if proj in fused:
|
||||||
|
shards = [prefix.replace(proj, s) for s in fused[proj]]
|
||||||
|
return all(
|
||||||
|
any(fl in sp for fl in self.fp8_layers) for sp in shards
|
||||||
|
)
|
||||||
|
return any(fl in prefix for fl in self.fp8_layers)
|
||||||
|
|
||||||
|
'''
|
||||||
|
assert src.count(a3) == 1, f"anchor3 count={src.count(a3)}"
|
||||||
|
src = src.replace(a3, methods + a3)
|
||||||
|
|
||||||
|
# 4. FP8 dispatch in the not-quantized blocks (awq/gptq/xpu/cpu are byte-identical;
|
||||||
|
# guard is inert unless fp8_config set, so patching all is safe)
|
||||||
|
a4 = (
|
||||||
|
" if not self.check_quantized(weight_bits):\n"
|
||||||
|
" if isinstance(layer, (LinearBase, ParallelLMHead)):\n"
|
||||||
|
" return UnquantizedLinearMethod()\n"
|
||||||
|
" else:\n"
|
||||||
|
" return None\n"
|
||||||
|
)
|
||||||
|
b4 = (
|
||||||
|
" if not self.check_quantized(weight_bits):\n"
|
||||||
|
" if self.fp8_config and self._is_layer_fp8(prefix): # spark-dflash-hybrid-fp8\n"
|
||||||
|
" from vllm.model_executor.layers.quantization.fp8 import (\n"
|
||||||
|
" Fp8LinearMethod,\n"
|
||||||
|
" )\n"
|
||||||
|
" return Fp8LinearMethod(self.fp8_config)\n"
|
||||||
|
" if isinstance(layer, (LinearBase, ParallelLMHead)):\n"
|
||||||
|
" return UnquantizedLinearMethod()\n"
|
||||||
|
" else:\n"
|
||||||
|
" return None\n"
|
||||||
|
)
|
||||||
|
n4 = src.count(a4)
|
||||||
|
assert n4 >= 2, f"anchor4 count={n4}"
|
||||||
|
src = src.replace(a4, b4)
|
||||||
|
|
||||||
|
# 5. FP8 dispatch in get_quant_method's extra_config (bits>=16) override
|
||||||
|
a5 = (
|
||||||
|
' ) and self.extra_config[layer_name].get("bits", 16) >= 16:\n'
|
||||||
|
" return UnquantizedLinearMethod()\n"
|
||||||
|
)
|
||||||
|
b5 = (
|
||||||
|
' ) and self.extra_config[layer_name].get("bits", 16) >= 16:\n'
|
||||||
|
" if self.fp8_config and self._is_layer_fp8(prefix): # spark-dflash-hybrid-fp8\n"
|
||||||
|
" from vllm.model_executor.layers.quantization.fp8 import (\n"
|
||||||
|
" Fp8LinearMethod,\n"
|
||||||
|
" )\n"
|
||||||
|
" return Fp8LinearMethod(self.fp8_config)\n"
|
||||||
|
" return UnquantizedLinearMethod()\n"
|
||||||
|
)
|
||||||
|
assert src.count(a5) == 1, f"anchor5 count={src.count(a5)}"
|
||||||
|
src = src.replace(a5, b5)
|
||||||
|
|
||||||
|
open(path, "w").write(src)
|
||||||
|
print(f"[patch_inc_hybrid] applied {SENT} to {path} (not-quant blocks x{n4})")
|
||||||
@@ -0,0 +1,124 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""INT8 W8A16 lm-head v3 for AEON vLLM 0.23 (sm121). Replaces ONLY the
|
||||||
|
`lm_head.quant_method.apply(...)` call inside LogitsProcessor._get_logits with a
|
||||||
|
batched int8 GEMV (one kernel launch for any batch), leaving the existing TP
|
||||||
|
gather + org_vocab_size trim untouched.
|
||||||
|
|
||||||
|
Fixes vs the broken v2 port:
|
||||||
|
* KEEPS the bf16 lm_head weight (DFlash drafter shares it) — does NOT zero it.
|
||||||
|
Trades the memory saving for correctness; the speed win is the int8 read in
|
||||||
|
_get_logits, independent of keeping bf16 around.
|
||||||
|
* Single BATCHED kernel (dot-based, pad B->16) for ALL B — no per-row Python
|
||||||
|
loop (the v2 B>4 loop was what made spec decode SLOWER).
|
||||||
|
* Fixed proven config (N128/K128/w4/s3, ~227 GB/s, argmax-exact vs bf16 on the
|
||||||
|
real [248320,3072] shape) — no autotune (avoids sm121 bad-config miscompiles).
|
||||||
|
|
||||||
|
Sentinel DGX_SPARK_INT8_LMHEAD_V3. Verified standalone: int8 3.35ms vs bf16 8.8ms
|
||||||
|
(B=1) / 6.5ms (B=13); maxerr 3e-4 < quant floor, argmax 100%.
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
TARGET = "/usr/local/lib/python3.12/site-packages/vllm/model_executor/layers/logits_processor.py"
|
||||||
|
|
||||||
|
ANCHOR = (
|
||||||
|
" # Get the logits for the next tokens.\n"
|
||||||
|
" logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)\n"
|
||||||
|
)
|
||||||
|
REPLACE = (
|
||||||
|
" # DGX_SPARK_INT8_LMHEAD_V3: int8 w8a16 GEMV for the huge vocab projection\n"
|
||||||
|
" logits = _spark_int8_lmhead_apply(self, lm_head, hidden_states, embedding_bias)\n"
|
||||||
|
)
|
||||||
|
|
||||||
|
MODULE_CODE = '''
|
||||||
|
|
||||||
|
# ===================== DGX_SPARK_INT8_LMHEAD_V3 =====================
|
||||||
|
import triton as _spark_triton
|
||||||
|
import triton.language as _spark_tl
|
||||||
|
|
||||||
|
|
||||||
|
@_spark_triton.jit
|
||||||
|
def _spark_k_int8(x_ptr, w_ptr, s_ptr, o_ptr, B, N, K,
|
||||||
|
sxb, sxk, swn, swk, sob, son,
|
||||||
|
BLOCK_B: _spark_tl.constexpr, BLOCK_N: _spark_tl.constexpr,
|
||||||
|
BLOCK_K: _spark_tl.constexpr):
|
||||||
|
pid_n = _spark_tl.program_id(0)
|
||||||
|
offs_b = _spark_tl.arange(0, BLOCK_B)
|
||||||
|
offs_n = pid_n * BLOCK_N + _spark_tl.arange(0, BLOCK_N)
|
||||||
|
offs_k = _spark_tl.arange(0, BLOCK_K)
|
||||||
|
x_ptrs = x_ptr + offs_b[:, None] * sxb + offs_k[None, :] * sxk
|
||||||
|
w_ptrs = w_ptr + offs_n[:, None] * swn + offs_k[None, :] * swk
|
||||||
|
acc = _spark_tl.zeros((BLOCK_B, BLOCK_N), dtype=_spark_tl.float32)
|
||||||
|
for k in range(0, K, BLOCK_K):
|
||||||
|
km = (offs_k[None, :] + k) < K
|
||||||
|
x = _spark_tl.load(x_ptrs, mask=(offs_b[:, None] < B) & km, other=0.0).to(_spark_tl.float16)
|
||||||
|
w = _spark_tl.load(w_ptrs, mask=(offs_n[:, None] < N) & km, other=0).to(_spark_tl.float16)
|
||||||
|
acc += _spark_tl.dot(x, w.T)
|
||||||
|
x_ptrs += BLOCK_K * sxk
|
||||||
|
w_ptrs += BLOCK_K * swk
|
||||||
|
s = _spark_tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0).to(_spark_tl.float32)
|
||||||
|
acc = acc * s[None, :]
|
||||||
|
o_ptrs = o_ptr + offs_b[:, None] * sob + offs_n[None, :] * son
|
||||||
|
_spark_tl.store(o_ptrs, acc, mask=(offs_b[:, None] < B) & (offs_n[None, :] < N))
|
||||||
|
|
||||||
|
|
||||||
|
def _spark_int8_gemm(hidden, w_int8, w_scale):
|
||||||
|
import torch
|
||||||
|
N, K = w_int8.shape
|
||||||
|
x = hidden.reshape(-1, K)
|
||||||
|
B = x.shape[0]
|
||||||
|
BLOCK_B = max(16, _spark_triton.next_power_of_2(B))
|
||||||
|
out = torch.empty(B, N, dtype=torch.float32, device=x.device)
|
||||||
|
xf = x.to(torch.float16)
|
||||||
|
grid = ((N + 127) // 128,)
|
||||||
|
_spark_k_int8[grid](xf, w_int8, w_scale, out, B, N, K,
|
||||||
|
xf.stride(0), xf.stride(1), w_int8.stride(0), w_int8.stride(1),
|
||||||
|
out.stride(0), out.stride(1),
|
||||||
|
BLOCK_B=BLOCK_B, BLOCK_N=128, BLOCK_K=128,
|
||||||
|
num_warps=4, num_stages=3)
|
||||||
|
return out.reshape(hidden.shape[:-1] + (N,))
|
||||||
|
|
||||||
|
|
||||||
|
def _spark_int8_lmhead_apply(self, lm_head, hidden_states, embedding_bias):
|
||||||
|
import sys
|
||||||
|
import torch
|
||||||
|
if not getattr(lm_head, "_spark_int8_ready", None) is True and \\
|
||||||
|
not getattr(lm_head, "_spark_int8_disabled", False):
|
||||||
|
w = getattr(lm_head, "weight", None)
|
||||||
|
if (w is not None and w.dtype in (torch.bfloat16, torch.float16)
|
||||||
|
and w.dim() == 2 and w.shape[0] > 100000):
|
||||||
|
with torch.no_grad():
|
||||||
|
scales = (w.float().abs().amax(dim=1) / 127.0).clamp(min=1e-12)
|
||||||
|
w_int8 = (w.float() / scales.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)
|
||||||
|
lm_head._spark_w_int8 = w_int8.contiguous()
|
||||||
|
lm_head._spark_w_scale = scales.to(torch.float16)
|
||||||
|
lm_head._spark_int8_ready = True
|
||||||
|
print("DGX_SPARK_INT8_LMHEAD_V3: lm_head -> int8 (%s), bf16 kept for shared drafter"
|
||||||
|
% (list(w_int8.shape),), file=sys.stderr, flush=True)
|
||||||
|
else:
|
||||||
|
lm_head._spark_int8_disabled = True
|
||||||
|
if getattr(lm_head, "_spark_int8_ready", False) and embedding_bias is None:
|
||||||
|
return _spark_int8_gemm(hidden_states, lm_head._spark_w_int8, lm_head._spark_w_scale)
|
||||||
|
return lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)
|
||||||
|
# =================== end DGX_SPARK_INT8_LMHEAD_V3 ===================
|
||||||
|
'''
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
if not os.path.exists(TARGET):
|
||||||
|
print(f"FAIL: {TARGET} not found"); sys.exit(1)
|
||||||
|
src = open(TARGET).read()
|
||||||
|
if "DGX_SPARK_INT8_LMHEAD_V3" in src:
|
||||||
|
print("SKIP: int8 lm-head v3 already applied"); return
|
||||||
|
if ANCHOR not in src:
|
||||||
|
print("FAIL: _get_logits apply-anchor not found"); sys.exit(1)
|
||||||
|
if src.count(ANCHOR) != 1:
|
||||||
|
print(f"FAIL: anchor count={src.count(ANCHOR)} (expected 1)"); sys.exit(1)
|
||||||
|
src = src.replace(ANCHOR, REPLACE)
|
||||||
|
src = src + MODULE_CODE
|
||||||
|
open(TARGET, "w").write(src)
|
||||||
|
print("OK: int8 lm-head v3 applied")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,55 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""spark-dflash-unify (CORRECTED): let the DFlash drafter's attention KV spec unify
|
||||||
|
with the hybrid GDN+mamba target's *padded* page size WITHOUT corrupting the drafter.
|
||||||
|
|
||||||
|
Root cause (AEON 0.23.0+aeon.sm121a.dflash): with --mamba-block-size set, vLLM's own
|
||||||
|
hybrid alignment makes target mamba page == target attention page by PADDING the mamba
|
||||||
|
page (e.g. "+0.54%"). That padded value becomes max_page_size. The drafter's attention
|
||||||
|
page is smaller; unify_kv_cache_spec_page_size scales its block_size by
|
||||||
|
ratio = max_page_size // layer_page_size, but because max_page_size is a *padded* (non
|
||||||
|
block-size-linear) number, the scaled page lands just under it and
|
||||||
|
`assert new_spec.page_size_bytes == max_page_size` fires.
|
||||||
|
|
||||||
|
The OLD spark-pad-unify patch "fixed" this by `replace(layer_spec,
|
||||||
|
page_size_padded=max_page_size)` -- but it DROPPED the block_size scaling, leaving the
|
||||||
|
drafter at block_size=16 behind a max-sized physical page. That mis-strided the drafter
|
||||||
|
KV and pinned mean accept length at ~1.47 (garbage drafts), so DFlash net-lost.
|
||||||
|
|
||||||
|
This corrected patch mirrors what get_kv_cache_groups already does for
|
||||||
|
HiddenStateCacheSpec layers: keep the SCALED block_size AND pad the <1% remainder.
|
||||||
|
Strides stay correct (block_size is properly scaled); only the page tail is padded.
|
||||||
|
"""
|
||||||
|
import pathlib, sys
|
||||||
|
|
||||||
|
P = pathlib.Path("/usr/local/lib/python3.12/site-packages/vllm/v1/core/kv_cache_utils.py")
|
||||||
|
src = P.read_text()
|
||||||
|
|
||||||
|
if "spark-dflash-unify" in src:
|
||||||
|
print("[patch_unify2] already patched", flush=True)
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
OLD = ''' new_spec = replace(layer_spec, block_size=new_block_size)
|
||||||
|
assert new_spec.page_size_bytes == max_page_size
|
||||||
|
new_kv_cache_spec[layer_name] = new_spec'''
|
||||||
|
|
||||||
|
NEW = ''' new_spec = replace(layer_spec, block_size=new_block_size)
|
||||||
|
if new_spec.page_size_bytes != max_page_size:
|
||||||
|
# spark-dflash-unify: max_page_size is a *padded* hybrid page; the
|
||||||
|
# scaled attention page lands just under it. Pad the remainder while
|
||||||
|
# KEEPING the scaled block_size, exactly as get_kv_cache_groups does
|
||||||
|
# for HiddenStateCacheSpec. (The old patch dropped the scaling ->
|
||||||
|
# mis-strided the DFlash drafter -> accept len stuck ~1.47.)
|
||||||
|
new_spec = replace(
|
||||||
|
layer_spec,
|
||||||
|
block_size=new_block_size,
|
||||||
|
page_size_padded=max_page_size,
|
||||||
|
)
|
||||||
|
assert new_spec.page_size_bytes == max_page_size
|
||||||
|
new_kv_cache_spec[layer_name] = new_spec'''
|
||||||
|
|
||||||
|
if OLD not in src:
|
||||||
|
print("[patch_unify2] ERROR: target block not found — vLLM source differs", flush=True)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
P.write_text(src.replace(OLD, NEW))
|
||||||
|
print("[patch_unify2] patched OK (scaled-block + pad-remainder)", flush=True)
|
||||||
Executable
+59
@@ -0,0 +1,59 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
# serve.sh — runs INSIDE the sm121 vLLM container (mounted at /host). Applies the
|
||||||
|
# runtime monkeypatches, then `vllm serve`s the Qwen3.5-122B-A10B INT4 target with
|
||||||
|
# the DFlash drafter. Driven by install.sh; can also be run by hand.
|
||||||
|
#
|
||||||
|
# args: $1 = num_speculative_tokens (0 = no-spec baseline)
|
||||||
|
# $2 = target attention backend (flash_attn | FLASHINFER)
|
||||||
|
# env: MODEL target path/repo (default Intel INT4; /model for hybrid)
|
||||||
|
# INC_HYBRID=1 apply the hybrid INT4+FP8 dense-expert dispatch patch
|
||||||
|
# INT8_LMHEAD_V3=1 apply the int8 lm-head GEMV patch
|
||||||
|
# MAX_MODEL_LEN GPU_MEM PORT
|
||||||
|
#
|
||||||
|
# Stack rationale: the DFlash drafter is non-causal -> needs FLASH_ATTN (FA2). The
|
||||||
|
# hybrid GDN+mamba+MoE target's KV page geometry won't absorb the drafter's
|
||||||
|
# attention spec without patch_unify2 (scale-block unify) + prefix-caching OFF
|
||||||
|
# (NoPrefixCache coordinator, dodges the hash assert). See docs/FINDINGS.md.
|
||||||
|
set -euo pipefail
|
||||||
|
NSPEC="${1:-12}"
|
||||||
|
BACKEND="${2:-flash_attn}"
|
||||||
|
MODEL="${MODEL:-Intel/Qwen3.5-122B-A10B-int4-AutoRound}"
|
||||||
|
DRAFT="${DRAFT:-z-lab/Qwen3.5-122B-A10B-DFlash}"
|
||||||
|
MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
|
||||||
|
GPU_MEM="${GPU_MEM:-0.8}"
|
||||||
|
PORT="${PORT:-8000}"
|
||||||
|
|
||||||
|
# FLA sm121 big-tile shmem fix (prefill/TTFT only on sm121; harmless, free).
|
||||||
|
echo "[serve] FLA sm121 big-tile shmem patch"
|
||||||
|
python3 /host/patch_fla_shmem.py || true
|
||||||
|
|
||||||
|
if [ "${INC_HYBRID:-0}" = "1" ]; then
|
||||||
|
echo "[serve] hybrid INT4+FP8 dispatch patch (inc.py)"
|
||||||
|
python3 /host/patch_inc_hybrid.py
|
||||||
|
fi
|
||||||
|
if [ "${INT8_LMHEAD_V3:-0}" = "1" ]; then
|
||||||
|
echo "[serve] int8 lm-head v3 patch (batched w8a16 GEMV)"
|
||||||
|
python3 /host/patch_int8_lmhead_v3.py
|
||||||
|
fi
|
||||||
|
|
||||||
|
if [ "$NSPEC" = "0" ]; then
|
||||||
|
SPEC_ARG=()
|
||||||
|
echo "[serve] NO-SPEC baseline (identical flags, prefix-off)"
|
||||||
|
else
|
||||||
|
SPEC_ARG=(--speculative-config "{\"method\":\"dflash\",\"model\":\"$DRAFT\",\"num_speculative_tokens\":$NSPEC,\"attention_backend\":\"FLASH_ATTN\"}")
|
||||||
|
echo "[serve] DFlash n=$NSPEC, target-backend=$BACKEND, drafter=FLASH_ATTN, model=$MODEL"
|
||||||
|
fi
|
||||||
|
python3 /host/patch_unify2.py || { [ "$NSPEC" = "0" ] && true; }
|
||||||
|
|
||||||
|
exec vllm serve "$MODEL" \
|
||||||
|
--served-model-name qwen \
|
||||||
|
--host 0.0.0.0 --port "$PORT" \
|
||||||
|
--max-model-len "$MAX_MODEL_LEN" \
|
||||||
|
--max-num-seqs 16 \
|
||||||
|
--max-num-batched-tokens "$MAX_MODEL_LEN" \
|
||||||
|
--gpu-memory-utilization "$GPU_MEM" \
|
||||||
|
--no-enable-prefix-caching \
|
||||||
|
--enable-chunked-prefill \
|
||||||
|
--trust-remote-code \
|
||||||
|
--attention-backend "$BACKEND" \
|
||||||
|
"${SPEC_ARG[@]}"
|
||||||
@@ -0,0 +1,77 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Faithful reproduction of albond's bench_qwen35.sh methodology so our DFlash/MTP
|
||||||
|
numbers are directly comparable to his reported 51.58 tok/s.
|
||||||
|
|
||||||
|
His method (verbatim): non-streaming /v1/chat/completions, time the WHOLE request,
|
||||||
|
tok/s = completion_tokens / wall_time (INCLUDES prefill + overhead = END-TO-END).
|
||||||
|
5 prompts (Q&A 256 / Code 512 / JSON 1024 / Math 64 / LongCode 2048), temp 0, run 1
|
||||||
|
discarded as JIT warmup. We add: decode-only tok/s isn't measured here on purpose
|
||||||
|
(his isn't either) + an aggregate spec-accept-len from /metrics deltas for context.
|
||||||
|
"""
|
||||||
|
import json, sys, time, urllib.request, statistics
|
||||||
|
|
||||||
|
BASE = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:8000"
|
||||||
|
MODEL = "qwen"
|
||||||
|
PROMPTS = [
|
||||||
|
("Q&A", "What are the main differences between TCP and UDP? Be concise.", 256),
|
||||||
|
("Code", "Write a Python function that implements binary search on a sorted list. Include type hints and docstring.", 512),
|
||||||
|
("JSON", "Generate a JSON array of 10 fictional employees with fields: name, age, department, salary, email, skills (array of 3). Output ONLY valid JSON, no explanation.", 1024),
|
||||||
|
("Math", "What is 7823 * 4519? Show only the answer.", 64),
|
||||||
|
("LongCode", "Write a complete Python implementation of a red-black tree with insert, delete, search, and in-order traversal. Include all rotation methods.", 2048),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def scrape():
|
||||||
|
acc = dr = 0.0
|
||||||
|
try:
|
||||||
|
txt = urllib.request.urlopen(BASE + "/metrics", timeout=10).read().decode()
|
||||||
|
for ln in txt.splitlines():
|
||||||
|
if ln.startswith("#") or not ln.split():
|
||||||
|
continue
|
||||||
|
v = float(ln.split()[-1])
|
||||||
|
if "spec_decode_num_accepted_tokens_total" in ln:
|
||||||
|
acc += v
|
||||||
|
elif "spec_decode_num_drafts_total" in ln:
|
||||||
|
dr += v
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return acc, dr
|
||||||
|
|
||||||
|
|
||||||
|
def chat_e2e(prompt, max_tokens):
|
||||||
|
body = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": prompt}],
|
||||||
|
"max_tokens": max_tokens, "temperature": 0.0}).encode()
|
||||||
|
req = urllib.request.Request(BASE + "/v1/chat/completions", data=body,
|
||||||
|
headers={"Content-Type": "application/json", "Authorization": "Bearer x"})
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
r = json.loads(urllib.request.urlopen(req, timeout=600).read())
|
||||||
|
elapsed = time.perf_counter() - t0
|
||||||
|
ct = r["usage"]["completion_tokens"]
|
||||||
|
return ct, elapsed, ct / elapsed if elapsed > 0 else 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
label = sys.argv[2] if len(sys.argv) > 2 else "server"
|
||||||
|
print(f"=== albond-method e2e bench :: {label} ===")
|
||||||
|
a0, d0 = scrape()
|
||||||
|
run2 = {}
|
||||||
|
for run in (1, 2):
|
||||||
|
tag = "WARMUP(discard)" if run == 1 else "RUN2"
|
||||||
|
for name, prompt, mt in PROMPTS:
|
||||||
|
try:
|
||||||
|
ct, el, tps = chat_e2e(prompt, mt)
|
||||||
|
except Exception as e:
|
||||||
|
print(f" [{name}] FAILED: {type(e).__name__}: {e}"); continue
|
||||||
|
if run == 2:
|
||||||
|
run2[name] = tps
|
||||||
|
print(f" {tag:16s} [{name:8s}] {ct:4d} tok in {el:6.2f}s = {tps:6.1f} tok/s (e2e)")
|
||||||
|
a1, d1 = scrape()
|
||||||
|
print(f"\n{label} RUN2 e2e tok/s: " + " ".join(f"{k}={v:.1f}" for k, v in run2.items()))
|
||||||
|
if run2:
|
||||||
|
print(f" cross-prompt mean (RUN2) = {statistics.mean(run2.values()):.1f} tok/s (albond reports 51.58)")
|
||||||
|
if d1 - d0 > 0:
|
||||||
|
print(f" aggregate spec accept length over bench = {1 + (a1-a0)/(d1-d0):.2f}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,138 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Single-stream (c=1) decode-rate benchmark for an OpenAI-compatible vLLM server.
|
||||||
|
|
||||||
|
Measures *decode* tok/s (excludes TTFT/prefill) over N sequential requests, and
|
||||||
|
optionally scrapes vLLM /metrics for speculative-decode acceptance length.
|
||||||
|
|
||||||
|
stdlib only — runs anywhere with python3, no pip installs.
|
||||||
|
|
||||||
|
python3 bench_decode.py --base-url http://127.0.0.1:8000 \
|
||||||
|
--model Intel/Qwen3.5-122B-A10B-int4-AutoRound \
|
||||||
|
--max-tokens 256 --runs 5 --label "int4 baseline"
|
||||||
|
"""
|
||||||
|
import argparse, json, statistics, sys, time, urllib.request, urllib.error
|
||||||
|
|
||||||
|
PROMPT = ("You are a careful writer. Write a detailed, continuous explanation of how "
|
||||||
|
"a modern mixture-of-experts transformer performs autoregressive decoding, "
|
||||||
|
"covering routing, KV cache, and memory bandwidth. Begin now:\n\n")
|
||||||
|
|
||||||
|
|
||||||
|
def post_stream(base_url, model, prompt, max_tokens, timeout):
|
||||||
|
"""Stream /v1/completions; return (completion_tokens, t_first, t_last)."""
|
||||||
|
body = json.dumps({
|
||||||
|
"model": model, "prompt": prompt, "max_tokens": max_tokens,
|
||||||
|
"temperature": 0.0, "stream": True,
|
||||||
|
"stream_options": {"include_usage": True},
|
||||||
|
# force the full token budget so we measure steady-state decode
|
||||||
|
"ignore_eos": True, "min_tokens": max_tokens,
|
||||||
|
}).encode()
|
||||||
|
req = urllib.request.Request(base_url.rstrip("/") + "/v1/completions",
|
||||||
|
data=body, headers={"Content-Type": "application/json",
|
||||||
|
"Authorization": "Bearer x"})
|
||||||
|
t_first = t_last = None
|
||||||
|
completion_tokens = 0
|
||||||
|
chunks = 0
|
||||||
|
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||||||
|
for raw in r:
|
||||||
|
line = raw.decode("utf-8", "replace").strip()
|
||||||
|
if not line.startswith("data:"):
|
||||||
|
continue
|
||||||
|
data = line[5:].strip()
|
||||||
|
if data == "[DONE]":
|
||||||
|
break
|
||||||
|
try:
|
||||||
|
obj = json.loads(data)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
continue
|
||||||
|
usage = obj.get("usage")
|
||||||
|
if usage and usage.get("completion_tokens"):
|
||||||
|
completion_tokens = usage["completion_tokens"]
|
||||||
|
choices = obj.get("choices") or []
|
||||||
|
if choices and choices[0].get("text"):
|
||||||
|
now = time.perf_counter()
|
||||||
|
if t_first is None:
|
||||||
|
t_first = now
|
||||||
|
t_last = now
|
||||||
|
chunks += 1
|
||||||
|
if completion_tokens == 0:
|
||||||
|
completion_tokens = chunks # fallback: 1 chunk ~= 1 token
|
||||||
|
return completion_tokens, t_first, t_last
|
||||||
|
|
||||||
|
|
||||||
|
def scrape_metrics(base_url, timeout=10):
|
||||||
|
"""Return dict of spec-decode counters from vLLM /metrics, if present."""
|
||||||
|
keys = ("num_accepted_tokens", "num_draft_tokens", "num_drafts",
|
||||||
|
"accepted_tokens", "draft_tokens")
|
||||||
|
out = {}
|
||||||
|
try:
|
||||||
|
with urllib.request.urlopen(base_url.rstrip("/") + "/metrics", timeout=timeout) as r:
|
||||||
|
for line in r.read().decode("utf-8", "replace").splitlines():
|
||||||
|
if line.startswith("#"):
|
||||||
|
continue
|
||||||
|
if "spec_decode" in line or "speculat" in line:
|
||||||
|
name, _, val = line.partition(" ")
|
||||||
|
try:
|
||||||
|
out[name] = out.get(name, 0.0) + float(val)
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
except (urllib.error.URLError, OSError):
|
||||||
|
pass
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
ap = argparse.ArgumentParser()
|
||||||
|
ap.add_argument("--base-url", default="http://127.0.0.1:8000")
|
||||||
|
ap.add_argument("--model", required=True)
|
||||||
|
ap.add_argument("--max-tokens", type=int, default=256)
|
||||||
|
ap.add_argument("--runs", type=int, default=5)
|
||||||
|
ap.add_argument("--warmup", type=int, default=1)
|
||||||
|
ap.add_argument("--timeout", type=float, default=600)
|
||||||
|
ap.add_argument("--label", default="")
|
||||||
|
ap.add_argument("--prompt", default=PROMPT)
|
||||||
|
args = ap.parse_args()
|
||||||
|
|
||||||
|
m_before = scrape_metrics(args.base_url)
|
||||||
|
for _ in range(args.warmup):
|
||||||
|
try:
|
||||||
|
post_stream(args.base_url, args.model, args.prompt, 32, args.timeout)
|
||||||
|
except Exception as e:
|
||||||
|
print(f"warmup failed: {e}", file=sys.stderr); sys.exit(2)
|
||||||
|
|
||||||
|
decode_tps, ttfts, e2e_tps = [], [], []
|
||||||
|
for i in range(args.runs):
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
toks, tf, tl = post_stream(args.base_url, args.model, args.prompt, args.max_tokens, args.timeout)
|
||||||
|
t1 = time.perf_counter()
|
||||||
|
if not toks or tf is None or tl is None or tl <= tf:
|
||||||
|
print(f" run {i}: degenerate (toks={toks})", file=sys.stderr); continue
|
||||||
|
dec = (toks - 1) / (tl - tf)
|
||||||
|
decode_tps.append(dec); ttfts.append((tf - t0) * 1000); e2e_tps.append(toks / (t1 - t0))
|
||||||
|
print(f" run {i}: {toks} tok decode={dec:6.1f} tok/s ttft={ (tf-t0)*1000:6.0f} ms")
|
||||||
|
m_after = scrape_metrics(args.base_url)
|
||||||
|
|
||||||
|
if not decode_tps:
|
||||||
|
print("no successful runs", file=sys.stderr); sys.exit(1)
|
||||||
|
print(f"\n=== {args.label or args.model} ===")
|
||||||
|
print(f"decode tok/s : median {statistics.median(decode_tps):.1f} "
|
||||||
|
f"mean {statistics.mean(decode_tps):.1f} min {min(decode_tps):.1f} max {max(decode_tps):.1f}")
|
||||||
|
print(f"ttft ms : median {statistics.median(ttfts):.0f}")
|
||||||
|
print(f"e2e tok/s : median {statistics.median(e2e_tps):.1f}")
|
||||||
|
|
||||||
|
# spec-decode acceptance length, if counters moved
|
||||||
|
def delta(k):
|
||||||
|
return m_after.get(k, 0.0) - m_before.get(k, 0.0)
|
||||||
|
acc = next((delta(k) for k in m_after if "accepted" in k), 0.0)
|
||||||
|
drafts = next((delta(k) for k in m_after if "num_drafts" in k or ("draft" in k and "tokens" not in k)), 0.0)
|
||||||
|
dtoks = next((delta(k) for k in m_after if "draft_tokens" in k or "num_draft_tokens" in k), 0.0)
|
||||||
|
if acc or dtoks:
|
||||||
|
al = (acc / drafts) if drafts else float("nan")
|
||||||
|
rate = (acc / dtoks) if dtoks else float("nan")
|
||||||
|
print(f"spec accept : +{acc:.0f} accepted, +{dtoks:.0f} drafted, "
|
||||||
|
f"mean accept len ~{al:.2f}, accept rate ~{rate:.1%}")
|
||||||
|
else:
|
||||||
|
print("spec accept : (no spec-decode counters — baseline/no drafter)")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,168 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Real-world DFlash/MTP acceptance bench on the actual Hermes agent workload.
|
||||||
|
|
||||||
|
Reconstructs real conversation contexts from ~/.hermes/state.db and has the loaded
|
||||||
|
server regenerate the *next assistant turn* (so the model that produced the history is
|
||||||
|
irrelevant — only the realistic context matters). Measures mean acceptance length
|
||||||
|
(authoritative vLLM /metrics deltas) + decode tok/s, aggregated over N real turns.
|
||||||
|
Privacy: runs entirely on the box; prints only metrics, not conversation content.
|
||||||
|
"""
|
||||||
|
import argparse, json, sqlite3, sys, time, urllib.request, urllib.error
|
||||||
|
|
||||||
|
DB = "/home/ent/.hermes/state.db"
|
||||||
|
|
||||||
|
|
||||||
|
def scrape(base):
|
||||||
|
out = {"acc": 0.0, "drafts": 0.0, "dtoks": 0.0}
|
||||||
|
try:
|
||||||
|
txt = urllib.request.urlopen(base.rstrip("/") + "/metrics", timeout=10).read().decode("utf-8", "replace")
|
||||||
|
except Exception:
|
||||||
|
return out
|
||||||
|
for ln in txt.splitlines():
|
||||||
|
if ln.startswith("#"):
|
||||||
|
continue
|
||||||
|
v = ln.split()[-1] if ln.split() else "0"
|
||||||
|
try:
|
||||||
|
val = float(v)
|
||||||
|
except ValueError:
|
||||||
|
continue
|
||||||
|
if "spec_decode_num_accepted_tokens_total" in ln:
|
||||||
|
out["acc"] += val
|
||||||
|
elif "spec_decode_num_drafts_total" in ln:
|
||||||
|
out["drafts"] += val
|
||||||
|
elif "spec_decode_num_draft_tokens_total" in ln:
|
||||||
|
out["dtoks"] += val
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def build_messages(cur, session_id, system_prompt, char_budget=24000):
|
||||||
|
rows = list(cur.execute(
|
||||||
|
"SELECT role, content, tool_calls, tool_call_id, tool_name FROM messages "
|
||||||
|
"WHERE session_id=? ORDER BY id", (session_id,)))
|
||||||
|
# find the LAST assistant turn -> generate it; prompt = everything before it
|
||||||
|
last_asst = None
|
||||||
|
for i, r in enumerate(rows):
|
||||||
|
if r[0] == "assistant":
|
||||||
|
last_asst = i
|
||||||
|
if last_asst is None or last_asst == 0:
|
||||||
|
return None
|
||||||
|
pre = rows[:last_asst]
|
||||||
|
msgs = []
|
||||||
|
for role, content, tool_calls, tool_call_id, tool_name in pre:
|
||||||
|
content = content or ""
|
||||||
|
if role == "user":
|
||||||
|
msgs.append({"role": "user", "content": content})
|
||||||
|
elif role == "assistant":
|
||||||
|
m = {"role": "assistant", "content": content}
|
||||||
|
if tool_calls:
|
||||||
|
try:
|
||||||
|
tc = json.loads(tool_calls)
|
||||||
|
if isinstance(tc, list) and tc:
|
||||||
|
m["tool_calls"] = tc
|
||||||
|
if not content:
|
||||||
|
m["content"] = ""
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
msgs.append(m)
|
||||||
|
elif role == "tool":
|
||||||
|
msgs.append({"role": "tool", "content": content,
|
||||||
|
"tool_call_id": tool_call_id or "call_0"})
|
||||||
|
# skip session_meta
|
||||||
|
if not msgs:
|
||||||
|
return None
|
||||||
|
# truncate oldest non-system messages to fit budget
|
||||||
|
sys_msg = [{"role": "system", "content": system_prompt}] if system_prompt else []
|
||||||
|
while sum(len(json.dumps(m)) for m in msgs) > char_budget and len(msgs) > 1:
|
||||||
|
msgs.pop(0)
|
||||||
|
return sys_msg + msgs
|
||||||
|
|
||||||
|
|
||||||
|
def stream_chat(base, model, messages, max_tokens, timeout):
|
||||||
|
body = json.dumps({"model": model, "messages": messages, "max_tokens": max_tokens,
|
||||||
|
"temperature": 0.0, "stream": True,
|
||||||
|
"stream_options": {"include_usage": True}}).encode()
|
||||||
|
req = urllib.request.Request(base.rstrip("/") + "/v1/chat/completions", data=body,
|
||||||
|
headers={"Content-Type": "application/json", "Authorization": "Bearer x"})
|
||||||
|
t_first = t_last = None
|
||||||
|
toks = 0
|
||||||
|
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||||||
|
for raw in r:
|
||||||
|
ln = raw.decode("utf-8", "replace").strip()
|
||||||
|
if not ln.startswith("data:"):
|
||||||
|
continue
|
||||||
|
d = ln[5:].strip()
|
||||||
|
if d == "[DONE]":
|
||||||
|
break
|
||||||
|
try:
|
||||||
|
o = json.loads(d)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
continue
|
||||||
|
u = o.get("usage")
|
||||||
|
if u and u.get("completion_tokens"):
|
||||||
|
toks = u["completion_tokens"]
|
||||||
|
ch = o.get("choices") or []
|
||||||
|
if ch and (ch[0].get("delta") or {}).get("content"):
|
||||||
|
now = time.perf_counter()
|
||||||
|
if t_first is None:
|
||||||
|
t_first = now
|
||||||
|
t_last = now
|
||||||
|
return toks, t_first, t_last
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
ap = argparse.ArgumentParser()
|
||||||
|
ap.add_argument("--base-url", default="http://127.0.0.1:8000")
|
||||||
|
ap.add_argument("--model", default="qwen")
|
||||||
|
ap.add_argument("--max-tokens", type=int, default=200)
|
||||||
|
ap.add_argument("--n-samples", type=int, default=10)
|
||||||
|
ap.add_argument("--min-msgs", type=int, default=6)
|
||||||
|
ap.add_argument("--timeout", type=float, default=300)
|
||||||
|
ap.add_argument("--label", default="HERMES real")
|
||||||
|
args = ap.parse_args()
|
||||||
|
|
||||||
|
db = sqlite3.connect(DB)
|
||||||
|
cur = db.cursor()
|
||||||
|
sess = list(cur.execute(
|
||||||
|
"SELECT s.id, s.system_prompt FROM sessions s "
|
||||||
|
"WHERE s.message_count >= ? ORDER BY s.started_at DESC LIMIT 40", (args.min_msgs,)))
|
||||||
|
|
||||||
|
m0 = scrape(args.base_url)
|
||||||
|
decode_tps, used = [], 0
|
||||||
|
for sid, sysp in sess:
|
||||||
|
if used >= args.n_samples:
|
||||||
|
break
|
||||||
|
try:
|
||||||
|
msgs = build_messages(cur, sid, sysp)
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
if not msgs:
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
toks, tf, tl = stream_chat(args.base_url, args.model, msgs, args.max_tokens, args.timeout)
|
||||||
|
except (urllib.error.HTTPError, urllib.error.URLError, OSError) as e:
|
||||||
|
print(f" skip {sid[:24]}: {type(e).__name__}", file=sys.stderr)
|
||||||
|
continue
|
||||||
|
if toks and tf and tl and tl > tf:
|
||||||
|
dec = (toks - 1) / (tl - tf)
|
||||||
|
decode_tps.append(dec)
|
||||||
|
used += 1
|
||||||
|
print(f" sample {used}: {toks} tok decode={dec:5.1f} tok/s ({sid[:28]})")
|
||||||
|
m1 = scrape(args.base_url)
|
||||||
|
|
||||||
|
dacc = m1["acc"] - m0["acc"]
|
||||||
|
ddr = m1["drafts"] - m0["drafts"]
|
||||||
|
ddt = m1["dtoks"] - m0["dtoks"]
|
||||||
|
import statistics
|
||||||
|
print(f"\n=== {args.label} (n={used} real turns) ===")
|
||||||
|
if decode_tps:
|
||||||
|
print(f"decode tok/s : median {statistics.median(decode_tps):.1f} mean {statistics.mean(decode_tps):.1f} "
|
||||||
|
f"min {min(decode_tps):.1f} max {max(decode_tps):.1f}")
|
||||||
|
if ddr > 0:
|
||||||
|
print(f"accept length: {1 + dacc/ddr:.2f} (accepted {dacc:.0f} / drafts {ddr:.0f}; "
|
||||||
|
f"draft tokens {ddt:.0f}; accept rate {dacc/ddt:.1%})")
|
||||||
|
else:
|
||||||
|
print("accept length: (no draft activity captured)")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Executable
+38
@@ -0,0 +1,38 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
# Monitor SGLang container startup with an OOM auto-kill guard.
|
||||||
|
# Breaks on: READY, ERROR (traceback), OOM-GUARD (avail mem too low), or timeout.
|
||||||
|
set -u
|
||||||
|
NAME="${1:-sglang-dflash}"
|
||||||
|
MAX_ITERS="${2:-60}" # 60 * 20s = 20 min
|
||||||
|
FLOOR_MB="${3:-4096}" # kill if available memory drops below this
|
||||||
|
|
||||||
|
for i in $(seq 1 "$MAX_ITERS"); do
|
||||||
|
if ! docker ps --filter "name=$NAME" --format '{{.Names}}' | grep -q "$NAME"; then
|
||||||
|
echo "STATE=EXITED iter=$i"
|
||||||
|
echo "--- last log ---"; docker logs "$NAME" 2>&1 | tail -25
|
||||||
|
exit 0
|
||||||
|
fi
|
||||||
|
AVAIL=$(free -m | awk '/^Mem:/{print $7}')
|
||||||
|
if [ "$AVAIL" -lt "$FLOOR_MB" ]; then
|
||||||
|
echo "STATE=OOM-GUARD iter=$i avail_mb=$AVAIL -> killing $NAME"
|
||||||
|
docker kill "$NAME" >/dev/null 2>&1
|
||||||
|
echo "--- last log ---"; docker logs "$NAME" 2>&1 | tail -25
|
||||||
|
exit 0
|
||||||
|
fi
|
||||||
|
LOG=$(docker logs "$NAME" 2>&1)
|
||||||
|
if echo "$LOG" | grep -qiE "server is fired up|Application startup complete|The server is ready"; then
|
||||||
|
echo "STATE=READY iter=$i avail_mb=$AVAIL"
|
||||||
|
echo "--- tail ---"; echo "$LOG" | tail -20
|
||||||
|
exit 0
|
||||||
|
fi
|
||||||
|
if echo "$LOG" | grep -qE "Traceback \(most recent call last\)|CUDA out of memory|RuntimeError|AssertionError|ValueError|raise NotImplementedError"; then
|
||||||
|
echo "STATE=ERROR iter=$i avail_mb=$AVAIL"
|
||||||
|
echo "--- tail ---"; echo "$LOG" | tail -40
|
||||||
|
exit 0
|
||||||
|
fi
|
||||||
|
LAST=$(echo "$LOG" | tail -1 | cut -c1-110)
|
||||||
|
echo "iter=$i avail_mb=$AVAIL :: $LAST"
|
||||||
|
sleep 20
|
||||||
|
done
|
||||||
|
echo "STATE=TIMEOUT after $MAX_ITERS iters"
|
||||||
|
docker logs "$NAME" 2>&1 | tail -20
|
||||||
Executable
+16
@@ -0,0 +1,16 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
# Run the 4-workload acceptance/throughput bank on whichever server is up.
|
||||||
|
# $1 = container name (for the authoritative server-log accept_len). Uses the SAME
|
||||||
|
# prompts as the DFlash run so DFlash vs MTP is apples-to-apples.
|
||||||
|
CN="${1:-vllm-mtp}"
|
||||||
|
BU=http://127.0.0.1:8000; M=qwen
|
||||||
|
acc_tail(){ docker logs "$CN" 2>&1 | grep "Mean acceptance length" | tail -1 \
|
||||||
|
| sed -E "s/.*Mean acceptance length: ([0-9.]+).*Per-position acceptance rate: ([0-9., ]+), Avg.*/accept_len=\1 per-pos=[\2]/"; }
|
||||||
|
run(){ L="$1"; P="$2"; echo "########## $L ##########"
|
||||||
|
python3 ~/bench_decode.py --base-url $BU --model $M --max-tokens 256 --runs 4 --warmup 1 --label "$L" --prompt "$P" 2>&1 | grep -E "decode tok/s"
|
||||||
|
echo " server $(acc_tail)"; }
|
||||||
|
run "PROSE" "You are a careful writer. Write a long flowing essay about the history and philosophy of science, with no lists. Begin: "
|
||||||
|
run "CODE" "Write a complete Python implementation of a binary search tree class with insert, search, delete, and inorder traversal. Include docstrings and type hints. Begin:\n\nclass BSTNode:"
|
||||||
|
run "COUNTING" "Write all the whole numbers from 1 to 400, separated by commas, with no other text. Begin: 1, 2, 3, "
|
||||||
|
echo "########## HERMES (real agent turns) ##########"
|
||||||
|
python3 ~/hermes_bench.py --base-url $BU --model $M --label "HERMES" --n-samples 10 --max-tokens 200 2>&1 | grep -E "decode tok/s|accept length"
|
||||||
@@ -0,0 +1,463 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Build a hybrid GPTQ-INT4 + FP8 checkpoint for Qwen3.5-122B-A10B.
|
||||||
|
|
||||||
|
Takes MoE expert weights from the GPTQ-INT4 checkpoint (0.5 bytes/param),
|
||||||
|
and dense layers (attention, shared experts, embeddings) from the official
|
||||||
|
FP8 checkpoint (1 byte/param + calibrated block scales).
|
||||||
|
|
||||||
|
Result: a checkpoint that is ~9 GB smaller than the GPTQ-INT4 original
|
||||||
|
while using properly calibrated FP8 scales (not naive cast), yielding
|
||||||
|
better decode throughput on bandwidth-limited hardware.
|
||||||
|
|
||||||
|
NOTE: Requires the hybrid FP8 dispatch patch from https://github.com/rmstxrx/vllm/tree/v0.17.1-hybrid-fp8
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python build-hybrid-checkpoint.py \
|
||||||
|
--gptq-dir ~/inference/models/hf/qwen3.5-122b-a10b-gptq-int4 \
|
||||||
|
--fp8-repo Qwen/Qwen3.5-122B-A10B-FP8 \
|
||||||
|
--output ~/inference/models/hf/qwen3.5-122b-a10b-fp8hybrid
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import shutil
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from huggingface_hub import hf_hub_download
|
||||||
|
from safetensors import safe_open
|
||||||
|
from safetensors.torch import load_file, save_file
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def get_fp8_non_expert_manifest(fp8_repo: str) -> dict[str, str]:
|
||||||
|
"""Get the weight map for non-expert tensors from the FP8 checkpoint.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
fp8_repo: Hugging Face repo ID for the FP8 checkpoint.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Mapping of tensor names to shard filenames for non-expert tensors.
|
||||||
|
"""
|
||||||
|
idx_path = hf_hub_download(fp8_repo, "model.safetensors.index.json")
|
||||||
|
with open(idx_path, encoding="utf-8") as f:
|
||||||
|
idx = json.load(f)
|
||||||
|
|
||||||
|
wm = idx["weight_map"]
|
||||||
|
return {k: v for k, v in wm.items() if ".experts." not in k}
|
||||||
|
|
||||||
|
|
||||||
|
def download_fp8_shards(fp8_repo: str, shards: set[str], cache_dir: Path) -> dict[str, Path]:
|
||||||
|
"""Download only the needed FP8 shards.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
fp8_repo: Hugging Face repo ID for the FP8 checkpoint.
|
||||||
|
shards: Shard filenames that contain non-expert tensors.
|
||||||
|
cache_dir: Local cache directory for downloaded shards.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Mapping of shard filename to downloaded local path.
|
||||||
|
"""
|
||||||
|
shard_paths: dict[str, Path] = {}
|
||||||
|
for shard in sorted(shards):
|
||||||
|
logger.info(" Downloading %s...", shard)
|
||||||
|
path = hf_hub_download(fp8_repo, shard, local_dir=str(cache_dir))
|
||||||
|
shard_paths[shard] = Path(path)
|
||||||
|
logger.info(" -> %s", path)
|
||||||
|
return shard_paths
|
||||||
|
|
||||||
|
|
||||||
|
def extract_fp8_tensors(shard_paths: dict[str, Path], wanted: dict[str, str]) -> dict[str, torch.Tensor]:
|
||||||
|
"""Extract the requested FP8 tensors from downloaded shards.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
shard_paths: Mapping of shard filename to local path.
|
||||||
|
wanted: Mapping of tensor name to shard filename.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Mapping of tensor name to loaded FP8 tensor.
|
||||||
|
"""
|
||||||
|
tensors: dict[str, torch.Tensor] = {}
|
||||||
|
for shard_name, shard_path in shard_paths.items():
|
||||||
|
keys_in_shard = [k for k, v in wanted.items() if v == shard_name]
|
||||||
|
if not keys_in_shard:
|
||||||
|
continue
|
||||||
|
|
||||||
|
logger.info(" Extracting %d tensors from %s...", len(keys_in_shard), shard_name)
|
||||||
|
with safe_open(str(shard_path), framework="pt") as f:
|
||||||
|
for k in keys_in_shard:
|
||||||
|
tensors[k] = f.get_tensor(k)
|
||||||
|
|
||||||
|
return tensors
|
||||||
|
|
||||||
|
|
||||||
|
def find_model_safetensors_files(directory: Path) -> list[Path]:
|
||||||
|
"""Find model safetensors files in a GPTQ checkpoint directory.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
directory: Directory containing GPTQ checkpoint files.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Sorted list of shard paths, or a single-element list for single-file checkpoints.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
FileNotFoundError: If no supported model safetensors files are found.
|
||||||
|
"""
|
||||||
|
gptq_shards = sorted(directory.glob("model.safetensors-*"))
|
||||||
|
if gptq_shards:
|
||||||
|
return gptq_shards
|
||||||
|
|
||||||
|
# AutoRound / HF naming: model-NNNNN-of-NNNNN.safetensors
|
||||||
|
gptq_shards = sorted(directory.glob("model-*-of-*.safetensors"))
|
||||||
|
if gptq_shards:
|
||||||
|
return gptq_shards
|
||||||
|
|
||||||
|
single_file = directory / "model.safetensors"
|
||||||
|
if single_file.is_file():
|
||||||
|
return [single_file]
|
||||||
|
|
||||||
|
raise FileNotFoundError(
|
||||||
|
f"No model.safetensors files found in {directory}. Expected a sharded "
|
||||||
|
"checkpoint (model.safetensors-NNNNN-of-NNNNN)."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def validate_gptq_input(gptq_dir: Path) -> None:
|
||||||
|
"""Validate the GPTQ checkpoint directory before any downloads.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
gptq_dir: Path to the local GPTQ checkpoint directory.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
FileNotFoundError: If the directory or required files do not exist.
|
||||||
|
NotADirectoryError: If the path exists but is not a directory.
|
||||||
|
"""
|
||||||
|
if not gptq_dir.exists():
|
||||||
|
raise FileNotFoundError(f"GPTQ directory does not exist: {gptq_dir}")
|
||||||
|
if not gptq_dir.is_dir():
|
||||||
|
raise NotADirectoryError(f"GPTQ path is not a directory: {gptq_dir}")
|
||||||
|
if not any(path.is_file() for path in gptq_dir.glob("*.safetensors*")):
|
||||||
|
raise FileNotFoundError(f"No .safetensors files found in {gptq_dir}")
|
||||||
|
if not (gptq_dir / "config.json").is_file():
|
||||||
|
raise FileNotFoundError(f"Missing config.json in {gptq_dir}")
|
||||||
|
|
||||||
|
|
||||||
|
def validate_output_dir(output_dir: Path, force: bool) -> None:
|
||||||
|
"""Validate the output directory before building.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
output_dir: Path where the hybrid checkpoint will be written.
|
||||||
|
force: Whether destructive cleanup is allowed.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
FileExistsError: If the output directory is non-empty without `force`.
|
||||||
|
NotADirectoryError: If the output path exists but is not a directory.
|
||||||
|
"""
|
||||||
|
if output_dir.exists() and not output_dir.is_dir():
|
||||||
|
raise NotADirectoryError(f"Output path is not a directory: {output_dir}")
|
||||||
|
|
||||||
|
if output_dir.exists() and any(output_dir.iterdir()) and not force:
|
||||||
|
raise FileExistsError(
|
||||||
|
f"Output directory {output_dir} exists and is not empty. Use --force "
|
||||||
|
"to remove existing model.safetensors* and config.json files before building."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def prepare_output_dir(output_dir: Path, force: bool) -> None:
|
||||||
|
"""Create or clean the output directory before writing files.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
output_dir: Path where the hybrid checkpoint will be written.
|
||||||
|
force: Whether destructive cleanup is allowed.
|
||||||
|
"""
|
||||||
|
if not output_dir.exists():
|
||||||
|
output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
return
|
||||||
|
|
||||||
|
if not force:
|
||||||
|
return
|
||||||
|
|
||||||
|
delete_targets = sorted(
|
||||||
|
{
|
||||||
|
path
|
||||||
|
for pattern in ("model.safetensors*", "config.json")
|
||||||
|
for path in output_dir.glob(pattern)
|
||||||
|
if path.is_file()
|
||||||
|
}
|
||||||
|
)
|
||||||
|
for path in delete_targets:
|
||||||
|
path.unlink()
|
||||||
|
logger.info(" Deleted %s", path)
|
||||||
|
|
||||||
|
|
||||||
|
def build_hybrid_checkpoint(
|
||||||
|
gptq_dir: Path,
|
||||||
|
fp8_tensors: dict[str, torch.Tensor],
|
||||||
|
output_dir: Path,
|
||||||
|
force: bool,
|
||||||
|
) -> tuple[int, int, int]:
|
||||||
|
"""Build the hybrid checkpoint from GPTQ and FP8 tensors.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
gptq_dir: Path to the local GPTQ checkpoint directory.
|
||||||
|
fp8_tensors: FP8 tensors keyed by tensor name.
|
||||||
|
output_dir: Output directory for rewritten checkpoint shards.
|
||||||
|
force: Whether to continue if many unexpected FP8 tensors are unmatched.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tuple of replaced tensor count, added scale tensor count, and bytes saved.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
FileNotFoundError: If no supported GPTQ model files are found.
|
||||||
|
RuntimeError: If too many FP8 tensors cannot be matched and `force` is not set.
|
||||||
|
ValueError: If a matched FP8 tensor has an incompatible shape.
|
||||||
|
"""
|
||||||
|
# 1. Copy non-safetensors files
|
||||||
|
for f in gptq_dir.iterdir():
|
||||||
|
if f.suffix != ".safetensors" and not f.name.startswith("."):
|
||||||
|
if f.is_file():
|
||||||
|
shutil.copy2(f, output_dir / f.name)
|
||||||
|
|
||||||
|
# 2. Process each GPTQ shard
|
||||||
|
gptq_shards = find_model_safetensors_files(gptq_dir)
|
||||||
|
total_replaced = 0
|
||||||
|
total_added = 0
|
||||||
|
total_saved_bytes = 0
|
||||||
|
|
||||||
|
# Track which FP8 tensors have been placed
|
||||||
|
placed_fp8: set[str] = set()
|
||||||
|
|
||||||
|
for i, shard_path in enumerate(gptq_shards):
|
||||||
|
shard_name = shard_path.name
|
||||||
|
logger.info(" [%d/%d] %s...", i + 1, len(gptq_shards), shard_name)
|
||||||
|
|
||||||
|
gptq_tensors = load_file(str(shard_path))
|
||||||
|
output_tensors: dict[str, torch.Tensor] = {}
|
||||||
|
replaced = 0
|
||||||
|
|
||||||
|
for name, tensor in gptq_tensors.items():
|
||||||
|
if name in fp8_tensors:
|
||||||
|
# Replace BF16 tensor with FP8 version
|
||||||
|
fp8_tensor = fp8_tensors[name]
|
||||||
|
if tensor.shape != fp8_tensor.shape:
|
||||||
|
raise ValueError(
|
||||||
|
f"Shape mismatch for {name}: GPTQ={tensor.shape}, FP8={fp8_tensor.shape}"
|
||||||
|
)
|
||||||
|
old_bytes = tensor.numel() * tensor.element_size()
|
||||||
|
new_bytes = fp8_tensor.numel() * fp8_tensor.element_size()
|
||||||
|
total_saved_bytes += old_bytes - new_bytes
|
||||||
|
output_tensors[name] = fp8_tensor
|
||||||
|
placed_fp8.add(name)
|
||||||
|
replaced += 1
|
||||||
|
|
||||||
|
# Also add the scale tensor if it exists
|
||||||
|
scale_name = name.replace(".weight", ".weight_scale_inv")
|
||||||
|
if scale_name in fp8_tensors and scale_name != name:
|
||||||
|
output_tensors[scale_name] = fp8_tensors[scale_name]
|
||||||
|
placed_fp8.add(scale_name)
|
||||||
|
total_added += 1
|
||||||
|
else:
|
||||||
|
output_tensors[name] = tensor
|
||||||
|
|
||||||
|
total_replaced += replaced
|
||||||
|
|
||||||
|
save_file(output_tensors, str(output_dir / shard_name))
|
||||||
|
logger.info(" replaced=%d, tensors=%d", replaced, len(output_tensors))
|
||||||
|
|
||||||
|
# 3. Check for any FP8 tensors not yet placed (e.g. scale tensors
|
||||||
|
# for weights that exist in shards we already processed)
|
||||||
|
unplaced = set(fp8_tensors.keys()) - placed_fp8
|
||||||
|
if unplaced:
|
||||||
|
expected_unplaced = {
|
||||||
|
name
|
||||||
|
for name in unplaced
|
||||||
|
if name.endswith(".weight_scale_inv")
|
||||||
|
and f"{name.removesuffix('.weight_scale_inv')}.weight" in placed_fp8
|
||||||
|
}
|
||||||
|
unexpected_unplaced = sorted(unplaced - expected_unplaced)
|
||||||
|
|
||||||
|
for name in unexpected_unplaced:
|
||||||
|
logger.warning("WARNING: unexpected unmatched FP8 tensor %s", name)
|
||||||
|
|
||||||
|
if len(unexpected_unplaced) > 5 and not force:
|
||||||
|
message = (
|
||||||
|
f"{len(unexpected_unplaced)} FP8 tensors could not be matched to GPTQ tensors. "
|
||||||
|
"This likely indicates a naming mismatch between the GPTQ and FP8 checkpoints. "
|
||||||
|
"Use --force to proceed anyway."
|
||||||
|
)
|
||||||
|
logger.error("ERROR: %s", message)
|
||||||
|
raise RuntimeError(message)
|
||||||
|
if len(unexpected_unplaced) > 5 and force:
|
||||||
|
logger.warning(
|
||||||
|
"WARNING: proceeding despite %d unexpected unmatched FP8 tensors because --force was provided",
|
||||||
|
len(unexpected_unplaced),
|
||||||
|
)
|
||||||
|
|
||||||
|
if expected_unplaced:
|
||||||
|
logger.info(" Adding %d expected unplaced FP8 scale tensors to final shard...", len(expected_unplaced))
|
||||||
|
|
||||||
|
if expected_unplaced:
|
||||||
|
# Load last shard, add expected scale tensors, re-save
|
||||||
|
last_shard = output_dir / gptq_shards[-1].name
|
||||||
|
existing = load_file(str(last_shard))
|
||||||
|
for name in sorted(expected_unplaced):
|
||||||
|
existing[name] = fp8_tensors[name]
|
||||||
|
total_added += 1
|
||||||
|
save_file(existing, str(last_shard))
|
||||||
|
|
||||||
|
return total_replaced, total_added, total_saved_bytes
|
||||||
|
|
||||||
|
|
||||||
|
def update_safetensors_index(output_dir: Path) -> None:
|
||||||
|
"""Rebuild `model.safetensors.index.json` from actual shard contents.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
output_dir: Directory containing rewritten model shards.
|
||||||
|
"""
|
||||||
|
weight_map: dict[str, str] = {}
|
||||||
|
total_size = 0
|
||||||
|
|
||||||
|
for shard_path in find_model_safetensors_files(output_dir):
|
||||||
|
with safe_open(str(shard_path), framework="pt") as f:
|
||||||
|
for key in f.keys():
|
||||||
|
weight_map[key] = shard_path.name
|
||||||
|
tensor = f.get_tensor(key)
|
||||||
|
total_size += tensor.numel() * tensor.element_size()
|
||||||
|
|
||||||
|
index = {
|
||||||
|
"metadata": {"total_size": total_size},
|
||||||
|
"weight_map": weight_map
|
||||||
|
}
|
||||||
|
|
||||||
|
with open(output_dir / "model.safetensors.index.json", "w", encoding="utf-8") as f:
|
||||||
|
json.dump(index, f, indent=2, sort_keys=True)
|
||||||
|
|
||||||
|
logger.info(" Index rebuilt: %d tensors, %.2f GB", len(weight_map), total_size / 1e9)
|
||||||
|
|
||||||
|
|
||||||
|
def update_config(output_dir: Path) -> None:
|
||||||
|
"""Update `config.json` with hybrid quantization metadata.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
output_dir: Directory containing the hybrid checkpoint.
|
||||||
|
"""
|
||||||
|
config_path = output_dir / "config.json"
|
||||||
|
with open(config_path, encoding="utf-8") as f:
|
||||||
|
config = json.load(f)
|
||||||
|
|
||||||
|
config["_hybrid_quant_info"] = {
|
||||||
|
"description": "Hybrid GPTQ-INT4 + FP8 checkpoint for single-GPU deployment",
|
||||||
|
"moe_experts": "GPTQ INT4 (group_size=128, sym=True, Marlin kernels)",
|
||||||
|
"dense_layers": "FP8 E4M3 block-128 (from official Qwen/Qwen3.5-122B-A10B-FP8, calibrated scales)",
|
||||||
|
"norms_gates_embeddings": "Preserved at source dtype (BF16 for norms/gates, FP8 for others)",
|
||||||
|
"source_gptq": "Qwen/Qwen3.5-122B-A10B-GPTQ-Int4",
|
||||||
|
"source_fp8": "Qwen/Qwen3.5-122B-A10B-FP8",
|
||||||
|
"vllm_patch": "https://github.com/rmstxrx/vllm/tree/v0.17.1-hybrid-fp8",
|
||||||
|
"target_hardware": "NVIDIA DGX Spark (GB10, 128GB unified, 273 GB/s)",
|
||||||
|
"converter": "build-hybrid-checkpoint.py"
|
||||||
|
}
|
||||||
|
|
||||||
|
with open(config_path, "w", encoding="utf-8") as f:
|
||||||
|
json.dump(config, f, indent=2)
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> int:
|
||||||
|
"""Build a hybrid GPTQ-INT4 + FP8 checkpoint.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Process exit code.
|
||||||
|
"""
|
||||||
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||||
|
parser = argparse.ArgumentParser(description="Build hybrid GPTQ-INT4 + FP8 checkpoint")
|
||||||
|
parser.add_argument("--gptq-dir", required=True, help="Path to GPTQ-INT4 model")
|
||||||
|
parser.add_argument("--fp8-repo", default="Qwen/Qwen3.5-122B-A10B-FP8", help="HF repo for FP8 model")
|
||||||
|
parser.add_argument("--output", required=True, help="Output directory")
|
||||||
|
parser.add_argument("--dry-run", action="store_true")
|
||||||
|
parser.add_argument(
|
||||||
|
"--force",
|
||||||
|
action="store_true",
|
||||||
|
help="Allow output cleanup and continue despite many unmatched FP8 tensors",
|
||||||
|
)
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
gptq_dir = Path(args.gptq_dir)
|
||||||
|
output_dir = Path(args.output)
|
||||||
|
|
||||||
|
validate_gptq_input(gptq_dir)
|
||||||
|
find_model_safetensors_files(gptq_dir)
|
||||||
|
validate_output_dir(output_dir, args.force)
|
||||||
|
|
||||||
|
logger.info("=== Hybrid GPTQ-INT4 + FP8 Checkpoint Builder ===")
|
||||||
|
logger.info(" GPTQ source: %s", gptq_dir)
|
||||||
|
logger.info(" FP8 source: %s", args.fp8_repo)
|
||||||
|
logger.info(" Output: %s", output_dir)
|
||||||
|
logger.info("")
|
||||||
|
|
||||||
|
# Step 1: Get manifest of non-expert tensors from FP8 checkpoint
|
||||||
|
logger.info("[1/5] Fetching FP8 tensor manifest...")
|
||||||
|
fp8_manifest = get_fp8_non_expert_manifest(args.fp8_repo)
|
||||||
|
shards_needed = set(fp8_manifest.values())
|
||||||
|
logger.info(" Non-expert tensors: %d", len(fp8_manifest))
|
||||||
|
logger.info(" Shards to download: %s", sorted(shards_needed))
|
||||||
|
|
||||||
|
if args.dry_run:
|
||||||
|
logger.info("")
|
||||||
|
logger.info("[DRY RUN] Would download shards and build hybrid. Exiting.")
|
||||||
|
return 0
|
||||||
|
|
||||||
|
prepare_output_dir(output_dir, args.force)
|
||||||
|
cache_dir = output_dir / ".fp8_cache"
|
||||||
|
cache_dir.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
# Step 2: Download needed FP8 shards
|
||||||
|
logger.info("")
|
||||||
|
logger.info("[2/5] Downloading %d FP8 shards...", len(shards_needed))
|
||||||
|
shard_paths = download_fp8_shards(args.fp8_repo, shards_needed, cache_dir)
|
||||||
|
|
||||||
|
# Step 3: Extract non-expert FP8 tensors
|
||||||
|
logger.info("")
|
||||||
|
logger.info("[3/5] Extracting FP8 tensors...")
|
||||||
|
fp8_tensors = extract_fp8_tensors(shard_paths, fp8_manifest)
|
||||||
|
logger.info(" Extracted: %d tensors", len(fp8_tensors))
|
||||||
|
|
||||||
|
# Show dtype breakdown
|
||||||
|
dtypes: dict[str, int] = {}
|
||||||
|
for name, t in fp8_tensors.items():
|
||||||
|
d = str(t.dtype)
|
||||||
|
dtypes[d] = dtypes.get(d, 0) + 1
|
||||||
|
logger.info(" Dtypes: %s", dtypes)
|
||||||
|
|
||||||
|
# Step 4: Build hybrid checkpoint
|
||||||
|
logger.info("")
|
||||||
|
logger.info("[4/5] Building hybrid checkpoint...")
|
||||||
|
replaced, added, saved = build_hybrid_checkpoint(
|
||||||
|
gptq_dir,
|
||||||
|
fp8_tensors,
|
||||||
|
output_dir,
|
||||||
|
args.force,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Step 5: Update index and config
|
||||||
|
logger.info("")
|
||||||
|
logger.info("[5/5] Updating index and config...")
|
||||||
|
update_safetensors_index(output_dir)
|
||||||
|
update_config(output_dir)
|
||||||
|
|
||||||
|
# Cleanup downloaded FP8 shards
|
||||||
|
shutil.rmtree(cache_dir, ignore_errors=True)
|
||||||
|
|
||||||
|
logger.info("")
|
||||||
|
logger.info("=== Complete ===")
|
||||||
|
logger.info(" Tensors replaced (BF16→FP8): %d", replaced)
|
||||||
|
logger.info(" Scale tensors added: %d", added)
|
||||||
|
logger.info(" Bytes saved: %.2f GB", saved / 1e9)
|
||||||
|
logger.info(" Output: %s", output_dir)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
raise SystemExit(main())
|
||||||
@@ -0,0 +1,115 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Inspect the Intel INT4 AutoRound checkpoint to decide if albond's hybrid
|
||||||
|
INT4+FP8 build will work: it swaps BF16 *dense* (non-expert) tensors for FP8 by
|
||||||
|
name. If AutoRound quantized the dense linears to INT4 (.qweight), the swap is a
|
||||||
|
near no-op. We need attention/shared_expert dense weights stored as BF16 (.weight).
|
||||||
|
|
||||||
|
Also pulls the FP8 repo's index to confirm it exists and that names line up.
|
||||||
|
"""
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
from collections import Counter
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
INT4_DIR = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(
|
||||||
|
"/home/ent/.cache/huggingface/hub/models--Intel--Qwen3.5-122B-A10B-int4-AutoRound/"
|
||||||
|
"snapshots/3045d02bb737effc4581da91bddbad3be02934e4")
|
||||||
|
FP8_REPO = sys.argv[2] if len(sys.argv) > 2 else "Qwen/Qwen3.5-122B-A10B-FP8"
|
||||||
|
|
||||||
|
idx = json.loads((INT4_DIR / "model.safetensors.index.json").read_text())
|
||||||
|
wm = idx["weight_map"]
|
||||||
|
names = list(wm.keys())
|
||||||
|
|
||||||
|
# Non-expert = no ".experts." (routed experts stay INT4 in the hybrid)
|
||||||
|
nonexp = [n for n in names if ".experts." not in n]
|
||||||
|
|
||||||
|
def bucket(n):
|
||||||
|
if ".experts." in n:
|
||||||
|
return "ROUTED-EXPERT"
|
||||||
|
if "shared_expert" in n:
|
||||||
|
return "shared_expert"
|
||||||
|
if "self_attn" in n or "linear_attn" in n or ".attn" in n:
|
||||||
|
return "attention"
|
||||||
|
if "embed_tokens" in n or "lm_head" in n:
|
||||||
|
return "embed/head"
|
||||||
|
if "mtp" in n:
|
||||||
|
return "mtp"
|
||||||
|
return "other"
|
||||||
|
|
||||||
|
print(f"=== Intel INT4 checkpoint: {INT4_DIR.name} ===")
|
||||||
|
print(f"total tensors: {len(names)} non-expert: {len(nonexp)}")
|
||||||
|
|
||||||
|
# Suffix histogram tells us quant scheme: .qweight/.scales/.qzeros => INT4; .weight => dense
|
||||||
|
suffix = Counter(n.rsplit(".", 1)[-1] for n in nonexp)
|
||||||
|
print("\nnon-expert tensor SUFFIX histogram (qweight/scales/qzeros = INT4-packed; weight = dense):")
|
||||||
|
for s, c in suffix.most_common():
|
||||||
|
print(f" .{s:20s} {c}")
|
||||||
|
|
||||||
|
# For each functional group, does it have .weight (BF16 dense) or .qweight (INT4)?
|
||||||
|
print("\nper-group quant scheme (sample names):")
|
||||||
|
groups = {}
|
||||||
|
for n in nonexp:
|
||||||
|
g = bucket(n)
|
||||||
|
groups.setdefault(g, {"weight": 0, "qweight": 0, "scale": 0, "other": 0, "ex": None})
|
||||||
|
suf = n.rsplit(".", 1)[-1]
|
||||||
|
if suf == "weight":
|
||||||
|
groups[g]["weight"] += 1
|
||||||
|
elif suf == "qweight":
|
||||||
|
groups[g]["qweight"] += 1
|
||||||
|
elif "scale" in suf or suf in ("qzeros",):
|
||||||
|
groups[g]["scale"] += 1
|
||||||
|
else:
|
||||||
|
groups[g]["other"] += 1
|
||||||
|
if groups[g]["ex"] is None and suf in ("weight", "qweight"):
|
||||||
|
groups[g]["ex"] = n
|
||||||
|
for g, d in sorted(groups.items()):
|
||||||
|
scheme = "INT4(.qweight)" if d["qweight"] else ("DENSE(.weight)" if d["weight"] else "?")
|
||||||
|
print(f" {g:16s} weight={d['weight']:4d} qweight={d['qweight']:4d} scale={d['scale']:4d} -> {scheme}")
|
||||||
|
print(f" e.g. {d['ex']}")
|
||||||
|
|
||||||
|
# Dtypes of a few non-expert .weight tensors (open the shard header only).
|
||||||
|
print("\ndtypes of sample non-expert '.weight' tensors (FP8 swap needs BF16 here):")
|
||||||
|
from safetensors import safe_open # noqa: E402
|
||||||
|
sample = [n for n in nonexp if n.endswith(".weight")
|
||||||
|
and ("self_attn" in n or "shared_expert" in n or "embed" in n or "lm_head" in n)]
|
||||||
|
seen_shards = {}
|
||||||
|
shown = 0
|
||||||
|
for n in sample:
|
||||||
|
shard = wm[n]
|
||||||
|
f = seen_shards.get(shard)
|
||||||
|
if f is None:
|
||||||
|
f = safe_open(str(INT4_DIR / shard), framework="pt")
|
||||||
|
seen_shards[shard] = f
|
||||||
|
try:
|
||||||
|
t = f.get_slice(n)
|
||||||
|
print(f" {n:60s} {t.get_dtype()} {tuple(t.get_shape())}")
|
||||||
|
except Exception as e:
|
||||||
|
print(f" {n:60s} <err {e}>")
|
||||||
|
shown += 1
|
||||||
|
if shown >= 12:
|
||||||
|
break
|
||||||
|
|
||||||
|
# FP8 repo: confirm exists + list its non-expert names/dtypes for name-match sanity.
|
||||||
|
print(f"\n=== FP8 repo manifest: {FP8_REPO} ===")
|
||||||
|
try:
|
||||||
|
from huggingface_hub import hf_hub_download
|
||||||
|
p = hf_hub_download(FP8_REPO, "model.safetensors.index.json")
|
||||||
|
fidx = json.loads(Path(p).read_text())
|
||||||
|
fwm = fidx["weight_map"]
|
||||||
|
fnon = [n for n in fwm if ".experts." not in n]
|
||||||
|
print(f"FP8 total tensors: {len(fwm)} non-expert: {len(fnon)}")
|
||||||
|
fsuf = Counter(n.rsplit('.', 1)[-1] for n in fnon)
|
||||||
|
print("FP8 non-expert suffix histogram:")
|
||||||
|
for s, c in fsuf.most_common(10):
|
||||||
|
print(f" .{s:20s} {c}")
|
||||||
|
# how many FP8 non-expert .weight names also exist in INT4 as .weight?
|
||||||
|
int4_weight = {n for n in nonexp if n.endswith('.weight')}
|
||||||
|
fp8_weight = {n for n in fnon if n.endswith('.weight')}
|
||||||
|
match = int4_weight & fp8_weight
|
||||||
|
print(f"\nname overlap (FP8 '.weight' that also exist as '.weight' in INT4): {len(match)} / {len(fp8_weight)} FP8 weights")
|
||||||
|
only_fp8 = sorted(fp8_weight - int4_weight)[:8]
|
||||||
|
print(f"FP8 '.weight' NOT present as '.weight' in INT4 (would not swap): {len(fp8_weight - int4_weight)}")
|
||||||
|
for n in only_fp8:
|
||||||
|
print(f" {n} (INT4 has: {'qweight' if n[:-7]+'.qweight' in set(names) else 'MISSING'})")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"<FP8 repo fetch failed: {e}>")
|
||||||
@@ -0,0 +1,96 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""v2 harness: proper bf16 GEMV baseline + a single BATCHED int8 w8a16 GEMM kernel
|
||||||
|
(one launch for any B, dot-based, pads B to >=16) vs albond's per-row loop, at the
|
||||||
|
real lm-head shape and B in {1,5,13} (base + DFlash verify-batch sizes).
|
||||||
|
"""
|
||||||
|
import time
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import triton
|
||||||
|
import triton.language as tl
|
||||||
|
|
||||||
|
V, H = 248320, 3072
|
||||||
|
DEV = "cuda"
|
||||||
|
torch.manual_seed(0)
|
||||||
|
|
||||||
|
W_bf16 = (torch.randn(V, H, device=DEV, dtype=torch.float32) * 0.02).to(torch.bfloat16)
|
||||||
|
scales = (W_bf16.float().abs().amax(dim=1) / 127.0).clamp(min=1e-12)
|
||||||
|
W_int8 = (W_bf16.float() / scales.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)
|
||||||
|
scales_f16 = scales.to(torch.float16)
|
||||||
|
INT8_BYTES = V * H
|
||||||
|
BF16_BYTES = V * H * 2
|
||||||
|
|
||||||
|
|
||||||
|
def bench(fn, iters=30, warmup=5):
|
||||||
|
for _ in range(warmup):
|
||||||
|
fn()
|
||||||
|
torch.cuda.synchronize()
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
for _ in range(iters):
|
||||||
|
fn()
|
||||||
|
torch.cuda.synchronize()
|
||||||
|
return (time.perf_counter() - t0) / iters
|
||||||
|
|
||||||
|
|
||||||
|
# ---- batched int8 w8a16 GEMM: out[B,N] = (x[B,K] @ (W_int8[N,K]*s[N]).T) ----
|
||||||
|
@triton.jit
|
||||||
|
def _k_batched(x_ptr, w_ptr, s_ptr, o_ptr, B, N, K,
|
||||||
|
sxb, sxk, swn, swk, sob, son,
|
||||||
|
BLOCK_B: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
|
||||||
|
pid_n = tl.program_id(0)
|
||||||
|
offs_b = tl.arange(0, BLOCK_B)
|
||||||
|
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||||
|
offs_k = tl.arange(0, BLOCK_K)
|
||||||
|
x_ptrs = x_ptr + offs_b[:, None] * sxb + offs_k[None, :] * sxk
|
||||||
|
w_ptrs = w_ptr + offs_n[:, None] * swn + offs_k[None, :] * swk
|
||||||
|
acc = tl.zeros((BLOCK_B, BLOCK_N), dtype=tl.float32)
|
||||||
|
for k in range(0, K, BLOCK_K):
|
||||||
|
km = (offs_k[None, :] + k) < K
|
||||||
|
x = tl.load(x_ptrs, mask=(offs_b[:, None] < B) & km, other=0.0).to(tl.float16)
|
||||||
|
w = tl.load(w_ptrs, mask=(offs_n[:, None] < N) & km, other=0).to(tl.float16)
|
||||||
|
acc += tl.dot(x, w.T) # [BB,BK] @ [BK,BN] -> [BB,BN]
|
||||||
|
x_ptrs += BLOCK_K * sxk
|
||||||
|
w_ptrs += BLOCK_K * swk
|
||||||
|
s = tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32)
|
||||||
|
acc = acc * s[None, :]
|
||||||
|
o_ptrs = o_ptr + offs_b[:, None] * sob + offs_n[None, :] * son
|
||||||
|
tl.store(o_ptrs, acc.to(tl.float16), mask=(offs_b[:, None] < B) & (offs_n[None, :] < N))
|
||||||
|
|
||||||
|
|
||||||
|
def run_batched(x, BLOCK_N=128, BLOCK_K=64, num_warps=4, num_stages=3):
|
||||||
|
B = x.shape[0]
|
||||||
|
BLOCK_B = max(16, triton.next_power_of_2(B))
|
||||||
|
out = torch.empty(B, V, dtype=torch.float16, device=DEV)
|
||||||
|
xf = x.to(torch.float16)
|
||||||
|
grid = ((V + BLOCK_N - 1) // BLOCK_N,)
|
||||||
|
_k_batched[grid](xf, W_int8, scales_f16, out, B, V, H,
|
||||||
|
xf.stride(0), xf.stride(1), W_int8.stride(0), W_int8.stride(1),
|
||||||
|
out.stride(0), out.stride(1),
|
||||||
|
BLOCK_B=BLOCK_B, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
|
||||||
|
num_warps=num_warps, num_stages=num_stages)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
for B in (1, 5, 13):
|
||||||
|
print(f"\n=== B={B} ===")
|
||||||
|
x = torch.randn(B, H, device=DEV, dtype=torch.bfloat16) * 0.1
|
||||||
|
ref_bf16 = (x.float() @ W_bf16.float().T)
|
||||||
|
ref_deq = (x.float() @ (W_int8.float() * scales.unsqueeze(1)).T)
|
||||||
|
am_floor = (ref_bf16.argmax(-1) == ref_deq.argmax(-1)).float().mean().item()
|
||||||
|
print(f" quant floor: argmax_match={am_floor*100:.1f}% maxerr={(ref_deq-ref_bf16).abs().max():.4f}")
|
||||||
|
|
||||||
|
# real bf16 GEMV baseline (stays bf16, reads 1.5GB)
|
||||||
|
dt = bench(lambda: torch.matmul(x, W_bf16.t()))
|
||||||
|
print(f" {'bf16 GEMV (real baseline)':30s} {'':33s}{dt*1e3:7.3f} ms {BF16_BYTES/dt/1e9:6.1f} GB/s")
|
||||||
|
|
||||||
|
# batched int8 kernel, a few configs
|
||||||
|
for (bn, bk, nw, ns) in [(128, 64, 4, 3), (256, 64, 8, 3), (128, 128, 4, 3), (64, 128, 4, 3)]:
|
||||||
|
try:
|
||||||
|
out = run_batched(x, bn, bk, nw, ns)
|
||||||
|
err = (out.float() - ref_deq.float()).abs().max().item()
|
||||||
|
am = (out.float().argmax(-1) == ref_bf16.argmax(-1)).float().mean().item()
|
||||||
|
dt = bench(lambda: run_batched(x, bn, bk, nw, ns))
|
||||||
|
print(f" batched N{bn}/K{bk}/w{nw}/s{ns:<2d} maxerr={err:8.4f} argmax={am*100:5.1f}% "
|
||||||
|
f"{dt*1e3:7.3f} ms {INT8_BYTES/dt/1e9:6.1f} GB/s")
|
||||||
|
except Exception as e:
|
||||||
|
print(f" batched N{bn}/K{bk}/w{nw}/s{ns}: ERR {str(e)[:70]}")
|
||||||
@@ -0,0 +1,27 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Validate patch_inc_hybrid.py landed on AEON 0.23's inc.py. Run AFTER the patch
|
||||||
|
in the same container: python3 /host/patch_inc_hybrid.py && python3 /host/validate_inc_patch.py
|
||||||
|
"""
|
||||||
|
import inspect
|
||||||
|
|
||||||
|
import vllm.model_executor.layers.quantization.inc as m
|
||||||
|
|
||||||
|
C = m.INCConfig
|
||||||
|
src = inspect.getsource(m)
|
||||||
|
print("import OK")
|
||||||
|
print("sentinel in source :", "spark-dflash-hybrid-fp8" in src)
|
||||||
|
print("maybe_update_config OWN :", "maybe_update_config" in C.__dict__)
|
||||||
|
print("_is_layer_fp8 OWN :", "_is_layer_fp8" in C.__dict__)
|
||||||
|
print("maybe_update_config hasattr:", hasattr(C, "maybe_update_config"))
|
||||||
|
print("_is_layer_fp8 hasattr :", hasattr(C, "_is_layer_fp8"))
|
||||||
|
# signature must accept hf_config kw (config/vllm.py calls it that way)
|
||||||
|
try:
|
||||||
|
sig = inspect.signature(C.maybe_update_config)
|
||||||
|
print("maybe_update_config sig :", str(sig))
|
||||||
|
print("accepts hf_config kw :", "hf_config" in sig.parameters)
|
||||||
|
except Exception as e:
|
||||||
|
print("sig err:", e)
|
||||||
|
# count FP8 dispatch sites
|
||||||
|
print("Fp8LinearMethod dispatch ct:", src.count("return Fp8LinearMethod(self.fp8_config)"))
|
||||||
|
# byte-compile sanity already implied by import; show line count
|
||||||
|
print("inc.py lines :", len(src.splitlines()))
|
||||||
@@ -0,0 +1,45 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Validate patch_int8_lmhead_v3 landed + the helper runs on GPU. Run AFTER the
|
||||||
|
patch in the same container."""
|
||||||
|
import torch
|
||||||
|
|
||||||
|
import vllm.model_executor.layers.logits_processor as lp
|
||||||
|
|
||||||
|
print("import OK | helpers:",
|
||||||
|
hasattr(lp, "_spark_int8_gemm"),
|
||||||
|
hasattr(lp, "_spark_int8_lmhead_apply"),
|
||||||
|
hasattr(lp, "_spark_k_int8"))
|
||||||
|
src = open(lp.__file__).read()
|
||||||
|
print("sentinel in _get_logits:", "DGX_SPARK_INT8_LMHEAD_V3: int8 w8a16" in src)
|
||||||
|
|
||||||
|
V, H = 4096, 512
|
||||||
|
torch.manual_seed(0)
|
||||||
|
W = torch.randn(V, H, device="cuda") * 0.02
|
||||||
|
s = (W.abs().amax(1) / 127).clamp(min=1e-12)
|
||||||
|
wi = (W / s.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8).contiguous()
|
||||||
|
sf = s.to(torch.float16)
|
||||||
|
for B in (1, 5, 13):
|
||||||
|
x = torch.randn(B, H, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||||
|
out = lp._spark_int8_gemm(x, wi, sf)
|
||||||
|
ref = x.float() @ (wi.float() * s.unsqueeze(1)).T
|
||||||
|
am = (out.argmax(-1) == ref.argmax(-1)).float().mean().item()
|
||||||
|
print(f" B={B}: out{tuple(out.shape)} {out.dtype} argmax={am*100:.0f}% "
|
||||||
|
f"maxerr={(out - ref).abs().max().item():.4f}")
|
||||||
|
|
||||||
|
# Exercise the FULL apply path (print + quantize-once + gemm) with a mock lm_head
|
||||||
|
# (vocab > 100k to trigger the int8 path). self is unused -> None.
|
||||||
|
print("--- full _spark_int8_lmhead_apply path (mock lm_head, V=131072) ---")
|
||||||
|
Vbig = 131072
|
||||||
|
|
||||||
|
class _MockLMHead:
|
||||||
|
pass
|
||||||
|
|
||||||
|
mh = _MockLMHead()
|
||||||
|
mh.weight = (torch.randn(Vbig, H, device="cuda", dtype=torch.float32) * 0.02).to(torch.bfloat16)
|
||||||
|
hs = torch.randn(3, H, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||||
|
o1 = lp._spark_int8_lmhead_apply(None, mh, hs, None) # first call: quantizes + prints
|
||||||
|
o2 = lp._spark_int8_lmhead_apply(None, mh, hs, None) # second call: reuses int8
|
||||||
|
refb = hs.float() @ mh.weight.float().T
|
||||||
|
am = (o2.argmax(-1) == refb.argmax(-1)).float().mean().item()
|
||||||
|
print(f" apply: out{tuple(o2.shape)} {o2.dtype} argmax_vs_bf16={am*100:.0f}% "
|
||||||
|
f"int8_ready={getattr(mh, '_spark_int8_ready', None)} weight_kept={mh.weight.numel() > 0}")
|
||||||
Reference in New Issue
Block a user