Rename the project to edge-infer, and complete the packaging audit
Renamed from vllm-embedding: nothing here is about embeddings, and the name had come from the parent working directory. edge-infer says what it does. The GitHub repository is renamed to match; nothing had been pushed yet. The audit found the packaging was incomplete in two ways. Scripts: 56 of them produced results that are published in this repo but were not included, so the numbers could be read and not checked. Added under benchmarks/edge_harness/experiments/ (one-off drivers: the AI Hub on-device runs, the 4-bit GEMM microbenchmark, fidelity and parity checks, checkpoint build and round-trip, activation capture, and the load generator used to prove the contention guard actually refuses) and benchmarks/edge_harness/runs/ (the 24 shell drivers behind published grids). Both directories carry a README saying plainly that they are a record of what was run rather than a supported interface -- most hardcode absolute paths from the machine they ran on. The supported entry points remain the tools one level up. One broken relative link fixed (comparison.md is under docs/engines/, not docs/analysis/). All other links and every script reference in the docs now resolve inside the repo. Secret scan: the token in ~/.qai_hub/client.ini is absent from the worktree and from all 61 commits. Worth noting that an earlier version of my own check reported a false positive, because `grep -rl ... | head` takes its exit status from head, which always succeeds -- the same class of mistake as the measurement errors this project keeps recording. The AI Hub scripts authenticate by reading that file at runtime and contain no credential. README reframed for the new name and the wider scope: the repository is the engine changes, the harness, the raw results, the baselines and the analyses, not just a patch set. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ
This commit is contained in:
tzhouamandClaude Opus 5 committed
1 parent
99ebfbbd04
commit
3834de3105
62 files changed
+3003
-12
No files matched your search
+1
-1
@@ -282,7 +282,7 @@ lance_compare/
|
||||
# Linting binaries downloaded by local pre-commit hooks
|
||||
shellcheck*/
|
||||
|
||||
# --- vllm-embedding ---------------------------------------------------------
|
||||
# --- edge-infer -------------------------------------------------------------
|
||||
# Measurement artifacts are the substance of this fork, not build output: every
|
||||
# published figure is traceable to a file below. They are tracked despite the
|
||||
# blanket `*.json` and `results/` rules above. Binary captures (.npz, .onnx,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
vllm-embedding
|
||||
edge-infer
|
||||
Copyright 2026 tzhouam
|
||||
|
||||
This product is a derivative work of vLLM-Omni
|
||||
|
||||
@@ -1,15 +1,22 @@
|
||||
# vllm-embedding
|
||||
# edge-infer
|
||||
|
||||
A fork of [vLLM-Omni](https://github.com/vllm-project/vllm-omni) for running
|
||||
multimodal models on **edge CPUs** — machines with tens of gigabytes of RAM and
|
||||
no GPU, rather than servers.
|
||||
Running multimodal models on **edge CPUs** — a mini-PC or a gateway box with
|
||||
tens of gigabytes of RAM and no GPU — instead of on servers.
|
||||
|
||||
The upstream project targets datacentre inference. Getting the same models onto
|
||||
a mini-PC or a gateway box turns out to be less about kernels than about three
|
||||
things nobody profiles: how much memory is resident while the model *loads*,
|
||||
how much is resident that nothing ever reads, and whether the number you just
|
||||
measured came from the code path you thought it did.
|
||||
Built as a fork of [vLLM-Omni](https://github.com/vllm-project/vllm-omni),
|
||||
because it is the only engine with a multi-stage multimodal pipeline,
|
||||
continuous batching and an OpenAI-compatible server; the alternatives were
|
||||
evaluated and the reasoning is in [docs/engines/](docs/engines/).
|
||||
|
||||
The work turned out to be less about kernels than about three things nobody
|
||||
profiles: how much memory is resident while the model *loads*, how much is
|
||||
resident that nothing ever reads, and whether the number you just measured came
|
||||
from the code path you thought it did. The largest single win here — 843 MB —
|
||||
was a setting that had been marked "no effect, do not retry" because the code
|
||||
path under test discarded it.
|
||||
|
||||
This repository holds the engine changes, the measurement harness, the raw
|
||||
results, the baselines they are compared against, and the written analyses.
|
||||
Upstream's own README is preserved at [docs/README.vllm-omni.md](docs/README.vllm-omni.md).
|
||||
|
||||
## What's here
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
# One-off experiment drivers
|
||||
|
||||
Scripts that produced a specific result rather than tooling meant to be reused.
|
||||
Kept because the results in `../results/` are only checkable if the script that
|
||||
made them is here.
|
||||
|
||||
| group | what |
|
||||
|---|---|
|
||||
| `aihub_*.py` | Qualcomm AI Hub on-device runs — Snapdragon layer timings, calibration, w8a16/w4a16 parity. Authentication comes from `~/.qai_hub/client.ini`; no credential is in these files |
|
||||
| `bench_int4_kernels.py` | the 4-bit GEMM microbenchmark — kernel-level shapes, used to size decisions before building anything |
|
||||
| `quality_int4.py`, `ppl_vllm.py`, `parity_hf_vs_vllm.py`, `ref_hf.py` | fidelity: greedy divergence, perplexity, and parity against a HuggingFace reference |
|
||||
| `quantize_int8.py`, `export_gptq_cpu.py`, `dequantize_int4.py`, `verify_export.py` | building and round-tripping checkpoints |
|
||||
| `capture_real_acts.py`, `capture_lm_head.py` | activation capture for calibrated quantization |
|
||||
| `cpu_busy.py` | a load generator, used to test that the harness's contention guard actually refuses |
|
||||
|
||||
**They are not portable as written.** Most carry absolute paths from the machine
|
||||
they ran on (`/data/<user>/embedding_infer/...`) and expect checkpoints under
|
||||
`models/`. They are a record of what was run, not a supported interface — the
|
||||
supported entry points are the tools one level up, described in
|
||||
[docs/edge/harness.md](../../../docs/edge/harness.md).
|
||||
@@ -0,0 +1,50 @@
|
||||
"""The architecture's repeating 4-layer block (3 sliding + 1 full) on a real S25.
|
||||
|
||||
Profiling one layer at a time pays the graph's Input/Output boundary once per
|
||||
layer. A real decode step pays it once for all 28. Measuring 1 layer and 4
|
||||
layers at the same context separates the fixed per-graph cost from the true
|
||||
per-layer cost, so the 28-layer step can be extrapolated honestly.
|
||||
"""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
OPT = {
|
||||
"fp16": "--target_runtime qnn_context_binary --quantize_full_type float16",
|
||||
"w4a16": "--target_runtime qnn_context_binary --quantize_full_type w4a16 --quantize_io",
|
||||
}
|
||||
|
||||
res, pend = {}, {}
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/ring_block4_ctx1024.onnx", name="spark25_block4")
|
||||
print("uploaded", flush=True)
|
||||
for q in ("w4a16", "fp16"):
|
||||
try:
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark block4 {q}", options=OPT[q])
|
||||
pend[q] = cj
|
||||
print("compile", q, cj.job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", q, str(e)[:200], flush=True)
|
||||
|
||||
prof = {}
|
||||
for q, cj in pend.items():
|
||||
st = cj.wait()
|
||||
print("compile", q, st.code, (st.message or "")[:200], flush=True)
|
||||
if st.success:
|
||||
prof[q] = hub.submit_profile_job(model=cj.get_target_model(),
|
||||
device=hub.Device(DEV),
|
||||
name=f"spark block4 {q}",
|
||||
options="--compute_unit npu")
|
||||
for q, pj in prof.items():
|
||||
st = pj.wait()
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[q] = p["execution_summary"]["estimated_inference_time"] / 1e3
|
||||
json.dump(p, open(f"{HERE}/profile_block4_{q}.json", "w"))
|
||||
print(f" block4 {q}: {res[q]:.3f} ms", flush=True)
|
||||
else:
|
||||
res[q] = None
|
||||
print(f" block4 {q} FAILED: {(st.message or '')[:250]}", flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_block4.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Calibrated quantisation + on-device parity for any exported graph."""
|
||||
import argparse, json
|
||||
import numpy as np
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
SCHEMES = {
|
||||
"w8a16": (hub.QuantizeDtype.INT8, hub.QuantizeDtype.INT16),
|
||||
"w16a16": (hub.QuantizeDtype.INT16, hub.QuantizeDtype.INT16),
|
||||
}
|
||||
|
||||
|
||||
def snr(a, b):
|
||||
a, b = np.asarray(a).ravel(), np.asarray(b).ravel()
|
||||
return float(20 * np.log10(np.linalg.norm(a) / max(np.linalg.norm(a - b), 1e-12)))
|
||||
|
||||
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--onnx", required=True)
|
||||
ap.add_argument("--acts", required=True)
|
||||
ap.add_argument("--map", default="", help="onnx_name=npz_name,...")
|
||||
ap.add_argument("--tag", required=True)
|
||||
ap.add_argument("--out", required=True)
|
||||
a = ap.parse_args()
|
||||
|
||||
import onnx
|
||||
g = onnx.load(a.onnx, load_external_data=False)
|
||||
live = [t.name for t in g.graph.input]
|
||||
mapping = dict(p.split("=") for p in a.map.split(",") if p)
|
||||
d = np.load(a.acts)
|
||||
n = d["x"].shape[0]
|
||||
calib = {name: [d[mapping.get(name, name)][i] for i in range(n)] for name in live}
|
||||
probe = {name: [d[mapping.get(name, name)][0]] for name in live}
|
||||
ref_out = d["hidden_out"][0]
|
||||
|
||||
src = hub.upload_model(a.onnx, name=f"spark25_{a.tag}_src")
|
||||
qj = {}
|
||||
for tag, (w, act) in SCHEMES.items():
|
||||
qj[tag] = hub.submit_quantize_job(model=src, calibration_data=calib,
|
||||
weights_dtype=w, activations_dtype=act,
|
||||
name=f"spark {a.tag} {tag}")
|
||||
print("quantize", tag, qj[tag].job_id, flush=True)
|
||||
|
||||
res, cj = {}, {}
|
||||
for tag, j in qj.items():
|
||||
st = j.wait()
|
||||
print("quantize", tag, st.code, (st.message or "")[:140], flush=True)
|
||||
if st.success:
|
||||
cj[tag] = hub.submit_compile_job(
|
||||
model=j.get_target_model(), device=hub.Device(DEV),
|
||||
name=f"spark {a.tag} {tag}",
|
||||
options="--target_runtime qnn_context_binary --quantize_io")
|
||||
ij, pj = {}, {}
|
||||
for tag, j in cj.items():
|
||||
st = j.wait()
|
||||
print("compile", tag, st.code, (st.message or "")[:140], flush=True)
|
||||
if st.success:
|
||||
t = j.get_target_model()
|
||||
ij[tag] = hub.submit_inference_job(model=t, device=hub.Device(DEV),
|
||||
inputs=probe, name=f"spark {a.tag} {tag}")
|
||||
pj[tag] = hub.submit_profile_job(model=t, device=hub.Device(DEV),
|
||||
name=f"spark {a.tag} {tag}",
|
||||
options="--compute_unit npu")
|
||||
for tag in list(ij):
|
||||
e = res.setdefault(tag, {})
|
||||
st = ij[tag].wait()
|
||||
if st.success:
|
||||
out = ij[tag].download_output_data()
|
||||
e["snr_db"] = round(snr(ref_out, np.asarray(list(out.values())[0][0])), 2)
|
||||
else:
|
||||
e["snr_db"] = None
|
||||
st = pj[tag].wait()
|
||||
e["ms"] = (pj[tag].download_profile()["execution_summary"]
|
||||
["estimated_inference_time"] / 1e3) if st.success else None
|
||||
print(f" {a.tag} {tag}: {e}", flush=True)
|
||||
json.dump(res, open(a.out, "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Calibrated quantisation of the sliding layer, judged on real activations.
|
||||
|
||||
Uncalibrated w4a16/w4a8 scored -1.6 / -1.28 dB against the reference -- output
|
||||
uncorrelated with the truth. This feeds the quantiser real activations
|
||||
captured from the reference model on real text, then measures the same way.
|
||||
"""
|
||||
import json
|
||||
import numpy as np
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
NAMES = ["x", "cos_sw", "sin_sw", "mask_sw", "k_cache_0", "v_cache_0"]
|
||||
|
||||
d = np.load(f"{HERE}/ring3_real_acts.npz")
|
||||
n_samples = d["x"].shape[0]
|
||||
calib = {n: [d[n][i] for i in range(n_samples)] for n in NAMES}
|
||||
probe = {n: [d[n][0]] for n in NAMES}
|
||||
ref_out = d["hidden_out"][0]
|
||||
|
||||
SCHEMES = {
|
||||
"w4a16": (hub.QuantizeDtype.INT4, hub.QuantizeDtype.INT16),
|
||||
"w8a16": (hub.QuantizeDtype.INT8, hub.QuantizeDtype.INT16),
|
||||
"w4a8": (hub.QuantizeDtype.INT4, hub.QuantizeDtype.INT8),
|
||||
"w8a8": (hub.QuantizeDtype.INT8, hub.QuantizeDtype.INT8),
|
||||
}
|
||||
|
||||
|
||||
def snr(a, b):
|
||||
a, b = np.asarray(a).ravel(), np.asarray(b).ravel()
|
||||
return float(20 * np.log10(np.linalg.norm(a) / max(np.linalg.norm(a - b), 1e-12)))
|
||||
|
||||
|
||||
src = hub.upload_model(f"{HERE}/onnx/ring3_sliding.onnx", name="spark25_calib_src")
|
||||
qjobs = {}
|
||||
for tag, (w, a) in SCHEMES.items():
|
||||
qj = hub.submit_quantize_job(model=src, calibration_data=calib,
|
||||
weights_dtype=w, activations_dtype=a,
|
||||
name=f"spark calib {tag}")
|
||||
qjobs[tag] = qj
|
||||
print("quantize", tag, qj.job_id, flush=True)
|
||||
|
||||
res = {}
|
||||
cjobs = {}
|
||||
for tag, qj in qjobs.items():
|
||||
st = qj.wait()
|
||||
print("quantize", tag, st.code, (st.message or "")[:150], flush=True)
|
||||
if st.success:
|
||||
cjobs[tag] = hub.submit_compile_job(
|
||||
model=qj.get_target_model(), device=hub.Device(DEV),
|
||||
name=f"spark calib {tag}",
|
||||
options="--target_runtime qnn_context_binary --quantize_io")
|
||||
|
||||
ijobs, pjobs = {}, {}
|
||||
for tag, cj in cjobs.items():
|
||||
st = cj.wait()
|
||||
print("compile", tag, st.code, (st.message or "")[:150], flush=True)
|
||||
if st.success:
|
||||
tgt = cj.get_target_model()
|
||||
ijobs[tag] = hub.submit_inference_job(model=tgt, device=hub.Device(DEV),
|
||||
inputs=probe, name=f"spark calib {tag}")
|
||||
pjobs[tag] = hub.submit_profile_job(model=tgt, device=hub.Device(DEV),
|
||||
name=f"spark calib {tag}",
|
||||
options="--compute_unit npu")
|
||||
|
||||
for tag in list(ijobs):
|
||||
st = ijobs[tag].wait()
|
||||
entry = res.setdefault(tag, {})
|
||||
if st.success:
|
||||
out = ijobs[tag].download_output_data()
|
||||
got = np.asarray(list(out.values())[0][0])
|
||||
entry["snr_db"] = round(snr(ref_out, got), 2)
|
||||
else:
|
||||
entry["snr_db"] = None
|
||||
print(tag, "inference FAILED", (st.message or "")[:150], flush=True)
|
||||
st = pjobs[tag].wait()
|
||||
if st.success:
|
||||
p = pjobs[tag].download_profile()
|
||||
entry["ms"] = p["execution_summary"]["estimated_inference_time"] / 1e3
|
||||
json.dump(p, open(f"{HERE}/profile_calib_{tag}.json", "w"))
|
||||
else:
|
||||
entry["ms"] = None
|
||||
print(f" {tag}: {entry}", flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_calibrated.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,40 @@
|
||||
"""A quantized phone-CPU reference for the same layer.
|
||||
|
||||
llama.cpp cannot be run on an AI Hub device, and it would use 4-bit weights on
|
||||
the CPU, so comparing our NPU path against an fp16 CPU run would flatter us.
|
||||
This asks for the closest thing AI Hub offers: the same graph, int8, on the
|
||||
S25's CPU.
|
||||
"""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
CAND = [
|
||||
("tflite_int8_cpu", "--target_runtime tflite --quantize_full_type int8 --quantize_io", "cpu"),
|
||||
("tflite_w8a8_cpu", "--target_runtime tflite --quantize_full_type w8a8 --quantize_io", "cpu"),
|
||||
]
|
||||
res, pend = {}, {}
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/ring3_sliding.onnx", name="spark25_ring3_cpu")
|
||||
for name, opt, unit in CAND:
|
||||
try:
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark {name}", options=opt)
|
||||
pend[(name, unit)] = cj; print("compile", name, cj.job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", name, str(e)[:180], flush=True)
|
||||
prof = {}
|
||||
for key, cj in pend.items():
|
||||
st = cj.wait(); print("compile", key[0], st.code, (st.message or "")[:180], flush=True)
|
||||
if st.success:
|
||||
prof[key] = hub.submit_profile_job(model=cj.get_target_model(), device=hub.Device(DEV),
|
||||
name=f"spark {key[0]}", options=f"--compute_unit {key[1]}")
|
||||
for key, pj in prof.items():
|
||||
st = pj.wait()
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[key[0]] = p["execution_summary"]["estimated_inference_time"]/1e3
|
||||
print(f" {key[0]}: {res[key[0]]:.3f} ms", flush=True)
|
||||
else:
|
||||
res[key[0]] = None; print(f" {key[0]} FAILED:", (st.message or "")[:180], flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_cpu_int8.json","w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Is the cheaper GELU actually cheaper on the Hexagon NPU?
|
||||
|
||||
The erf form is 11-13% of every layer in the device profile. tanh holds 65.2 dB
|
||||
against the reference with top-1/top-5 intact, so it is adoptable if it buys
|
||||
time. This measures whether it does.
|
||||
"""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
OPT = "--target_runtime qnn_context_binary --quantize_full_type float16"
|
||||
res = {}
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/ring_sliding_tanh.onnx", name="spark25_sliding_tanh")
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name="spark sliding ring tanh fp16", options=OPT)
|
||||
print("compile", cj.job_id, flush=True)
|
||||
st = cj.wait(); print("compile", st.code, (st.message or "")[:160], flush=True)
|
||||
if st.success:
|
||||
pj = hub.submit_profile_job(model=cj.get_target_model(), device=hub.Device(DEV),
|
||||
name="spark sliding ring tanh fp16",
|
||||
options="--compute_unit npu")
|
||||
st2 = pj.wait()
|
||||
if st2.success:
|
||||
p = pj.download_profile()
|
||||
res["ring_sliding_tanh_fp16"] = p["execution_summary"]["estimated_inference_time"] / 1e3
|
||||
json.dump(p, open(f"{HERE}/profile_ring_sliding_tanh_fp16.json", "w"))
|
||||
print(f" tanh: {res['ring_sliding_tanh_fp16']:.3f} ms", flush=True)
|
||||
else:
|
||||
print(" profile FAILED:", (st2.message or "")[:200], flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_gelu.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,86 @@
|
||||
"""int4 weights on the NPU, calibrated.
|
||||
|
||||
int4 halves the weight traffic that dominates every matmul in the step, so it
|
||||
is the largest remaining lever. The compile-path form (--quantize_full_type
|
||||
w4a16) runs at 0.838 ms but is uncalibrated and scores -1.6 dB. The
|
||||
quantize-job form carries calibration but failed to compile. This retries that
|
||||
compile without --quantize_io, which is the one thing the failing jobs had in
|
||||
common with an already-quantized input model.
|
||||
"""
|
||||
import argparse, json
|
||||
import numpy as np
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--onnx", default=f"{HERE}/onnx/ring3_sliding.onnx")
|
||||
ap.add_argument("--acts", default=f"{HERE}/ring3_real_acts.npz")
|
||||
ap.add_argument("--map", default="")
|
||||
ap.add_argument("--tag", default="sliding")
|
||||
ap.add_argument("--out", default=f"{HERE}/aihub_int4.json")
|
||||
a = ap.parse_args()
|
||||
|
||||
import onnx
|
||||
live = [t.name for t in onnx.load(a.onnx, load_external_data=False).graph.input]
|
||||
mapping = dict(p.split("=") for p in a.map.split(",") if p)
|
||||
d = np.load(a.acts)
|
||||
n = d["x"].shape[0]
|
||||
calib = {x: [d[mapping.get(x, x)][i] for i in range(n)] for x in live}
|
||||
probe = {x: [d[mapping.get(x, x)][0]] for x in live}
|
||||
ref = d["hidden_out"][0]
|
||||
|
||||
|
||||
def snr(u, v):
|
||||
u, v = np.asarray(u).ravel(), np.asarray(v).ravel()
|
||||
return float(20 * np.log10(np.linalg.norm(u) / max(np.linalg.norm(u - v), 1e-12)))
|
||||
|
||||
|
||||
src = hub.upload_model(a.onnx, name=f"spark25_{a.tag}_int4")
|
||||
qj = hub.submit_quantize_job(model=src, calibration_data=calib,
|
||||
weights_dtype=hub.QuantizeDtype.INT4,
|
||||
activations_dtype=hub.QuantizeDtype.INT16,
|
||||
name=f"spark {a.tag} int4 calib")
|
||||
print("quantize", qj.job_id, flush=True)
|
||||
st = qj.wait()
|
||||
print("quantize", st.code, (st.message or "")[:200], flush=True)
|
||||
res = {}
|
||||
if st.success:
|
||||
qm = qj.get_target_model()
|
||||
# The failing attempts all passed --quantize_io on an already-quantized
|
||||
# model; try without it, and with the plain QNN DLC target as a fallback.
|
||||
VARIANTS = {
|
||||
"ctx_no_qio": "--target_runtime qnn_context_binary",
|
||||
"ctx_qio": "--target_runtime qnn_context_binary --quantize_io",
|
||||
}
|
||||
cj = {}
|
||||
for tag, opt in VARIANTS.items():
|
||||
try:
|
||||
cj[tag] = hub.submit_compile_job(model=qm, device=hub.Device(DEV),
|
||||
name=f"spark {a.tag} int4 {tag}",
|
||||
options=opt)
|
||||
print("compile", tag, cj[tag].job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", tag, str(e)[:160], flush=True)
|
||||
for tag, j in cj.items():
|
||||
s2 = j.wait()
|
||||
print("compile", tag, s2.code, (s2.message or "")[:200], flush=True)
|
||||
e = res.setdefault(tag, {})
|
||||
if not s2.success:
|
||||
e["ms"] = e["snr_db"] = None
|
||||
json.dump(res, open(a.out, "w"), indent=1)
|
||||
continue
|
||||
t = j.get_target_model()
|
||||
ij = hub.submit_inference_job(model=t, device=hub.Device(DEV), inputs=probe,
|
||||
name=f"spark {a.tag} int4 {tag}")
|
||||
pj = hub.submit_profile_job(model=t, device=hub.Device(DEV),
|
||||
name=f"spark {a.tag} int4 {tag}",
|
||||
options="--compute_unit npu")
|
||||
s3 = ij.wait()
|
||||
e["snr_db"] = round(snr(ref, np.asarray(list(ij.download_output_data().values())[0][0])), 2) if s3.success else None
|
||||
s4 = pj.wait()
|
||||
e["ms"] = pj.download_profile()["execution_summary"]["estimated_inference_time"] / 1e3 if s4.success else None
|
||||
print(f" {a.tag} int4 {tag}: {e}", flush=True)
|
||||
json.dump(res, open(a.out, "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,60 @@
|
||||
"""On-device fidelity of the quantized sliding layer.
|
||||
|
||||
Speed is only half the answer: this repo has already found automatic int8
|
||||
numerically broken on Qwen3-TTS and MiniCPM-o. This runs the real layer on a
|
||||
real S25 at each quantisation and compares against the fp32 reference.
|
||||
"""
|
||||
import json
|
||||
import numpy as np
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
OPT = {
|
||||
"fp16": "--target_runtime qnn_context_binary --quantize_full_type float16",
|
||||
"w4a16": "--target_runtime qnn_context_binary --quantize_full_type w4a16 --quantize_io",
|
||||
"w4a8": "--target_runtime qnn_context_binary --quantize_full_type w4a8 --quantize_io",
|
||||
}
|
||||
|
||||
ref = np.load(f"{HERE}/ring3_ref_inputs.npz")
|
||||
# The graph's live inputs; ONNX prunes the full-attention rotary/mask tensors
|
||||
# a sliding layer never reads.
|
||||
names = ["x", "cos_sw", "sin_sw", "mask_sw", "k_cache_0", "v_cache_0"]
|
||||
inputs = {n: [ref[n].astype(np.float32)] for n in names}
|
||||
ref_out = ref["hidden_out"]
|
||||
|
||||
|
||||
def snr(a, b):
|
||||
a, b = a.ravel(), b.ravel()
|
||||
return float(20 * np.log10(np.linalg.norm(a) / max(np.linalg.norm(a - b), 1e-12)))
|
||||
|
||||
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/ring3_sliding.onnx", name="spark25_ring3_parity")
|
||||
res, pend = {}, {}
|
||||
for q, o in OPT.items():
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark parity {q}", options=o)
|
||||
pend[q] = cj
|
||||
print("compile", q, cj.job_id, flush=True)
|
||||
infj = {}
|
||||
for q, cj in pend.items():
|
||||
st = cj.wait()
|
||||
if st.success:
|
||||
infj[q] = hub.submit_inference_job(model=cj.get_target_model(),
|
||||
device=hub.Device(DEV), inputs=inputs,
|
||||
name=f"spark parity {q}")
|
||||
else:
|
||||
print("compile", q, "FAILED", (st.message or "")[:160], flush=True)
|
||||
for q, ij in infj.items():
|
||||
st = ij.wait()
|
||||
if st.success:
|
||||
out = ij.download_output_data()
|
||||
key = [k for k in out if "hidden" in k or "output" in k.lower()] or list(out)
|
||||
got = np.asarray(out[key[0]][0])
|
||||
res[q] = round(snr(ref_out, got), 2)
|
||||
print(f" {q}: {res[q]} dB", flush=True)
|
||||
else:
|
||||
res[q] = None
|
||||
print(f" {q} inference FAILED:", (st.message or "")[:160], flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_parity_a8.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,54 @@
|
||||
"""The ring decode layers across every processor the S25 exposes.
|
||||
|
||||
llama.cpp cannot be run on an AI Hub device, so the honest phone-CPU
|
||||
reference is the same graph compiled for TFLite and run on the CPU. That
|
||||
gives an NPU-vs-CPU ratio on identical work, on identical silicon.
|
||||
"""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
GRAPHS = ["ring_sliding_ctx1024", "ring_full_ctx1024"]
|
||||
TARGETS = [
|
||||
("qnn_npu", "--target_runtime qnn_context_binary --quantize_full_type float16", "npu"),
|
||||
("tflite_npu", "--target_runtime tflite", "npu"),
|
||||
("tflite_cpu", "--target_runtime tflite", "cpu"),
|
||||
("tflite_gpu", "--target_runtime tflite", "gpu"),
|
||||
("onnx_npu", "--target_runtime onnx", "npu"),
|
||||
]
|
||||
|
||||
pend, res = {}, {}
|
||||
for g in GRAPHS:
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/{g}.onnx", name=f"spark25_{g}_proc")
|
||||
for name, opt, unit in TARGETS:
|
||||
try:
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark {g} {name}", options=opt)
|
||||
pend[(g, name, unit)] = cj
|
||||
print("compile", g, name, cj.job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", g, name, str(e)[:140], flush=True)
|
||||
|
||||
prof = {}
|
||||
for key, cj in pend.items():
|
||||
st = cj.wait()
|
||||
if st.success:
|
||||
prof[key] = hub.submit_profile_job(model=cj.get_target_model(),
|
||||
device=hub.Device(DEV),
|
||||
name=f"spark {key[0]} {key[1]}",
|
||||
options=f"--compute_unit {key[2]}")
|
||||
else:
|
||||
print("compile", key, "FAILED", (st.message or "")[:140], flush=True)
|
||||
for key, pj in prof.items():
|
||||
st = pj.wait()
|
||||
tag = f"{key[0]}|{key[1]}"
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[tag] = p["execution_summary"]["estimated_inference_time"] / 1e3
|
||||
print(f" {tag}: {res[tag]:.3f} ms", flush=True)
|
||||
else:
|
||||
res[tag] = None
|
||||
print(f" {tag} FAILED: {(st.message or '')[:160]}", flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_processors.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Ring-buffer decode step on a real Galaxy S25, plus the output head.
|
||||
|
||||
The roll layout spent ~20% of a sliding layer writing the recomputed window
|
||||
back out. The ring layout returns one K/V entry instead, so this measures
|
||||
what that bought. The output head (2048 x 131072, tied to the embedding) was
|
||||
missing from the earlier per-layer numbers entirely and is measured here.
|
||||
"""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
|
||||
JOBS = [
|
||||
("ring_sliding_ctx1024", ["fp16", "w4a16"]),
|
||||
("ring_full_ctx1024", ["fp16", "w4a16"]),
|
||||
("ring_full_ctx4096", ["w4a16"]),
|
||||
("lm_head", ["w4a16", "w8a16"]),
|
||||
]
|
||||
OPT = {
|
||||
"fp16": "--target_runtime qnn_context_binary --quantize_full_type float16",
|
||||
"w8a16": "--target_runtime qnn_context_binary --quantize_full_type w8a16 --quantize_io",
|
||||
"w4a16": "--target_runtime qnn_context_binary --quantize_full_type w4a16 --quantize_io",
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
pend = {}
|
||||
for name, quants in JOBS:
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/{name}.onnx", name=f"spark25_{name}")
|
||||
print("uploaded", name, flush=True)
|
||||
for q in quants:
|
||||
try:
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark {name} {q}", options=OPT[q])
|
||||
pend[(name, q)] = cj
|
||||
print("compile", name, q, cj.job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", name, q, str(e)[:160], flush=True)
|
||||
|
||||
prof, res = {}, {}
|
||||
for key, cj in pend.items():
|
||||
st = cj.wait()
|
||||
print("compile", key, st.code, (st.message or "")[:150], flush=True)
|
||||
if st.success:
|
||||
prof[key] = hub.submit_profile_job(
|
||||
model=cj.get_target_model(), device=hub.Device(DEV),
|
||||
name=f"spark {key[0]} {key[1]}", options="--compute_unit npu")
|
||||
for key, pj in prof.items():
|
||||
st = pj.wait()
|
||||
tag = f"{key[0]}|{key[1]}"
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[tag] = p["execution_summary"]["estimated_inference_time"] / 1e3
|
||||
json.dump(p, open(f"{HERE}/profile_{key[0]}_{key[1]}.json", "w"))
|
||||
print(f" {tag}: {res[tag]:.3f} ms", flush=True)
|
||||
else:
|
||||
res[tag] = None
|
||||
print(f" {tag} FAILED: {(st.message or '')[:200]}", flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_ring.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Ring sliding layer, concat-free (online softmax), at w4a16 and fp16."""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
OPT = {
|
||||
"w4a16": "--target_runtime qnn_context_binary --quantize_full_type w4a16 --quantize_io",
|
||||
"fp16": "--target_runtime qnn_context_binary --quantize_full_type float16",
|
||||
}
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/ring2_sliding_ctx1024.onnx", name="spark25_ring2_sliding")
|
||||
pend, res = {}, {}
|
||||
for q, o in OPT.items():
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark ring2 sliding {q}", options=o)
|
||||
pend[q] = cj; print("compile", q, cj.job_id, flush=True)
|
||||
prof = {}
|
||||
for q, cj in pend.items():
|
||||
st = cj.wait(); print("compile", q, st.code, (st.message or "")[:180], flush=True)
|
||||
if st.success:
|
||||
prof[q] = hub.submit_profile_job(model=cj.get_target_model(), device=hub.Device(DEV),
|
||||
name=f"spark ring2 sliding {q}",
|
||||
options="--compute_unit npu")
|
||||
for q, pj in prof.items():
|
||||
st = pj.wait()
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[q] = p["execution_summary"]["estimated_inference_time"]/1e3
|
||||
json.dump(p, open(f"{HERE}/profile_ring2_sliding_{q}.json","w"))
|
||||
print(f" ring2 sliding {q}: {res[q]:.3f} ms", flush=True)
|
||||
else:
|
||||
res[q] = None; print(f" {q} FAILED:", (st.message or "")[:200], flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_ring2.json","w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Ring sliding layer, concat-free (online softmax), at w4a16 and fp16."""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
OPT = {
|
||||
"w4a16": "--target_runtime qnn_context_binary --quantize_full_type w4a16 --quantize_io",
|
||||
"fp16": "--target_runtime qnn_context_binary --quantize_full_type float16",
|
||||
}
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/ring3_sliding.onnx", name="spark25_ring3_sliding")
|
||||
pend, res = {}, {}
|
||||
for q, o in OPT.items():
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark ring3 sliding {q}", options=o)
|
||||
pend[q] = cj; print("compile", q, cj.job_id, flush=True)
|
||||
prof = {}
|
||||
for q, cj in pend.items():
|
||||
st = cj.wait(); print("compile", q, st.code, (st.message or "")[:180], flush=True)
|
||||
if st.success:
|
||||
prof[q] = hub.submit_profile_job(model=cj.get_target_model(), device=hub.Device(DEV),
|
||||
name=f"spark ring3 sliding {q}",
|
||||
options="--compute_unit npu")
|
||||
for q, pj in prof.items():
|
||||
st = pj.wait()
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[q] = p["execution_summary"]["estimated_inference_time"]/1e3
|
||||
json.dump(p, open(f"{HERE}/profile_ring3_sliding_{q}.json","w"))
|
||||
print(f" ring3 sliding {q}: {res[q]:.3f} ms", flush=True)
|
||||
else:
|
||||
res[q] = None; print(f" {q} FAILED:", (st.message or "")[:200], flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_ring3.json","w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,44 @@
|
||||
"""w4a16 retry for the ring layout, after expressing the new-token term as a
|
||||
matmul instead of a two-axis broadcast multiply (which QNN's context-binary
|
||||
converter rejected at w4a16 while accepting it at fp16)."""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
OPT = "--target_runtime qnn_context_binary --quantize_full_type w4a16 --quantize_io"
|
||||
NAMES = ["ring_sliding_ctx1024", "ring_full_ctx1024", "ring_full_ctx4096",
|
||||
]
|
||||
|
||||
pend, res = {}, {}
|
||||
for n in NAMES:
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/{n}.onnx", name=f"spark25_{n}_v3")
|
||||
try:
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark {n} w4a16 v3", options=OPT)
|
||||
pend[n] = cj
|
||||
print("compile", n, cj.job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", n, str(e)[:200], flush=True)
|
||||
|
||||
prof = {}
|
||||
for n, cj in pend.items():
|
||||
st = cj.wait()
|
||||
print("compile", n, st.code, (st.message or "")[:200], flush=True)
|
||||
if st.success:
|
||||
prof[n] = hub.submit_profile_job(model=cj.get_target_model(),
|
||||
device=hub.Device(DEV),
|
||||
name=f"spark {n} w4a16 v3",
|
||||
options="--compute_unit npu")
|
||||
for n, pj in prof.items():
|
||||
st = pj.wait()
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[n] = p["execution_summary"]["estimated_inference_time"] / 1e3
|
||||
json.dump(p, open(f"{HERE}/profile_{n}_w4a16.json", "w"))
|
||||
print(f" {n}: {res[n]:.3f} ms", flush=True)
|
||||
else:
|
||||
res[n] = None
|
||||
print(f" {n} FAILED: {(st.message or '')[:250]}", flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_ring_w4_v3.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Profile Spark-X2.5 decoder layers on real Snapdragon hardware.
|
||||
|
||||
One sliding layer and one full-attention layer are profiled separately. The
|
||||
sliding layer's cache is a fixed 512 entries at any context; the full layer's
|
||||
grows, so it is profiled at 1024 and 4096 to show the slope. A whole decode
|
||||
step is 21 sliding + 7 full layers plus the head.
|
||||
"""
|
||||
import json, sys
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
HID, HEADS, KV, D = 2048, 8, 2, 256
|
||||
SW = 512
|
||||
|
||||
GRAPHS = [
|
||||
("sliding", 1024, SW),
|
||||
("full", 1024, 1024),
|
||||
("full", 4096, 4096),
|
||||
]
|
||||
QUANTS = [
|
||||
("fp16", "--target_runtime qnn_context_binary --quantize_full_type float16"),
|
||||
("w8a16", "--target_runtime qnn_context_binary --quantize_full_type w8a16 --quantize_io"),
|
||||
("w4a16", "--target_runtime qnn_context_binary --quantize_full_type w4a16 --quantize_io"),
|
||||
]
|
||||
DEVICES = ["Samsung Galaxy S25"]
|
||||
|
||||
|
||||
def spec_for(ctx, cache_len):
|
||||
s = {
|
||||
"x": ((1, 1, HID), "float32"),
|
||||
"cos_sw": ((1, 1, 1, D), "float32"),
|
||||
"sin_sw": ((1, 1, 1, D), "float32"),
|
||||
"cos_full": ((1, 1, 1, D // 4), "float32"),
|
||||
"sin_full": ((1, 1, 1, D // 4), "float32"),
|
||||
"mask_sw": ((1, 1, 1, SW), "float32"),
|
||||
"mask_full": ((1, 1, 1, ctx + 1), "float32"),
|
||||
"k_cache_0": ((1, KV, cache_len, D), "float32"),
|
||||
"v_cache_0": ((1, KV, cache_len, D), "float32"),
|
||||
}
|
||||
return s
|
||||
|
||||
|
||||
def main():
|
||||
results, pending = {}, {}
|
||||
for kind, ctx, cache_len in GRAPHS:
|
||||
path = f"{HERE}/onnx/spark_{kind}_ctx{ctx}.onnx"
|
||||
mdl = hub.upload_model(path, name=f"spark25_{kind}_ctx{ctx}")
|
||||
print("uploaded", kind, ctx, flush=True)
|
||||
for dev in DEVICES:
|
||||
for q, opt in QUANTS:
|
||||
key = (kind, ctx, dev, q)
|
||||
try:
|
||||
# No input_specs: every shape is static, and the ONNX
|
||||
# export prunes the inputs a given layer type never
|
||||
# reads (a full layer has no sliding mask), so a
|
||||
# hand-written spec would not match the graph.
|
||||
cj = hub.submit_compile_job(
|
||||
model=mdl, device=hub.Device(dev),
|
||||
name=f"spark {kind} ctx{ctx} {q}", options=opt)
|
||||
pending[key] = cj
|
||||
print("compile submitted", key, cj.job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", key, str(e)[:160], flush=True)
|
||||
|
||||
prof = {}
|
||||
for key, cj in pending.items():
|
||||
st = cj.wait()
|
||||
print("compile", key, st.code, (st.message or "")[:160], flush=True)
|
||||
if st.success:
|
||||
prof[key] = hub.submit_profile_job(
|
||||
model=cj.get_target_model(), device=hub.Device(key[2]),
|
||||
name=f"spark {key[0]} ctx{key[1]} {key[3]}",
|
||||
options="--compute_unit npu")
|
||||
|
||||
for key, pj in prof.items():
|
||||
st = pj.wait()
|
||||
kind, ctx, dev, q = key
|
||||
tag = f"{kind}|ctx{ctx}|{dev}|{q}"
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
ms = p["execution_summary"]["estimated_inference_time"] / 1e3
|
||||
results[tag] = ms
|
||||
json.dump(p, open(f"{HERE}/profile_spark_{kind}_{ctx}_{q}.json", "w"))
|
||||
print(f" {tag}: {ms:.3f} ms/layer", flush=True)
|
||||
else:
|
||||
results[tag] = None
|
||||
print(f" {tag} FAILED: {(st.message or '')[:200]}", flush=True)
|
||||
json.dump(results, open(f"{HERE}/aihub_spark_layers.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Same Spark layer, three phone runtimes, one device.
|
||||
|
||||
Isolates runtime choice from export choice: if QNN, TFLite and ONNX Runtime
|
||||
land close together, the export layout is what sets on-device speed, not the
|
||||
engine. Run on the sliding layer, which is 21 of Spark's 28 layers.
|
||||
"""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
PATH = f"{HERE}/onnx/spark_sliding_ctx1024.onnx"
|
||||
|
||||
# Each runtime spells precision differently: only QNN takes
|
||||
# --quantize_full_type, TFLite and ONNX reject it outright, and a QNN context
|
||||
# binary is NPU-only so the GPU leg has to go through TFLite.
|
||||
RUNTIMES = [
|
||||
("qnn_npu", "--target_runtime qnn_context_binary --quantize_full_type float16", "npu"),
|
||||
("tflite_npu", "--target_runtime tflite", "npu"),
|
||||
("onnx_npu", "--target_runtime onnx", "npu"),
|
||||
("tflite_gpu", "--target_runtime tflite", "gpu"),
|
||||
("tflite_cpu", "--target_runtime tflite", "cpu"),
|
||||
]
|
||||
|
||||
|
||||
def main():
|
||||
mdl = hub.upload_model(PATH, name="spark25_sliding_runtimes")
|
||||
pend = {}
|
||||
for name, opt, unit in RUNTIMES:
|
||||
try:
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark sliding {name}", options=opt)
|
||||
pend[name] = (cj, unit)
|
||||
print("compile", name, cj.job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", name, str(e)[:160], flush=True)
|
||||
|
||||
prof, res = {}, {}
|
||||
for name, (cj, unit) in pend.items():
|
||||
st = cj.wait()
|
||||
print("compile", name, st.code, (st.message or "")[:140], flush=True)
|
||||
if st.success:
|
||||
prof[name] = hub.submit_profile_job(
|
||||
model=cj.get_target_model(), device=hub.Device(DEV),
|
||||
name=f"spark sliding {name}", options=f"--compute_unit {unit}")
|
||||
for name, pj in prof.items():
|
||||
st = pj.wait()
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[name] = p["execution_summary"]["estimated_inference_time"] / 1e3
|
||||
print(f" {name}: {res[name]:.3f} ms/layer", flush=True)
|
||||
else:
|
||||
res[name] = None
|
||||
print(f" {name} FAILED: {(st.message or '')[:180]}", flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_spark_runtimes.json", "w"), indent=1)
|
||||
print("done", flush=True)
|
||||
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,42 @@
|
||||
"""int8 activations: does halving KV-cache traffic pay?
|
||||
|
||||
Input -- reading the layer's KV cache -- is 13.4% of a sliding layer, and at
|
||||
w4a16 the cache crosses the boundary as int16. w8a8 halves that and uses int8
|
||||
MACs. Prior work in this repo found automatic int8 numerically broken on two
|
||||
other model families, so this measures speed first and only then asks whether
|
||||
the fidelity is survivable.
|
||||
"""
|
||||
import json
|
||||
import qai_hub as hub
|
||||
HERE = "/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge"
|
||||
DEV = "Samsung Galaxy S25"
|
||||
CAND = {
|
||||
"w8a8": "--target_runtime qnn_context_binary --quantize_full_type w8a8 --quantize_io",
|
||||
"w4a8": "--target_runtime qnn_context_binary --quantize_full_type w4a8 --quantize_io",
|
||||
}
|
||||
mdl = hub.upload_model(f"{HERE}/onnx/ring3_sliding.onnx", name="spark25_ring3_a8")
|
||||
pend, res = {}, {}
|
||||
for q, o in CAND.items():
|
||||
try:
|
||||
cj = hub.submit_compile_job(model=mdl, device=hub.Device(DEV),
|
||||
name=f"spark ring3 {q}", options=o)
|
||||
pend[q] = cj; print("compile", q, cj.job_id, flush=True)
|
||||
except Exception as e:
|
||||
print("submit failed", q, str(e)[:180], flush=True)
|
||||
prof = {}
|
||||
for q, cj in pend.items():
|
||||
st = cj.wait(); print("compile", q, st.code, (st.message or "")[:180], flush=True)
|
||||
if st.success:
|
||||
prof[q] = hub.submit_profile_job(model=cj.get_target_model(), device=hub.Device(DEV),
|
||||
name=f"spark ring3 {q}", options="--compute_unit npu")
|
||||
for q, pj in prof.items():
|
||||
st = pj.wait()
|
||||
if st.success:
|
||||
p = pj.download_profile()
|
||||
res[q] = p["execution_summary"]["estimated_inference_time"]/1e3
|
||||
json.dump(p, open(f"{HERE}/profile_ring3_{q}.json","w"))
|
||||
print(f" ring3 {q}: {res[q]:.3f} ms", flush=True)
|
||||
else:
|
||||
res[q] = None; print(f" {q} FAILED:", (st.message or "")[:180], flush=True)
|
||||
json.dump(res, open(f"{HERE}/aihub_w8a8.json","w"), indent=1)
|
||||
print("done", flush=True)
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Aggregate decode throughput across concurrent sequences (vLLM, CPU).
|
||||
|
||||
Mirrors `llama-batched-bench`: B sequences each with a short prompt generate
|
||||
N tokens; the figure is total generated tokens per second.
|
||||
"""
|
||||
import argparse, json, os, time
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
|
||||
ap.add_argument("--dtype", default="bfloat16")
|
||||
ap.add_argument("--batches", default="1,2,4,8,16")
|
||||
ap.add_argument("--gen", type=int, default=128)
|
||||
ap.add_argument("--prompt", type=int, default=128)
|
||||
ap.add_argument("--tag", default="vllm")
|
||||
ap.add_argument("--out", default="batch_bench.jsonl")
|
||||
a = ap.parse_args()
|
||||
|
||||
import vllm_omni # noqa: F401
|
||||
from vllm import LLM, SamplingParams
|
||||
batches = [int(b) for b in a.batches.split(",")]
|
||||
llm = LLM(model=a.model, dtype=a.dtype, max_model_len=1024,
|
||||
max_num_seqs=max(batches), enable_prefix_caching=False,
|
||||
trust_remote_code=False, disable_log_stats=True)
|
||||
sp = SamplingParams(temperature=0.0, max_tokens=a.gen, ignore_eos=True)
|
||||
res = {}
|
||||
for b in batches:
|
||||
prompts = [{"prompt_token_ids": list(range(10 + i * 7, 10 + i * 7 + a.prompt))}
|
||||
for i in range(b)]
|
||||
llm.generate(prompts[:1], SamplingParams(temperature=0.0, max_tokens=1,
|
||||
ignore_eos=True), use_tqdm=False)
|
||||
t0 = time.perf_counter()
|
||||
llm.generate(prompts, sp, use_tqdm=False)
|
||||
dt = time.perf_counter() - t0
|
||||
res[b] = round(b * a.gen / dt, 2)
|
||||
print(f"batch {b:>3}: {res[b]:>8.2f} gen tok/s total", flush=True)
|
||||
row = {"tag": a.tag, "threads": os.environ.get("OMP_NUM_THREADS", ""),
|
||||
"gen": a.gen, "prompt": a.prompt, "tok_s_by_batch": res}
|
||||
with open(a.out, "a") as f:
|
||||
f.write(json.dumps(row) + "\n")
|
||||
print(json.dumps(row))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,88 @@
|
||||
"""pp512 / tg128 for Spark-X2.5 on one GPU, for vLLM or SGLang.
|
||||
|
||||
Same protocol as the CPU grid and llama-bench: a 512-token prefill measured on
|
||||
its own, then 128 generated tokens with the one-token prefill subtracted.
|
||||
Prefix caching is disabled on both engines so a replayed prompt is actually
|
||||
recomputed.
|
||||
"""
|
||||
import argparse, json, os, time
|
||||
|
||||
|
||||
def bench_vllm(a):
|
||||
import vllm_omni # noqa: F401 registers Spark2_5ForCausalLM
|
||||
from vllm import LLM, SamplingParams
|
||||
t0 = time.perf_counter()
|
||||
llm = LLM(model=a.model, dtype=a.dtype, max_model_len=a.max_model_len,
|
||||
max_num_seqs=1, enable_prefix_caching=False,
|
||||
gpu_memory_utilization=0.85, trust_remote_code=False,
|
||||
disable_log_stats=True)
|
||||
init = time.perf_counter() - t0
|
||||
sp1 = SamplingParams(temperature=0.0, max_tokens=1, ignore_eos=True)
|
||||
spg = SamplingParams(temperature=0.0, max_tokens=a.tg, ignore_eos=True)
|
||||
pp_prompt = {"prompt_token_ids": list(range(10, 10 + a.pp))}
|
||||
tg_prompt = {"prompt_token_ids": [10]}
|
||||
|
||||
def run(p, sp):
|
||||
s = time.perf_counter()
|
||||
llm.generate([p], sp, use_tqdm=False)
|
||||
return time.perf_counter() - s
|
||||
|
||||
run(pp_prompt, sp1)
|
||||
pp = min(run(pp_prompt, sp1) for _ in range(a.reps))
|
||||
tg = min(run(tg_prompt, spg) for _ in range(a.reps))
|
||||
tg1 = min(run(tg_prompt, sp1) for _ in range(a.reps))
|
||||
return init, pp, tg, tg1
|
||||
|
||||
|
||||
def bench_sglang(a):
|
||||
import sglang as sgl
|
||||
t0 = time.perf_counter()
|
||||
llm = sgl.Engine(model_path=a.model, dtype=a.dtype,
|
||||
context_length=a.max_model_len,
|
||||
disable_radix_cache=True, mem_fraction_static=0.85,
|
||||
disable_cuda_graph=False, log_level="error")
|
||||
init = time.perf_counter() - t0
|
||||
pp_ids = list(range(10, 10 + a.pp))
|
||||
|
||||
def run(ids, n):
|
||||
s = time.perf_counter()
|
||||
llm.generate(input_ids=[ids],
|
||||
sampling_params={"temperature": 0.0, "max_new_tokens": n,
|
||||
"ignore_eos": True})
|
||||
return time.perf_counter() - s
|
||||
|
||||
run(pp_ids, 1)
|
||||
pp = min(run(pp_ids, 1) for _ in range(a.reps))
|
||||
tg = min(run([10], a.tg) for _ in range(a.reps))
|
||||
tg1 = min(run([10], 1) for _ in range(a.reps))
|
||||
llm.shutdown()
|
||||
return init, pp, tg, tg1
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--engine", choices=["vllm", "sglang"], required=True)
|
||||
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
|
||||
ap.add_argument("--dtype", default="bfloat16")
|
||||
ap.add_argument("--pp", type=int, default=512)
|
||||
ap.add_argument("--tg", type=int, default=128)
|
||||
ap.add_argument("--reps", type=int, default=3)
|
||||
ap.add_argument("--max-model-len", type=int, default=2048)
|
||||
ap.add_argument("--out", default="gpu_bench.jsonl")
|
||||
a = ap.parse_args()
|
||||
|
||||
init, pp, tg, tg1 = (bench_vllm if a.engine == "vllm" else bench_sglang)(a)
|
||||
res = {
|
||||
"engine": a.engine, "dtype": a.dtype, "gpu": os.environ.get("GPU_NAME", ""),
|
||||
"init_s": round(init, 2),
|
||||
"pp_tokens": a.pp, "pp_s": round(pp, 4), "pp_tps": round(a.pp / pp, 2),
|
||||
"tg_tokens": a.tg, "tg_s": round(tg - tg1, 4),
|
||||
"tg_tps": round((a.tg - 1) / max(tg - tg1, 1e-9), 2),
|
||||
}
|
||||
print(json.dumps(res, indent=2))
|
||||
with open(a.out, "a") as f:
|
||||
f.write(json.dumps(res) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Microbenchmark of candidate CPU 4-bit weight kernels at Spark-X2.5 shapes.
|
||||
|
||||
Decode is one token, so every linear is a GEMV and the only thing that matters
|
||||
is how fast the weight bytes can be streamed. llama.cpp's Q4_K_M moves ~4.5
|
||||
bits per weight; anything we pick has to get close to that at a comparable
|
||||
fraction of peak bandwidth.
|
||||
"""
|
||||
import argparse, json, time
|
||||
import torch
|
||||
|
||||
SHAPES = [
|
||||
("qkv", 3072, 2048),
|
||||
("o", 2048, 2048),
|
||||
("gate_up", 13312, 2048),
|
||||
("down", 2048, 6656),
|
||||
("lm_head", 131072, 2048),
|
||||
]
|
||||
|
||||
|
||||
def timeit(fn, iters, warmup=5):
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
best = float("inf")
|
||||
for _ in range(iters):
|
||||
t = time.perf_counter()
|
||||
fn()
|
||||
best = min(best, time.perf_counter() - t)
|
||||
return best
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--m", type=int, default=1)
|
||||
ap.add_argument("--iters", type=int, default=30)
|
||||
ap.add_argument("--group", type=int, default=128)
|
||||
ap.add_argument("--out", default="int4_kernels.json")
|
||||
a = ap.parse_args()
|
||||
torch.set_num_threads(torch.get_num_threads())
|
||||
print("threads", torch.get_num_threads())
|
||||
|
||||
rows = []
|
||||
for name, N, K in SHAPES:
|
||||
x16 = torch.randn(a.m, K, dtype=torch.bfloat16)
|
||||
x32 = x16.float()
|
||||
w16 = torch.randn(N, K, dtype=torch.bfloat16)
|
||||
rec = {"name": name, "N": N, "K": K, "M": a.m}
|
||||
|
||||
# bf16 reference (oneDNN through aten)
|
||||
t = timeit(lambda: torch.nn.functional.linear(x16, w16), a.iters)
|
||||
rec["bf16_ms"] = t * 1e3
|
||||
rec["bf16_GBps"] = N * K * 2 / t / 1e9
|
||||
|
||||
# int8 weight-only (aten _weight_int8pack_mm)
|
||||
try:
|
||||
w8 = torch.randint(-127, 127, (N, K), dtype=torch.int8)
|
||||
s8 = torch.rand(N, dtype=torch.bfloat16) * 0.01
|
||||
t = timeit(lambda: torch.ops.aten._weight_int8pack_mm(x16, w8, s8), a.iters)
|
||||
rec["int8_ms"] = t * 1e3
|
||||
rec["int8_GBps"] = N * K / t / 1e9
|
||||
except Exception as e:
|
||||
rec["int8_err"] = f"{type(e).__name__}: {e}"
|
||||
|
||||
# int4 group-wise (aten _weight_int4pack_mm_for_cpu)
|
||||
for dt, xx in (("bf16", x16), ("fp32", x32)):
|
||||
try:
|
||||
wint = torch.randint(0, 15, (N, K), dtype=torch.int32)
|
||||
packed = torch.ops.aten._convert_weight_to_int4pack_for_cpu(wint, 1)
|
||||
sz = torch.rand(K // a.group, N, 2, dtype=xx.dtype)
|
||||
f = lambda: torch.ops.aten._weight_int4pack_mm_for_cpu(xx, packed, a.group, sz)
|
||||
out = f()
|
||||
t = timeit(f, a.iters)
|
||||
rec[f"int4_{dt}_ms"] = t * 1e3
|
||||
rec[f"int4_{dt}_GBps"] = N * K * 0.5 / t / 1e9
|
||||
rec[f"int4_{dt}_out"] = str(tuple(out.shape)) + " " + str(out.dtype)
|
||||
except Exception as e:
|
||||
rec[f"int4_{dt}_err"] = f"{type(e).__name__}: {e}"
|
||||
|
||||
rows.append(rec)
|
||||
print(json.dumps(rec))
|
||||
|
||||
with open(a.out, "w") as f:
|
||||
json.dump(rows, f, indent=2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Is vLLM's own CPU 4-bit kernel faster than the tinygemm one?
|
||||
|
||||
vLLM 0.28's CPU wheel does ship a 4-bit dense GEMM -- `_C::int4_scaled_mm_cpu`,
|
||||
an AMX W4A8 kernel reached by loading a GPTQ checkpoint with
|
||||
`VLLM_CPU_INT4_W4A8=1`. It is not named like the ops an earlier search looked
|
||||
for, which is why it was missed. This measures it against
|
||||
`aten::_weight_int4pack_mm_for_cpu` at Spark-X2.5's shapes.
|
||||
"""
|
||||
import json, time
|
||||
|
||||
import torch
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
pack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
|
||||
_to_tinygemm, quantize_groupwise,
|
||||
)
|
||||
|
||||
SHAPES = [("qkv", 3072, 2048), ("o", 2048, 2048), ("gate_up", 13312, 2048),
|
||||
("down", 2048, 6656), ("lm_head", 131072, 2048)]
|
||||
GROUP = 128
|
||||
|
||||
|
||||
def timeit(fn, iters=20):
|
||||
for _ in range(5):
|
||||
fn()
|
||||
best = float("inf")
|
||||
for _ in range(iters):
|
||||
t = time.perf_counter()
|
||||
fn()
|
||||
best = min(best, time.perf_counter() - t)
|
||||
return best
|
||||
|
||||
|
||||
def build_w4a8(w, group):
|
||||
"""[N, K] float -> the blocked tensors `int4_scaled_mm_cpu` wants."""
|
||||
n, k = w.shape
|
||||
gnum = k // group
|
||||
wf = w.float().reshape(n, gnum, group)
|
||||
wmax, wmin = wf.amax(-1), wf.amin(-1)
|
||||
scale = ((wmax - wmin).clamp(min=1e-9) / 15.0)
|
||||
zp = (-wmin / scale).round().clamp(0, 15)
|
||||
q = ((wf - wmin.unsqueeze(-1)) / scale.unsqueeze(-1)).round().clamp_(0, 15)
|
||||
codes = q.reshape(n, k).to(torch.int32).t().contiguous() # [K, N]
|
||||
zpi = zp.to(torch.int32).t().contiguous() # [G, N]
|
||||
|
||||
# AWQ interleave the kernel's repack expects
|
||||
codes = codes.view(k, n // 8, 8)[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(k, n)
|
||||
zpi = zpi.view(gnum, n // 8, 8)[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(gnum, n)
|
||||
packed_w = pack_quantized_values_into_int32(codes, scalar_types.uint4, 1).contiguous()
|
||||
packed_zp = pack_quantized_values_into_int32(zpi, scalar_types.uint4, 1).contiguous()
|
||||
scales = scale.t().contiguous().to(torch.bfloat16) # [G, N]
|
||||
return ops.convert_weight_packed_scale_zp(
|
||||
packed_w, packed_zp, scales, ops.CPUQuantAlgo.AWQ)
|
||||
|
||||
|
||||
def main():
|
||||
rows = []
|
||||
for name, N, K in SHAPES:
|
||||
w = torch.randn(N, K)
|
||||
packed, s, z = quantize_groupwise(w, GROUP)
|
||||
inner, sz = _to_tinygemm(packed, s, z)
|
||||
try:
|
||||
bw, bzp, bs = build_w4a8(w, GROUP)
|
||||
except Exception as e:
|
||||
print(json.dumps({"name": name, "w4a8_build_error": f"{type(e).__name__}: {e}"}))
|
||||
continue
|
||||
rec = {"name": name, "N": N, "K": K, "group": GROUP}
|
||||
for M in (1, 8, 512):
|
||||
x = torch.randn(M, K, dtype=torch.bfloat16)
|
||||
t_tg = timeit(lambda: torch.ops.aten._weight_int4pack_mm_for_cpu(x, inner, GROUP, sz))
|
||||
try:
|
||||
out = ops.int4_scaled_mm_cpu(x, bw, bzp, bs, None)
|
||||
t_w4 = timeit(lambda: ops.int4_scaled_mm_cpu(x, bw, bzp, bs, None))
|
||||
ref = torch.ops.aten._weight_int4pack_mm_for_cpu(x, inner, GROUP, sz)
|
||||
rel = ((out.float() - ref.float()).norm() / ref.float().norm()).item()
|
||||
rec[f"m{M}"] = {"tinygemm_ms": round(t_tg * 1e3, 4),
|
||||
"w4a8_ms": round(t_w4 * 1e3, 4),
|
||||
"speedup": round(t_tg / t_w4, 2),
|
||||
"rel_diff_vs_tinygemm": round(rel, 4)}
|
||||
except Exception as e:
|
||||
rec[f"m{M}"] = {"tinygemm_ms": round(t_tg * 1e3, 4),
|
||||
"w4a8_error": f"{type(e).__name__}: {str(e)[:90]}"}
|
||||
rows.append(rec)
|
||||
print(json.dumps(rec))
|
||||
with open("w4a8_vs_tinygemm.json", "w") as f:
|
||||
json.dump(rows, f, indent=2)
|
||||
for M in (1, 8, 512):
|
||||
tg = sum(r[f"m{M}"].get("tinygemm_ms", 0) for r in rows[:4])
|
||||
w4 = sum(r[f"m{M}"].get("w4a8_ms", 0) for r in rows[:4])
|
||||
if w4:
|
||||
head_tg = rows[4][f"m{M}"].get("tinygemm_ms", 0)
|
||||
head_w4 = rows[4][f"m{M}"].get("w4a8_ms", 0)
|
||||
print(f"M={M:4d} 28 layers + head: tinygemm {tg*28+head_tg:8.3f} ms "
|
||||
f"w4a8 {w4*28+head_w4:8.3f} ms")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,49 @@
|
||||
"""Real inputs for the output head: the last layer's output, pre-final-norm."""
|
||||
import importlib.util, json, sys, types
|
||||
from pathlib import Path
|
||||
import numpy as np, torch
|
||||
|
||||
ROOT = Path("/data/zhoutaichang/embedding_infer/vllm-omni")
|
||||
HERE = Path("/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge")
|
||||
def _load(n, p):
|
||||
sp = importlib.util.spec_from_file_location(n, p); m = importlib.util.module_from_spec(sp)
|
||||
sys.modules[n] = m; sp.loader.exec_module(m); return m
|
||||
for pkg in ("vllm_omni", "vllm_omni.edge"):
|
||||
mod = types.ModuleType(pkg); mod.__path__ = []; sys.modules[pkg] = mod
|
||||
_load("vllm_omni.edge.qnn_export", ROOT / "vllm_omni/edge/qnn_export.py")
|
||||
se = _load("vllm_omni.edge.spark_export", ROOT / "vllm_omni/edge/spark_export.py")
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
MODEL = "XHToken/Spark-X2.5-1.7B"
|
||||
SNAP = ("/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/"
|
||||
"snapshots/448e61eb392c00f2c403185c5b56d5e0665bfaab")
|
||||
TEXTS = ["The quick brown fox jumps over the lazy dog. " * 60,
|
||||
"Engineers argued about memory bandwidth on the train home. " * 60,
|
||||
"def forward(self, x):\n return self.down(self.act(self.gate(x)))\n" * 60,
|
||||
"移动端推理的瓶颈通常是内存带宽而不是算力。" * 80]
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||||
m = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True,
|
||||
dtype=torch.float32).eval()
|
||||
hf = json.loads(Path(f"{SNAP}/config.json").read_text())
|
||||
cfg = se.config_from_spark(hf, slice(0, 0), include_lm_head=True)
|
||||
step = se.SparkDecodeStep(cfg); se.load_spark_weights(step, SNAP); step.eval()
|
||||
|
||||
grabbed = {}
|
||||
h = m.model.layers[-1].register_forward_hook(
|
||||
lambda mod, inp, out: grabbed.__setitem__("h", out[0] if isinstance(out, tuple) else out))
|
||||
xs, outs = [], []
|
||||
for t in TEXTS:
|
||||
ids = tok(t, add_special_tokens=False)["input_ids"][:256]
|
||||
with torch.no_grad():
|
||||
m(torch.tensor([ids]))
|
||||
x = grabbed["h"][:, -1:, :].float()
|
||||
logits = step(x, *[torch.zeros(1)] * 0) if False else step(
|
||||
x, torch.zeros(1,1,1,4), torch.zeros(1,1,1,4),
|
||||
torch.zeros(1,1,1,4), torch.zeros(1,1,1,4),
|
||||
torch.zeros(1,1,1,1), torch.zeros(1,1,1,1))[0]
|
||||
xs.append(x.numpy().astype(np.float32)); outs.append(logits.numpy().astype(np.float32))
|
||||
h.remove()
|
||||
np.savez(HERE / "real_acts_lmhead.npz", x=np.stack(xs), hidden_out=np.stack(outs))
|
||||
print("captured", len(TEXTS), "| x absmax", np.abs(np.stack(xs)).max(),
|
||||
"| logits absmax", np.abs(np.stack(outs)).max())
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Capture real layer inputs for calibration and parity.
|
||||
|
||||
Quantisation ranges inferred from random tensors describe a distribution the
|
||||
model never sees. This runs the reference on real text and records what
|
||||
actually enters a sliding layer at several positions, together with the fp32
|
||||
output of that layer, so the same tensors can calibrate the quantiser and
|
||||
then judge it.
|
||||
"""
|
||||
import importlib.util, json, sys, types
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
ROOT = Path("/data/zhoutaichang/embedding_infer/vllm-omni")
|
||||
HERE = Path("/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge")
|
||||
|
||||
|
||||
def _load(name, path):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
m = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = m
|
||||
spec.loader.exec_module(m)
|
||||
return m
|
||||
|
||||
|
||||
for pkg in ("vllm_omni", "vllm_omni.edge"):
|
||||
mod = types.ModuleType(pkg); mod.__path__ = []; sys.modules[pkg] = mod
|
||||
_load("vllm_omni.edge.qnn_export", ROOT / "vllm_omni/edge/qnn_export.py")
|
||||
se = _load("vllm_omni.edge.spark_export", ROOT / "vllm_omni/edge/spark_export.py")
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
MODEL = "XHToken/Spark-X2.5-1.7B"
|
||||
SNAP = ("/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/"
|
||||
"snapshots/448e61eb392c00f2c403185c5b56d5e0665bfaab")
|
||||
TEXTS = [
|
||||
"The quick brown fox jumps over the lazy dog. " * 300,
|
||||
"In a distant galaxy, engineers argued about memory bandwidth and cache lines. " * 200,
|
||||
"def forward(self, x):\n return self.down_proj(self.act_fn(self.gate_proj(x)))\n" * 200,
|
||||
"人工智能模型在移动设备上的推理速度取决于内存带宽和算子调度。" * 300,
|
||||
]
|
||||
C = 1024
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||||
m = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True,
|
||||
dtype=torch.float32).eval()
|
||||
hf = json.loads(Path(f"{SNAP}/config.json").read_text())
|
||||
import os
|
||||
WHICH = os.environ.get("CAPTURE", "sliding")
|
||||
li = hf["layer_types"].index(se.SLIDING if WHICH == "sliding" else se.FULL)
|
||||
cfg = se.config_from_spark(hf, slice(li, li + 1), include_lm_head=False,
|
||||
cache_layout="ring" if WHICH == "sliding" else "roll")
|
||||
step = se.SparkDecodeStep(cfg)
|
||||
se.load_spark_weights(step, SNAP)
|
||||
step.eval()
|
||||
|
||||
LT = se.SLIDING if WHICH == "sliding" else se.FULL
|
||||
sw = cfg.cache_len(LT, C)
|
||||
cos, sin = se._angles(C, cfg.rotary_dim(LT), cfg.rope_theta[LT], torch.float32)
|
||||
samples = {k: [] for k in ("x", "cos_sw", "sin_sw", "mask_sw", "k_cache_0",
|
||||
"v_cache_0", "hidden_out")}
|
||||
|
||||
for text in TEXTS:
|
||||
ids = tok(text, add_special_tokens=False)["input_ids"][: C + 1]
|
||||
assert len(ids) == C + 1
|
||||
with torch.no_grad():
|
||||
pre = m(torch.tensor([ids[:-1]]), use_cache=True)
|
||||
cache = pre.past_key_values
|
||||
try:
|
||||
k, v = cache.layers[li].keys, cache.layers[li].values
|
||||
except AttributeError:
|
||||
k, v = cache.key_cache[li], cache.value_cache[li]
|
||||
k, v = k[:, :, -sw:].float(), v[:, :, -sw:].float()
|
||||
if k.shape[2] < sw: # pad the ring, mask it out
|
||||
pad = sw - k.shape[2]
|
||||
z = torch.zeros(k.shape[0], k.shape[1], pad, k.shape[3])
|
||||
k, v = torch.cat([z, k], 2), torch.cat([z, v], 2)
|
||||
mask = torch.zeros(1, 1, 1, sw + 1)
|
||||
# Real hidden state entering this layer, not a random vector.
|
||||
# Real hidden state entering this layer: run the prefix through the
|
||||
# layers below it rather than assuming the embedding.
|
||||
x = m.model.embedding(torch.tensor([[ids[-1]]])).float()
|
||||
if li:
|
||||
pos = torch.tensor([[C]])
|
||||
hs = x
|
||||
for j in range(li):
|
||||
lt = hf["layer_types"][j]
|
||||
prf = m.config.get_partial_rotary_factor(lt)
|
||||
cj_, sj_ = __import__("importlib").import_module(
|
||||
m.__class__.__module__).compute_rope_cos_sin(
|
||||
pos[0], hf["head_dim"], m.config.get_rope_theta(lt),
|
||||
partial_rotary_factor=prf)
|
||||
hs = m.model.layers[j](hs, position_embeddings=(cj_, sj_),
|
||||
past_key_values=cache, cache_position=pos[0])
|
||||
hs = hs[0] if isinstance(hs, tuple) else hs
|
||||
x = hs.float()
|
||||
out = step(x, cos, sin, cos, sin, mask, mask, k, v)[0]
|
||||
for name, arr in (("x", x), ("cos_sw", cos), ("sin_sw", sin),
|
||||
("mask_sw", mask), ("k_cache_0", k), ("v_cache_0", v),
|
||||
("hidden_out", out)):
|
||||
samples[name].append(arr.numpy().astype(np.float32))
|
||||
|
||||
np.savez(HERE / f"real_acts_{WHICH}.npz",
|
||||
**{k: np.stack(v) for k, v in samples.items()})
|
||||
print("captured", len(TEXTS), "real samples for", WHICH, "layer", li)
|
||||
for k, v in samples.items():
|
||||
a = np.stack(v)
|
||||
print(f" {k:<12} {str(a.shape):<26} absmax {np.abs(a).max():.3f}")
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Compare the vLLM-Omni Spark2_5 implementation against ref.json."""
|
||||
import argparse, json
|
||||
import numpy as np
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--ref", default="ref.json")
|
||||
ap.add_argument("--ref-logits", default="ref_logits.npy")
|
||||
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
|
||||
ap.add_argument("--dtype", default="bfloat16")
|
||||
ap.add_argument("--max-model-len", type=int, default=2048)
|
||||
ap.add_argument("--out", default="parity.json")
|
||||
a = ap.parse_args()
|
||||
|
||||
ref = json.load(open(a.ref))
|
||||
ref_logits = np.load(a.ref_logits)
|
||||
|
||||
import vllm_omni # noqa: F401 registers Spark2_5ForCausalLM
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
llm = LLM(model=a.model, dtype=a.dtype,
|
||||
max_model_len=a.max_model_len, enforce_eager=True,
|
||||
max_num_seqs=1, trust_remote_code=False)
|
||||
n_new = len(ref["new_ids"])
|
||||
sp = SamplingParams(temperature=0.0, max_tokens=n_new, logprobs=20)
|
||||
out = llm.generate([{"prompt_token_ids": ref["prompt_ids"]}], sp)[0]
|
||||
v_ids = list(out.outputs[0].token_ids)
|
||||
|
||||
n = min(len(ref["new_ids"]), len(v_ids))
|
||||
match = sum(int(x == y) for x, y in zip(ref["new_ids"][:n], v_ids[:n]))
|
||||
# first-step logprob agreement on the reference top-20
|
||||
lp0 = out.outputs[0].logprobs[0]
|
||||
ref_top = np.argsort(-ref_logits)[:20]
|
||||
ref_lse = np.log(np.exp(ref_logits - ref_logits.max()).sum()) + ref_logits.max()
|
||||
pairs = [(int(t), float(ref_logits[t] - ref_lse), lp0[t].logprob)
|
||||
for t in ref_top if t in lp0]
|
||||
d = [abs(r - v) for _, r, v in pairs]
|
||||
rep = {
|
||||
"ref_text": ref["text"], "vllm_text": out.outputs[0].text,
|
||||
"greedy_token_match": f"{match}/{n}",
|
||||
"exact_match": ref["new_ids"][:n] == v_ids[:n],
|
||||
"ref_top1": ref["top1"], "vllm_first_token": v_ids[0] if v_ids else None,
|
||||
"first_token_agree": bool(v_ids and v_ids[0] == ref["top1"]),
|
||||
"top20_logprob_overlap": f"{len(pairs)}/20",
|
||||
"top20_logprob_max_abs_diff": max(d) if d else None,
|
||||
"top20_logprob_mean_abs_diff": float(np.mean(d)) if d else None,
|
||||
"dtype": a.dtype,
|
||||
}
|
||||
json.dump(rep, open(a.out, "w"), indent=2)
|
||||
print(json.dumps(rep, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Percent-busy of a core range over a short window.
|
||||
|
||||
Two engines sharing cores measure the scheduler, not the engine, so every run
|
||||
in this directory records what else was on its cores before it started.
|
||||
"""
|
||||
import sys, time
|
||||
|
||||
|
||||
def snap():
|
||||
d = {}
|
||||
for line in open("/proc/stat"):
|
||||
if line.startswith("cpu") and line[3].isdigit():
|
||||
p = line.split()
|
||||
d[int(p[0][3:])] = (sum(map(int, p[1:])), int(p[4]))
|
||||
return d
|
||||
|
||||
|
||||
def busy(lo, hi, window=2.0):
|
||||
a = snap(); time.sleep(window); b = snap()
|
||||
tot = idle = 0
|
||||
for c in range(lo, hi + 1):
|
||||
tot += b[c][0] - a[c][0]
|
||||
idle += b[c][1] - a[c][1]
|
||||
return 100.0 * (1 - idle / max(tot, 1))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
lo, hi = (int(x) for x in sys.argv[1].split("-"))
|
||||
print(f"{busy(lo, hi):.1f}")
|
||||
@@ -0,0 +1,58 @@
|
||||
"""Expand a `cpu_int4` checkpoint back to a plain bf16 one.
|
||||
|
||||
Used as a control: if the engine is applying the quantized weights faithfully,
|
||||
the expanded checkpoint must score exactly the same perplexity as the
|
||||
quantized one, and any difference is the runtime's, not the quantizer's.
|
||||
"""
|
||||
import argparse, json, shutil
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
from vllm_omni.model_executor.layers.quantization.cpu_int4 import unpack_nibbles
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--src", required=True)
|
||||
ap.add_argument("--dst", required=True)
|
||||
a = ap.parse_args()
|
||||
src, dst = Path(a.src), Path(a.dst)
|
||||
dst.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
cfg = json.loads((src / "config.json").read_text())
|
||||
group = cfg["quantization_config"]["group_size"]
|
||||
cfg.pop("quantization_config", None)
|
||||
|
||||
state = {}
|
||||
for shard in sorted(src.glob("*.safetensors")):
|
||||
state.update(load_file(str(shard)))
|
||||
|
||||
out = {}
|
||||
for name, t in state.items():
|
||||
if name.endswith(".weight_packed"):
|
||||
base = name[: -len(".weight_packed")]
|
||||
codes = unpack_nibbles(t).float()
|
||||
n, k = codes.shape
|
||||
s = state[base + ".weight_scale"].float()
|
||||
z = state[base + ".weight_zero"].float()
|
||||
deq = ((codes.reshape(n, k // group, group) - 8.0) * s.unsqueeze(-1)
|
||||
+ z.unsqueeze(-1)).reshape(n, k)
|
||||
out[base + ".weight"] = deq.to(torch.bfloat16)
|
||||
elif name.endswith((".weight_scale", ".weight_zero")):
|
||||
continue
|
||||
else:
|
||||
out[name] = t
|
||||
|
||||
save_file(out, str(dst / "model.safetensors"), metadata={"format": "pt"})
|
||||
(dst / "config.json").write_text(json.dumps(cfg, indent=2))
|
||||
for f in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
|
||||
"vocab.json", "merges.txt", "generation_config.json", "chat_template.jinja"):
|
||||
if (src / f).exists():
|
||||
shutil.copy(src / f, dst / f)
|
||||
print(json.dumps({"tensors": len(out), "group": group, "dst": str(dst)}))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,190 @@
|
||||
"""Write Spark-X2.5 as a symmetric 4-bit GPTQ checkpoint.
|
||||
|
||||
vLLM 0.28's CPU backend already has a 4-bit dense GEMM -- `int4_scaled_mm_cpu`,
|
||||
an AMX W4A8 kernel, plus `cpu_gemm_wna16` as the W4A16 fallback -- reached by
|
||||
loading a GPTQ checkpoint. It is not named like the ops an earlier search
|
||||
looked for (`woq_int4_linear`, `int4_scaled_mm_with_quant`, `da8w4_linear`),
|
||||
which is how it was missed. Measured at Spark's shapes it beats the tinygemm
|
||||
kernel 1.7x at one row and 8.4x at 512, so this is the path that matters.
|
||||
|
||||
`sym=True` selects vLLM's `uint4b8` weight type, whose dequantization is
|
||||
|
||||
w = (code - 8) * scale, code in [0, 15]
|
||||
|
||||
and whose zero points are ignored. Getting that convention wrong does not
|
||||
fail loudly -- the checkpoint loads and generates, just badly -- so
|
||||
`--verify` re-reads what was written and reports the reconstruction SNR.
|
||||
"""
|
||||
import argparse, json, shutil
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
QUANT_SUFFIXES = (
|
||||
"self_attn.q_proj.weight",
|
||||
"self_attn.k_proj.weight",
|
||||
"self_attn.v_proj.weight",
|
||||
"self_attn.out_proj.weight",
|
||||
"mlp.gate_proj.weight",
|
||||
"mlp.up_proj.weight",
|
||||
"mlp.down_proj.weight",
|
||||
)
|
||||
EMBED_NAME = "model.embedding.weight"
|
||||
|
||||
|
||||
def pack_int32(codes: torch.Tensor, dim: int) -> torch.Tensor:
|
||||
"""Pack 8 nibbles per int32 along `dim`, low-order nibble first."""
|
||||
codes = codes.to(torch.int32)
|
||||
if dim == 0:
|
||||
k, n = codes.shape
|
||||
out = torch.zeros(k // 8, n, dtype=torch.int32)
|
||||
for i in range(8):
|
||||
out |= codes[i::8, :] << (4 * i)
|
||||
else:
|
||||
k, n = codes.shape
|
||||
out = torch.zeros(k, n // 8, dtype=torch.int32)
|
||||
for i in range(8):
|
||||
out |= codes[:, i::8] << (4 * i)
|
||||
return out.contiguous()
|
||||
|
||||
|
||||
def quantize_asym(w: torch.Tensor, group: int):
|
||||
"""[N, K] -> GPTQ tensors with a real per-group zero point.
|
||||
|
||||
`w = (code - zp) * scale` over the group's true min/max, which does not
|
||||
waste a level the way a symmetric grid does on an off-centre group. The
|
||||
CPU W4A8 kernel's stand-in zero point for symmetric checkpoints is
|
||||
0x88888888 -- every nibble 8, used directly -- so `qzeros` here holds the
|
||||
zero point itself, with no AutoGPTQ-style -1.
|
||||
"""
|
||||
n, k = w.shape
|
||||
wt = w.float().t().contiguous()
|
||||
g = k // group
|
||||
blocks = wt.reshape(g, group, n)
|
||||
wmax, wmin = blocks.amax(dim=1), blocks.amin(dim=1)
|
||||
scale = ((wmax - wmin).clamp(min=1e-9) / 15.0)
|
||||
zp = (-wmin / scale).round().clamp_(0, 15)
|
||||
codes = ((blocks / scale.unsqueeze(1)).round() + zp.unsqueeze(1)).clamp_(0, 15)
|
||||
deq = ((codes - zp.unsqueeze(1)) * scale.unsqueeze(1)).reshape(k, n)
|
||||
codes = codes.reshape(k, n)
|
||||
snr = float(20 * torch.log10(wt.norm() / (wt - deq).norm().clamp(min=1e-12)))
|
||||
qweight = pack_int32(codes, dim=0)
|
||||
qzeros = pack_int32(zp.to(torch.int32), dim=1)
|
||||
scales = scale.to(torch.float16)
|
||||
g_idx = torch.arange(k, dtype=torch.int32) // group
|
||||
return qweight, qzeros, scales, g_idx, snr
|
||||
|
||||
|
||||
def quantize_sym(w: torch.Tensor, group: int):
|
||||
"""[N, K] -> GPTQ tensors. Symmetric about code 8, group-wise along K."""
|
||||
n, k = w.shape
|
||||
wt = w.float().t().contiguous() # GPTQ stores [K, N]
|
||||
g = k // group
|
||||
blocks = wt.reshape(g, group, n)
|
||||
# code 8 is zero and codes run 0..15, so the representable range is
|
||||
# [-8, 7] * scale; size it off the negative side to keep it symmetric.
|
||||
amax = blocks.abs().amax(dim=1) # [g, n]
|
||||
scale = (amax / 8.0).clamp(min=1e-9)
|
||||
codes = (blocks / scale.unsqueeze(1)).round().add_(8).clamp_(0, 15)
|
||||
codes = codes.reshape(k, n)
|
||||
deq = ((codes.reshape(g, group, n) - 8.0) * scale.unsqueeze(1)).reshape(k, n)
|
||||
snr = float(20 * torch.log10(wt.norm() / (wt - deq).norm().clamp(min=1e-12)))
|
||||
qweight = pack_int32(codes, dim=0) # [K/8, N]
|
||||
# sym GPTQ still carries qzeros; the classic convention stores zp - 1 = 7.
|
||||
qzeros = pack_int32(torch.full((g, n), 7, dtype=torch.int32), dim=1) # [g, N/8]
|
||||
scales = scale.to(torch.float16) # [g, N]
|
||||
g_idx = torch.arange(k, dtype=torch.int32) // group
|
||||
return qweight, qzeros, scales, g_idx, snr
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--src", required=True)
|
||||
ap.add_argument("--dst", required=True)
|
||||
ap.add_argument("--group", type=int, default=128)
|
||||
ap.add_argument("--quantize-embedding", action="store_true")
|
||||
ap.add_argument("--asym", action="store_true",
|
||||
help="per-group zero points instead of a symmetric grid")
|
||||
ap.add_argument("--quantize-lm-head", action="store_true",
|
||||
help="untie the head and quantize it; the lm_head GEMM is "
|
||||
"the single largest read in a decode step")
|
||||
a = ap.parse_args()
|
||||
src, dst = Path(a.src), Path(a.dst)
|
||||
dst.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
state = {}
|
||||
for shard in sorted(src.glob("*.safetensors")):
|
||||
state.update(load_file(str(shard)))
|
||||
|
||||
# Split the fused q_k_v_proj here rather than at load time: the model's
|
||||
# splitter cuts dim 0, and GPTQ tensors are stored [K, N], so the output
|
||||
# dim moves. Splitting along output columns is exactly equivalent for a
|
||||
# grouping that runs along K, so do it before quantizing and let vLLM's
|
||||
# ordinary stacked-parameter mapping take over.
|
||||
cfg_src = json.loads((src / "config.json").read_text())
|
||||
q_dim = cfg_src["num_attention_heads"] * cfg_src["head_dim"]
|
||||
kv_dim = cfg_src["num_key_value_heads"] * cfg_src["head_dim"]
|
||||
for name in [k for k in state if ".self_attn.q_k_v_proj." in k]:
|
||||
q, k, v = state.pop(name).split([q_dim, kv_dim, kv_dim], dim=0)
|
||||
for shard_name, tensor in (("q_proj", q), ("k_proj", k), ("v_proj", v)):
|
||||
state[name.replace("q_k_v_proj", shard_name)] = tensor.contiguous()
|
||||
|
||||
quantize = quantize_asym if a.asym else quantize_sym
|
||||
out, report = {}, []
|
||||
for name, t in sorted(state.items()):
|
||||
is_embed = name == EMBED_NAME and a.quantize_embedding
|
||||
if name.endswith(QUANT_SUFFIXES) or is_embed:
|
||||
qw, qz, sc, gi, snr = quantize(t, a.group)
|
||||
base = name[: -len(".weight")]
|
||||
out[f"{base}.qweight"] = qw
|
||||
out[f"{base}.qzeros"] = qz
|
||||
out[f"{base}.scales"] = sc
|
||||
out[f"{base}.g_idx"] = gi
|
||||
report.append({"name": name, "snr_db": round(snr, 2)})
|
||||
else:
|
||||
out[name] = t
|
||||
|
||||
if a.quantize_lm_head:
|
||||
# Spark ties the head to the vocab table. Quantizing it means untying:
|
||||
# the table stays bf16 for the lookup (a few hundred bytes a token)
|
||||
# and the head gets its own 4-bit copy, which is what a decode step
|
||||
# actually streams -- 131072 x 2048, the largest read in the step.
|
||||
qw, qz, sc, gi, snr = quantize(state[EMBED_NAME], a.group)
|
||||
out["lm_head.qweight"] = qw
|
||||
out["lm_head.qzeros"] = qz
|
||||
out["lm_head.scales"] = sc
|
||||
out["lm_head.g_idx"] = gi
|
||||
report.append({"name": "lm_head", "snr_db": round(snr, 2)})
|
||||
|
||||
save_file(out, str(dst / "model.safetensors"), metadata={"format": "pt"})
|
||||
for f in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
|
||||
"vocab.json", "merges.txt", "generation_config.json", "chat_template.jinja"):
|
||||
if (src / f).exists():
|
||||
shutil.copy(src / f, dst / f)
|
||||
|
||||
cfg = json.loads((src / "config.json").read_text())
|
||||
cfg["quantization_config"] = {
|
||||
"quant_method": "gptq", "bits": 4, "group_size": a.group,
|
||||
"desc_act": False, "sym": not a.asym, "true_sequential": True,
|
||||
"checkpoint_format": "gptq",
|
||||
"lm_head": bool(a.quantize_lm_head),
|
||||
}
|
||||
if a.quantize_lm_head:
|
||||
cfg["tie_word_embeddings"] = False
|
||||
(dst / "config.json").write_text(json.dumps(cfg, indent=2))
|
||||
|
||||
snrs = [r["snr_db"] for r in report]
|
||||
summary = {"group": a.group, "tensors": len(report),
|
||||
"snr_db_mean": round(sum(snrs) / len(snrs), 2),
|
||||
"snr_db_min": round(min(snrs), 2),
|
||||
"bits_per_weight": 4 + (16 + 4) / a.group,
|
||||
"quantized_embedding": a.quantize_embedding,
|
||||
"quantized_lm_head": a.quantize_lm_head}
|
||||
(dst / "quantization_report.json").write_text(
|
||||
json.dumps({**summary, "per_tensor": report}, indent=2))
|
||||
print(json.dumps(summary, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Export the sliding layer with F.rms_norm at a higher opset.
|
||||
|
||||
The calibrated layer spends ~11% of its cycles in an unfused RMSNorm
|
||||
(Pow/ReduceMean/Div/Mul), where the earlier build had a single fused `rms`
|
||||
op at 1%. ONNX opset 23 has a real RMSNormalization node; if the exporter
|
||||
emits it and QNN maps it to its fused kernel, that 11% largely goes away.
|
||||
"""
|
||||
import json, sys
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
sys.path.insert(0, "/data/zhoutaichang/embedding_infer/vllm-omni")
|
||||
import vllm_omni.edge.spark_export as se
|
||||
|
||||
SNAP = ("/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/"
|
||||
"snapshots/448e61eb392c00f2c403185c5b56d5e0665bfaab")
|
||||
HERE = Path("/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge")
|
||||
|
||||
_orig = se._rms
|
||||
|
||||
|
||||
def _rms_fused(x, weight, eps):
|
||||
return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps)
|
||||
|
||||
|
||||
hf = json.loads(Path(f"{SNAP}/config.json").read_text())
|
||||
i = hf["layer_types"].index(se.SLIDING)
|
||||
|
||||
for tag, fn, opset in (("rmsfused_op23", _rms_fused, 23),
|
||||
("rmsfused_op18", _rms_fused, 18)):
|
||||
se._rms = fn
|
||||
cfg = se.config_from_spark(hf, slice(i, i + 1), include_lm_head=False,
|
||||
cache_layout="ring")
|
||||
step = se.SparkDecodeStep(cfg)
|
||||
se.load_spark_weights(step, SNAP)
|
||||
out = HERE / f"onnx/{tag}_sliding.onnx"
|
||||
try:
|
||||
info = se.export_onnx(step, 1024, out, opset=opset)
|
||||
import onnx
|
||||
g = onnx.load(str(out), load_external_data=False)
|
||||
ops = {n.op_type for n in g.graph.node}
|
||||
print(tag, "OK ->", out.name,
|
||||
"| RMSNormalization" if "RMSNormalization" in ops else "| decomposed",
|
||||
"| ops:", len(g.graph.node), flush=True)
|
||||
except Exception as e:
|
||||
print(tag, "export FAILED:", str(e)[:200], flush=True)
|
||||
se._rms = _orig
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Numerical parity: HF reference Spark-X2.5 vs the vLLM-Omni implementation.
|
||||
|
||||
Runs both on CPU, greedy, and compares generated ids plus the first-step
|
||||
logit vector.
|
||||
"""
|
||||
import argparse, json, os, sys
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
MODEL = os.environ.get("SPARK_MODEL", "XHToken/Spark-X2.5-1.7B")
|
||||
PROMPT_MSGS = [
|
||||
{"role": "user", "content": "What is the capital of France? Answer in one word."},
|
||||
]
|
||||
|
||||
|
||||
def hf_reference(model_path, prompt_ids, n_new, dtype):
|
||||
from transformers import AutoConfig, AutoModelForCausalLM
|
||||
from transformers.dynamic_module_utils import get_class_from_dynamic_module
|
||||
|
||||
# modeling_spark.py targets transformers 4.57; v5 expects the tied-weight
|
||||
# map to be a dict rather than a list of keys.
|
||||
cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
|
||||
cls = get_class_from_dynamic_module(
|
||||
cfg.auto_map["AutoModelForCausalLM"], model_path
|
||||
)
|
||||
if isinstance(getattr(cls, "_tied_weights_keys", None), list):
|
||||
cls._tied_weights_keys = {"lm_head.weight": "model.embedding.weight"}
|
||||
m = cls.from_pretrained(model_path, dtype=dtype).eval()
|
||||
ids = torch.tensor([prompt_ids])
|
||||
with torch.no_grad():
|
||||
out = m(ids)
|
||||
first_logits = out.logits[0, -1].float().numpy()
|
||||
gen = m.generate(ids, max_new_tokens=n_new, do_sample=False,
|
||||
pad_token_id=2)
|
||||
return first_logits, gen[0, len(prompt_ids):].tolist()
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--n-new", type=int, default=24)
|
||||
ap.add_argument("--dtype", default="bfloat16")
|
||||
ap.add_argument("--max-model-len", type=int, default=2048)
|
||||
ap.add_argument("--out", default="parity.json")
|
||||
args = ap.parse_args()
|
||||
dtype = getattr(torch, args.dtype)
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||||
prompt = tok.apply_chat_template(PROMPT_MSGS, tokenize=False,
|
||||
add_generation_prompt=True)
|
||||
prompt_ids = tok(prompt, add_special_tokens=False)["input_ids"]
|
||||
print(f"prompt tokens: {len(prompt_ids)}")
|
||||
|
||||
hf_logits, hf_ids = hf_reference(MODEL, prompt_ids, args.n_new, dtype)
|
||||
print("HF :", repr(tok.decode(hf_ids)))
|
||||
|
||||
import vllm_omni # noqa: F401 (registers Spark2_5ForCausalLM)
|
||||
from vllm import LLM, SamplingParams
|
||||
llm = LLM(model=MODEL, dtype=args.dtype, max_model_len=args.max_model_len,
|
||||
enforce_eager=True, gpu_memory_utilization=0.30,
|
||||
max_num_seqs=1, trust_remote_code=False)
|
||||
sp = SamplingParams(temperature=0.0, max_tokens=args.n_new, logprobs=0,
|
||||
prompt_logprobs=0)
|
||||
res = llm.generate([{"prompt_token_ids": prompt_ids}], sp)[0]
|
||||
v_ids = list(res.outputs[0].token_ids)
|
||||
print("vLLM:", repr(res.outputs[0].text))
|
||||
|
||||
n = min(len(hf_ids), len(v_ids))
|
||||
match = sum(int(a == b) for a, b in zip(hf_ids[:n], v_ids[:n]))
|
||||
# first generated token's logit agreement
|
||||
first_hf_top = int(np.argmax(hf_logits))
|
||||
report = {
|
||||
"prompt_len": len(prompt_ids),
|
||||
"dtype": args.dtype,
|
||||
"hf_text": tok.decode(hf_ids),
|
||||
"vllm_text": res.outputs[0].text,
|
||||
"hf_ids": hf_ids, "vllm_ids": v_ids,
|
||||
"greedy_token_match": f"{match}/{n}",
|
||||
"exact_match": hf_ids[:n] == v_ids[:n],
|
||||
"hf_first_top1": first_hf_top,
|
||||
"vllm_first_token": v_ids[0] if v_ids else None,
|
||||
"first_token_agree": bool(v_ids and v_ids[0] == first_hf_top),
|
||||
}
|
||||
json.dump(report, open(args.out, "w"), indent=2)
|
||||
print(json.dumps({k: v for k, v in report.items()
|
||||
if k not in ("hf_ids", "vllm_ids")}, indent=2))
|
||||
return 0 if report["exact_match"] else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,57 @@
|
||||
"""Teacher-forced perplexity of a text file under a vLLM model.
|
||||
|
||||
Quantization is scored against the same engine's own bf16 run, not against
|
||||
llama.cpp's: each engine tokenizes and chunks its own way, and comparing the
|
||||
*increase* over an in-engine baseline cancels that, where comparing absolute
|
||||
perplexity across engines would not.
|
||||
"""
|
||||
import argparse, json, math
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", required=True)
|
||||
ap.add_argument("--corpus", required=True)
|
||||
ap.add_argument("--window", type=int, default=512)
|
||||
ap.add_argument("--max-chunks", type=int, default=0)
|
||||
ap.add_argument("--tag", default="")
|
||||
ap.add_argument("--out", default="ppl.jsonl")
|
||||
a = ap.parse_args()
|
||||
|
||||
import vllm_omni # noqa: F401
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
llm = LLM(model=a.model, dtype="bfloat16", max_model_len=a.window + 8,
|
||||
max_num_seqs=1, enable_prefix_caching=False,
|
||||
trust_remote_code=False, disable_log_stats=True)
|
||||
tok = llm.get_tokenizer()
|
||||
ids = tok(open(a.corpus).read(), add_special_tokens=False)["input_ids"]
|
||||
|
||||
chunks = [ids[i:i + a.window] for i in range(0, len(ids) - a.window + 1, a.window)]
|
||||
if a.max_chunks:
|
||||
chunks = chunks[: a.max_chunks]
|
||||
|
||||
sp = SamplingParams(temperature=0.0, max_tokens=1, prompt_logprobs=0,
|
||||
ignore_eos=True, detokenize=False)
|
||||
outs = llm.generate([{"prompt_token_ids": c} for c in chunks], sp, use_tqdm=False)
|
||||
|
||||
total_nll, n = 0.0, 0
|
||||
for o, c in zip(outs, chunks):
|
||||
pl = o.prompt_logprobs
|
||||
# position 0 has no prediction behind it
|
||||
for pos in range(1, len(c)):
|
||||
entry = pl[pos]
|
||||
lp = entry[c[pos]].logprob
|
||||
total_nll -= lp
|
||||
n += 1
|
||||
ppl = math.exp(total_nll / n)
|
||||
rec = {"tag": a.tag or a.model, "model": a.model, "window": a.window,
|
||||
"chunks": len(chunks), "scored_tokens": n,
|
||||
"nll": total_nll / n, "ppl": ppl}
|
||||
print(json.dumps(rec, indent=2))
|
||||
with open(a.out, "a") as f:
|
||||
f.write(json.dumps(rec) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,163 @@
|
||||
"""How much does 4-bit cost, measured the same way on both engines.
|
||||
|
||||
Absolute quality numbers are not comparable across engines -- different
|
||||
tokenizers' edge cases, different chat handling. What *is* comparable is how
|
||||
far each engine's own 4-bit build drifts from its own bf16 build on the same
|
||||
prompts under greedy decoding. That isolates the quantization, which is the
|
||||
question: is our W4A16 as faithful as llama.cpp's Q4_K_M?
|
||||
"""
|
||||
import argparse, json, subprocess, time
|
||||
from pathlib import Path
|
||||
|
||||
HERE = Path(__file__).parent
|
||||
|
||||
PLAIN = [
|
||||
"The capital of Australia is",
|
||||
"In 1969, the Apollo 11 mission",
|
||||
"def fibonacci(n):\n \"\"\"Return the nth Fibonacci number.\"\"\"\n",
|
||||
"The three laws of thermodynamics state that",
|
||||
"To make a proper espresso you need",
|
||||
"The difference between TCP and UDP is",
|
||||
"Photosynthesis converts",
|
||||
"A binary search tree is a data structure that",
|
||||
"The Treaty of Westphalia was signed in",
|
||||
]
|
||||
|
||||
|
||||
def prompts():
|
||||
out = []
|
||||
for f, name in (("ref.json", "chat_short"), ("ref_long.json", "chat_long"),
|
||||
("ref_tools.json", "chat_tools")):
|
||||
p = HERE / f
|
||||
if p.exists():
|
||||
out.append((name, json.loads(p.read_text())["prompt"]))
|
||||
out += [(f"plain_{i}", t) for i, t in enumerate(PLAIN)]
|
||||
return out
|
||||
|
||||
|
||||
def run_vllm(model, n_predict):
|
||||
import vllm_omni # noqa: F401
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
llm = LLM(model=model, dtype="bfloat16", max_model_len=4096,
|
||||
max_num_seqs=1, enable_prefix_caching=False,
|
||||
trust_remote_code=False, disable_log_stats=True)
|
||||
sp = SamplingParams(temperature=0.0, max_tokens=n_predict, ignore_eos=False)
|
||||
names, texts = zip(*prompts())
|
||||
outs = llm.generate(list(texts), sp, use_tqdm=False)
|
||||
return {n: {"text": o.outputs[0].text,
|
||||
"token_ids": list(o.outputs[0].token_ids)}
|
||||
for n, o in zip(names, outs)}
|
||||
|
||||
|
||||
def run_llamacpp(binary, gguf, n_predict, threads, port=8399):
|
||||
"""Drive llama.cpp through its server's /completion endpoint.
|
||||
|
||||
llama-cli in this build defaults to interactive conversation mode and
|
||||
refuses Spark's chat template without --jinja, so it never emits a plain
|
||||
completion. The server takes the prompt as raw text, which is what a
|
||||
like-for-like comparison needs, and loads the model once instead of once
|
||||
per prompt.
|
||||
"""
|
||||
import json as _json
|
||||
import urllib.request
|
||||
|
||||
log = open("/tmp/llama_server_quality.log", "w")
|
||||
proc = subprocess.Popen(
|
||||
[binary, "-m", gguf, "-t", str(threads), "-c", "4096",
|
||||
"--port", str(port), "--host", "127.0.0.1", "-np", "1"],
|
||||
stdout=log, stderr=subprocess.STDOUT,
|
||||
)
|
||||
try:
|
||||
base = f"http://127.0.0.1:{port}"
|
||||
for _ in range(600):
|
||||
try:
|
||||
urllib.request.urlopen(base + "/health", timeout=2).read()
|
||||
break
|
||||
except Exception:
|
||||
if proc.poll() is not None:
|
||||
raise RuntimeError("llama-server exited; see /tmp/llama_server_quality.log")
|
||||
time.sleep(1)
|
||||
else:
|
||||
raise RuntimeError("llama-server did not become healthy")
|
||||
|
||||
res = {}
|
||||
for name, text in prompts():
|
||||
body = _json.dumps({
|
||||
"prompt": text, "n_predict": n_predict,
|
||||
"temperature": 0.0, "top_k": 1, "seed": 0,
|
||||
"cache_prompt": False,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
base + "/completion", data=body,
|
||||
headers={"Content-Type": "application/json"})
|
||||
out = _json.loads(urllib.request.urlopen(req, timeout=1800).read())
|
||||
res[name] = {"text": out["content"], "token_ids": None}
|
||||
return res
|
||||
finally:
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
log.close()
|
||||
|
||||
|
||||
def common_prefix(a, b):
|
||||
n = 0
|
||||
for x, y in zip(a, b):
|
||||
if x != y:
|
||||
break
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
def compare(ref, got, label):
|
||||
rows, exact, tot_pref, tot_len = [], 0, 0, 0
|
||||
for name in ref:
|
||||
r, g = ref[name]["text"], got[name]["text"]
|
||||
cp = common_prefix(r, g)
|
||||
exact += int(r == g)
|
||||
tot_pref += cp
|
||||
tot_len += max(len(r), 1)
|
||||
rows.append({"prompt": name, "exact": r == g,
|
||||
"common_prefix_chars": cp, "ref_chars": len(r)})
|
||||
return {"label": label, "prompts": len(ref),
|
||||
"exact_match": f"{exact}/{len(ref)}",
|
||||
"mean_prefix_kept": round(tot_pref / max(tot_len, 1), 4),
|
||||
"detail": rows}
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
sub = ap.add_subparsers(dest="cmd", required=True)
|
||||
g = sub.add_parser("gen")
|
||||
g.add_argument("--engine", choices=["vllm", "llamacpp"], required=True)
|
||||
g.add_argument("--model", default=None)
|
||||
g.add_argument("--gguf", default=None)
|
||||
g.add_argument("--binary",
|
||||
default="/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-server")
|
||||
g.add_argument("--threads", type=int, default=32)
|
||||
g.add_argument("--n-predict", type=int, default=64)
|
||||
g.add_argument("--out", required=True)
|
||||
c = sub.add_parser("cmp")
|
||||
c.add_argument("--ref", required=True)
|
||||
c.add_argument("--got", required=True)
|
||||
c.add_argument("--label", required=True)
|
||||
c.add_argument("--out", required=True)
|
||||
a = ap.parse_args()
|
||||
|
||||
if a.cmd == "gen":
|
||||
res = (run_vllm(a.model, a.n_predict) if a.engine == "vllm"
|
||||
else run_llamacpp(a.binary, a.gguf, a.n_predict, a.threads))
|
||||
Path(a.out).write_text(json.dumps(res, indent=2))
|
||||
print(f"wrote {a.out} ({len(res)} prompts)")
|
||||
else:
|
||||
rep = compare(json.loads(Path(a.ref).read_text()),
|
||||
json.loads(Path(a.got).read_text()), a.label)
|
||||
Path(a.out).write_text(json.dumps(rep, indent=2))
|
||||
print(json.dumps({k: v for k, v in rep.items() if k != "detail"}, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Write a compressed-tensors W8A8-int8 checkpoint for Spark-X2.5.
|
||||
|
||||
Done directly from the safetensors rather than through llmcompressor: that
|
||||
package pulls transformers 5.x, on which Spark's remote modelling code does
|
||||
not import, and symmetric per-channel round-to-nearest needs no calibration
|
||||
data anyway.
|
||||
|
||||
The head-wise gate (g_proj) is deliberately left in full precision. It is a
|
||||
2048x8 projection whose eight outputs each scale a whole attention head, so
|
||||
quantisation error there is multiplied across 256 dims -- and it is 0.001% of
|
||||
the weights, so leaving it out costs nothing.
|
||||
"""
|
||||
import argparse, json, shutil
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
QUANT_SUFFIXES = (
|
||||
"self_attn.q_k_v_proj.weight",
|
||||
"self_attn.out_proj.weight",
|
||||
"mlp.gate_proj.weight",
|
||||
"mlp.up_proj.weight",
|
||||
"mlp.down_proj.weight",
|
||||
)
|
||||
|
||||
|
||||
def quantize_channelwise(w: torch.Tensor):
|
||||
"""Symmetric per-output-channel int8."""
|
||||
w = w.to(torch.float32)
|
||||
scale = w.abs().amax(dim=1, keepdim=True) / 127.0
|
||||
scale = torch.clamp(scale, min=1e-12)
|
||||
q = torch.clamp(torch.round(w / scale), -127, 127).to(torch.int8)
|
||||
return q, scale.to(torch.float32)
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--src", required=True)
|
||||
ap.add_argument("--dst", required=True)
|
||||
a = ap.parse_args()
|
||||
src, dst = Path(a.src), Path(a.dst)
|
||||
dst.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
state = {}
|
||||
for shard in sorted(src.glob("*.safetensors")):
|
||||
state.update(load_file(str(shard)))
|
||||
|
||||
out, n_q, err = {}, 0, []
|
||||
for name, t in state.items():
|
||||
if name.endswith(QUANT_SUFFIXES):
|
||||
q, s = quantize_channelwise(t)
|
||||
out[name] = q
|
||||
out[name + "_scale"] = s
|
||||
deq = q.to(torch.float32) * s
|
||||
ref = t.to(torch.float32)
|
||||
err.append(float(
|
||||
20 * torch.log10(ref.norm() / (ref - deq).norm().clamp(min=1e-12))
|
||||
))
|
||||
n_q += 1
|
||||
else:
|
||||
out[name] = t
|
||||
|
||||
save_file(out, str(dst / "model.safetensors"), metadata={"format": "pt"})
|
||||
|
||||
for f in ("config.json", "tokenizer.json", "tokenizer_config.json",
|
||||
"special_tokens_map.json", "vocab.json", "merges.txt",
|
||||
"generation_config.json", "chat_template.jinja"):
|
||||
if (src / f).exists():
|
||||
shutil.copy(src / f, dst / f)
|
||||
|
||||
cfg = json.loads((dst / "config.json").read_text())
|
||||
cfg.pop("auto_map", None) # the vendored vLLM config/model handle this arch
|
||||
cfg["quantization_config"] = {
|
||||
"quant_method": "compressed-tensors",
|
||||
"format": "int-quantized",
|
||||
"quantization_status": "compressed",
|
||||
"ignore": ["lm_head", "re:.*g_proj.*"],
|
||||
"config_groups": {
|
||||
"group_0": {
|
||||
"targets": ["Linear"],
|
||||
"weights": {
|
||||
"num_bits": 8, "type": "int", "symmetric": True,
|
||||
"strategy": "channel", "dynamic": False,
|
||||
"observer": "minmax",
|
||||
},
|
||||
"input_activations": {
|
||||
"num_bits": 8, "type": "int", "symmetric": True,
|
||||
"strategy": "token", "dynamic": True, "observer": None,
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
(dst / "config.json").write_text(json.dumps(cfg, indent=2))
|
||||
print(json.dumps({
|
||||
"quantized_tensors": n_q,
|
||||
"weight_snr_db_min": round(min(err), 2),
|
||||
"weight_snr_db_mean": round(sum(err) / len(err), 2),
|
||||
"dst": str(dst),
|
||||
}, indent=2))
|
||||
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Ground-truth Spark-X2.5 run using the official remote code on
|
||||
transformers 4.57.1 (the version config.json was written for)."""
|
||||
import argparse, json
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
MODEL = "XHToken/Spark-X2.5-1.7B"
|
||||
|
||||
def build_msgs(scenario):
|
||||
if scenario == "short":
|
||||
return [{"role": "user", "content": "What is the capital of France? Answer in one word."}]
|
||||
if scenario == "long":
|
||||
# >512 tokens so the 512-window sliding layers actually slide and
|
||||
# diverge from the full-attention layers.
|
||||
filler = " ".join(
|
||||
f"Item {i}: the code for warehouse district {i} is {1000 + i * 7}."
|
||||
for i in range(1, 121)
|
||||
)
|
||||
return [{"role": "user", "content":
|
||||
"Here is an inventory listing.\n" + filler +
|
||||
"\nWhat is the code for warehouse district 73? Answer with just the number."}]
|
||||
if scenario == "tools":
|
||||
return [
|
||||
{"role": "system", "content": "You are a helpful assistant with tools."},
|
||||
{"role": "user", "content": "What is the weather in Shenzhen right now, in celsius?"},
|
||||
]
|
||||
raise ValueError(scenario)
|
||||
|
||||
|
||||
TOOLS = [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather for a city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string", "description": "City name"},
|
||||
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}]
|
||||
|
||||
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--n-new", type=int, default=24)
|
||||
ap.add_argument("--dtype", default="bfloat16")
|
||||
ap.add_argument("--scenario", default="short", choices=["short","long","tools"])
|
||||
ap.add_argument("--out", default="ref.json")
|
||||
ap.add_argument("--logits-out", default="ref_logits.npy")
|
||||
a = ap.parse_args()
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||||
kw = {"tools": TOOLS} if a.scenario == "tools" else {}
|
||||
prompt = tok.apply_chat_template(build_msgs(a.scenario), tokenize=False,
|
||||
add_generation_prompt=True, **kw)
|
||||
ids = tok(prompt, add_special_tokens=False)["input_ids"]
|
||||
m = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL, trust_remote_code=True, dtype=getattr(torch, a.dtype)).eval()
|
||||
t = torch.tensor([ids])
|
||||
with torch.no_grad():
|
||||
logits = m(t).logits[0, -1].float().numpy()
|
||||
gen = m.generate(t, max_new_tokens=a.n_new, do_sample=False, pad_token_id=2)
|
||||
new = gen[0, len(ids):].tolist()
|
||||
np.save(a.logits_out, logits)
|
||||
json.dump({"prompt": prompt, "prompt_ids": ids, "new_ids": new,
|
||||
"text": tok.decode(new), "dtype": a.dtype,
|
||||
"top1": int(np.argmax(logits)), "scenario": a.scenario}, open(a.out, "w"), indent=2)
|
||||
print("prompt_len", len(ids), "| top1", int(np.argmax(logits)))
|
||||
print("text:", repr(tok.decode(new)))
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Check the exported Spark decode step against the HF reference.
|
||||
|
||||
Runs the reference over a C-token prompt, hands its KV cache to
|
||||
SparkDecodeStep (sliding layers get only the last 512 entries, which is the
|
||||
whole point of the hybrid layout) and compares next-token logits.
|
||||
"""
|
||||
import argparse, importlib.util, json, sys, types
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
ROOT = Path("/data/zhoutaichang/embedding_infer/vllm-omni")
|
||||
|
||||
|
||||
def _load(name, path):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
m = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = m
|
||||
spec.loader.exec_module(m)
|
||||
return m
|
||||
|
||||
|
||||
for pkg_name in ("vllm_omni", "vllm_omni.edge"):
|
||||
mod = types.ModuleType(pkg_name)
|
||||
mod.__path__ = []
|
||||
sys.modules[pkg_name] = mod
|
||||
_load("vllm_omni.edge.qnn_export", ROOT / "vllm_omni/edge/qnn_export.py")
|
||||
se = _load("vllm_omni.edge.spark_export", ROOT / "vllm_omni/edge/spark_export.py")
|
||||
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
|
||||
ap.add_argument("--context", type=int, default=1024)
|
||||
ap.add_argument("--cache-layout", default="auto", choices=["auto", "ring", "roll"])
|
||||
ap.add_argument("--gelu", default="exact", choices=["exact", "tanh", "sigmoid"])
|
||||
ap.add_argument("--out", default="export_parity.json")
|
||||
a = ap.parse_args()
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True)
|
||||
text = ("The quick brown fox jumps over the lazy dog. " * 400)
|
||||
ids = tok(text, add_special_tokens=False)["input_ids"][: a.context + 1]
|
||||
assert len(ids) == a.context + 1, f"only got {len(ids)} tokens"
|
||||
|
||||
m = AutoModelForCausalLM.from_pretrained(a.model, trust_remote_code=True,
|
||||
dtype=torch.float32).eval()
|
||||
import copy
|
||||
with torch.no_grad():
|
||||
pre = m(torch.tensor([ids[:-1]]), use_cache=True)
|
||||
# The decode call appends to the cache in place, so snapshot it first.
|
||||
cache = copy.deepcopy(pre.past_key_values)
|
||||
ref = m(torch.tensor([[ids[-1]]]), past_key_values=pre.past_key_values,
|
||||
use_cache=True)
|
||||
ref_logits = ref.logits[0, -1].float().numpy()
|
||||
|
||||
def kv(i):
|
||||
try:
|
||||
return cache.layers[i].keys, cache.layers[i].values
|
||||
except AttributeError:
|
||||
return cache.key_cache[i], cache.value_cache[i]
|
||||
|
||||
hf_cfg = json.loads((Path(m.config._name_or_path) / "config.json").read_text()) \
|
||||
if Path(str(m.config._name_or_path)).exists() else m.config.to_dict()
|
||||
cfg = se.config_from_spark(hf_cfg, None, include_lm_head=True,
|
||||
cache_layout=a.cache_layout, gelu_mode=a.gelu)
|
||||
step = se.SparkDecodeStep(cfg)
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
wdir = snapshot_download(a.model, allow_patterns=["*.safetensors"])
|
||||
print("loaded", se.load_spark_weights(step, wdir), "tensors")
|
||||
step.eval()
|
||||
|
||||
C = a.context
|
||||
pos = C
|
||||
emb = m.model.embedding(torch.tensor([[ids[-1]]])).float()
|
||||
cos_sw, sin_sw = se._angles(pos, cfg.rotary_dim(se.SLIDING), cfg.rope_theta[se.SLIDING], torch.float32)
|
||||
cos_f, sin_f = se._angles(pos, cfg.rotary_dim(se.FULL), cfg.rope_theta[se.FULL], torch.float32)
|
||||
sw_len = cfg.cache_len(se.SLIDING, C)
|
||||
# Ring masks only cache slots; roll masks the rebuilt window (cache + new).
|
||||
mask_sw = torch.zeros(1, 1, 1, sw_len + (1 if cfg.layout_for(se.SLIDING) == "ring" else 0))
|
||||
mask_full = torch.zeros(1, 1, 1, C + 1)
|
||||
|
||||
shapes = {}
|
||||
caches = []
|
||||
for i, t in enumerate(cfg.layer_types):
|
||||
k, v = kv(i)
|
||||
shapes.setdefault(t, k.shape[2])
|
||||
if t == se.SLIDING:
|
||||
# transformers already truncates sliding layers to the window; take
|
||||
# the newest sw_len entries and left-pad (masked out) if it holds fewer.
|
||||
k, v = k[:, :, -sw_len:], v[:, :, -sw_len:]
|
||||
have = k.shape[2]
|
||||
if have < sw_len:
|
||||
pad = sw_len - have
|
||||
z = torch.zeros(k.shape[0], k.shape[1], pad, k.shape[3])
|
||||
# No mask needed: the window roll drops exactly these pad slots.
|
||||
k, v = torch.cat([z, k.float()], 2), torch.cat([z, v.float()], 2)
|
||||
caches += [k.float(), v.float()]
|
||||
print("hf cache lens by layer type:", shapes, "| sw_len used:", sw_len)
|
||||
|
||||
with torch.no_grad():
|
||||
out = step(emb, cos_sw, sin_sw, cos_f, sin_f, mask_sw, mask_full, *caches)
|
||||
got = out[0][0].float().numpy()
|
||||
|
||||
def snr(ref, got):
|
||||
return 20 * np.log10(np.linalg.norm(ref) / max(np.linalg.norm(ref - got), 1e-12))
|
||||
|
||||
rep = {
|
||||
"context": C,
|
||||
"cache_layout": a.cache_layout,
|
||||
"gelu_mode": a.gelu,
|
||||
"logits_snr_db": round(float(snr(ref_logits, got)), 2),
|
||||
"max_abs_diff": float(np.max(np.abs(ref_logits - got))),
|
||||
"ref_top1": int(np.argmax(ref_logits)),
|
||||
"export_top1": int(np.argmax(got)),
|
||||
"top1_agree": int(np.argmax(ref_logits)) == int(np.argmax(got)),
|
||||
"top5_agree": np.argsort(-ref_logits)[:5].tolist() == np.argsort(-got)[:5].tolist(),
|
||||
"kv_bytes": se.kv_cache_bytes(cfg, C),
|
||||
}
|
||||
json.dump(rep, open(a.out, "w"), indent=2)
|
||||
print(json.dumps(rep, indent=2))
|
||||
@@ -0,0 +1,11 @@
|
||||
# Shell drivers
|
||||
|
||||
The exact invocations behind published grids — thread sweeps, interleaved A/Bs,
|
||||
NUMA-node comparisons, fusion ablations.
|
||||
|
||||
They are kept verbatim, absolute paths and all, because a grid's *command* is
|
||||
part of its result: which cores, which environment variables, which order. The
|
||||
harness now records all of that into each artifact's `provenance` block, so new
|
||||
work does not need a script like these; these document the runs that predate it.
|
||||
|
||||
Read them for what was done, not as a portable interface.
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
#!/bin/bash
|
||||
# Attribute the gate delta per-op rather than by throughput: the ablation's
|
||||
# 4.1 ms/token is six times what the aten::mm it removes was measured to cost,
|
||||
# so either the profile or the throughput number is misleading.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
M=/data/zhoutaichang/embedding_infer/models/spark2_5-int4-g64
|
||||
cd $HERE
|
||||
for g in 0 1; do
|
||||
echo "#### gate_reduction=$g"
|
||||
VLLM_OMNI_SPARK_GATE_REDUCTION=$g KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=24 VLLM_CPU_OMP_THREADS_BIND=32-55 \
|
||||
taskset -c 32-55 timeout 2400 $V/bin/python profile_cpu_step.py --model $M \
|
||||
--out $HERE/profile_gate_$g.json 2>&1 | sed -n '/^{/,$p' | head -30
|
||||
done
|
||||
echo "#### GATE PROFILE DONE"
|
||||
Executable
+48
@@ -0,0 +1,48 @@
|
||||
#!/bin/bash
|
||||
# Decode-speed ablation on cores 32-55 (node 0; another session holds 0-31).
|
||||
# One change at a time, each engine start back to back on the same cores, so
|
||||
# the deltas are attributable even though the machine is not idle.
|
||||
# gate head-gate projection as an eager reduction instead of aten::mm
|
||||
# chunk inductor keeps small elementwise kernels single-threaded
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
G=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B-GGUF/snapshots/23e1fcac55e7dd71e4c12a23723cc228ba0e5e85
|
||||
B=/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-bench
|
||||
M=/data/zhoutaichang/embedding_infer/models
|
||||
CHUNK='{"inductor_compile_config":{"cpp.min_chunk_size":65536}}'
|
||||
OUT=$HERE/ab_grid.jsonl
|
||||
CORES=32-55
|
||||
T=24
|
||||
cd $HERE
|
||||
|
||||
run () { # tag model gate_reduction [--cc json]
|
||||
local tag=$1 model=$2 gate=$3; shift 3
|
||||
local cc=""
|
||||
[ "${1:-}" = "--cc" ] && { cc="$2"; shift 2; }
|
||||
VLLM_OMNI_SPARK_GATE_REDUCTION=$gate \
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$T VLLM_CPU_OMP_THREADS_BIND=$CORES \
|
||||
taskset -c $CORES timeout 2400 $V/bin/python bench_cpu_vllm.py --model "$model" \
|
||||
--tag "$tag" --reps 3 --out $OUT ${cc:+--compilation-config "$cc"} 2>&1 \
|
||||
| grep -E '"(tag|pp_tps|tg_tps)"'
|
||||
}
|
||||
|
||||
llamacpp () {
|
||||
taskset -c $CORES $B -m $G/Spark-X2.5-1.7B-Q4_K_M.gguf -p 512 -n 128 -t $T -o json 2>/dev/null \
|
||||
> $HERE/ab_llamacpp_$1.json
|
||||
python3 -c "
|
||||
import json;d=json.load(open('$HERE/ab_llamacpp_$1.json'))
|
||||
print(' llama.cpp Q4_K_M', {('pp512' if r['n_prompt'] else 'tg128'): round(1e9*(r['n_prompt'] or r['n_gen'])/r['avg_ns'],2) for r in d})"
|
||||
}
|
||||
|
||||
echo "#### llama.cpp Q4_K_M t=$T (before)"; llamacpp before
|
||||
run "base" $M/spark2_5-int4-g64 0
|
||||
run "gate" $M/spark2_5-int4-g64 1
|
||||
run "gate+chunk" $M/spark2_5-int4-g64 1 --cc "$CHUNK"
|
||||
run "gate+chunk+g128" $M/spark2_5-int4-g128 1 --cc "$CHUNK"
|
||||
run "w4a8+gate+chunk" $M/spark2_5-gptq-g64 1 --cc "$CHUNK"
|
||||
echo "#### llama.cpp Q4_K_M t=$T (after, drift check)"; llamacpp after
|
||||
echo "#### AB GRID DONE"
|
||||
Executable
+24
@@ -0,0 +1,24 @@
|
||||
#!/bin/bash
|
||||
# Round two, with cpp.min_chunk_size dropped: it measured -6% decode and -19%
|
||||
# prefill, so the remaining levers get tested against `gate` alone.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
M=/data/zhoutaichang/embedding_infer/models
|
||||
OUT=$HERE/ab_grid2.jsonl
|
||||
cd $HERE
|
||||
run () {
|
||||
local tag=$1 model=$2
|
||||
VLLM_OMNI_SPARK_GATE_REDUCTION=1 KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=24 VLLM_CPU_OMP_THREADS_BIND=32-55 \
|
||||
taskset -c 32-55 timeout 2400 $V/bin/python bench_cpu_vllm.py --model "$model" \
|
||||
--tag "$tag" --reps 5 --out $OUT 2>&1 | grep -E '"(tag|pp_tps|tg_tps)"'
|
||||
}
|
||||
run "gate-g64" $M/spark2_5-int4-g64
|
||||
run "gate-g128" $M/spark2_5-int4-g128
|
||||
run "gate-g64-noop" $M/spark2_5-int4-g64-nocustomop
|
||||
run "gate-w4a8" $M/spark2_5-gptq-g64
|
||||
run "gate-g64-again" $M/spark2_5-int4-g64
|
||||
echo "#### AB GRID 2 DONE"
|
||||
Executable
+23
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
# Strictly sequential: both engines get the same pinned cores, one at a time,
|
||||
# so neither steals the other's cache or memory bandwidth.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
G=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B-GGUF/snapshots/23e1fcac55e7dd71e4c12a23723cc228ba0e5e85
|
||||
B=/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-bench
|
||||
cd $HERE
|
||||
for t in 4 8 16 32; do
|
||||
last=$((t-1))
|
||||
echo "######## threads=$t : llama.cpp"
|
||||
taskset -c 0-$last $B -m $G/Spark-X2.5-1.7B.gguf -m $G/Spark-X2.5-1.7B-Q8_0.gguf \
|
||||
-m $G/Spark-X2.5-1.7B-Q4_K_M.gguf -p 512 -n 128 -t $t -o json 2>/dev/null \
|
||||
> $HERE/llamacpp_grid_t$t.json
|
||||
echo "######## threads=$t : vllm"
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$t VLLM_CPU_OMP_THREADS_BIND=0-$last \
|
||||
taskset -c 0-$last timeout 2400 $V/bin/python bench_cpu_vllm.py \
|
||||
--tag "vllm-bf16-t$t" --out vllm_cpu_bench.jsonl 2>&1 | grep -E '"(tag|pp_tps|tg_tps|init_s)"'
|
||||
done
|
||||
echo "######## GRID DONE"
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
#!/bin/bash
|
||||
# Final matched comparison, strictly sequential on the same pinned cores with
|
||||
# nothing else running: llama.cpp Q4_K_M against both vLLM 4-bit paths.
|
||||
# W4A16 group-wise weights, bf16 activations, tinygemm kernel (cpu_int4)
|
||||
# W4A8 group-wise weights, int8 activations, AMX kernel (gptq)
|
||||
# W4A8 is the arithmetic llama.cpp's Q4_K_M actually does (Q4_K x Q8_K).
|
||||
#
|
||||
# The OMP team is kept hot between regions (KMP_BLOCKTIME/OMP_WAIT_POLICY),
|
||||
# which is the one thread-fan-out knob that survived measurement. Forcing
|
||||
# inductor's small elementwise kernels single-threaded (cpp.min_chunk_size)
|
||||
# was tried and dropped: -6% decode and -19% prefill, because it also
|
||||
# serialises the kernels that have real work at prefill.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
G=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B-GGUF/snapshots/23e1fcac55e7dd71e4c12a23723cc228ba0e5e85
|
||||
B=/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-bench
|
||||
M=/data/zhoutaichang/embedding_infer/models
|
||||
W4A16=$M/spark2_5-int4-g64
|
||||
W4A8=$M/spark2_5-gptq-g64
|
||||
OUT=$HERE/final_grid.jsonl
|
||||
cd $HERE
|
||||
|
||||
run_vllm () { # tag model threads [--cc json] [env...]
|
||||
local tag=$1 model=$2 t=$3; shift 3
|
||||
local cc=""
|
||||
if [ "${1:-}" = "--cc" ]; then cc="$2"; shift 2; fi
|
||||
local last=$((t-1))
|
||||
env "$@" LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$t VLLM_CPU_OMP_THREADS_BIND=0-$last \
|
||||
taskset -c 0-$last timeout 2400 $V/bin/python bench_cpu_vllm.py --model "$model" \
|
||||
--tag "$tag" --reps 3 --out $OUT ${cc:+--compilation-config "$cc"} 2>&1 \
|
||||
| grep -E '"(tag|pp_tps|tg_tps|rss_mb)"'
|
||||
}
|
||||
|
||||
llamacpp () {
|
||||
local t=$1 last=$((t-1))
|
||||
taskset -c 0-$last $B -m $G/Spark-X2.5-1.7B-Q4_K_M.gguf -p 512 -n 128 -t $t -o json 2>/dev/null \
|
||||
> $HERE/final_llamacpp_t$t.json
|
||||
python3 -c "
|
||||
import json;d=json.load(open('$HERE/final_llamacpp_t$t.json'))
|
||||
print(' llama.cpp Q4_K_M t=$t ', {('pp512' if r['n_prompt'] else 'tg128'): round(1e9*(r['n_prompt'] or r['n_gen'])/r['avg_ns'],2) for r in d})"
|
||||
}
|
||||
|
||||
for t in 32 16; do
|
||||
echo "######## threads=$t"
|
||||
llamacpp $t
|
||||
run_vllm "w4a16-t$t" $W4A16 $t KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active
|
||||
run_vllm "w4a8-t$t" $W4A8 $t KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active
|
||||
run_vllm "w4a16-g128-t$t" $M/spark2_5-int4-g128 $t KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active
|
||||
done
|
||||
echo "######## bf16 control t=32"
|
||||
run_vllm "bf16-t32" XHToken/Spark-X2.5-1.7B 32 KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active
|
||||
echo "######## FINAL GRID DONE"
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
# Fused CPU RMSNorm kernels on/off, interleaved, on both 4-bit paths.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
M=/data/zhoutaichang/embedding_infer/models
|
||||
OUT=$HERE/fusion_ab.jsonl
|
||||
CORES=32-55; T=24
|
||||
cd $HERE
|
||||
arm () { # tag model fused
|
||||
VLLM_OMNI_CPU_FUSED_NORMS=$3 VLLM_OMNI_SPARK_GATE_REDUCTION=1 \
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$T VLLM_CPU_OMP_THREADS_BIND=$CORES \
|
||||
taskset -c $CORES timeout 2400 $V/bin/python bench_cpu_vllm.py --model "$2" \
|
||||
--tag "$1" --reps 3 --out $OUT 2>&1 | grep -E '"(tag|pp_tps|tg_tps|init_s|rss_mb)"'
|
||||
}
|
||||
for pass in 1 2 3; do
|
||||
echo "======== pass $pass"
|
||||
arm "w4a8-g128-nofuse" $M/spark2_5-gptq-g128 0
|
||||
arm "w4a8-g128-fused" $M/spark2_5-gptq-g128 1
|
||||
arm "w4a16-g128-fused" $M/spark2_5-int4-g128 1
|
||||
done
|
||||
echo "#### FUSION AB DONE"
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
#!/bin/bash
|
||||
# Init time and memory, the two things that decide whether this can be served
|
||||
# from an edge box. Sequential; each row changes exactly one thing.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
M=/data/zhoutaichang/embedding_infer/models/spark2_5-gptq-g128
|
||||
SCRATCH=/tmp/claude-2005/-data-zhoutaichang-embedding-infer/e695f57b-0163-4287-b1be-1d0ec9330768/scratchpad
|
||||
OUT=$HERE/init_memory.jsonl
|
||||
CORES=32-55; T=24
|
||||
cd $HERE
|
||||
|
||||
run () { # tag cache_root shim kv_gib max_len
|
||||
VLLM_CACHE_ROOT=$2 VLLM_OMNI_AOT_CACHE_SHIM=$3 VLLM_CPU_KVCACHE_SPACE=$4 \
|
||||
VLLM_OMNI_SPARK_GATE_REDUCTION=1 VLLM_OMNI_CPU_FUSED_NORMS=0 \
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$T VLLM_CPU_OMP_THREADS_BIND=$CORES \
|
||||
taskset -c $CORES timeout 1800 $V/bin/python bench_init_memory.py --model $M \
|
||||
--tag "$1" --max-model-len $5 --out $OUT 2>&1 \
|
||||
| grep -E '"(tag|init_s|rss_mb|rss_peak_mb|tokens_generated|kvcache|aot_shim)"'
|
||||
}
|
||||
|
||||
FRESH=$SCRATCH/vllm_cache_fresh_$$
|
||||
rm -rf $FRESH
|
||||
echo "#### 1. cold cache, shim on"
|
||||
run "cold-shim-on" $FRESH 1 4 2048
|
||||
echo "#### 2. warm cache, shim on (the point of the shim)"
|
||||
run "warm-shim-on" $FRESH 1 4 2048
|
||||
echo "#### 3. warm cache, shim OFF (vLLM recompiles: the bug)"
|
||||
run "warm-shim-off" $FRESH 0 4 2048
|
||||
echo "#### 4. warm, shim on, KV sized to the context instead of 4 GiB"
|
||||
run "warm-kv-0.25" $FRESH 1 0.25 2048
|
||||
echo "#### 5. warm, shim on, KV 0.25 GiB and 1k context"
|
||||
run "warm-kv-0.25-1k" $FRESH 1 0.25 1024
|
||||
rm -rf $FRESH
|
||||
echo "#### INIT MEMORY DONE"
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
#!/bin/bash
|
||||
# Strictly sequential: one engine at a time on the same pinned cores, so
|
||||
# neither steals the other's memory bandwidth. Same rules as run_cpu_grid.sh.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
G=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B-GGUF/snapshots/23e1fcac55e7dd71e4c12a23723cc228ba0e5e85
|
||||
B=/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-bench
|
||||
INT4=/data/zhoutaichang/embedding_infer/models/spark2_5-int4-g64
|
||||
OUT=${1:-$HERE/int4_grid.jsonl}
|
||||
cd $HERE
|
||||
for t in 4 8 16 32; do
|
||||
last=$((t-1))
|
||||
echo "######## threads=$t : llama.cpp Q4_K_M"
|
||||
taskset -c 0-$last $B -m $G/Spark-X2.5-1.7B-Q4_K_M.gguf -p 512 -n 128 -t $t -o json 2>/dev/null \
|
||||
> $HERE/llamacpp_q4km_t$t.json
|
||||
python3 -c "
|
||||
import json;d=json.load(open('$HERE/llamacpp_q4km_t$t.json'))
|
||||
for r in d: print(' ', r['n_prompt'] or 'tg', round(1e9*(r['n_prompt'] or r['n_gen'])/r['avg_ns'],2), 'tok/s')"
|
||||
echo "######## threads=$t : vLLM int4 g64"
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$t VLLM_CPU_OMP_THREADS_BIND=0-$last \
|
||||
taskset -c 0-$last timeout 2400 $V/bin/python bench_cpu_vllm.py --model $INT4 \
|
||||
--tag "vllm-int4-g64-t$t" --reps 3 --out $OUT 2>&1 | grep -E '"(tag|pp_tps|tg_tps|init_s)"'
|
||||
done
|
||||
echo "######## threads=32 : vLLM bf16 control"
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=32 VLLM_CPU_OMP_THREADS_BIND=0-31 \
|
||||
taskset -c 0-31 timeout 2400 $V/bin/python bench_cpu_vllm.py \
|
||||
--tag "vllm-bf16-t32-control" --reps 3 --out $OUT 2>&1 | grep -E '"(tag|pp_tps|tg_tps|init_s)"'
|
||||
echo "######## GRID DONE"
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
#!/bin/bash
|
||||
# Both engines, strictly one at a time, on NUMA node 1's physical cores.
|
||||
# Node 0 is in use by another job, and a shared-core benchmark measures the
|
||||
# scheduler, not the engine.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
G=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B-GGUF/snapshots/23e1fcac55e7dd71e4c12a23723cc228ba0e5e85
|
||||
B=/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-bench
|
||||
BASE=56
|
||||
cd $HERE
|
||||
for t in 4 8 16 32; do
|
||||
lo=$BASE; hi=$((BASE + t - 1))
|
||||
echo "######## threads=$t cores $lo-$hi : llama.cpp"
|
||||
taskset -c $lo-$hi $B -m $G/Spark-X2.5-1.7B-Q4_K_M.gguf -p 512 -n 128 -t $t -o json \
|
||||
2>/dev/null > $HERE/llamacpp_node1_t$t.json
|
||||
echo "######## threads=$t cores $lo-$hi : vllm int4"
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$t VLLM_CPU_OMP_THREADS_BIND=$lo-$hi \
|
||||
taskset -c $lo-$hi timeout 2400 $V/bin/python bench_cpu_vllm.py \
|
||||
--model ${MODEL:-/data/zhoutaichang/embedding_infer/models/spark2_5-int4-g64} \
|
||||
--tag "${TAG:-int4-g64}-t$t" --reps 3 --out ${OUT:-$HERE/vllm_cpu_int4_node1.jsonl} 2>&1 \
|
||||
| grep -E '"(tag|pp_tps|tg_tps|init_s)"'
|
||||
done
|
||||
echo "#### GRID DONE"
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
#!/bin/bash
|
||||
# Interleaved, not sequential.
|
||||
#
|
||||
# A sequential grid on this machine is not trustworthy: another tenant on the
|
||||
# same socket swung the llama.cpp control from 86.7 to 60.4 tok/s across one
|
||||
# pass, and the vLLM baseline alone varied 57.3 -> 68.0 between two identical
|
||||
# runs. So arms are rotated and each is run three times; the figure taken for
|
||||
# each arm is its best pass, which is the least-contaminated estimate of what
|
||||
# that configuration does when the machine cooperates.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
G=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B-GGUF/snapshots/23e1fcac55e7dd71e4c12a23723cc228ba0e5e85
|
||||
B=/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-bench
|
||||
M=/data/zhoutaichang/embedding_infer/models
|
||||
OUT=$HERE/interleaved.jsonl
|
||||
CORES=32-55; T=24
|
||||
cd $HERE
|
||||
|
||||
vllm_arm () {
|
||||
VLLM_OMNI_SPARK_GATE_REDUCTION=$3 KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$T VLLM_CPU_OMP_THREADS_BIND=$CORES \
|
||||
taskset -c $CORES timeout 2400 $V/bin/python bench_cpu_vllm.py --model "$2" \
|
||||
--tag "$1" --reps 3 --out $OUT 2>&1 | grep -E '"(tag|pp_tps|tg_tps)"'
|
||||
}
|
||||
|
||||
for pass in 1 2 3; do
|
||||
echo "======== pass $pass"
|
||||
taskset -c $CORES $B -m $G/Spark-X2.5-1.7B-Q4_K_M.gguf -p 512 -n 128 -t $T -o json 2>/dev/null \
|
||||
> $HERE/il_llamacpp_p$pass.json
|
||||
python3 -c "
|
||||
import json;d=json.load(open('$HERE/il_llamacpp_p$pass.json'))
|
||||
print(' llamacpp-q4km', {('pp512' if r['n_prompt'] else 'tg128'): round(1e9*(r['n_prompt'] or r['n_gen'])/r['avg_ns'],2) for r in d})"
|
||||
vllm_arm "w4a16-g64-nogate" $M/spark2_5-int4-g64 0
|
||||
vllm_arm "w4a16-g64" $M/spark2_5-int4-g64 1
|
||||
vllm_arm "w4a16-g128" $M/spark2_5-int4-g128 1
|
||||
vllm_arm "w4a8-g64" $M/spark2_5-gptq-g64 1
|
||||
vllm_arm "w4a8-g128" $M/spark2_5-gptq-g128 1
|
||||
done
|
||||
echo "#### INTERLEAVED DONE"
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
#!/bin/bash
|
||||
# Two untried overhead knobs, applied to the leading config (W4A8 g128).
|
||||
# split_kv vLLM splits the KV range across threads; at batch 1 with two KV
|
||||
# heads that may be fan-out for its own sake
|
||||
# nodetok skip detokenization, which is pure host work per step
|
||||
# Interleaved for the same reason as run_interleaved.sh.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
MODEL=/data/zhoutaichang/embedding_infer/models/spark2_5-gptq-g128
|
||||
OUT=$HERE/leader_knobs.jsonl
|
||||
CORES=32-55; T=24
|
||||
cd $HERE
|
||||
arm () { # tag, split_kv, extra-flags...
|
||||
local tag=$1 split=$2; shift 2
|
||||
VLLM_OMNI_SPARK_GATE_REDUCTION=1 VLLM_CPU_ATTN_SPLIT_KV=$split \
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=$T VLLM_CPU_OMP_THREADS_BIND=$CORES \
|
||||
taskset -c $CORES timeout 2400 $V/bin/python bench_cpu_vllm.py --model $MODEL \
|
||||
--tag "$tag" --reps 3 --out $OUT "$@" 2>&1 | grep -E '"(tag|pp_tps|tg_tps)"'
|
||||
}
|
||||
for pass in 1 2 3; do
|
||||
echo "======== pass $pass"
|
||||
arm "leader" 1
|
||||
arm "leader-nosplitkv" 0
|
||||
arm "leader-nodetok" 1 --no-detokenize
|
||||
done
|
||||
echo "#### LEADER KNOBS DONE"
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
#!/bin/bash
|
||||
G=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B-GGUF/snapshots/23e1fcac55e7dd71e4c12a23723cc228ba0e5e85
|
||||
B=/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-bench
|
||||
for t in 4 8 16 32; do
|
||||
last=$((t-1))
|
||||
taskset -c 0-$last $B -m $G/Spark-X2.5-1.7B.gguf -m $G/Spark-X2.5-1.7B-Q8_0.gguf -m $G/Spark-X2.5-1.7B-Q4_K_M.gguf \
|
||||
-p 512 -n 128 -t $t -o json 2>/dev/null
|
||||
done
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
#!/bin/bash
|
||||
# Memory, sized from the model instead of from a server default.
|
||||
# Spark-X2.5 is 28 layers x 2 KV heads x 256 head_dim x 2 (K,V) x 2 B
|
||||
# = 56 KiB/token, so a 2048-token single-stream deployment needs 112 MiB of KV
|
||||
# cache. The CPU default asks for 4 GiB.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
M=/data/zhoutaichang/embedding_infer/models/spark2_5-gptq-g128
|
||||
OUT=$HERE/init_memory.jsonl
|
||||
cd $HERE
|
||||
run () { # tag kv_gib max_len [kv_bytes]
|
||||
local extra=""
|
||||
[ -n "${4:-}" ] && extra="--kv-cache-memory-bytes $4"
|
||||
VLLM_CPU_KVCACHE_SPACE=$2 VLLM_OMNI_SPARK_GATE_REDUCTION=1 VLLM_OMNI_CPU_FUSED_NORMS=0 \
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=24 VLLM_CPU_OMP_THREADS_BIND=32-55 \
|
||||
taskset -c 32-55 timeout 1800 $V/bin/python bench_init_memory.py --model $M \
|
||||
--tag "$1" --max-model-len $3 $extra --out $OUT 2>&1 \
|
||||
| grep -E '"(tag|init_s|rss_mb|rss_peak_mb|tokens_generated|kv_cache_memory_bytes)"'
|
||||
}
|
||||
echo "#### server default: 4 GiB of KV for a 112 MiB need"
|
||||
run "mem-kv4gib" 4 2048
|
||||
echo "#### KV sized to a 2048-token context"
|
||||
run "mem-kv128mib" 1 2048 $((128*1024*1024))
|
||||
echo "#### KV sized to a 1024-token context"
|
||||
run "mem-kv64mib" 1 1024 $((64*1024*1024))
|
||||
echo "#### MEMORY SWEEP DONE"
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
#!/bin/bash
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
cd $HERE
|
||||
for M in ${MODELS:-g64 g128}; do
|
||||
for S in "short:ref.json:ref_logits.npy:2048" "long:ref_long.json:ref_long.npy:4096" "tools:ref_tools.json:ref_tools.npy:2048"; do
|
||||
name=${S%%:*}; r=${S#*:}; rj=${r%%:*}; r2=${r#*:}; rn=${r2%%:*}; ml=${r2##*:}
|
||||
echo "######## int4-$M $name"
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=32 VLLM_CPU_OMP_THREADS_BIND=0-31 \
|
||||
taskset -c 0-31 timeout 1800 $V/bin/python cmp_vllm.py \
|
||||
--model /data/zhoutaichang/embedding_infer/models/spark2_5-int4-$M \
|
||||
--ref $rj --ref-logits $rn --max-model-len $ml \
|
||||
--out $HERE/parity_int4_${M}_${name}.json 2>&1 \
|
||||
| grep -E '"(greedy_token_match|first_token_agree|top20_logprob_overlap|top20_logprob_max_abs_diff)"'
|
||||
done
|
||||
done
|
||||
echo "#### PARITY DONE"
|
||||
Executable
+18
@@ -0,0 +1,18 @@
|
||||
#!/bin/bash
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
cd $HERE
|
||||
run () {
|
||||
echo "######## $1"
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=ACTIVE \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=32 VLLM_CPU_OMP_THREADS_BIND=0-31 \
|
||||
taskset -c 0-31 timeout 2400 $V/bin/python ppl_vllm.py --model "$2" --corpus ppl_corpus.txt \
|
||||
--window 512 --tag "$1" --out $HERE/ppl_vllm.jsonl 2>&1 | grep -E '"(tag|ppl|scored_tokens|chunks)"'
|
||||
}
|
||||
run "vllm-bf16" XHToken/Spark-X2.5-1.7B
|
||||
run "vllm-int4-g64" /data/zhoutaichang/embedding_infer/models/spark2_5-int4-g64
|
||||
run "vllm-int4-g128" /data/zhoutaichang/embedding_infer/models/spark2_5-int4-g128
|
||||
echo "#### PPL DONE"
|
||||
Executable
+31
@@ -0,0 +1,31 @@
|
||||
#!/bin/bash
|
||||
# Output comparison only: greedy text is load-independent, so this can run
|
||||
# while the machine is busy.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
G=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B-GGUF/snapshots/23e1fcac55e7dd71e4c12a23723cc228ba0e5e85
|
||||
BF16=/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots/448e61eb392c00f2c403185c5b56d5e0665bfaab
|
||||
cd $HERE
|
||||
gen_vllm () {
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=24 VLLM_CPU_OMP_THREADS_BIND=88-111 \
|
||||
taskset -c 88-111 timeout 3600 $V/bin/python quality_int4.py gen --engine vllm \
|
||||
--model "$1" --out "$2" 2>&1 | tail -2
|
||||
}
|
||||
gen_vllm "$BF16" quality_vllm_bf16.json
|
||||
gen_vllm /data/zhoutaichang/embedding_infer/models/spark2_5-int4-g64 quality_vllm_w4a16.json
|
||||
gen_vllm /data/zhoutaichang/embedding_infer/models/spark2_5-gptq-g64 quality_vllm_w4a8.json
|
||||
taskset -c 88-111 $V/bin/python quality_int4.py gen --engine llamacpp --threads 24 \
|
||||
--gguf $G/Spark-X2.5-1.7B.gguf --out quality_llamacpp_bf16.json 2>&1 | tail -1
|
||||
taskset -c 88-111 $V/bin/python quality_int4.py gen --engine llamacpp --threads 24 \
|
||||
--gguf $G/Spark-X2.5-1.7B-Q4_K_M.gguf --out quality_llamacpp_q4km.json 2>&1 | tail -1
|
||||
echo "### generations done"
|
||||
for pair in "quality_vllm_bf16.json quality_vllm_w4a16.json vllm-W4A16-tinygemm" \
|
||||
"quality_vllm_bf16.json quality_vllm_w4a8.json vllm-W4A8-gptq" \
|
||||
"quality_llamacpp_bf16.json quality_llamacpp_q4km.json llamacpp-Q4_K_M"; do
|
||||
set -- $pair
|
||||
$V/bin/python quality_int4.py cmp --ref $1 --got $2 --label $3 --out quality_cmp_$3.json 2>&1 | tail -6
|
||||
done
|
||||
echo "### QUALITY DONE"
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
#!/bin/bash
|
||||
# Per-op profile instead of throughput: the two 4-bit paths differ by ~10% and
|
||||
# this machine's throughput noise is ~30%, so the GEMM self time is the only
|
||||
# way to rank them.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
M=/data/zhoutaichang/embedding_infer/models
|
||||
cd $HERE
|
||||
prof () {
|
||||
VLLM_OMNI_SPARK_GATE_REDUCTION=1 VLLM_OMNI_CPU_FUSED_NORMS=$3 \
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=active \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=24 VLLM_CPU_OMP_THREADS_BIND=32-55 \
|
||||
taskset -c 32-55 timeout 1800 $V/bin/python profile_cpu_step.py --model "$2" \
|
||||
--out $HERE/profile_$1.json 2>&1 | sed -n '/^{/,$p' | head -12
|
||||
}
|
||||
for pass in 1 2; do
|
||||
echo "==== pass $pass"
|
||||
prof "w4a16_g128_p$pass" $M/spark2_5-int4-g128 0
|
||||
prof "w4a8_g128_p$pass" $M/spark2_5-gptq-g128 0
|
||||
done
|
||||
echo "#### QUANT PROFILE DONE"
|
||||
Executable
+30
@@ -0,0 +1,30 @@
|
||||
#!/bin/bash
|
||||
# Does the AOT cache-load fix generalise past Spark? The shim sits in the
|
||||
# shared torch.compiler path, so any CPU model vLLM compiles should stop
|
||||
# recompiling on every start -- Qwen3-TTS is the harder case because the
|
||||
# pipeline has two engine cores, each with its own compiled artifact.
|
||||
set -u
|
||||
OMNI=/data/zhoutaichang/embedding_infer/vllm-omni
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
SCRATCH=/tmp/claude-2005/-data-zhoutaichang-embedding-infer/e695f57b-0163-4287-b1be-1d0ec9330768/scratchpad
|
||||
FRESH=$SCRATCH/vllm_cache_tts_$$
|
||||
rm -rf $FRESH
|
||||
cd $OMNI
|
||||
run () { # tag shim
|
||||
echo "#### $1 (shim=$2)"
|
||||
VLLM_CACHE_ROOT=$FRESH VLLM_OMNI_AOT_CACHE_SHIM=$2 \
|
||||
VLLM_CPU_KVCACHE_SPACE=4 HF_HUB_OFFLINE=1 \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_ENABLE_V1_MULTIPROCESSING=0 \
|
||||
OMP_NUM_THREADS=24 VLLM_CPU_OMP_THREADS_BIND=32-55 \
|
||||
taskset -c 32-55 timeout 2400 $V/bin/python benchmarks/tts/stream_latency_bench.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --query-type CustomVoice \
|
||||
--txt-prompts benchmarks/tts/prompts_12.txt --warmup 1 --init-only \
|
||||
--tag "$1" --init-timeout 1800 --stage-init-timeout 1200 2>&1 \
|
||||
| grep -iE "init engine|ready|took|init_s|Compiling model again|Directly load AOT|error|Traceback" | tail -8
|
||||
}
|
||||
run tts-cold-shim-on 1
|
||||
run tts-warm-shim-on 1
|
||||
run tts-warm-shim-off 0
|
||||
rm -rf $FRESH
|
||||
echo "#### TTS INIT DONE"
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
# Pin to physical cores 0..N-1 on socket 0 (cores 0-55 physical, 112+ are HT
|
||||
# siblings) so the thread count is enforced, not merely requested.
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
export LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so"
|
||||
export CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0
|
||||
for t in 4 8 16 32; do
|
||||
last=$((t-1))
|
||||
echo "### threads=$t"
|
||||
OMP_NUM_THREADS=$t VLLM_CPU_OMP_THREADS_BIND=0-$last \
|
||||
taskset -c 0-$last timeout 2400 $V/bin/python bench_cpu_vllm.py \
|
||||
--tag "vllm-bf16-t$t" --out vllm_cpu_bench.jsonl 2>&1 | grep -E '"(tag|pp_tps|tg_tps|init_s)"'
|
||||
done
|
||||
+9
@@ -0,0 +1,9 @@
|
||||
#!/bin/bash
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
export LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so"
|
||||
export CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=8
|
||||
export OMP_NUM_THREADS=16 VLLM_CPU_OMP_THREADS_BIND=64-79
|
||||
exec taskset -c 64-79 $V/bin/vllm serve XHToken/Spark-X2.5-1.7B \
|
||||
--dtype bfloat16 --max-model-len 8192 --max-num-seqs 4 \
|
||||
--tool-call-parser glm45 --reasoning-parser glm45 --enable-auto-tool-choice \
|
||||
--host 127.0.0.1 --port 8077
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
#!/bin/bash
|
||||
# vLLM compiles one dynamic-shape graph, so inductor cannot see that a decode
|
||||
# step's tensors are a single row and emits a parallel loop -- and a barrier --
|
||||
# for every norm, rope and GEGLU. compile_sizes pins a static graph for the
|
||||
# shape decode actually runs at.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
M=${MODEL:-/data/zhoutaichang/embedding_infer/models/spark2_5-int4-g64}
|
||||
OUT=${OUT:-$HERE/vllm_cpu_int4_compile.jsonl}
|
||||
CORES=${CORES:-0-31}
|
||||
cd $HERE
|
||||
run () {
|
||||
local name=$1; local cc=$2; shift 2
|
||||
echo "######## $name cc=$cc cores_busy_before=$(python3 cpu_busy.py $CORES)%"
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=ACTIVE KMP_LIBRARY=turnaround \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 \
|
||||
VLLM_ENABLE_V1_MULTIPROCESSING=0 OMP_NUM_THREADS=${THREADS:-32} VLLM_CPU_OMP_THREADS_BIND=$CORES \
|
||||
taskset -c $CORES timeout 2400 $V/bin/python bench_cpu_vllm.py \
|
||||
--model $M --tag "$name" --reps 3 --out $OUT ${cc:+--compilation-config "$cc"} "$@" 2>&1 \
|
||||
| grep -E '"(tag|pp_tps|tg_tps)"|Error|Traceback' | head -8
|
||||
}
|
||||
run "cmp-baseline" ""
|
||||
run "cmp-size1" '{"compile_sizes":[1]}'
|
||||
run "cmp-eager" "" --enforce-eager
|
||||
run "cmp-baseline2" ""
|
||||
echo "#### COMPILE SWEEP DONE"
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
OUT=${OUT:-$HERE/vllm_cpu_int4_decode.jsonl}
|
||||
CORES=${CORES:-0-31}
|
||||
cd $HERE
|
||||
run () {
|
||||
local name=$1; local model=$2; shift 2
|
||||
echo "######## $name busy_before=$(python3 cpu_busy.py $CORES)%"
|
||||
KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=ACTIVE KMP_LIBRARY=turnaround \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 \
|
||||
VLLM_ENABLE_V1_MULTIPROCESSING=0 OMP_NUM_THREADS=32 VLLM_CPU_OMP_THREADS_BIND=$CORES \
|
||||
taskset -c $CORES timeout 2400 $V/bin/python bench_cpu_vllm.py \
|
||||
--model $model --tag "$name" --reps 3 --out $OUT "$@" 2>&1 \
|
||||
| grep -E '"(tag|pp_tps|tg_tps)"|Error|Traceback' | head -6
|
||||
}
|
||||
G64=/data/zhoutaichang/embedding_infer/models/spark2_5-int4-g64
|
||||
G128=/data/zhoutaichang/embedding_infer/models/spark2_5-int4-g128
|
||||
run "g64" $G64
|
||||
run "g64-nodetok" $G64 --no-detokenize
|
||||
run "g128" $G128
|
||||
run "g128-nodetok" $G128 --no-detokenize
|
||||
echo "#### DECODE SWEEP DONE"
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
# Decode at batch 1 is dominated by ~250 OpenMP parallel regions per step, so
|
||||
# how the runtime parks its threads between them matters as much as the GEMM.
|
||||
set -u
|
||||
HERE=/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
M=${MODEL:-/data/zhoutaichang/embedding_infer/models/spark2_5-int4-g64}
|
||||
OUT=${OUT:-$HERE/vllm_cpu_int4_env.jsonl}
|
||||
cd $HERE
|
||||
run () { # name, extra env assignments...
|
||||
local name=$1; shift
|
||||
echo "######## $name"
|
||||
env "$@" \
|
||||
LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so" \
|
||||
CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 \
|
||||
VLLM_ENABLE_V1_MULTIPROCESSING=0 OMP_NUM_THREADS=32 VLLM_CPU_OMP_THREADS_BIND=56-87 \
|
||||
taskset -c 56-87 timeout 2400 $V/bin/python bench_cpu_vllm.py \
|
||||
--model $M --tag "$name" --reps 3 --out $OUT 2>&1 \
|
||||
| grep -E '"(tag|pp_tps|tg_tps)"'
|
||||
}
|
||||
run "env-baseline" DUMMY=1
|
||||
run "env-blocktime-inf" KMP_BLOCKTIME=infinite
|
||||
run "env-active-wait" KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=ACTIVE KMP_AFFINITY=granularity=fine,compact,1,0
|
||||
run "env-spin-hard" KMP_BLOCKTIME=infinite OMP_WAIT_POLICY=ACTIVE KMP_LIBRARY=turnaround
|
||||
echo "#### ENV SWEEP DONE"
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
#!/bin/bash
|
||||
V=/data/zhoutaichang/embedding_infer/.venvs/omni-cpu
|
||||
export LD_PRELOAD="$V/lib/python3.12/site-packages/vllm/libs/libtcmalloc_minimal.so.4:$V/lib/libiomp5.so"
|
||||
export CUDA_VISIBLE_DEVICES= VLLM_TARGET_DEVICE=cpu VLLM_CPU_KVCACHE_SPACE=4 VLLM_ENABLE_V1_MULTIPROCESSING=0
|
||||
for t in 4 8 16 32; do
|
||||
last=$((t-1))
|
||||
echo "### threads=$t"
|
||||
VLLM_CPU_OMP_THREADS_BIND=0-$last timeout 2400 $V/bin/python bench_cpu_vllm.py \
|
||||
--tag "vllm-bf16-t$t" --out vllm_cpu_bench.jsonl 2>&1 | grep -E '"(tag|pp_tps|tg_tps|init_s)"'
|
||||
done
|
||||
@@ -76,7 +76,7 @@ Full analyses in [docs/engines/](../engines/). In brief, for this workload:
|
||||
| **ncnn** | mobile-first, no LLM serving story of the kind needed here |
|
||||
| **vLLM / vLLM-Omni** | the only one with a multi-stage multimodal pipeline, continuous batching and an OpenAI-compatible server — which is why the edge work is a fork of it rather than a port away from it |
|
||||
|
||||
The decision recorded in [docs/analysis/comparison.md](../analysis/comparison.md)
|
||||
The decision recorded in [docs/analysis/comparison.md](../engines/comparison.md)
|
||||
is that a phone NPU cannot run vLLM at all: there the deliverable is an exported
|
||||
graph, and the quantization scheme is decided by the export toolchain. On device
|
||||
(Snapdragon via QNN) only calibrated **w8a16** was measured correct — 36.9 dB;
|
||||
|
||||
Reference in new issue
Block a user