electron-rare/turboquant
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Strip to bare minimum, switch license to GPL-3.0

Browse Source 
Remove 30 scratch/experimental files: HTML reports, profiling scripts,
validation scripts, autoresearch harness, tests, and benchmarks.

Keep only: turboquant/ package, proof.py, benchmark.py, setup.py, README.
License changed from MIT to GPL-3.0.
This commit is contained in:
seroxdesign
2026-03-27 18:28:25 -04:00
parent dbf85683e6
commit 7ac9b8d165
34 changed files with 771 additions and 9467 deletions
Download Patch File Download Diff File Expand all files Collapse all files
HANDOFF_RESEARCH_5090.html
-239
View File
@@ -1,239 +0,0 @@
<!doctype html>
<html><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1">
<title>TurboQuant Research Handoff</title>
<style>body{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,Arial,sans-serif;max-width:980px;margin:40px auto;padding:0 20px;line-height:1.5}pre{white-space:pre-wrap;background:#f6f8fa;padding:12px;border-radius:8px;overflow:auto}code{background:#f6f8fa;padding:2px 4px;border-radius:4px}table{border-collapse:collapse}th,td{border:1px solid #ddd;padding:6px 8px}</style></head><body><h1 id="turboquant-research-handoff-rtx-5090-titpod">TurboQuant Research Handoff (RTX 5090 / TitPod)</h1>
<h2 id="1-scope-and-current-status">1) Scope and current status</h2>
<p>This handoff captures:
- how to access the remote RTX 5090 box,
- where code/data/scripts live (local + remote),
- what has been implemented and validated,
- exact run commands + known working configs,
- known blockers and next research steps.</p>
<p>Current state:
- <strong>TurboQuant modular serving stack is implemented and working</strong> (capture/store/score/integration adapter).
- <strong>No-alloc + shared-KV path for long-context experiments is implemented</strong>.
- <strong>30k context A/B telemetry collected</strong> (baseline vs TurboQuant, separately).
- <strong>200k context with TurboQuant has completed successfully in prior runs</strong> (single-shot completion done), but full dual-case telemetry at 200k is unstable due long-run process kills / baseline stall.</p>
<hr />
<h2 id="2-access-to-device">2) Access to device</h2>
<h2 id="ssh-details">SSH details</h2>
<ul>
<li>Host: <code>your-gpu-host.example.com</code></li>
<li>Port: <code>3003</code></li>
<li>User: <code>root</code></li>
<li>Key (local): <code>~/.ssh/id_rsa</code></li>
</ul>
<p>Use:</p>
<pre><code class="language-bash">ssh -F /dev/null \
-o StrictHostKeyChecking=no \
-o ConnectTimeout=12 \
-o IdentityFile=&quot;~/.ssh/id_rsa&quot; \
-o IdentitiesOnly=yes \
root@your-gpu-host.example.com -p 3003
</code></pre>
<hr />
<h2 id="3-directory-map">3) Directory map</h2>
<h2 id="local">Local</h2>
<ul>
<li>Workspace root: <code>.</code></li>
<li>Main repo used in this work: <code>./turboquant_pkg</code></li>
</ul>
<h2 id="remote">Remote</h2>
<ul>
<li>Work root: <code>/5090-qwen3.5-27b</code></li>
<li>Venv: <code>/5090-qwen3.5-27b/.venv</code></li>
<li>Models: <code>/5090-qwen3.5-27b/models</code></li>
<li>Scripts: <code>/5090-qwen3.5-27b/scripts</code></li>
<li>TurboQuant repo mirror: <code>/5090-qwen3.5-27b/turboquant</code></li>
<li>Runtime package path: <code>/5090-qwen3.5-27b/turboquant/turboquant</code></li>
<li>Logs: <code>/5090-qwen3.5-27b/logs</code></li>
</ul>
<hr />
<h2 id="4-environment-and-versions-remote">4) Environment and versions (remote)</h2>
<ul>
<li>Python: <code>3.12.3</code></li>
<li>Torch: <code>2.10.0+cu128</code></li>
<li>vLLM: <code>0.18.0</code></li>
<li>transformers: <code>5.3.0.dev0</code></li>
</ul>
<p>GPU:
- RTX 5090 32GB (single GPU used in all runs shown here).</p>
<hr />
<h2 id="5-models-on-remote">5) Models on remote</h2>
<ul>
<li><code>/5090-qwen3.5-27b/models/QuantTrio-Qwen3.5-27B-AWQ</code> (~21 GB) ← primary model used</li>
<li><code>/5090-qwen3.5-27b/models/Qwen-Qwen3.5-27B-FP8</code> (~29 GB) ← FP8 baseline experiments only</li>
</ul>
<hr />
<h2 id="6-code-architecture-implemented">6) Code architecture implemented</h2>
<p>Core modular files (in <code>turboquant/</code>):</p>
<ul>
<li><code>capture.py</code></li>
<li><code>RingBuffer</code> + <code>KVCaptureEngine</code></li>
<li>bulk prefill capture + decode buffering</li>
<li><code>store.py</code></li>
<li><code>CompressedKVStore</code>, chunked quantized storage, lazy flattening</li>
<li><code>score.py</code></li>
<li><code>compute_hybrid_attention()</code> for compressed-history + exact-recent merge</li>
<li>memory-safe GQA path (removed large head repeat blowups)</li>
<li><code>integration/vllm.py</code></li>
<li>clean adapter (<code>off | capture_only | hybrid | full_tq</code>)</li>
<li>no-alloc support path added</li>
<li><code>vllm_attn_backend.py</code></li>
<li>compatibility shim + no-alloc hook enablement + KV sharing patching</li>
</ul>
<p>Important 200k-related implementation details:
- <code>enable_no_alloc(...)</code> used before <code>LLM(...)</code>.
- Hook install occurs during executor KV spec phase.
- Flash-attn layers can share KV target layer to reduce paged KV pressure.
- Added layout patch to avoid missing-layer KeyError in hybrid attn+mamba layout update.
- Added no-alloc prefill attention path in modular integration.</p>
<hr />
<h2 id="7-remote-scripts-already-present">7) Remote scripts already present</h2>
<p>Key scripts in <code>/5090-qwen3.5-27b/scripts</code>:
- <code>final_test3.py</code> (baseline + TQ + coherence + 24k needle + KV free)
- <code>long_ctx_runner.py</code>, <code>long_ctx_test.py</code> (long context harness)
- <code>full_bench.py</code>, <code>profile_overhead.py</code>, <code>verify_runner.py</code>, etc.</p>
<hr />
<h2 id="8-validation-completed">8) Validation completed</h2>
<h2 id="test-suites">Test suites</h2>
<p>Local and remote both passed:
- <code>test_modular.py</code> → <code>19/19 pass</code>
- <code>test_turboquant.py</code> → pass</p>
<p>Remote <code>final_test3.py</code> passed historically with:
- coherence suite: pass
- 24K needle retrieval: pass
- KV freed: ~2.9 GB</p>
<hr />
<h2 id="9-measured-results-latest-reliable-set">9) Measured results (latest reliable set)</h2>
<h2 id="a-30k-context-same-prompt-hash-both-runs">A) 30k context (same prompt hash both runs)</h2>
<p>Prompt:
- <code>max_model_len=32768</code>
- prompt tokens: <code>32720</code>
- prompt hash: <code>6f5f3d4d1971c67a2a0ffff7023626a282cdd5aea300c80b0748f1367f5bd510</code></p>
<h3 id="baseline-30k-separate-ttft-and-full-runs">Baseline 30k (separate TTFT and full runs)</h3>
<ul>
<li>Init: <code>40.160s</code></li>
<li>TTFT: <code>18.138s</code></li>
<li>Prefill throughput est: <code>1803.92 tok/s</code></li>
<li>Full run (24 tok): <code>18.992s</code></li>
<li>End-to-end output rate: <code>1.264 tok/s</code></li>
<li>VRAM used:</li>
<li>after init: <code>27833 MB</code></li>
<li>after TTFT/full: <code>~28393 MB</code></li>
<li>Power/temp (after full): <code>313.26W</code>, <code>55C</code></li>
<li>KV cache reserved (unique): <code>2.724 GB</code></li>
<li>Activation est (full pass): <code>644.61 MB</code></li>
</ul>
<h3 id="turboquant-30k-no-alloc-stack">TurboQuant 30k (no-alloc stack)</h3>
<ul>
<li>Init: <code>43.749s</code></li>
<li>TTFT: <code>17.162s</code></li>
<li>Prefill throughput est: <code>1906.52 tok/s</code></li>
<li>Full run (24 tok): <code>18.415s</code></li>
<li>End-to-end output rate: <code>1.303 tok/s</code></li>
<li>VRAM used:</li>
<li>after init: <code>27823 MB</code></li>
<li>after TTFT: <code>28411 MB</code></li>
<li>after full: <code>28419 MB</code></li>
<li>Power/temp (after full): <code>210.64W</code>, <code>53C</code></li>
<li>KV cache reserved (unique): <code>2.724 GB</code></li>
<li>Shared KV layers: <code>15</code></li>
<li>TQ stats in runtime: <code>num_layers=16</code>, mode <code>hybrid</code></li>
<li>Activation est (full pass): <code>599.20 MB</code></li>
</ul>
<h3 id="30k-delta-tq-vs-baseline">30k delta (TQ vs baseline)</h3>
<ul>
<li>TTFT improved by ~<code>0.98s</code></li>
<li>Prefill throughput ~<code>+5.7%</code></li>
<li>Full 24-token latency ~<code>0.58s</code> faster</li>
<li>Full-pass activation estimate ~<code>45 MB</code> lower</li>
</ul>
<h2 id="b-200k-status">B) 200k status</h2>
<p>Reliable observed facts:
- TurboQuant run completed at ~200k prompt:
- prompt tokens: <code>199952</code>
- output tokens: <code>24</code>
- elapsed: <code>58.34s</code>
- gpu mem near end: <code>~31.9 GB</code>
- tq hooked layers: <code>16</code>
- Baseline 200k repeatedly stalls/times out in this setup.
- Full dual-case 200k telemetry scripts were repeatedly interrupted by remote process kills (<code>exit 137</code>) or cancellation side effects.</p>
<hr />
<h2 id="10-known-blockers-failure-modes">10) Known blockers / failure modes</h2>
<ol>
<li><strong>Interrupted runs leave GPU processes alive</strong></li>
<li>cancellation can leave a Python process attached to GPU using ~28 GB VRAM.</li>
<li><strong>Baseline 200k not practical in current config</strong></li>
<li>repeated timeout/stall behavior.</li>
<li><strong>Host/process kill (<code>exit 137</code>) during long scripted telemetry</strong></li>
<li>long single-process “collect everything” jobs can be killed before final JSON write.</li>
<li><strong>FP8 KV + TurboQuant stacking</strong></li>
<li>still incompatible in earlier testing path (FlashInfer metadata mismatch).</li>
</ol>
<hr />
<h2 id="11-operational-playbook-researcher">11) Operational playbook (researcher)</h2>
<h2 id="a-always-clean-gpu-before-runs">A) Always clean GPU before runs</h2>
<pre><code class="language-bash">python3 - &lt;&lt;'PY'
import os, subprocess
try:
out=subprocess.check_output(['nvidia-smi','--query-compute-apps=pid','--format=csv,noheader,nounits'], text=True)
for x in out.splitlines():
x=x.strip()
if x:
try: os.kill(int(x),9)
except: pass
except: pass
print(subprocess.check_output(['nvidia-smi','--query-gpu=memory.used,memory.free,utilization.gpu','--format=csv,noheader,nounits'], text=True))
PY
</code></pre>
<h2 id="b-use-remote-venv-explicitly">B) Use remote venv explicitly</h2>
<pre><code class="language-bash">/5090-qwen3.5-27b/.venv/bin/python ...
</code></pre>
<h2 id="c-split-experiments-by-phase">C) Split experiments by phase</h2>
<p>Do <strong>separate invocations</strong> for:
- init telemetry,
- TTFT/prefill telemetry,
- full generation telemetry,
instead of one giant all-in-one script.</p>
<h2 id="d-suggested-stable-env-flags">D) Suggested stable env flags</h2>
<pre><code class="language-bash">CUDA_VISIBLE_DEVICES=0
VLLM_ENABLE_V1_MULTIPROCESSING=0
TOKENIZERS_PARALLELISM=false
</code></pre>
<h2 id="e-turboquant-enablement">E) TurboQuant enablement</h2>
<p>Before <code>LLM(...)</code>:</p>
<pre><code class="language-python">from turboquant.vllm_attn_backend import enable_no_alloc
enable_no_alloc(key_bits=3, value_bits=2, buffer_size=128, initial_layers_count=4)
</code></pre>
<hr />
<h2 id="12-where-outputslogs-should-be-written">12) Where outputs/logs should be written</h2>
<p>Recommended:
- <code>/5090-qwen3.5-27b/logs/telemetry_30k_baseline.json</code>
- <code>/5090-qwen3.5-27b/logs/telemetry_30k_tq.json</code>
- <code>/5090-qwen3.5-27b/logs/telemetry_200k_tq.json</code></p>
<p>Use one file per phase/case to avoid losing everything on kill.</p>
<hr />
<h2 id="13-what-has-been-accomplished-research-summary">13) What has been accomplished (research summary)</h2>
<ul>
<li>Modular TurboQuant runtime refactor delivered (capture/store/score/integration).</li>
<li>Compatibility shim maintained for older backend entrypoint.</li>
<li>Runtime bug fixes done for real inference paths (shape handling, no-alloc plumbing, KV sharing wiring).</li>
<li>Real remote validations completed on RTX 5090:</li>
<li>test suites pass,</li>
<li>coherence and 24k retrieval pass in established validation scripts,</li>
<li>~2.9 GB KV freeing confirmed in previous runs,</li>
<li>30k A/B telemetry shows small but consistent TQ advantage.</li>
<li>200k with TQ has been demonstrated as runnable; robust full baseline-vs-TQ 200k telemetry remains a follow-up task due baseline instability + long-run process kills.</li>
</ul>
<hr />
<h2 id="14-recommended-next-research-steps">14) Recommended next research steps</h2>
<ol>
<li>Build a <strong>phase-isolated telemetry harness</strong> (each phase separate process).</li>
<li>Add auto-timeout and partial JSON checkpointing after each phase.</li>
<li>Benchmark 30k/50k/80k/120k for baseline+TQ with identical prompt hashes.</li>
<li>For 200k:</li>
<li>prioritize TQ-only telemetry first,</li>
<li>then baseline with reduced scope and strict timeout budget.</li>
<li>Investigate FP8+TQ compatibility path only after stable 200k telemetry pipeline.</li>
</ol></body></html>
HANDOFF_RESEARCH_5090.md
-280
View File
@@ -1,280 +0,0 @@
# TurboQuant Research Handoff (RTX 5090 / TitPod)
## 1) Scope and current status
This handoff captures:
- how to access the remote RTX 5090 box,
- where code/data/scripts live (local + remote),
- what has been implemented and validated,
- exact run commands + known working configs,
- known blockers and next research steps.
Current state:
- **TurboQuant modular serving stack is implemented and working** (capture/store/score/integration adapter).
- **No-alloc + shared-KV path for long-context experiments is implemented**.
- **30k context A/B telemetry collected** (baseline vs TurboQuant, separately).
- **200k context with TurboQuant has completed successfully in prior runs** (single-shot completion done), but full dual-case telemetry at 200k is unstable due long-run process kills / baseline stall.
---
## 2) Access to device
## SSH details
*Redacted for public repo. Configure your own remote GPU host.*
```bash
ssh -o StrictHostKeyChecking=no -o ConnectTimeout=12 user@your-gpu-host
```
---
## 3) Directory map
## Local
- Workspace root: `.`
- Main repo used in this work: `./turboquant_pkg`
## Remote
- Work root: `/5090-qwen3.5-27b`
- Venv: `/5090-qwen3.5-27b/.venv`
- Models: `/5090-qwen3.5-27b/models`
- Scripts: `/5090-qwen3.5-27b/scripts`
- TurboQuant repo mirror: `/5090-qwen3.5-27b/turboquant`
- Runtime package path: `/5090-qwen3.5-27b/turboquant/turboquant`
- Logs: `/5090-qwen3.5-27b/logs`
---
## 4) Environment and versions (remote)
- Python: `3.12.3`
- Torch: `2.10.0+cu128`
- vLLM: `0.18.0`
- transformers: `5.3.0.dev0`
GPU:
- RTX 5090 32GB (single GPU used in all runs shown here).
---
## 5) Models on remote
- `/5090-qwen3.5-27b/models/QuantTrio-Qwen3.5-27B-AWQ` (~21 GB) ← primary model used
- `/5090-qwen3.5-27b/models/Qwen-Qwen3.5-27B-FP8` (~29 GB) ← FP8 baseline experiments only
---
## 6) Code architecture implemented
Core modular files (in `turboquant/`):
- `capture.py`
- `RingBuffer` + `KVCaptureEngine`
- bulk prefill capture + decode buffering
- `store.py`
- `CompressedKVStore`, chunked quantized storage, lazy flattening
- `score.py`
- `compute_hybrid_attention()` for compressed-history + exact-recent merge
- memory-safe GQA path (removed large head repeat blowups)
- `integration/vllm.py`
- clean adapter (`off | capture_only | hybrid | full_tq`)
- no-alloc support path added
- `vllm_attn_backend.py`
- compatibility shim + no-alloc hook enablement + KV sharing patching
Important 200k-related implementation details:
- `enable_no_alloc(...)` used before `LLM(...)`.
- Hook install occurs during executor KV spec phase.
- Flash-attn layers can share KV target layer to reduce paged KV pressure.
- Added layout patch to avoid missing-layer KeyError in hybrid attn+mamba layout update.
- Added no-alloc prefill attention path in modular integration.
---
## 7) Remote scripts already present
Key scripts in `/5090-qwen3.5-27b/scripts`:
- `final_test3.py` (baseline + TQ + coherence + 24k needle + KV free)
- `long_ctx_runner.py`, `long_ctx_test.py` (long context harness)
- `full_bench.py`, `profile_overhead.py`, `verify_runner.py`, etc.
---
## 8) Validation completed
## Test suites
Local and remote both passed:
- `test_modular.py` → `19/19 pass`
- `test_turboquant.py` → pass
Remote `final_test3.py` passed historically with:
- coherence suite: pass
- 24K needle retrieval: pass
- KV freed: ~2.9 GB
---
## 9) Measured results (latest reliable set)
## A) 30k context (same prompt hash both runs)
Prompt:
- `max_model_len=32768`
- prompt tokens: `32720`
- prompt hash: `6f5f3d4d1971c67a2a0ffff7023626a282cdd5aea300c80b0748f1367f5bd510`
### Baseline 30k (separate TTFT and full runs)
- Init: `40.160s`
- TTFT: `18.138s`
- Prefill throughput est: `1803.92 tok/s`
- Full run (24 tok): `18.992s`
- End-to-end output rate: `1.264 tok/s`
- VRAM used:
- after init: `27833 MB`
- after TTFT/full: `~28393 MB`
- Power/temp (after full): `313.26W`, `55C`
- KV cache reserved (unique): `2.724 GB`
- Activation est (full pass): `644.61 MB`
### TurboQuant 30k (no-alloc stack)
- Init: `43.749s`
- TTFT: `17.162s`
- Prefill throughput est: `1906.52 tok/s`
- Full run (24 tok): `18.415s`
- End-to-end output rate: `1.303 tok/s`
- VRAM used:
- after init: `27823 MB`
- after TTFT: `28411 MB`
- after full: `28419 MB`
- Power/temp (after full): `210.64W`, `53C`
- KV cache reserved (unique): `2.724 GB`
- Shared KV layers: `15`
- TQ stats in runtime: `num_layers=16`, mode `hybrid`
- Activation est (full pass): `599.20 MB`
### 30k delta (TQ vs baseline)
- TTFT improved by ~`0.98s`
- Prefill throughput ~`+5.7%`
- Full 24-token latency ~`0.58s` faster
- Full-pass activation estimate ~`45 MB` lower
## B) 200k status
Reliable observed facts:
- TurboQuant run completed at ~200k prompt:
- prompt tokens: `199952`
- output tokens: `24`
- elapsed: `58.34s`
- gpu mem near end: `~31.9 GB`
- tq hooked layers: `16`
- Baseline 200k repeatedly stalls/times out in this setup.
- Full dual-case 200k telemetry scripts were repeatedly interrupted by remote process kills (`exit 137`) or cancellation side effects.
---
## 10) Known blockers / failure modes
1. **Interrupted runs leave GPU processes alive**
- cancellation can leave a Python process attached to GPU using ~28 GB VRAM.
2. **Baseline 200k not practical in current config**
- repeated timeout/stall behavior.
3. **Host/process kill (`exit 137`) during long scripted telemetry**
- long single-process “collect everything” jobs can be killed before final JSON write.
4. **FP8 KV + TurboQuant stacking**
- still incompatible in earlier testing path (FlashInfer metadata mismatch).
---
## 11) Operational playbook (researcher)
## A) Always clean GPU before runs
```bash
python3 - <<'PY'
import os, subprocess
try:
out=subprocess.check_output(['nvidia-smi','--query-compute-apps=pid','--format=csv,noheader,nounits'], text=True)
for x in out.splitlines():
x=x.strip()
if x:
try: os.kill(int(x),9)
except: pass
except: pass
print(subprocess.check_output(['nvidia-smi','--query-gpu=memory.used,memory.free,utilization.gpu','--format=csv,noheader,nounits'], text=True))
PY
```
## B) Use remote venv explicitly
```bash
/5090-qwen3.5-27b/.venv/bin/python ...
```
## C) Split experiments by phase
Do **separate invocations** for:
- init telemetry,
- TTFT/prefill telemetry,
- full generation telemetry,
instead of one giant all-in-one script.
## D) Suggested stable env flags
```bash
CUDA_VISIBLE_DEVICES=0
VLLM_ENABLE_V1_MULTIPROCESSING=0
TOKENIZERS_PARALLELISM=false
```
## E) TurboQuant enablement
Before `LLM(...)`:
```python
from turboquant.vllm_attn_backend import enable_no_alloc
enable_no_alloc(key_bits=3, value_bits=2, buffer_size=128, initial_layers_count=4)
```
---
## 12) Where outputs/logs should be written
Recommended:
- `/5090-qwen3.5-27b/logs/telemetry_30k_baseline.json`
- `/5090-qwen3.5-27b/logs/telemetry_30k_tq.json`
- `/5090-qwen3.5-27b/logs/telemetry_200k_tq.json`
Use one file per phase/case to avoid losing everything on kill.
---
## 13) What has been accomplished (research summary)
- Modular TurboQuant runtime refactor delivered (capture/store/score/integration).
- Compatibility shim maintained for older backend entrypoint.
- Runtime bug fixes done for real inference paths (shape handling, no-alloc plumbing, KV sharing wiring).
- Real remote validations completed on RTX 5090:
- test suites pass,
- coherence and 24k retrieval pass in established validation scripts,
- ~2.9 GB KV freeing confirmed in previous runs,
- 30k A/B telemetry shows small but consistent TQ advantage.
- 200k with TQ has been demonstrated as runnable; robust full baseline-vs-TQ 200k telemetry remains a follow-up task due baseline instability + long-run process kills.
---
## 14) Recommended next research steps
1. Build a **phase-isolated telemetry harness** (each phase separate process).
2. Add auto-timeout and partial JSON checkpointing after each phase.
3. Benchmark 30k/50k/80k/120k for baseline+TQ with identical prompt hashes.
4. For 200k:
- prioritize TQ-only telemetry first,
- then baseline with reduced scope and strict timeout budget.
5. Investigate FP8+TQ compatibility path only after stable 200k telemetry pipeline.
LICENSE
+670 -17
View File
@@ -1,21 +1,674 @@
MIT License
GNU GENERAL PUBLIC LICENSE
Version 3, 29 June 2007
Copyright (c) 2025
Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
Everyone is permitted to copy and distribute verbatim copies
of this license document, but changing it is not allowed.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
Preamble
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
The GNU General Public License is a free, copyleft license for
software and other kinds of works.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
The licenses for most software and other practical works are designed
to take away your freedom to share and change the works. By contrast,
the GNU General Public License is intended to guarantee your freedom to
share and change all versions of a program--to make sure it remains free
software for all its users. We, the Free Software Foundation, use the
GNU General Public License for most of our software; it applies also to
any other work released this way by its authors. You can apply it to
your programs, too.
When we speak of free software, we are referring to freedom, not
price. Our General Public Licenses are designed to make sure that you
have the freedom to distribute copies of free software (and charge for
them if you wish), that you receive source code or can get it if you
want it, that you can change the software or use pieces of it in new
free programs, and that you know you can do these things.
To protect your rights, we need to prevent others from denying you
these rights or asking you to surrender the rights. Therefore, you have
certain responsibilities if you distribute copies of the software, or if
you modify it: responsibilities to respect the freedom of others.
For example, if you distribute copies of such a program, whether
gratis or for a fee, you must pass on to the recipients the same
freedoms that you received. You must make sure that they, too, receive
or can get the source code. And you must show them these terms so they
know their rights.
Developers that use the GNU GPL protect your rights with two steps:
(1) assert copyright on the software, and (2) offer you this License
giving you legal permission to copy, distribute and/or modify it.
For the developers' and authors' protection, the GPL clearly explains
that there is no warranty for this free software. For both users' and
authors' sake, the GPL requires that modified versions be marked as
changed, so that their problems will not be attributed erroneously to
authors of previous versions.
Some devices are designed to deny users access to install or run
modified versions of the software inside them, although the manufacturer
can do so. This is fundamentally incompatible with the aim of
protecting users' freedom to change the software. The systematic
pattern of such abuse occurs in the area of products for individuals to
use, which is precisely where it is most unacceptable. Therefore, we
have designed this version of the GPL to prohibit the practice for those
products. If such problems arise substantially in other domains, we
stand ready to extend this provision to those domains in future versions
of the GPL, as needed to protect the freedom of users.
Finally, every program is threatened constantly by software patents.
States should not allow patents to restrict development and use of
software on general-purpose computers, but in those that do, we wish to
avoid the special danger that patents applied to a free program could
make it effectively proprietary. To prevent this, the GPL assures that
patents cannot be used to render the program non-free.
The precise terms and conditions for copying, distribution and
modification follow.
TERMS AND CONDITIONS
0. Definitions.
"This License" refers to version 3 of the GNU General Public License.
"Copyright" also means copyright-like laws that apply to other kinds of
works, such as semiconductor masks.
"The Program" refers to any copyrightable work licensed under this
License. Each licensee is addressed as "you". "Licensees" and
"recipients" may be individuals or organizations.
To "modify" a work means to copy from or adapt all or part of the work
in a fashion requiring copyright permission, other than the making of an
exact copy. The resulting work is called a "modified version" of the
earlier work or a work "based on" the earlier work.
A "covered work" means either the unmodified Program or a work based
on the Program.
To "propagate" a work means to do anything with it that, without
permission, would make you directly or secondarily liable for
infringement under applicable copyright law, except executing it on a
computer or modifying a private copy. Propagation includes copying,
distribution (with or without modification), making available to the
public, and in some countries other activities as well.
To "convey" a work means any kind of propagation that enables other
parties to make or receive copies. Mere interaction with a user through
a computer network, with no transfer of a copy, is not conveying.
An interactive user interface displays "Appropriate Legal Notices"
to the extent that it includes a convenient and prominently visible
feature that (1) displays an appropriate copyright notice, and (2)
tells the user that there is no warranty for the work (except to the
extent that warranties are provided), that licensees may convey the
work under this License, and how to view a copy of this License. If
the interface presents a list of user commands or options, such as a
menu, a prominent item in the list meets this criterion.
1. Source Code.
The "source code" for a work means the preferred form of the work
for making modifications to it. "Object code" means any non-source
form of a work.
A "Standard Interface" means an interface that either is an official
standard defined by a recognized standards body, or, in the case of
interfaces specified for a particular programming language, one that
is widely used among developers working in that language.
The "System Libraries" of an executable work include anything, other
than the work as a whole, that (a) is included in the normal form of
packaging a Major Component, but which is not part of that Major
Component, and (b) serves only to enable use of the work with that
Major Component, or to implement a Standard Interface for which an
implementation is available to the public in source code form. A
"Major Component", in this context, means a major essential component
(kernel, window system, and so on) of the specific operating system
(if any) on which the executable work runs, or a compiler used to
produce the work, or an object code interpreter used to run it.
The "Corresponding Source" for a work in object code form means all
the source code needed to generate, install, and (for an executable
work) run the object code and to modify the work, including scripts to
control those activities. However, it does not include the work's
System Libraries, or general-purpose tools or generally available free
programs which are used unmodified in performing those activities but
which are not part of the work. For example, Corresponding Source
includes interface definition files associated with source files for
the work, and the source code for shared libraries and dynamically
linked subprograms that the work is specifically designed to require,
such as by intimate data communication or control flow between those
subprograms and other parts of the work.
The Corresponding Source need not include anything that users
can regenerate automatically from other parts of the Corresponding
Source.
The Corresponding Source for a work in source code form is that
same work.
2. Basic Permissions.
All rights granted under this License are granted for the term of
copyright on the Program, and are irrevocable provided the stated
conditions are met. This License explicitly affirms your unlimited
permission to run the unmodified Program. The output from running a
covered work is covered by this License only if the output, given its
content, constitutes a covered work. This License acknowledges your
rights of fair use or other equivalent, as provided by copyright law.
You may make, run and propagate covered works that you do not
convey, without conditions so long as your license otherwise remains
in force. You may convey covered works to others for the sole purpose
of having them make modifications exclusively for you, or provide you
with facilities for running those works, provided that you comply with
the terms of this License in conveying all material for which you do
not control copyright. Those thus making or running the covered works
for you must do so exclusively on your behalf, under your direction
and control, on terms that prohibit them from making any copies of
your copyrighted material outside their relationship with you.
Conveying under any other circumstances is permitted solely under
the conditions stated below. Sublicensing is not allowed; section 10
makes it unnecessary.
3. Protecting Users' Legal Rights From Anti-Circumvention Law.
No covered work shall be deemed part of an effective technological
measure under any applicable law fulfilling obligations under article
11 of the WIPO copyright treaty adopted on 20 December 1996, or
similar laws prohibiting or restricting circumvention of such
measures.
When you convey a covered work, you waive any legal power to forbid
circumvention of technological measures to the extent such circumvention
is effected by exercising rights under this License with respect to
the covered work, and you disclaim any intention to limit operation or
modification of the work as a means of enforcing, against the work's
users, your or third parties' legal rights to forbid circumvention of
technological measures.
4. Conveying Verbatim Copies.
You may convey verbatim copies of the Program's source code as you
receive it, in any medium, provided that you conspicuously and
appropriately publish on each copy an appropriate copyright notice;
keep intact all notices stating that this License and any
non-permissive terms added in accord with section 7 apply to the code;
keep intact all notices of the absence of any warranty; and give all
recipients a copy of this License along with the Program.
You may charge any price or no price for each copy that you convey,
and you may offer support or warranty protection for a fee.
5. Conveying Modified Source Versions.
You may convey a work based on the Program, or the modifications to
produce it from the Program, in the form of source code under the
terms of section 4, provided that you also meet all of these conditions:
a) The work must carry prominent notices stating that you modified
it, and giving a relevant date.
b) The work must carry prominent notices stating that it is
released under this License and any conditions added under section
7. This requirement modifies the requirement in section 4 to
"keep intact all notices".
c) You must license the entire work, as a whole, under this
License to anyone who comes into possession of a copy. This
License will therefore apply, along with any applicable section 7
additional terms, to the whole of the work, and all its parts,
regardless of how they are packaged. This License gives no
permission to license the work in any other way, but it does not
invalidate such permission if you have separately received it.
d) If the work has interactive user interfaces, each must display
Appropriate Legal Notices; however, if the Program has interactive
interfaces that do not display Appropriate Legal Notices, your
work need not make them do so.
A compilation of a covered work with other separate and independent
works, which are not by their nature extensions of the covered work,
and which are not combined with it such as to form a larger program,
in or on a volume of a storage or distribution medium, is called an
"aggregate" if the compilation and its resulting copyright are not
used to limit the access or legal rights of the compilation's users
beyond what the individual works permit. Inclusion of a covered work
in an aggregate does not cause this License to apply to the other
parts of the aggregate.
6. Conveying Non-Source Forms.
You may convey a covered work in object code form under the terms
of sections 4 and 5, provided that you also convey the
machine-readable Corresponding Source under the terms of this License,
in one of these ways:
a) Convey the object code in, or embodied in, a physical product
(including a physical distribution medium), accompanied by the
Corresponding Source fixed on a durable physical medium
customarily used for software interchange.
b) Convey the object code in, or embodied in, a physical product
(including a physical distribution medium), accompanied by a
written offer, valid for at least three years and valid for as
long as you offer spare parts or customer support for that product
model, to give anyone who possesses the object code either (1) a
copy of the Corresponding Source for all the software in the
product that is covered by this License, on a durable physical
medium customarily used for software interchange, for a price no
more than your reasonable cost of physically performing this
conveying of source, or (2) access to copy the
Corresponding Source from a network server at no charge.
c) Convey individual copies of the object code with a copy of the
written offer to provide the Corresponding Source. This
alternative is allowed only occasionally and noncommercially, and
only if you received the object code with such an offer, in accord
with subsection 6b.
d) Convey the object code by offering access from a designated
place (gratis or for a charge), and offer equivalent access to the
Corresponding Source in the same way through the same place at no
further charge. You need not require recipients to copy the
Corresponding Source along with the object code. If the place to
copy the object code is a network server, the Corresponding Source
may be on a different server (operated by you or a third party)
that supports equivalent copying facilities, provided you maintain
clear directions next to the object code saying where to find the
Corresponding Source. Regardless of what server hosts the
Corresponding Source, you remain obligated to ensure that it is
available for as long as needed to satisfy these requirements.
e) Convey the object code using peer-to-peer transmission, provided
you inform other peers where the object code and Corresponding
Source of the work are being offered to the general public at no
charge under subsection 6d.
A separable portion of the object code, whose source code is excluded
from the Corresponding Source as a System Library, need not be
included in conveying the object code work.
A "User Product" is either (1) a "consumer product", which means any
tangible personal property which is normally used for personal, family,
or household purposes, or (2) anything designed or sold for incorporation
into a dwelling. In determining whether a product is a consumer product,
doubtful cases shall be resolved in favor of coverage. For a particular
product received by a particular user, "normally used" refers to a
typical or common use of that class of product, regardless of the status
of the particular user or of the way in which the particular user
actually uses, or expects or is expected to use, the product. A product
is a consumer product regardless of whether the product has substantial
commercial, industrial or non-consumer uses, unless such uses represent
the only significant mode of use of the product.
"Installation Information" for a User Product means any methods,
procedures, authorization keys, or other information required to install
and execute modified versions of a covered work in that User Product from
a modified version of its Corresponding Source. The information must
suffice to ensure that the continued functioning of the modified object
code is in no case prevented or interfered with solely because
modification has been made.
If you convey an object code work under this section in, or with, or
specifically for use in, a User Product, and the conveying occurs as
part of a transaction in which the right of possession and use of the
User Product is transferred to the recipient in perpetuity or for a
fixed term (regardless of how the transaction is characterized), the
Corresponding Source conveyed under this section must be accompanied
by the Installation Information. But this requirement does not apply
if neither you nor any third party retains the ability to install
modified object code on the User Product (for example, the work has
been installed in ROM).
The requirement to provide Installation Information does not include a
requirement to continue to provide support service, warranty, or updates
for a work that has been modified or installed by the recipient, or for
the User Product in which it has been modified or installed. Access to a
network may be denied when the modification itself materially and
adversely affects the operation of the network or violates the rules and
protocols for communication across the network.
Corresponding Source conveyed, and Installation Information provided,
in accord with this section must be in a format that is publicly
documented (and with an implementation available to the public in
source code form), and must require no special password or key for
unpacking, reading or copying.
7. Additional Terms.
"Additional permissions" are terms that supplement the terms of this
License by making exceptions from one or more of its conditions.
Additional permissions that are applicable to the entire Program shall
be treated as though they were included in this License, to the extent
that they are valid under applicable law. If additional permissions
apply only to part of the Program, that part may be used separately
under those permissions, but the entire Program remains governed by
this License without regard to the additional permissions.
When you convey a copy of a covered work, you may at your option
remove any additional permissions from that copy, or from any part of
it. (Additional permissions may be written to require their own
removal in certain cases when you modify the work.) You may place
additional permissions on material, added by you to a covered work,
for which you have or can give appropriate copyright permission.
Notwithstanding any other provision of this License, for material you
add to a covered work, you may (if authorized by the copyright holders of
that material) supplement the terms of this License with terms:
a) Disclaiming warranty or limiting liability differently from the
terms of sections 15 and 16 of this License; or
b) Requiring preservation of specified reasonable legal notices or
author attributions in that material or in the Appropriate Legal
Notices displayed by works containing it; or
c) Prohibiting misrepresentation of the origin of that material, or
requiring that modified versions of such material be marked in
reasonable ways as different from the original version; or
d) Limiting the use for publicity purposes of names of licensors or
authors of the material; or
e) Declining to grant rights under trademark law for use of some
trade names, trademarks, or service marks; or
f) Requiring indemnification of licensors and authors of that
material by anyone who conveys the material (or modified versions of
it) with contractual assumptions of liability to the recipient, for
any liability that these contractual assumptions directly impose on
those licensors and authors.
All other non-permissive additional terms are considered "further
restrictions" within the meaning of section 10. If the Program as you
received it, or any part of it, contains a notice stating that it is
governed by this License along with a term that is a further
restriction, you may remove that term. If a license document contains
a further restriction but permits relicensing or conveying under this
License, you may add to a covered work material governed by the terms
of that license document, provided that the further restriction does
not survive such relicensing or conveying.
If you add terms to a covered work in accord with this section, you
must place, in the relevant source files, a statement of the
additional terms that apply to those files, or a notice indicating
where to find the applicable terms.
Additional terms, permissive or non-permissive, may be stated in the
form of a separately written license, or stated as exceptions;
the above requirements apply either way.
8. Termination.
You may not propagate or modify a covered work except as expressly
provided under this License. Any attempt otherwise to propagate or
modify it is void, and will automatically terminate your rights under
this License (including any patent licenses granted under the third
paragraph of section 11).
However, if you cease all violation of this License, then your
license from a particular copyright holder is reinstated (a)
provisionally, unless and until the copyright holder explicitly and
finally terminates your license, and (b) permanently, if the copyright
holder fails to notify you of the violation by some reasonable means
prior to 60 days after the cessation.
Moreover, your license from a particular copyright holder is
reinstated permanently if the copyright holder notifies you of the
violation by some reasonable means, this is the first time you have
received notice of violation of this License (for any work) from that
copyright holder, and you cure the violation prior to 30 days after
your receipt of the notice.
Termination of your rights under this section does not terminate the
licenses of parties who have received copies or rights from you under
this License. If your rights have been terminated and not permanently
reinstated, you do not qualify to receive new licenses for the same
material under section 10.
9. Acceptance Not Required for Having Copies.
You are not required to accept this License in order to receive or
run a copy of the Program. Ancillary propagation of a covered work
occurring solely as a consequence of using peer-to-peer transmission
to receive a copy likewise does not require acceptance. However,
nothing other than this License grants you permission to propagate or
modify any covered work. These actions infringe copyright if you do
not accept this License. Therefore, by modifying or propagating a
covered work, you indicate your acceptance of this License to do so.
10. Automatic Licensing of Downstream Recipients.
Each time you convey a covered work, the recipient automatically
receives a license from the original licensors, to run, modify and
propagate that work, subject to this License. You are not responsible
for enforcing compliance by third parties with this License.
An "entity transaction" is a transaction transferring control of an
organization, or substantially all assets of one, or subdividing an
organization, or merging organizations. If propagation of a covered
work results from an entity transaction, each party to that
transaction who receives a copy of the work also receives whatever
licenses to the work the party's predecessor in interest had or could
give under the previous paragraph, plus a right to possession of the
Corresponding Source of the work from the predecessor in interest, if
the predecessor has it or can get it with reasonable efforts.
You may not impose any further restrictions on the exercise of the
rights granted or affirmed under this License. For example, you may
not impose a license fee, royalty, or other charge for exercise of
rights granted under this License, and you may not initiate litigation
(including a cross-claim or counterclaim in a lawsuit) alleging that
any patent claim is infringed by making, using, selling, offering for
sale, or importing the Program or any portion of it.
11. Patents.
A "contributor" is a copyright holder who authorizes use under this
License of the Program or a work on which the Program is based. The
work thus licensed is called the contributor's "contributor version".
A contributor's "essential patent claims" are all patent claims
owned or controlled by the contributor, whether already acquired or
hereafter acquired, that would be infringed by some manner, permitted
by this License, of making, using, or selling its contributor version,
but do not include claims that would be infringed only as a
consequence of further modification of the contributor version. For
purposes of this definition, "control" includes the right to grant
patent sublicenses in a manner consistent with the requirements of
this License.
Each contributor grants you a non-exclusive, worldwide, royalty-free
patent license under the contributor's essential patent claims, to
make, use, sell, offer for sale, import and otherwise run, modify and
propagate the contents of its contributor version.
In the following three paragraphs, a "patent license" is any express
agreement or commitment, however denominated, not to enforce a patent
(such as an express permission to practice a patent or covenant not to
sue for patent infringement). To "grant" such a patent license to a
party means to make such an agreement or commitment not to enforce a
patent against the party.
If you convey a covered work, knowingly relying on a patent license,
and the Corresponding Source of the work is not available for anyone
to copy, free of charge and under the terms of this License, through a
publicly available network server or other readily accessible means,
then you must either (1) cause the Corresponding Source to be so
available, or (2) arrange to deprive yourself of the benefit of the
patent license for this particular work, or (3) arrange, in a manner
consistent with the requirements of this License, to extend the patent
license to downstream recipients. "Knowingly relying" means you have
actual knowledge that, but for the patent license, your conveying the
covered work in a country, or your recipient's use of the covered work
in a country, would infringe one or more identifiable patents in that
country that you have reason to believe are valid.
If, pursuant to or in connection with a single transaction or
arrangement, you convey, or propagate by procuring conveyance of, a
covered work, and grant a patent license to some of the parties
receiving the covered work authorizing them to use, propagate, modify
or convey a specific copy of the covered work, then the patent license
you grant is automatically extended to all recipients of the covered
work and works based on it.
A patent license is "discriminatory" if it does not include within
the scope of its coverage, prohibits the exercise of, or is
conditioned on the non-exercise of one or more of the rights that are
specifically granted under this License. You may not convey a covered
work if you are a party to an arrangement with a third party that is
in the business of distributing software, under which you make payment
to the third party based on the extent of your activity of conveying
the work, and under which the third party grants, to any of the
parties who would receive the covered work from you, a discriminatory
patent license (a) in connection with copies of the covered work
conveyed by you (or copies made from those copies), or (b) primarily
for and in connection with specific products or compilations that
contain the covered work, unless you entered into that arrangement,
or that patent license was granted, prior to 28 March 2007.
Nothing in this License shall be construed as excluding or limiting
any implied license or other defenses to infringement that may
otherwise be available to you under applicable patent law.
12. No Surrender of Others' Freedom.
If conditions are imposed on you (whether by court order, agreement or
otherwise) that contradict the conditions of this License, they do not
excuse you from the conditions of this License. If you cannot convey a
covered work so as to satisfy simultaneously your obligations under this
License and any other pertinent obligations, then as a consequence you may
not convey it at all. For example, if you agree to terms that obligate you
to collect a royalty for further conveying from those to whom you convey
the Program, the only way you could satisfy both those terms and this
License would be to refrain entirely from conveying the Program.
13. Use with the GNU Affero General Public License.
Notwithstanding any other provision of this License, you have
permission to link or combine any covered work with a work licensed
under version 3 of the GNU Affero General Public License into a single
combined work, and to convey the resulting work. The terms of this
License will continue to apply to the part which is the covered work,
but the special requirements of the GNU Affero General Public License,
section 13, concerning interaction through a network will apply to the
combination as such.
14. Revised Versions of this License.
The Free Software Foundation may publish revised and/or new versions of
the GNU General Public License from time to time. Such new versions will
be similar in spirit to the present version, but may differ in detail to
address new problems or concerns.
Each version is given a distinguishing version number. If the
Program specifies that a certain numbered version of the GNU General
Public License "or any later version" applies to it, you have the
option of following the terms and conditions either of that numbered
version or of any later version published by the Free Software
Foundation. If the Program does not specify a version number of the
GNU General Public License, you may choose any version ever published
by the Free Software Foundation.
If the Program specifies that a proxy can decide which future
versions of the GNU General Public License can be used, that proxy's
public statement of acceptance of a version permanently authorizes you
to choose that version for the Program.
Later license versions may give you additional or different
permissions. However, no additional obligations are imposed on any
author or copyright holder as a result of your choosing to follow a
later version.
15. Disclaimer of Warranty.
THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
16. Limitation of Liability.
IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
SUCH DAMAGES.
17. Interpretation of Sections 15 and 16.
If the disclaimer of warranty and limitation of liability provided
above cannot be given local legal effect according to their terms,
reviewing courts shall apply local law that most closely approximates
an absolute waiver of all civil liability in connection with the
Program, unless a warranty or assumption of liability accompanies a
copy of the Program in return for a fee.
END OF TERMS AND CONDITIONS
How to Apply These Terms to Your New Programs
If you develop a new program, and you want it to be of the greatest
possible use to the public, the best way to achieve this is to make it
free software which everyone can redistribute and change under these terms.
To do so, attach the following notices to the program. It is safest
to attach them to the start of each source file to most effectively
state the exclusion of warranty; and each file should have at least
the "copyright" line and a pointer to where the full notice is found.
<one line to give the program's name and a brief idea of what it does.>
Copyright (C) <year> <name of author>
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
Also add information on how to contact you by electronic and paper mail.
If the program does terminal interaction, make it output a short
notice like this when it starts in an interactive mode:
<program> Copyright (C) <year> <name of author>
This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
This is free software, and you are welcome to redistribute it
under certain conditions; type `show c' for details.
The hypothetical commands `show w' and `show c' should show the appropriate
parts of the General Public License. Of course, your program's commands
might be different; for a GUI interface, you would use an "about box".
You should also get your employer (if you work as a programmer) or school,
if any, to sign a "copyright disclaimer" for the program, if necessary.
For more information on this, and how to apply and follow the GNU GPL, see
<https://www.gnu.org/licenses/>.
The GNU General Public License does not permit incorporating your program
into proprietary programs. If your program is a subroutine library, you
may consider it more useful to permit linking proprietary applications with
the library. If this is what you want to do, use the GNU Lesser General
Public License instead of this License. But first, please read
<https://www.gnu.org/licenses/why-not-lgpl.html>.
REVIEW_PAPER_VS_IMPL.html
-294
View File
@@ -1,294 +0,0 @@
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>TurboQuant: Paper vs Implementation</title>
<script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>
<style>
:root {
--bg: #0d1117; --surface: #161b22; --raised: #1c2129;
--border: #30363d; --text: #e6edf3; --muted: #8b949e;
--accent: #58a6ff; --green: #3fb950; --red: #f85149; --yellow: #d29922;
--green-bg: rgba(63,185,80,0.10); --red-bg: rgba(248,81,73,0.10); --yellow-bg: rgba(210,153,34,0.10);
}
* { margin: 0; padding: 0; box-sizing: border-box; }
body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; background: var(--bg); color: var(--text); line-height: 1.55; }
.hdr { border-bottom: 1px solid var(--border); padding: 40px 24px 32px; text-align: center; }
.hdr h1 { font-size: 28px; font-weight: 700; letter-spacing: -0.4px; }
.hdr h1 span { color: var(--accent); }
.hdr .sub { color: var(--muted); font-size: 14px; margin-top: 6px; }
.wrap { max-width: 960px; margin: 0 auto; padding: 24px 20px 48px; }
.s { background: var(--surface); border: 1px solid var(--border); border-radius: 10px; padding: 24px; margin-bottom: 20px; }
.s h2 { font-size: 19px; font-weight: 700; margin-bottom: 16px; padding-bottom: 10px; border-bottom: 1px solid var(--border); }
h3 { font-size: 15px; font-weight: 600; margin: 16px 0 8px; }
p, li { font-size: 14px; } p { margin-bottom: 10px; }
ul, ol { padding-left: 18px; margin-bottom: 10px; } li { margin-bottom: 3px; }
strong { color: #fff; }
code { font-family: 'SF Mono', monospace; font-size: 12px; background: var(--raised); padding: 1px 5px; border-radius: 3px; border: 1px solid var(--border); }
table { width: 100%; border-collapse: collapse; margin: 12px 0; font-size: 13px; }
th { text-align: left; padding: 8px 10px; background: var(--raised); border: 1px solid var(--border); font-weight: 600; font-size: 12px; text-transform: uppercase; letter-spacing: 0.3px; color: var(--muted); }
td { padding: 8px 10px; border: 1px solid var(--border); }
.st { display: inline-block; font-size: 11px; font-weight: 600; padding: 2px 6px; border-radius: 3px; }
.ok { background: var(--green-bg); color: var(--green); }
.warn { background: var(--yellow-bg); color: var(--yellow); }
.fail { background: var(--red-bg); color: var(--red); }
.badge { display: inline-block; font-size: 10px; font-weight: 600; text-transform: uppercase; padding: 2px 6px; border-radius: 4px; margin-left: 6px; vertical-align: middle; }
.mw { background: #fff; border-radius: 6px; padding: 16px 12px; margin: 12px 0; overflow-x: auto; }
.g2 { display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin: 12px 0; }
.g2 .c { background: var(--raised); border: 1px solid var(--border); border-radius: 6px; padding: 16px; }
.c h3 { margin-top: 0; }
.c ul { margin-bottom: 0; }
.cw { border-left: 3px solid; padding: 10px 14px; border-radius: 0 6px 6px 0; margin: 12px 0; font-size: 13px; }
.cw-r { border-color: var(--red); background: var(--red-bg); }
@media (max-width: 640px) { .g2 { grid-template-columns: 1fr; } .hdr h1 { font-size: 22px; } .s { padding: 16px; } }
</style>
</head>
<body>
<div class="hdr">
<h1><span>TurboQuant</span>: Paper vs Implementation</h1>
<p class="sub">arXiv:2504.19874 (ICLR 2026) &mdash; Qwen3.5-27B-AWQ on RTX 5090 &mdash; <a href="https://github.com/0xSero/turboquant" style="color:var(--accent)">github.com/0xSero/turboquant</a></p>
</div>
<div class="wrap">
<!-- A: Implementation Status -->
<div class="s">
<h2>Implementation Status vs Paper</h2>
<table>
<tr><th>Paper Component</th><th>Status</th><th>File</th></tr>
<tr><td>Random orthogonal rotation (QR)</td><td><span class="st ok">Faithful</span></td><td><code>rotation.py</code></td></tr>
<tr><td>Lloyd-Max codebook on Beta PDF</td><td><span class="st ok">Faithful</span></td><td><code>codebook.py</code></td></tr>
<tr><td>TurboQuant_MSE (Algorithm 1)</td><td><span class="st ok">Faithful</span></td><td><code>quantizer.py</code></td></tr>
<tr><td>TurboQuant_Prod (Algorithm 2)</td><td><span class="st ok">Faithful</span></td><td><code>quantizer.py</code></td></tr>
<tr><td>QJL projection + sign packing</td><td><span class="st ok">Faithful</span></td><td><code>quantizer.py</code>, <code>rotation.py</code></td></tr>
<tr><td>Value quantization</td><td><span class="st warn">Not from paper</span></td><td><code>kv_cache.py</code> (group quant)</td></tr>
<tr><td>Outlier channel splitting</td><td><span class="st fail">Removed</span></td><td>Was dead code, stripped</td></tr>
<tr><td>MLA support</td><td><span class="st fail">Stub</span></td><td><code>integration/vllm.py</code></td></tr>
<tr><td>vLLM integration + Triton kernels</td><td><span class="st ok">Novel eng.</span></td><td><code>integration/</code>, <code>triton_kernels.py</code></td></tr>
<tr><td>Modular capture/store/score</td><td><span class="st ok">Novel eng.</span></td><td><code>capture.py</code>, <code>store.py</code>, <code>score.py</code></td></tr>
</table>
</div>
<!-- B: Attention Backends -->
<div class="s">
<h2>Attention Backends</h2>
<table>
<tr><th>Backend</th><th>Status</th><th>Notes</th></tr>
<tr><td>Flash Attention (MHA/GQA)</td><td><span class="st ok">Working</span></td><td>Primary path, hooks on FlashAttentionImpl</td></tr>
<tr><td>Flash Attention (no-alloc prefill)</td><td><span class="st warn">Degraded</span></td><td>Falls back to F.scaled_dot_product_attention</td></tr>
<tr><td>FlashInfer</td><td><span class="st fail">Incompatible</span></td><td>Metadata mismatch with FP8 KV</td></tr>
<tr><td>MLA</td><td><span class="st fail">Stub</span></td><td>Passthrough only</td></tr>
<tr><td>GDN / Mamba</td><td><span class="st warn">N/A</span></td><td>No KV cache, TQ hooks on attention layers only</td></tr>
</table>
</div>
<!-- C: Known Issues -->
<div class="s">
<h2>Known Issues <span class="badge" style="background:var(--red-bg);color:var(--red);">After Trim</span></h2>
<ul>
<li><strong>No-alloc decode fallback returns zeros</strong> when TQ store has &lt;16 tokens</li>
<li><strong>No-alloc prefill uses <code>repeat_interleave</code></strong> for GQA, massive memory at 200k</li>
<li><strong>Value group_size hardcoded to 32</strong> (now explicit, was previously fragile math)</li>
<li><strong>MLA path is passthrough</strong> &mdash; no quantization for DeepSeek V3 etc.</li>
<li><strong>200k baseline stalls</strong> &mdash; no dual-case telemetry available</li>
<li><strong>Quality benchmarks are synthetic</strong> &mdash; no LongBench/RULER runs</li>
</ul>
</div>
<!-- D: Validation Results -->
<div class="s">
<h2>Validation Results (validate_paper.py) <span class="badge" style="background:var(--green-bg);color:var(--green);">9/9 pass</span></h2>
<table>
<tr><th>Test</th><th>What it validates</th><th>Result</th></tr>
<tr><td>MSE distortion bounds</td><td>Theorem 1: MSE &le; &radic;3&pi;/2 &middot; 1/4<sup>b</sup></td><td><span class="st ok">Pass</span> (b=1..4)</td></tr>
<tr><td>Codebook Table 1</td><td>Lloyd-Max MSE matches paper values</td><td><span class="st ok">Pass</span></td></tr>
<tr><td>Unbiasedness</td><td>Theorem 2: E[&langle;y, &tilde;x&rangle;] = &langle;y, x&rangle;</td><td><span class="st ok">Pass</span> (bits=2,3,4)</td></tr>
<tr><td>Distortion scaling</td><td>Theorem 3: 1/4<sup>b</sup> scaling</td><td><span class="st ok">Pass</span> (ratios &gt;2x per bit)</td></tr>
<tr><td>Recall@8 (N=4096)</td><td>Attention ranking quality at scale</td><td><span class="st ok">Pass</span></td></tr>
<tr><td>Rank correlation (N=2048)</td><td>Spearman corr 3-bit &gt;0.75, 4-bit &gt;0.90</td><td><span class="st ok">Pass</span></td></tr>
<tr><td>Needle@5 depths (N=4096)</td><td>Retrieval at 10/25/50/75/90% depth</td><td><span class="st ok">Pass</span> (3+4 bit)</td></tr>
<tr><td>Needle chunked (N=8192)</td><td>Multi-chunk retrieval</td><td><span class="st ok">Pass</span></td></tr>
<tr><td>Compression ratio</td><td>&gt;2x vs FP16</td><td><span class="st ok">Pass</span> (5.1x)</td></tr>
</table>
</div>
<!-- E: What was stripped -->
<div class="s">
<h2>What Was Stripped</h2>
<div class="g2">
<div class="c">
<h3 style="color:var(--red);">Removed</h3>
<ul>
<li><code>n_outlier_channels</code> / <code>outlier_key_bits</code> params (dead code)</li>
<li><code>_effective_bits()</code> method (never called)</li>
<li><code>_expand_query_for_gqa()</code> (dead function)</li>
<li><code>try_fused_decode()</code> (never called from working paths)</li>
<li>750+ lines of legacy classes in <code>vllm_attn_backend.py</code></li>
<li>Dead <code>console_scripts</code> entry point in <code>setup.py</code></li>
<li>Unused imports (<code>F</code>, <code>Tuple</code>, <code>MSEQuantized</code>, etc.)</li>
</ul>
</div>
<div class="c">
<h3 style="color:var(--green);">Fixed</h3>
<ul>
<li>Missing <code>import os</code> in <code>test_triton_kernels.py</code></li>
<li>Fragile group_size calc in <code>score.py</code> &rarr; explicit <code>32</code></li>
<li><code>vllm_attn_backend.py</code> reduced from 756 to ~190 lines (thin shim)</li>
</ul>
</div>
</div>
</div>
<!-- Diagram 1: System Architecture -->
<div class="s">
<h2>Architecture</h2>
<div class="mw">
<pre class="mermaid">
graph TB
subgraph "Core Quantization"
CB["codebook.py<br/>Lloyd-Max on Beta PDF"]
ROT["rotation.py<br/>QR rotation + QJL matrix"]
QMSE["TurboQuantMSE<br/>Rotate -> Quantize -> Pack"]
QPROD["TurboQuantProd<br/>MSE@(b-1) + QJL residual"]
end
subgraph "Serving Stack"
CAP["capture.py<br/>RingBuffer + KVCaptureEngine"]
STORE["store.py<br/>CompressedKVStore"]
SCORE["score.py<br/>compute_hybrid_attention"]
end
subgraph "vLLM Integration"
VLLM["integration/vllm.py<br/>off | capture_only | hybrid"]
SHIM["vllm_attn_backend.py<br/>enable_no_alloc shim"]
end
CB --> QMSE
ROT --> QMSE
ROT --> QPROD
QMSE --> QPROD
QPROD --> STORE
CAP --> STORE
STORE --> SCORE
VLLM --> CAP
VLLM --> SCORE
SHIM --> VLLM
</pre>
</div>
</div>
<!-- Diagram 2: Quantization Pipeline -->
<div class="s">
<h2>Quantization Pipeline</h2>
<div class="mw">
<pre class="mermaid">
graph LR
subgraph "Key: TurboQuant_Prod"
K["Key (H, T, D)"] --> NORM["||x||, normalize"]
NORM --> ROT["y = x @ Pi_T"]
ROT --> IDX["searchsorted -> indices"]
IDX --> PACK["bit-pack (b-1 bits)"]
ROT --> DQ["dequant MSE"]
DQ --> RES["r = x - x_mse"]
RES --> QJL["sign(r @ S_T)"]
QJL --> SIGNS["pack signs (1 bit/coord)"]
RES --> RNORM["||r||"]
end
subgraph "Value: Group Quant"
V["Value (H, T, D)"] --> GRP["reshape to groups"]
GRP --> MM["per-group min/max"]
MM --> VQ["round + bit-pack (2-bit)"]
end
</pre>
</div>
</div>
<!-- Diagram 3: Read/Write Paths -->
<div class="s">
<h2>Read / Write Paths</h2>
<div class="mw">
<pre class="mermaid">
graph TB
subgraph "Write"
PRE["Prefill N tokens"] --> SP{"N > ring?"}
SP -->|Yes| CMP["Compress first N-ring -> store"]
SP -->|Yes| BUF["Last ring tokens -> ring buffer"]
SP -->|No| BA["All -> ring buffer"]
DEC["Decode 1 token"] --> RW["ring.write"]
RW --> OV{"Overflow?"}
OV -->|Yes| FL["Drain -> store.append_chunk"]
end
subgraph "Read (Decode)"
Q["Query"] --> HH{"History >= 16?"}
HH -->|No| EX["Attend exact buffer only"]
HH -->|Yes| HY["Hybrid: dequant history + cat recent"]
HY --> ATT["einsum attention with GQA broadcast"]
end
</pre>
</div>
</div>
<!-- Diagram 4: vLLM Hook Chain -->
<div class="s">
<h2>vLLM Hook Chain</h2>
<div class="mw">
<pre class="mermaid">
sequenceDiagram
participant U as User
participant E as enable_no_alloc
participant X as Executor
participant I as FlashAttentionImpl
participant T as TQ LayerState
U->>E: enable_no_alloc(bits=3)
E->>X: patch get_kv_cache_specs
U->>X: LLM() init
X->>I: install_hooks (16 layers)
I->>T: create store + engine per layer
Note over I: Share KV cache across TQ layers
U->>X: generate(prompt)
Note over I: PREFILL
I->>T: SDPA attention + capture K/V
Note over I: DECODE
I->>T: compute_hybrid_attention
T-->>I: output
</pre>
</div>
</div>
<!-- Summary -->
<div class="s">
<h2>Summary</h2>
<div class="g2">
<div class="c">
<h3 style="color:var(--green);">What Works</h3>
<ul>
<li>Core TQ_MSE + TQ_Prod faithful to paper</li>
<li>All 3 theorems validated empirically</li>
<li>Needle retrieval passes at 8192 tokens</li>
<li>5.1x compression vs FP16</li>
<li>30k A/B: +5.7% prefill, +3.1% output speed</li>
<li>200k TQ completed (199,952 tokens)</li>
</ul>
</div>
<div class="c">
<h3 style="color:var(--red);">What Doesn't</h3>
<ul>
<li>MLA: stub only (DeepSeek V3 unsupported)</li>
<li>FlashInfer: incompatible</li>
<li>200k quality unvalidated (no LongBench)</li>
<li>Multi-sequence batching untested</li>
<li>No-alloc prefill slow (SDPA, not flash)</li>
</ul>
</div>
</div>
<div class="cw cw-r">
<strong>Key gap:</strong> The "identical output" claim holds only for trivial prompts. Long-context retrieval quality is unquantified against standard benchmarks (LongBench, RULER, Needle-in-Haystack at 200k). The 200k needle failure in the handoff doc is concerning.
</div>
</div>
</div>
<script>
mermaid.initialize({ startOnLoad: true, theme: 'default', securityLevel: 'loose',
flowchart: { useMaxWidth: true, curve: 'basis' },
sequence: { useMaxWidth: true, wrap: true }
});
</script>
</body>
</html>
REVIEW_PAPER_VS_IMPL.md
-496
View File
@@ -1,496 +0,0 @@
# TurboQuant: Paper vs Implementation — Comprehensive Review
> **Paper**: "TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate" (Zandieh, Daliri, Hadian, Mirrokni — Google Research, ICLR 2026, arXiv:2504.19874)
>
> **Implementation**: `./turboquant_pkg/`
---
## A) PAPER vs IMPLEMENTATION COMPARISON
### What the Paper Proposes
The TurboQuant paper introduces a **data-oblivious online vector quantization** framework with two algorithms:
1. **TurboQuant_MSE (Algorithm 1)** — Minimizes mean-squared error:
- Randomly rotate input vector via orthogonal matrix Π (QR of Gaussian)
- Each rotated coordinate follows Beta distribution → near-Gaussian for high d
- Apply **Lloyd-Max optimal scalar quantizer** per coordinate (1D continuous k-means on Beta PDF)
- Store norm separately, quantize unit-norm vector
- **Proven MSE**: ≤ (√3·π/2) · 1/4^b; for b=1,2,3,4: ≈ 0.36, 0.117, 0.03, 0.009
2. **TurboQuant_Prod (Algorithm 2)** — Unbiased inner-product estimator:
- Apply TurboQuant_MSE at **(b−1) bits** to get x̃_mse
- Compute residual r = x − x̃_mse
- Apply **QJL** (Quantized Johnson-Lindenstrauss): sign(S·r) where S ∈ ℝ^{d×d} is i.i.d. N(0,1)
- Store ‖r‖₂ for rescaling
- Dequant: x̃ = x̃_mse + (√(π/2)/d) · ‖r‖ · S^T · signs
- **Proven unbiased**: E[⟨y, x̃⟩] = ⟨y, x⟩
- **Proven distortion**: D_prod ≤ (√3·π²·‖y‖²/d) · 1/4^b
3. **Key claims**: Near-optimal (within ≈2.7× of information-theoretic lower bound), data-oblivious, accelerator-friendly, perfect needle-in-a-haystack retrieval at 3.5-bit, 5×+ KV cache compression.
4. **Paper also proposes** (Section 4.3): Outlier channel splitting — separate outlier channels quantized at higher bits for mixed-precision (e.g., 32 outlier channels at 3-bit + 96 at 2-bit = 2.5-bit effective).
### Which Components Are Implemented
| Paper Component | Implementation Status | Files |
|---|---|---|
| Random orthogonal rotation (Π via QR) | ✅ **Faithful** | `rotation.py` |
| Lloyd-Max codebook on Beta PDF | ✅ **Faithful** | `codebook.py` |
| TurboQuant_MSE (Algorithm 1) | ✅ **Faithful** | `quantizer.py::TurboQuantMSE` |
| TurboQuant_Prod (Algorithm 2) | ✅ **Faithful** | `quantizer.py::TurboQuantProd` |
| QJL projection matrix (i.i.d. Gaussian) | ✅ **Faithful** | `rotation.py::generate_qjl_matrix` |
| QJL sign packing/unpacking | ✅ **Faithful** | `quantizer.py::_pack_qjl_signs/_unpack_qjl_signs` |
| Bit-packed index storage | ✅ **Faithful** | `quantizer.py::_pack_indices/_unpack_indices` |
| Value group quantization | ⚠️ **Not from paper** — paper doesn't describe value quantization separately | `kv_cache.py::quantize_values` |
| Outlier channel splitting | ❌ **Not implemented** | Configured in `TurboQuantKVCache.__init__` but never used |
| Entropy coding of indices | ❌ **Not implemented** (paper also chose not to) | — |
| KV cache integration (vLLM) | ✅ **Novel engineering** (not in paper) | `integration/vllm.py`, `vllm_attn_backend.py` |
| Fused Triton kernels | ✅ **Novel engineering** (not in paper) | `triton_kernels.py` |
| Modular capture/store/score | ✅ **Novel engineering** | `capture.py`, `store.py`, `score.py` |
### Faithfulness Assessment
**The core quantization algorithms (MSE + Prod) are faithful to the paper.** Specifically:
1. **Rotation**: Uses QR decomposition of Gaussian matrix, exactly as Algorithm 1 specifies. Sign correction via diag(R) is a standard numerical refinement.
2. **Codebook**: Implements the exact Beta PDF from Lemma 1, solves Lloyd-Max via numerical integration (scipy). Precomputes and caches codebooks. The MSE values in `test_turboquant.py` validate against paper's Table (0.36, 0.117, 0.03, 0.009) with 30% tolerance.
3. **MSE Quantizer**: Normalize → rotate → searchsorted against decision boundaries → bit-pack. This matches Algorithm 1 exactly. Uses `searchsorted` instead of argmin for efficiency (equivalent for sorted centroids).
4. **Prod Quantizer**: MSE at (b−1) bits → residual → QJL sign(S·r) → store ‖r‖. Matches Algorithm 2 line-by-line. The `qjl_scale = sqrt(π/2)/d` matches the paper's dequantization constant.
5. **Attention score computation**: The `attention_score` method correctly implements the asymmetric estimator: `⟨y, x̃_mse⟩ + qjl_scale · ‖r‖ · ⟨S^T·y, signs⟩`, which is the unbiased estimator from Theorem 2.
### What Diverges from the Paper
1. **Value quantization**: The paper focuses exclusively on key quantization via TurboQuant. The implementation uses **asymmetric group quantization** (min-max, per-group scales/zeros) for values — a standard technique from KIVI/GEAR, not from TurboQuant.
2. **No outlier channel splitting**: The paper's Section 4.3 describes splitting channels into outlier vs non-outlier sets for mixed precision (e.g., 2.5-bit). The implementation has `n_outlier_channels` parameter in `TurboQuantKVCache` but **never actually uses it** — all channels get the same quantization.
3. **No PolarQuant integration**: The Google blog mentions PolarQuant as a companion method. Not implemented.
4. **QJL matrix is dense Gaussian**: Both paper and implementation use dense S ∈ ℝ^{d×d}. The paper notes this is i.i.d. N(0,1), which matches `generate_qjl_matrix`. However, for d=128 this is 64KB per layer — the implementation correctly makes this a module buffer.
5. **Paper benchmarks on Llama-3.1-8B and Ministral-7B with JAX**: The implementation targets Qwen3.5-27B with vLLM/PyTorch. Different model, different framework, different scale.
---
## B) ATTENTION MECHANISMS SUPPORTED
### Backend Support Matrix
| Backend | Status | Details |
|---|---|---|
| **Flash Attention (standard MHA/GQA)** | ✅ **Working (hybrid mode)** | Primary path. Hooks `do_kv_cache_update` and `forward` on `FlashAttentionImpl` |
| **Flash Attention (no-alloc prefill)** | ⚠️ **Functional but degraded** | Uses `F.scaled_dot_product_attention` as fallback — no flash kernel, materializes full attention matrix |
| **FlashInfer** | ❌ **Incompatible** | HANDOFF doc notes "FP8 KV + TurboQuant stacking still incompatible (FlashInfer metadata mismatch)" |
| **MLA (Multi-Latent Attention)** | ❌ **Stub only** | `_make_patched_mla_forward` is pure passthrough. Logs "MLA active decode is not implemented yet" |
| **GDN** | ❌ **Not addressed** | No code references GDN at all |
| **Mamba/Linear Attention** | ❌ **Not applicable** | These layers don't use KV cache. README notes "TQ hooks onto all 16 [full attention layers], rest are linear-attention/Mamba" |
### How Hybrid Attention Works
The hybrid attention mechanism (in `score.py::compute_hybrid_attention`) works as follows:
1. **Compressed history segment**: All tokens older than the ring buffer are stored in `CompressedKVStore` as quantized keys (TurboQuant_Prod) and quantized values (group quantization).
2. **Exact recent segment**: The most recent `ring_capacity` tokens (default 128) are kept in full precision in the `RingBuffer`.
3. **Merge**: During decode, both segments are dequantized/concatenated and standard matmul attention is computed over the concatenated KV. The log-sum-exp trick is **not** used for merging — instead, `_attend_hybrid` simply concatenates dequantized historical K/V with exact recent K/V and runs full attention over the combined set.
4. **Critical limitation**: The `_attend_hybrid` path **fully dequantizes** all historical tokens to float, meaning there's no memory savings during the attention computation itself. The savings come only from storage (compressed store uses ~3 bits/element vs 16 bits/element for bf16).
### The No-Alloc Path
The `enable_no_alloc()` function patches the vLLM executor to:
1. Install TQ hooks during `get_kv_cache_specs()`
2. Share paged KV cache across all TQ-hooked flash layers (only 1 layer's worth allocated)
3. Use `F.scaled_dot_product_attention` for prefill (instead of flash attention)
4. Use TQ compressed store for decode
**This path has significant issues** — see Section C.
---
## C) ERROR ANALYSIS
### Known Bugs and Incomplete Paths
1. **`n_outlier_channels` is dead code**: `TurboQuantKVCache.__init__` accepts it, creates `self.n_outlier_channels` and `self.outlier_key_bits`, but these are **never referenced** in any quantization/dequantization path. The paper's mixed-precision outlier strategy is entirely unimplemented.
2. **Value group size calculation in `score.py` is fragile**:
```python
gs = quantizer.mse_quantizer.dim // (quantizer.mse_quantizer.dim // 32)
```
This computes `dim / (dim / 32)` = 32 always. It's a roundabout way of hardcoding `group_size=32`. If dim isn't divisible by 32, this could silently produce wrong results.
3. **No-alloc decode fallback returns zeros**:
```python
# In _make_patched_forward, the no_alloc fallback:
return torch.zeros(num_actual, ...)
```
When the TQ store doesn't have enough data (< 16 tokens) and no_alloc=True, the decode path returns **literal zeros**. This means the first few decode tokens after a very short prefill would produce garbage.
4. **`test_triton_kernels.py` has a missing import**: Line 14 has `import os` missing — `os.path.dirname` is used but `os` is never imported. The test file would crash immediately.
5. **Fused decode kernel reads unpacked values**: In `turboquant_fused_decode`, value data is pre-unpacked in Python before being passed to the Triton kernel. The kernel reads `V_DATA_ptr` as float-like values but the docstring claims it handles bit-packed values. The Triton kernel itself doesn't unpack — it relies on Python-side unpacking, negating some of the "fused" benefit.
6. **GQA handling in `_attend_hybrid` allocates large intermediates**: While `_matmul_attend` uses broadcasting (`k.unsqueeze(1)`) to avoid `repeat_interleave`, the `_no_alloc_prefill_attention` in `vllm.py` **does** use `repeat_interleave`, which materializes the full expanded KV tensor. For a model with GQA ratio 4 and 200k tokens, this is a massive memory allocation.
### Failure Modes at Different Context Lengths
| Context | Status | Issue |
|---|---|---|
| **< 128 tokens** | ⚠️ Works but no compression | Everything stays in ring buffer, no quantization occurs |
| **128–32k tokens** | ✅ Working | Normal hybrid path, validated at 30k |
| **32k–100k** | ⚠️ Theoretical | No dual-case telemetry collected |
| **200k** | ⚠️ **Unstable** | "Needle retrieval failures at 200k" per handoff doc. TQ completed but baseline couldn't finish. Full dual-case comparison never obtained |
| **> 200k** | ❌ Unknown | Never tested |
The HANDOFF document explicitly states: "200k context with TurboQuant has completed successfully in prior runs (single-shot completion done), but full dual-case telemetry at 200k is unstable due to long-run process kills / baseline stall."
### TODO/Stub Paths That Don't Work
1. **MLA path**: `_make_patched_mla_forward` is pure passthrough. The log message says "TQ MLA path is deferred."
2. **`full_tq` mode**: Listed as `MODE_FULL_TQ = "full_tq"` but the `_make_patched_forward` never checks for it. The code comments say "(future) TQ handles everything including prefill."
3. **Zero-allocation path (Iteration 6)**: README says "attempted, not yet working" due to vLLM's compiled graph hardcoding `unified_kv_cache_update`. The current `enable_no_alloc` is a **different, partial workaround** that shares KV across layers rather than eliminating allocation.
4. **Entropy encoding**: Paper mentions it, implementation doesn't have it (paper also chose not to).
---
## D) BENCHMARKING METHODOLOGY
### Performance Benchmarks
| Script | What It Measures | How |
|---|---|---|
| `proof.py` | VRAM savings, context capacity | Two separate subprocesses (baseline vs TQ), same prompt, nvidia-smi memory readings |
| `benchmark.py` | VRAM, throughput (tok/s), quality | Two subprocesses per model, captures blocks/vram/tps/quality responses |
| `test_triton_kernels.py` | Kernel correctness + latency | Triton vs PyTorch reference, wall-clock timing |
### Quality Benchmarks
| Script | What It Measures |
|---|---|
| `test_turboquant.py` | Codebook MSE vs paper values, rotation properties, quantizer distortion, attention score correlation, memory savings |
| `test_modular.py` | Ring buffer, store, capture engine, hybrid attention, needle retrieval, rank correlation |
### Are the Benchmarks Fair?
**Partially fair, with significant caveats:**
1. **`proof.py` uses the same prompt** for both baseline and TQ — good. But it only generates 64 tokens, which is too short to test quality degradation over long generation.
2. **`benchmark.py` quality test is weak**: Tests 5 factual questions (capital of France, 17×23, etc.) — these are trivially answerable and don't test long-context comprehension. The paper uses LongBench, Needle-in-Haystack up to 104k, ZeroSCROLLS, RULER, and L-Eval. The implementation's quality testing is orders of magnitude less rigorous.
3. **Memory measurement via nvidia-smi**: This captures total GPU memory, not just KV cache. Differences in CUDA allocator behavior, fragmentation, and PyTorch memory pooling mean the reported numbers are approximate. The proof.py approach of measuring before/after `free_kv_cache()` is more reliable.
4. **No standardized quality benchmark**: The implementation never runs LongBench, Needle-in-Haystack (the paper's benchmark), or any established long-context evaluation. The `test_modular.py` needle test uses synthetic data with d=64 and N=200 — trivially easy compared to real LLM attention at d=128 with N=200k.
5. **TQ recall thresholds are very lenient**: `test_modular.py::test_attention_recall` requires only recall@8 ≥ 0.50 for 3-bit. The paper shows near-perfect recall. The low bar hides real quality issues.
### What Results Actually Show vs What Is Claimed
**README claims**:
- "2.0x context improvement" — **Legitimate**. Based on `free_kv_cache` releasing 30GB across 4 GPUs.
- "Both produce identical output for the same prompt" — **Only for trivial prompts**. For long contexts with precise retrieval, TQ is lossy by design.
- "2.6x compression per full-attention layer" — **Mathematically correct** calculation (198 bytes/token vs 512).
**What's missing**: No evidence of quality preservation at long contexts. The 200k needle retrieval failure mentioned in the handoff document is concerning and unreported in the README.
---
## E) ARCHITECTURE DIAGRAMS
### 1. Overall System Architecture
```mermaid
graph TB
subgraph "TurboQuant Package"
direction TB
subgraph "Core Quantization"
CB[codebook.py<br/>Lloyd-Max on Beta PDF]
ROT[rotation.py<br/>QR rotation + QJL matrix]
QMSE[quantizer.py::TurboQuantMSE<br/>Rotate → Scalar Quantize → Pack]
QPROD[quantizer.py::TurboQuantProd<br/>MSE@(b-1) + QJL@1-bit residual]
end
subgraph "KV Cache Layer"
KVC[kv_cache.py::TurboQuantKVCache<br/>Key: TQ_Prod, Value: GroupQuant<br/>Buffer: recent unquantized tokens]
VQ[kv_cache.py::quantize_values<br/>Asymmetric group quantization]
end
subgraph "Modular Serving Stack"
CAP[capture.py<br/>RingBuffer + KVCaptureEngine]
STORE[store.py<br/>CompressedKVStore<br/>Chunked, lazy flatten]
SCORE[score.py<br/>compute_hybrid_attention<br/>Compressed + Exact merge]
end
subgraph "Triton Kernels"
TK1[triton_kernels.py<br/>MSE score kernel]
TK2[triton_kernels.py<br/>QJL score kernel]
TK3[triton_kernels.py<br/>Fused decode kernel]
end
subgraph "vLLM Integration"
VLLM_NEW[integration/vllm.py<br/>New modular adapter<br/>Modes: off|capture_only|hybrid|full_tq]
VLLM_OLD[vllm_attn_backend.py<br/>Legacy compat shim<br/>Modes: accumulate|shadow|active]
NOALLOC[vllm_attn_backend.py<br/>enable_no_alloc<br/>Patches Executor + get_kv_cache_specs]
end
end
CB --> QMSE
ROT --> QMSE
ROT --> QPROD
QMSE --> QPROD
QPROD --> KVC
VQ --> KVC
QPROD --> STORE
VQ --> STORE
CAP --> STORE
STORE --> SCORE
TK1 --> SCORE
TK2 --> SCORE
TK3 --> SCORE
VLLM_NEW --> CAP
VLLM_NEW --> SCORE
VLLM_OLD --> VLLM_NEW
NOALLOC --> VLLM_OLD
```
### 2. Quantization Pipeline
```mermaid
graph LR
subgraph "Key Quantization (TurboQuant_Prod)"
direction TB
K_IN["Input Key<br/>(batch, heads, seq, d)"] --> NORM["Compute ‖x‖₂<br/>Normalize to unit"]
NORM --> ROTATE["y = x_unit @ Π^T<br/>(random rotation)"]
ROTATE --> SEARCH["searchsorted on<br/>decision boundaries"]
SEARCH --> PACK_MSE["Bit-pack indices<br/>(b-1 bits/coord)"]
ROTATE --> DEQUANT_MSE["Dequant MSE:<br/>centroid[idx] @ Π"]
DEQUANT_MSE --> RESIDUAL["r = x - x̃_mse"]
RESIDUAL --> QJL_PROJ["projected = r @ S^T"]
QJL_PROJ --> SIGN["signs = sign(projected)"]
SIGN --> PACK_QJL["Pack signs<br/>(8 per byte)"]
RESIDUAL --> RES_NORM["‖r‖₂"]
PACK_MSE --> K_OUT["ProdQuantized:<br/>mse_indices, qjl_signs,<br/>norms, residual_norms"]
PACK_QJL --> K_OUT
RES_NORM --> K_OUT
NORM --> K_OUT
end
subgraph "Value Quantization (Group Quant)"
direction TB
V_IN["Input Value<br/>(batch, heads, seq, d)"] --> GROUP["Reshape to groups<br/>(n_groups × group_size)"]
GROUP --> MINMAX["Per-group min/max"]
MINMAX --> SCALE["scale = (max-min) / (2^b - 1)<br/>zero = min"]
SCALE --> QUANT_V["v_q = round((v - zero) / scale)"]
QUANT_V --> PACK_V["Bit-pack<br/>(4 values/byte for 2-bit)"]
PACK_V --> V_OUT["ValueQuantized:<br/>data, scales, zeros"]
end
```
### 3. Capture / Store / Score Read/Write Paths
```mermaid
graph TB
subgraph "Write Path"
direction TB
PREFILL["Prefill tokens<br/>(N tokens)"] --> SPLIT{"N > ring_capacity?"}
SPLIT -->|Yes| COMPRESS["Compress first N-ring_cap tokens<br/>→ store.append_chunk()"]
SPLIT -->|Yes| BUFFER["Buffer last ring_cap tokens<br/>→ ring.write()"]
SPLIT -->|No| BUFFER_ALL["All to ring buffer<br/>→ ring.write()"]
DECODE["Decode token<br/>(1 token)"] --> RING_W["ring.write(k, v, 1)"]
RING_W --> OVERFLOW{"Ring overflow?"}
OVERFLOW -->|Yes| FLUSH["Drain ring → store.append_chunk()"]
OVERFLOW -->|No| DONE_W["Done"]
COMPRESS --> STORE_W["CompressedKVStore<br/>quantizer.quantize(k)<br/>quantize_values(v)<br/>→ append to chunk lists"]
FLUSH --> STORE_W
end
subgraph "Read Path"
direction TB
QUERY["Decode query<br/>(1, Q_heads, d)"] --> HAS_HIST{"History ≥ 16 tokens?"}
HAS_HIST -->|No| EXACT_ONLY["attend exact buffer only"]
HAS_HIST -->|Yes| HAS_RECENT{"Has recent buffer?"}
HAS_RECENT -->|No| COMP_ONLY["attend compressed only"]
HAS_RECENT -->|Yes| HYBRID["Hybrid attend"]
HYBRID --> FLAT["store.get_flat_cache()<br/>(lazy concatenate chunks)"]
FLAT --> DEQUANT_K["quantizer.dequantize(prod_q)<br/>→ k_hist (H_kv, N, D)"]
FLAT --> DEQUANT_V["dequantize_values(value_q)<br/>→ v_hist (H_kv, N, D)"]
DEQUANT_K --> CONCAT["cat(k_hist, k_recent)<br/>cat(v_hist, v_recent)"]
DEQUANT_V --> CONCAT
CONCAT --> MATMUL["einsum attention<br/>with GQA broadcast"]
MATMUL --> OUTPUT["attention output<br/>(T, Q_heads, D)"]
end
```
### 4. vLLM Integration Hook Chain
```mermaid
sequenceDiagram
participant User as User Code
participant EA as enable_no_alloc()
participant Exec as vLLM Executor
participant Runner as GPUModelRunner
participant Impl as FlashAttentionImpl
participant TQ as TurboQuant Layer State
User->>EA: enable_no_alloc(key_bits=3, ...)
EA->>Exec: Monkey-patch get_kv_cache_specs()
EA->>Runner: Monkey-patch _update_hybrid_attention_mamba_layout()
User->>Exec: LLM(...) creates engine
Exec->>Exec: patched get_kv_cache_specs()
Exec->>Runner: collective_rpc(_worker_install_tq)
Runner->>Impl: install_turboquant_hooks()
Note over Impl: For each attention layer:
Impl->>TQ: Create LayerState (store + engine)
Impl->>Impl: Monkey-patch do_kv_cache_update → capture K/V
Impl->>Impl: Monkey-patch forward → hybrid decode
Note over Impl: Share KV cache across TQ layers<br/>(only 1 paged alloc for all 16 layers)
User->>Exec: llm.generate(prompt)
Note over Impl: PREFILL PHASE
Impl->>TQ: patched forward (is_prefill=True, no_alloc=True)
TQ->>TQ: F.scaled_dot_product_attention(q, k, v)
TQ->>TQ: engine.ingest_prefill(k, v)
Note over Impl: DECODE PHASE
Impl->>TQ: patched forward (is_prefill=False)
TQ->>TQ: engine.ingest_decode(k, v, 1)
TQ->>TQ: compute_hybrid_attention(q, store, recent_k, recent_v)
TQ-->>Impl: return attention output
User->>Exec: collective_rpc(free_kv_cache)
Note over Runner: Replace paged KV tensors with 1-byte dummy<br/>torch.cuda.empty_cache()
```
### 5. No-Alloc Long-Context Path
```mermaid
graph TB
subgraph "Engine Initialization"
PATCH["enable_no_alloc()<br/>patches Executor.get_kv_cache_specs"]
LLM_INIT["LLM() constructor"]
SPECS["patched get_kv_cache_specs()"]
INSTALL["collective_rpc: install_turboquant_hooks<br/>mode=ACTIVE, no_alloc=True"]
SHARE["Share paged KV across TQ layers<br/>kv_sharing_target_layer_name = flash_layers[0]"]
LAYOUT["patched _update_hybrid_attention_mamba_layout<br/>handles shared layers in kv_caches dict"]
end
subgraph "Prefill (no_alloc=True)"
P_FWD["patched forward called"]
P_CHECK{"is_prefill?"}
P_CHECK -->|Yes| P_SDPA["F.scaled_dot_product_attention<br/>(no flash kernel, full materialization)"]
P_SDPA --> P_CAP["Capture K/V into TQ store<br/>ring_write + flush"]
P_CHECK -->|No| D_PATH["Decode path"]
end
subgraph "Decode (no_alloc=True)"
D_FWD["patched forward called"]
D_DATA{"Has TQ data ≥ 16?"}
D_DATA -->|Yes| D_HYBRID["compute_hybrid_attention<br/>(compressed + ring buffer)"]
D_DATA -->|No| D_ZERO["⚠️ Return torch.zeros(...)"]
end
subgraph "Post-Generation"
FREE["free_kv_cache()"]
REPLACE["Replace shared paged tensor with 1-byte dummy"]
EMPTY["torch.cuda.empty_cache()"]
RESULT["VRAM freed for new allocations"]
end
PATCH --> LLM_INIT --> SPECS --> INSTALL --> SHARE --> LAYOUT
P_FWD --> P_CHECK
D_FWD --> D_DATA
FREE --> REPLACE --> EMPTY --> RESULT
```
### 6. Attention Computation: Baseline vs TQ Hybrid
```mermaid
graph TB
subgraph "Baseline vLLM Attention"
direction TB
BQ["Query (1, Q, D)"] --> B_FLASH["Flash Attention<br/>reads from paged KV cache<br/>(all tokens in bf16)"]
B_FLASH --> B_OUT["Attention output (1, Q, D)"]
B_MEM["Memory: 2 × H_kv × N × D × 2 bytes<br/>= 512 bytes/token for d=128, bf16"]
end
subgraph "TurboQuant Hybrid Attention"
direction TB
TQ_Q["Query (1, Q, D)"]
subgraph "Compressed Segment (old tokens)"
TQ_FLAT["Flat cache: ProdQuantized + ValueQuantized"]
TQ_DQ_K["Dequantize keys:<br/>centroid[idx] @ Π · norms<br/>+ √(π/2)/d · ‖r‖ · S^T · signs"]
TQ_DQ_V["Dequantize values:<br/>data × scale + zero"]
TQ_FLAT --> TQ_DQ_K
TQ_FLAT --> TQ_DQ_V
end
subgraph "Exact Segment (recent 128 tokens)"
TQ_BUF_K["Ring buffer keys (bf16)"]
TQ_BUF_V["Ring buffer values (bf16)"]
end
TQ_DQ_K --> TQ_CAT_K["Concatenate K"]
TQ_BUF_K --> TQ_CAT_K
TQ_DQ_V --> TQ_CAT_V["Concatenate V"]
TQ_BUF_V --> TQ_CAT_V
TQ_Q --> TQ_SCORE["scores = Q @ K^T × scale"]
TQ_CAT_K --> TQ_SCORE
TQ_SCORE --> TQ_SOFT["weights = softmax(scores)"]
TQ_SOFT --> TQ_OUT["output = weights @ V"]
TQ_CAT_V --> TQ_OUT
TQ_MEM["Compressed memory:<br/>~198 bytes/token<br/>(3-bit keys + 2-bit values + overhead)<br/>= 2.6× savings"]
end
subgraph "Fused Triton Path (Alternative)"
direction TB
FQ["Query"] --> FROT["q_rot = q @ Π^T<br/>q_sketch = q @ S^T"]
FROT --> FKERNEL["Single Triton kernel:<br/>MSE scores + QJL scores<br/>+ online softmax + value aggregation"]
FKERNEL --> FOUT["Attention output"]
FMEM["Reads packed data directly<br/>Never materializes full KV"]
end
```
---
## Summary of Key Findings
### What Works Well
1. **Core algorithms are faithfully implemented** — TurboQuant_MSE and TurboQuant_Prod match the paper's Algorithms 1 and 2 precisely
2. **Codebook computation matches paper's theoretical MSE values**
3. **Modular architecture** (capture/store/score) is clean and well-separated
4. **Real VRAM savings demonstrated** — 30GB freed across 4 GPUs on Qwen3.5-27B
5. **Three Triton kernels provide a genuine fused decode path** (MSE score, QJL score, fused decode with online softmax)
### What Doesn't Work or Is Missing
1. **Outlier channel splitting** — Dead code despite being a key paper technique for achieving 2.5/3.5-bit mixed precision
2. **MLA backend** — Pure stub, no quantization support
3. **No-alloc prefill** — Falls back to `F.scaled_dot_product_attention` (slow, materializes full attention)
4. **No-alloc decode fallback** — Returns zeros when TQ data is insufficient
5. **200k context** — Unstable, no reliable dual-case comparison
6. **Quality benchmarks** — Far weaker than paper's (5 trivial questions vs LongBench/Needle-in-a-Haystack/RULER)
7. **test_triton_kernels.py** — Broken (missing `import os`)
8. **Value quantization** — Not from the paper; standard group quantization, no theoretical backing for optimality
### Risk Assessment
- **For production use at 30k context**: Reasonable confidence — validated with matching outputs
- **For 100k+ context**: Unvalidated — quality degradation likely, especially for retrieval tasks
- **For MLA-based models (DeepSeek V3, etc.)**: Completely unsupported
- **For multi-sequence batching**: Untested — all benchmarks use `max_num_seqs=1`
audit_claims.py
-447
View File
@@ -1,447 +0,0 @@
#!/usr/bin/env python3
"""
Adversarial audit of TurboQuant claims.
Goal: find every lie, exaggeration, or misleading result.
"""
import math
import torch
import numpy as np
torch.manual_seed(42)
np.random.seed(42)
def section(title):
print(f"\n{'='*60}")
print(f" {title}")
print(f"{'='*60}\n")
# ======================================================================
# LIE #1: "5.1x compression"
# The memory_bytes() method only counts the quantized data tensors.
# It does NOT count:
# - The rotation matrix Pi (D*D*4 bytes per layer)
# - The QJL matrix S (D*D*4 bytes per layer)
# - The codebook centroids + boundaries
# - The Python object overhead
# - The ring buffer (128 * H_kv * D * 2 bytes per layer, in bf16)
# ======================================================================
section("AUDIT 1: Is '5.1x compression' honest?")
from turboquant.store import CompressedKVStore
from turboquant.quantizer import TurboQuantProd
d = 128; H_kv = 8; N = 4096
store = CompressedKVStore(head_dim=d, num_kv_heads=H_kv, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"))
k = torch.randn(N, H_kv, d); v = torch.randn(N, H_kv, d)
store.append_chunk(k, v)
tq_data_bytes = store.memory_bytes()
fp16_bytes = N * H_kv * d * 2 * 2
naive_ratio = fp16_bytes / tq_data_bytes
# Now count what memory_bytes() DOESN'T count
quantizer = store.quantizer
Pi_bytes = quantizer.mse_quantizer.Pi.nelement() * quantizer.mse_quantizer.Pi.element_size()
S_bytes = quantizer.S.nelement() * quantizer.S.element_size()
centroids_bytes = quantizer.mse_quantizer.centroids.nelement() * 4
boundaries_bytes = quantizer.mse_quantizer.boundaries.nelement() * 4
decision_bytes = quantizer.mse_quantizer.decision_boundaries.nelement() * 4
# Per layer overhead
per_layer_overhead = Pi_bytes + S_bytes + centroids_bytes + boundaries_bytes + decision_bytes
total_honest = tq_data_bytes + per_layer_overhead
honest_ratio = fp16_bytes / total_honest
# Ring buffer (not counted either — it holds 128 tokens in full precision)
ring_bytes = 128 * H_kv * d * 2 # bf16
total_with_ring = total_honest + ring_bytes
ratio_with_ring = fp16_bytes / total_with_ring
print(f"FP16 KV for {N} tokens, {H_kv} heads, d={d}:")
print(f" FP16 total: {fp16_bytes:>12,} bytes ({fp16_bytes/1e6:.1f} MB)")
print(f" TQ data only: {tq_data_bytes:>12,} bytes -> {naive_ratio:.2f}x (CLAIMED)")
print(f" + Pi ({d}x{d}): {Pi_bytes:>12,} bytes")
print(f" + S ({d}x{d}): {S_bytes:>12,} bytes")
print(f" + codebook: {centroids_bytes + boundaries_bytes + decision_bytes:>12,} bytes")
print(f" Per-layer overhead: {per_layer_overhead:>12,} bytes")
print(f" TQ data + overhead: {total_honest:>12,} bytes -> {honest_ratio:.2f}x (HONEST)")
print(f" + ring buffer (128 tok): {ring_bytes:>12,} bytes")
print(f" TQ all-in: {total_with_ring:>12,} bytes -> {ratio_with_ring:.2f}x (ALL-IN)")
print()
# At what N does Pi+S overhead become negligible?
for test_n in [128, 512, 1024, 4096, 32768]:
data = tq_data_bytes * test_n / N
total = data + per_layer_overhead + ring_bytes
fp16 = test_n * H_kv * d * 2 * 2
r = fp16 / total
print(f" N={test_n:>6}: all-in ratio = {r:.2f}x")
# ======================================================================
# LIE #2: "Needle-in-haystack passes"
# Our needle tests use signal-to-noise ratio of 3.0 / 0.02 = 150x.
# Real LLM attention has MUCH more subtle signal. The needle is
# ridiculously loud.
# ======================================================================
section("AUDIT 2: Are needle tests realistic?")
from turboquant.store import CompressedKVStore
d = 128; H_kv = 4; N = 4096
def run_needle(needle_magnitude, noise_magnitude, N, bits):
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=bits, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
keys = torch.randn(N, H_kv, d) * noise_magnitude
values = torch.randn(N, H_kv, d)
needle_pos = N // 2
needle_key = torch.randn(1, H_kv, d) * needle_magnitude
keys[needle_pos] = needle_key.squeeze(0)
store.append_chunk(keys, values)
flat = store.get_flat_cache()
k_dequant = store.quantizer.dequantize(flat.prod_q)
query_vec = needle_key.squeeze(0).unsqueeze(0)
scores = torch.bmm(
query_vec.float().transpose(0, 1),
k_dequant.float().transpose(1, 2),
).squeeze(1)
correct = 0
for h in range(H_kv):
if scores[h].argmax().item() == needle_pos:
correct += 1
return correct / H_kv
print("Needle retrieval accuracy at different signal-to-noise ratios:")
print(f"{'SNR':>8} {'Needle':>8} {'Noise':>8} {'3-bit':>8} {'4-bit':>8}")
print("-" * 48)
for needle_mag, noise_mag in [(3.0, 0.02), (1.0, 0.1), (0.5, 0.1), (0.3, 0.1),
(0.2, 0.1), (0.15, 0.1), (0.1, 0.1)]:
snr = needle_mag / noise_mag
acc3 = run_needle(needle_mag, noise_mag, N, 3)
acc4 = run_needle(needle_mag, noise_mag, N, 4)
print(f"{snr:>8.1f} {needle_mag:>8.2f} {noise_mag:>8.2f} {acc3:>8.1%} {acc4:>8.1%}")
print()
print("Our tests use SNR=150x. Even at SNR=1.0, retrieval is 100%.")
print("This is because score-space SNR is amplified by sqrt(d)=11.3.")
print("The needle test proves: 'does argmax survive quantization?'")
print("Answer: yes, always, because dominant keys are always preserved.")
print("What it does NOT prove: ranking quality of non-dominant keys.")
# ======================================================================
# LIE #3: "Recall@8 passes (0.40 for 3-bit, 0.55 for 4-bit)"
# These thresholds are absurdly low. Random baseline recall@8 from
# N=4096 is 8/4096 = 0.002. The paper claims near-perfect retrieval.
# Our 0.40 threshold means we're losing 60% of the top-k keys.
# ======================================================================
section("AUDIT 3: How bad is recall@8 really?")
from turboquant.quantizer import TurboQuantProd
d = 128; N = 4096; n_queries = 64; device = "cpu"
print(f"Recall@k at d={d}, N={N}, averaged over {n_queries} queries:")
print(f"{'bits':>6} {'k=4':>8} {'k=8':>8} {'k=16':>8} {'k=32':>8} {'k=64':>8}")
print("-" * 48)
for bits in [2, 3, 4]:
q = TurboQuantProd(dim=d, bits=bits, device=device, seed=42)
keys = torch.randn(1, 1, N, d) * 0.1
queries = torch.randn(1, 1, n_queries, d) * 0.1
true_scores = torch.matmul(queries, keys.transpose(-2, -1)).squeeze(0).squeeze(0)
key_q = q.quantize(keys)
tq_scores = q.attention_score(queries, key_q).squeeze(0).squeeze(0)
row = f"{bits:>6}"
for k in [4, 8, 16, 32, 64]:
true_topk = set()
tq_topk_set = set()
total_recall = 0
for qi in range(n_queries):
t_set = set(true_scores[qi].topk(k).indices.tolist())
tq_set = set(tq_scores[qi].topk(k).indices.tolist())
total_recall += len(t_set & tq_set) / k
mean_recall = total_recall / n_queries
row += f" {mean_recall:>8.3f}"
print(row)
print()
print("Paper claims near-perfect Needle-in-a-Haystack at 3.5-bit.")
print("Our 3-bit recall@8 ~ 0.4-0.5 means we LOSE HALF the important keys.")
# ======================================================================
# LIE #4: "Hybrid attention works"
# The hybrid path fully dequantizes ALL compressed tokens to float32.
# This means during decode we allocate H_kv * N * D * 4 bytes.
# For 200k tokens, d=128, H_kv=8: that's 200000*8*128*4 = 819 MB
# per decode step. The "memory savings" only exist in storage, not
# during the actual attention computation.
# ======================================================================
section("AUDIT 4: Does hybrid decode actually save memory?")
print("Memory allocated during _matmul_attend (per decode step):")
print("All compressed tokens are dequantized to float32 for the matmul.\n")
for N_tokens in [1024, 4096, 30000, 100000, 200000]:
# k_dequant + v_dequant, both (H_kv, N, D) float32
kv_decompressed = 2 * 8 * N_tokens * 128 * 4 # H_kv=8, D=128, float32
# scores tensor: (H_kv, G, T, N) float32, T=1, G=6 for Qwen3.5 (48q/8kv)
scores_mem = 8 * 6 * 1 * N_tokens * 4
total = kv_decompressed + scores_mem
print(f" N={N_tokens:>7}: decompressed KV = {kv_decompressed/1e6:>8.1f} MB, "
f"scores = {scores_mem/1e6:>6.1f} MB, total = {total/1e6:>8.1f} MB")
print()
print("At 200k tokens, we allocate ~1 GB per decode step just for dequantized KV.")
print("The 'savings' are only in between decode steps (storage), not during compute.")
print("This is NOT what the paper describes - the paper uses fused kernels that")
print("never materialize the full dequantized KV.")
# ======================================================================
# AUDIT 5: "Distortion scales as 1/4^b"
# We test this with threshold ">2x per bit" but paper claims ~4x.
# CORRECTION: The paper bound is for UNIT NORM vectors.
# Our initial test used unnormalized randn vectors (||x|| ~ sqrt(d) ~ 11.3),
# which inflated the distortion by ~d = 128x.
# ======================================================================
section("AUDIT 5: Does distortion actually follow 1/4^b?")
from turboquant.quantizer import TurboQuantProd
d = 128; N = 10000
print(f"Inner-product distortion at d={d}, N={N}:")
print()
print("--- With UNNORMALIZED vectors (how we tested initially, WRONG comparison) ---")
print(f"{'bits':>6} {'raw MSE':>12} {'||x||^2':>10} {'||y||^2':>10}")
print("-" * 44)
for bits in [2, 3, 4]:
q = TurboQuantProd(dim=d, bits=bits, device="cpu", seed=42)
x = torch.randn(1, 1, N, d)
y = torch.randn(1, 1, 1, d)
true_ip = (y * x).sum(dim=-1).squeeze()
key_q = q.quantize(x)
est_ip = q.attention_score(y, key_q).squeeze()
raw_mse = ((est_ip - true_ip) ** 2).mean().item()
x_norm_sq = (x ** 2).sum(dim=-1).mean().item()
y_norm_sq = (y ** 2).sum().item()
print(f"{bits:>6} {raw_mse:>12.4f} {x_norm_sq:>10.2f} {y_norm_sq:>10.2f}")
print()
print("--- With UNIT NORM vectors (what the paper's bound assumes) ---")
print(f"{'bits':>6} {'MSE':>12} {'paper_bound':>14} {'ratio':>10}")
print("-" * 44)
for bits in [2, 3, 4]:
q = TurboQuantProd(dim=d, bits=bits, device="cpu", seed=42)
x = torch.randn(1, 1, N, d)
x = x / x.norm(dim=-1, keepdim=True)
y = torch.randn(1, 1, 1, d)
y = y / y.norm(dim=-1, keepdim=True)
true_ip = (y * x).sum(dim=-1).squeeze()
key_q = q.quantize(x)
est_ip = q.attention_score(y, key_q).squeeze()
mse = ((est_ip - true_ip) ** 2).mean().item()
paper_bound = math.sqrt(3) * math.pi**2 / d * (1.0 / 4**bits)
print(f"{bits:>6} {mse:>12.8f} {paper_bound:>14.8f} {mse/paper_bound:>10.2f}x")
print()
print("VERDICT: Distortion IS within paper bounds for unit-norm vectors.")
print("Our audit_v1 was wrong — we compared unnormalized MSE to a normalized bound.")
# ======================================================================
# LIE #6: "Unbiased estimator"
# We check with abs(bias) < 0.05 which sounds tight but the signal
# magnitude matters. Let's look at relative bias.
# ======================================================================
section("AUDIT 6: How unbiased is it really?")
from turboquant.quantizer import TurboQuantProd
d = 128; N = 5000
print(f"Bias analysis at d={d}, N={N}:")
print(f"{'bits':>6} {'mean_bias':>12} {'mean_|ip|':>12} {'rel_bias%':>12} {'max_|bias|':>12}")
print("-" * 60)
for bits in [2, 3, 4]:
all_biases = []
all_true_ips = []
for trial in range(50):
q = TurboQuantProd(dim=d, bits=bits, device="cpu", seed=trial * 100)
x = torch.randn(1, 1, N, d)
y = torch.randn(1, 1, 1, d)
true_ip = (y * x).sum(dim=-1).squeeze()
key_q = q.quantize(x)
est_ip = q.attention_score(y, key_q).squeeze()
per_sample_bias = (est_ip - true_ip)
all_biases.append(per_sample_bias.mean().item())
all_true_ips.append(true_ip.abs().mean().item())
mean_bias = np.mean(all_biases)
mean_abs_ip = np.mean(all_true_ips)
rel_bias = abs(mean_bias) / mean_abs_ip * 100
max_abs_bias = np.max(np.abs(all_biases))
print(f"{bits:>6} {mean_bias:>12.6f} {mean_abs_ip:>12.4f} {rel_bias:>11.2f}% {max_abs_bias:>12.6f}")
# ======================================================================
# LIE #7: "30k benchmark shows TQ is faster"
# TTFT: 17.162 vs 18.138 (~5.7% faster)
# But TQ init is SLOWER: 43.749 vs 40.160 (9% slower)
# And the speed "gain" could be noise — it's a SINGLE run per case.
# ======================================================================
section("AUDIT 7: Is the 30k speedup real?")
print("30k telemetry (SINGLE run each, no error bars):")
print()
print(" Metric Baseline TQ Delta")
print(" " + "-" * 55)
print(" Init time 40.160s 43.749s +3.589s (TQ 8.9% SLOWER)")
print(" TTFT 18.138s 17.162s -0.976s (TQ 5.4% faster)")
print(" Full 24-tok run 18.992s 18.415s -0.577s (TQ 3.0% faster)")
print(" Prefill tok/s 1803.92 1906.52 +102.6 (+5.7%)")
print(" End-to-end tok/s 1.264 1.303 +0.039 (+3.1%)")
print(" Activation est MB 644.61 599.20 -45.41 (-7.0%)")
print()
print("Problems with this data:")
print(" 1. N=1 per condition. No error bars. Could be noise.")
print(" 2. TQ init is 3.6s slower — this cost is IGNORED in speedup claims.")
print(" 3. Total wall time (init+gen): baseline=59.15s, TQ=62.16s -> TQ is SLOWER overall.")
print(" 4. 'Activation est' is peak_alloc - end_alloc, which includes allocator fragmentation.")
print(" 5. The 5.7% prefill speedup makes no sense — TQ uses SDPA which is slower than flash.")
print(" (Unless the measurement noise is >=6%, which for N=1 it certainly is.)")
# ======================================================================
# LIE #8: "200k context works"
# A single completion doesn't prove it works. What about output quality?
# Was the output checked at all? Was needle retrieval tested at 200k?
# ======================================================================
section("AUDIT 8: Does 200k context actually work?")
print("200k TQ completion facts:")
print(" - prompt tokens: 199,952")
print(" - output tokens: 24")
print(" - elapsed: 58.34s")
print(" - GPU mem: ~31.9 GB")
print()
print("What was NOT checked:")
print(" 1. Output quality — was the output coherent? Nobody verified.")
print(" 2. Needle retrieval — handoff doc says 'needle retrieval failures at 200k'")
print(" 3. Baseline comparison — baseline stalled, so there's ZERO comparison data")
print(" 4. Perplexity — no measurement of how much quality degraded")
print(" 5. The no-alloc SDPA prefill materializes full attention matrix at 200k:")
print(f" 48 heads * 200000 * 200000 * 4 bytes = {48 * 200000 * 200000 * 4 / 1e12:.1f} TB")
print(" This CANNOT work. The causal mask in SDPA may help, but it's")
print(" computed in chunks. Each chunk still materializes huge intermediates.")
print()
print(" The 200k 'success' likely means: the process didn't crash.")
print(" It does NOT mean: the output was correct, or even coherent.")
# ======================================================================
# LIE #9: Compression ratio in the REAL benchmark
# README claims "2.0x context improvement" based on free_kv_cache.
# But that 30GB freed is across 4 GPUs (proof.py uses 4 RTX 3090s).
# Per-GPU it's ~7.5 GB freed. And the freed memory is paged KV cache
# which vLLM may not reuse efficiently.
# ======================================================================
section("AUDIT 9: Is '2x context improvement' honest?")
print("README's proof.py result: ~30GB freed across 4 GPUs")
print(" Per-GPU: ~7.5 GB freed")
print(" Model weights: ~19.78 GB (AWQ 4-bit Qwen3.5-27B)")
print(" On single 5090 (32GB): 32 - 19.78 = 12.22 GB for KV cache")
print(" 7.5 GB freed out of 12.22 GB KV = 61% of KV freed")
print()
print(" But wait — can we actually USE the freed memory?")
print(" vLLM pre-allocates paged blocks. Freeing them doesn't mean")
print(" we can allocate more blocks. The freed memory goes back to")
print(" the CUDA allocator, not vLLM's block allocator.")
print()
print(" The '2x context' claim is theoretical. In practice, you'd need")
print(" to restart the engine with a higher max_model_len to actually")
print(" serve 2x more tokens.")
# ======================================================================
# SUMMARY
# ======================================================================
section("SUMMARY: What we're lying about (or at least misleading)")
lies = [
("5.1x compression", "MISLEADING",
"Doesn't count Pi/S matrices (128KB/layer), ring buffer, or Python overhead. "
"Real ratio ~4.6x at 4096 tokens, ~1.1x at 128 tokens. Only approaches 5x at 32k+."),
("Needle-in-haystack passes", "HONEST BUT MEANINGLESS",
"Needle test uses query=key which gives perfect match. This tests 'does argmax survive' "
"which is trivial. At SNR=1.0 it still passes because score-space SNR is 11x due to "
"d=128 dimensionality. Real LLM queries are NOT copies of keys."),
("Recall@8 >= 0.40 (3-bit)", "DECEPTIVELY LOW BAR",
"Paper claims near-perfect retrieval. Our 3-bit recall@1=38%, recall@8=55%. "
"BUT this only matters for FLAT attention. When attention is spiky (dominant tokens exist), "
"TQ preserves the important keys perfectly. The recall failure is on the unimportant tail."),
("Hybrid decode saves memory", "TRUE FOR STORAGE, FALSE FOR COMPUTE",
"Storage is compressed (~3 bits/element). But during compute, ALL history is "
"dequantized to float32. At 200k tokens that's ~1.6 GB per decode step (H_kv=8). "
"Paper uses fused Triton kernels that never materialize full KV. We have those kernels "
"but the hybrid path doesn't use them — it uses the PyTorch dequantize-then-matmul path."),
("Distortion follows 1/4^b", "TRUE (was falsely accused)",
"Initial audit showed 75x above bound, but that was comparing non-normalized vectors. "
"With unit-norm vectors as the paper specifies: 2-bit=0.70x, 3-bit=0.82x, 4-bit=0.97x. "
"All WITHIN the theoretical bound. Implementation is faithful."),
("30k TQ is faster than baseline", "WITHIN NOISE",
"Single run, no error bars. Total wall time (init+gen) TQ is actually slower (62.2s vs 59.2s). "
"The 5.7% prefill 'speedup' is within measurement noise for N=1."),
("200k context works", "UNVERIFIED",
"Process completed without crashing. Output quality never checked. "
"Needle retrieval reportedly fails at 200k. No perplexity measurement. "
"The SDPA prefill at 200k chunks internally but still allocates huge intermediates."),
("2x context improvement", "THEORETICAL ONLY",
"Freed memory returns to CUDA allocator, not vLLM block allocator. "
"Can't actually serve 2x more tokens without engine restart with higher max_model_len."),
]
for claim, verdict, detail in lies:
print(f" [{verdict}] \"{claim}\"")
print(f" {detail}")
print()
baseline_vs_tq.py
-334
View File
@@ -1,334 +0,0 @@
#!/usr/bin/env python3
"""
Baseline vs TQ: Actual VRAM measurement.
vLLM pre-allocates KV cache blocks at startup, so nvidia-smi won't show per-request
deltas. Instead we measure:
1. BASELINE (current server, bf16 KV):
- vLLM reports num_gpu_blocks=5190, block_size=272
- KV cache usage % during inference at different context lengths
- We can compute exact bytes from: blocks_used * block_size * per_token_kv_bytes
2. Calculate what TQ would save based on measured block usage
3. Also: restart server with LOWER gpu_memory_utilization to see where it OOMs,
which tells us the REAL memory pressure.
"""
import json
import time
import subprocess
import tempfile
import os
BASE_URL = "http://localhost:8000/v1"
MODEL = "Qwen/Qwen3.5-35B-A3B"
# From model config and vLLM metrics:
NUM_FULL_ATTN_LAYERS = 10
NUM_LINEAR_LAYERS = 30
HEAD_DIM = 256
NUM_KV_HEADS = 2
BLOCK_SIZE = 272 # tokens per block
NUM_GPU_BLOCKS = 5190 # total blocks allocated
TP_SIZE = 8
# Per-token KV size for full_attention layers (bf16)
# K + V per layer = 2 * kv_heads * head_dim * 2 bytes (bf16) = 2 * 2 * 256 * 2 = 2048 bytes
KV_PER_TOKEN_PER_FULL_LAYER = 2 * NUM_KV_HEADS * HEAD_DIM * 2 # 2048 bytes
KV_PER_TOKEN_ALL_FULL = KV_PER_TOKEN_PER_FULL_LAYER * NUM_FULL_ATTN_LAYERS # 20480 bytes
# Linear attention state per block (mamba_page_size_padded=278528 per layer)
# This is a FIXED cost per block, not per token
MAMBA_PAGE_SIZE = 278528 # bytes per layer per block (from vLLM config)
# Per block total:
# full_attn: BLOCK_SIZE * KV_PER_TOKEN_ALL_FULL = 272 * 20480 = 5,570,560 bytes
# linear_attn: MAMBA_PAGE_SIZE * NUM_LINEAR_LAYERS = 278528 * 30 = 8,355,840 bytes
# Total per block: ~13.9 MB
FULL_ATTN_PER_BLOCK = BLOCK_SIZE * KV_PER_TOKEN_ALL_FULL
LINEAR_PER_BLOCK = MAMBA_PAGE_SIZE * NUM_LINEAR_LAYERS
TOTAL_PER_BLOCK = FULL_ATTN_PER_BLOCK + LINEAR_PER_BLOCK
# TQ per-token (3-bit keys, 2-bit values) for full_attention layers:
# Key: mse_indices (packed) + qjl_signs (packed) + norms (float32) + residual_norms (float32)
# For head_dim=256, 3-bit TQ:
# mse_indices: 256 * 2 / 8 = 64 bytes (2-bit MSE, packed)
# qjl_signs: 256 / 8 = 32 bytes
# norms: 4 bytes, residual_norms: 4 bytes
# Total key per head: 104 bytes
# Value (2-bit group quant, group_size=32):
# data: 256 * 2 / 8 = 64 bytes (packed)
# scales + zeros: 2 * (256/32) * 4 = 64 bytes
# Total value per head: 128 bytes
# Per token per layer: (104 + 128) * 2 kv_heads = 464 bytes
TQ_PER_TOKEN_PER_FULL_LAYER = 464
TQ_PER_TOKEN_ALL_FULL = TQ_PER_TOKEN_PER_FULL_LAYER * NUM_FULL_ATTN_LAYERS # 4640 bytes
# TQ overhead per layer (Pi + S matrices): 256*256*4 * 2 = 524288 bytes = 0.5 MB
TQ_OVERHEAD_PER_LAYER = 2 * HEAD_DIM * HEAD_DIM * 4
TQ_OVERHEAD_TOTAL = TQ_OVERHEAD_PER_LAYER * NUM_FULL_ATTN_LAYERS # 5.24 MB
def curl_post_file(endpoint, data, timeout=1200):
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, dir='/tmp') as f:
json.dump(data, f)
tmppath = f.name
try:
cmd = ["curl", "-s", "-X", "POST", f"{BASE_URL}/{endpoint}",
"-H", "Content-Type: application/json", "-d", f"@{tmppath}",
"--max-time", str(timeout)]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 30)
return json.loads(result.stdout) if result.returncode == 0 else {"error": result.stderr[:500]}
except json.JSONDecodeError:
return {"error": "bad json"}
finally:
os.unlink(tmppath)
def get_kv_usage():
result = subprocess.run(["curl", "-s", "http://localhost:8000/metrics"],
capture_output=True, text=True, timeout=10)
for line in result.stdout.splitlines():
if line.startswith("vllm:kv_cache_usage_perc{"):
return float(line.split()[-1])
return 0
def build_filler_prompt(target_tokens):
filler = (
"In the vast digital landscape, information flows through networks. "
"Each node processes data according to its designated protocols. "
"Modern cloud infrastructure supports concurrent operations. "
"Load balancers distribute traffic among server clusters. "
)
target_chars = target_tokens * 4
text = ""
while len(text) < target_chars:
text += filler
return text[:target_chars] + "\n\nSay 'OK' and nothing else."
def measure_kv_during_request(target_tokens):
"""Send a request and measure KV cache usage DURING processing."""
prompt = build_filler_prompt(target_tokens)
# Start the request (non-blocking via streaming)
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, dir='/tmp') as f:
json.dump({
"model": MODEL,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 10,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
"stream": True,
}, f)
tmppath = f.name
proc = subprocess.Popen(
["curl", "-s", "-N", "-X", "POST", f"{BASE_URL}/chat/completions",
"-H", "Content-Type: application/json", "-d", f"@{tmppath}",
"--max-time", "600"],
stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
)
# Poll KV usage while request is in flight
max_kv = 0
samples = 0
ttft = None
t0 = time.time()
for line in proc.stdout:
line = line.strip()
# Sample KV usage
kv = get_kv_usage()
max_kv = max(max_kv, kv)
samples += 1
if line.startswith("data: "):
data_str = line[6:]
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
if content and ttft is None:
ttft = time.time() - t0
except json.JSONDecodeError:
pass
proc.wait(timeout=10)
elapsed = time.time() - t0
# One more sample after completion
kv_after = get_kv_usage()
os.unlink(tmppath)
return max_kv, kv_after, ttft, elapsed, samples
print("=" * 80)
print(" BASELINE vs TURBOQUANT: VRAM ANALYSIS")
print("=" * 80)
print()
# Total pre-allocated KV cache memory
total_kv_cache_bytes = NUM_GPU_BLOCKS * TOTAL_PER_BLOCK
print(f"vLLM KV cache pool:")
print(f" Blocks: {NUM_GPU_BLOCKS} x {BLOCK_SIZE} tokens = {NUM_GPU_BLOCKS * BLOCK_SIZE:,} token capacity")
print(f" Per block: full_attn={FULL_ATTN_PER_BLOCK/1e6:.2f}MB + linear={LINEAR_PER_BLOCK/1e6:.2f}MB = {TOTAL_PER_BLOCK/1e6:.2f}MB")
print(f" Total pre-allocated: {total_kv_cache_bytes/1e9:.2f} GB across {TP_SIZE} GPUs ({total_kv_cache_bytes/TP_SIZE/1e9:.2f} GB/GPU)")
print()
# How much of this is full_attention (compressible by TQ)?
full_attn_total = NUM_GPU_BLOCKS * FULL_ATTN_PER_BLOCK
linear_total = NUM_GPU_BLOCKS * LINEAR_PER_BLOCK
print(f" Full attention (TQ-compressible): {full_attn_total/1e9:.2f} GB ({full_attn_total/total_kv_cache_bytes*100:.1f}%)")
print(f" Linear attention (incompressible): {linear_total/1e9:.2f} GB ({linear_total/total_kv_cache_bytes*100:.1f}%)")
print()
print("=" * 80)
print(" BASELINE: KV cache usage at different context lengths")
print("=" * 80)
print()
print(f"{'Context':>10} {'Prompt':>8} {'KV peak':>10} {'KV bytes':>14} {'full_attn':>12} {'linear':>12} {'TTFT':>8}")
print(f"{'target':>10} {'tokens':>8} {'(%)':>10} {'used':>14} {'KV':>12} {'state':>12} {'(s)':>8}")
print("-" * 85)
results = []
for target in [1000, 4000, 8000, 16000, 32000, 64000, 100000, 131000]:
max_kv, kv_after, ttft, elapsed, samples = measure_kv_during_request(target)
# Estimate actual tokens from KV usage
blocks_used = max_kv * NUM_GPU_BLOCKS
tokens_used = blocks_used * BLOCK_SIZE
# Bytes used
full_attn_bytes = blocks_used * FULL_ATTN_PER_BLOCK
linear_bytes = blocks_used * LINEAR_PER_BLOCK
total_bytes = full_attn_bytes + linear_bytes
# With TQ: full_attn portion compressed
tq_full_attn_bytes = (tokens_used * TQ_PER_TOKEN_ALL_FULL) + TQ_OVERHEAD_TOTAL if tokens_used > 0 else 0
tq_total = tq_full_attn_bytes + linear_bytes
savings = total_bytes - tq_total
r = {
"target": target,
"kv_peak_pct": max_kv,
"blocks_used": blocks_used,
"tokens_est": tokens_used,
"total_bytes": total_bytes,
"full_attn_bytes": full_attn_bytes,
"linear_bytes": linear_bytes,
"tq_total": tq_total,
"savings_bytes": savings,
"ttft": ttft,
}
results.append(r)
print(f"{target:>10,} {'?':>8} {max_kv:>10.4%} {total_bytes/1e6:>12.1f}MB {full_attn_bytes/1e6:>10.1f}MB {linear_bytes/1e6:>10.1f}MB {ttft if ttft else '?':>8}")
print()
print("=" * 80)
print(" BASELINE vs TURBOQUANT COMPARISON")
print("=" * 80)
print()
print(f"{'Context':>10} {'Baseline':>14} {'With TQ':>14} {'Savings':>14} {'Savings':>10} {'Savings':>12}")
print(f"{'target':>10} {'KV total':>14} {'KV total':>14} {'total':>14} {'per GPU':>10} {'(%)':>12}")
print("-" * 80)
for r in results:
baseline = r["total_bytes"]
tq = r["tq_total"]
sav = r["savings_bytes"]
if baseline > 0:
pct = sav / baseline * 100
print(f"{r['target']:>10,} {baseline/1e6:>12.1f}MB {tq/1e6:>12.1f}MB {sav/1e6:>12.1f}MB {sav/TP_SIZE/1e6:>8.1f}MB {pct:>10.1f}%")
else:
print(f"{r['target']:>10,} {'N/A':>14}")
# Now do the key comparison: what if we gave TQ's freed memory BACK to context?
print()
print("=" * 80)
print(" CONTEXT EXTENSION: How much more context with TQ?")
print("=" * 80)
print()
# Current capacity: 5190 blocks * 272 = 1,411,680 tokens
# But max_model_len=131072
# If we compress full_attn KV with TQ, we free up memory that could be used for MORE blocks
# The freed memory per block = (FULL_ATTN_PER_BLOCK - TQ equivalent per block)
tq_full_attn_per_block = BLOCK_SIZE * TQ_PER_TOKEN_ALL_FULL
saved_per_block = FULL_ATTN_PER_BLOCK - tq_full_attn_per_block
# But we can't just add more blocks -- the linear_attn state per block stays the same
# New blocks would need: LINEAR_PER_BLOCK + tq_full_attn_per_block bytes each
new_block_cost = LINEAR_PER_BLOCK + tq_full_attn_per_block
# Total memory currently used for KV: NUM_GPU_BLOCKS * TOTAL_PER_BLOCK
# With TQ, each block costs: LINEAR_PER_BLOCK + tq_full_attn_per_block
# So we can fit: total_memory / new_block_cost blocks
total_kv_memory = NUM_GPU_BLOCKS * TOTAL_PER_BLOCK
new_num_blocks = total_kv_memory // new_block_cost
new_capacity = new_num_blocks * BLOCK_SIZE
print(f"Current KV cache budget: {total_kv_memory/1e9:.2f} GB")
print(f"Current block cost: {TOTAL_PER_BLOCK/1e6:.2f} MB (full_attn: {FULL_ATTN_PER_BLOCK/1e6:.2f} + linear: {LINEAR_PER_BLOCK/1e6:.2f})")
print(f"TQ block cost: {new_block_cost/1e6:.2f} MB (tq_full_attn: {tq_full_attn_per_block/1e6:.2f} + linear: {LINEAR_PER_BLOCK/1e6:.2f})")
print()
print(f"Current: {NUM_GPU_BLOCKS} blocks = {NUM_GPU_BLOCKS * BLOCK_SIZE:,} tokens")
print(f"With TQ: {new_num_blocks} blocks = {new_capacity:,} tokens")
print(f"Context extension: {new_capacity / (NUM_GPU_BLOCKS * BLOCK_SIZE):.2f}x")
print()
# What about a model with NO linear attention (pure transformer)?
print("--- For comparison: if this were a PURE transformer (no linear attention) ---")
pure_block_cost_baseline = FULL_ATTN_PER_BLOCK # no linear state
pure_block_cost_tq = tq_full_attn_per_block
pure_blocks_baseline = total_kv_memory // pure_block_cost_baseline
pure_blocks_tq = total_kv_memory // pure_block_cost_tq
print(f"Pure transformer baseline: {pure_blocks_baseline} blocks = {pure_blocks_baseline * BLOCK_SIZE:,} tokens")
print(f"Pure transformer with TQ: {pure_blocks_tq} blocks = {pure_blocks_tq * BLOCK_SIZE:,} tokens")
print(f"Context extension: {pure_blocks_tq / pure_blocks_baseline:.2f}x")
print()
print("=" * 80)
print(" ACTUAL VRAM BREAKDOWN (per GPU)")
print("=" * 80)
print()
# nvidia-smi shows ~22.3 GB used per GPU
# gpu_memory_utilization=0.92 means 24576 * 0.92 = 22610 MB reserved
# Model weights: 54 GB / 8 = 6.75 GB per GPU
# KV cache pool: total_kv_memory / 8
# CUDA overhead + graphs: remainder
model_weights_per_gpu = 54000 / TP_SIZE # approximate from disk size
kv_cache_per_gpu = total_kv_memory / TP_SIZE
reserved_per_gpu = 24576 * 0.92
overhead = reserved_per_gpu - model_weights_per_gpu - kv_cache_per_gpu / 1e6
print(f"Total VRAM per GPU: 24,576 MB")
print(f"Reserved (0.92): {reserved_per_gpu:.0f} MB")
print(f"Model weights: ~{model_weights_per_gpu:.0f} MB")
print(f"KV cache pool: {kv_cache_per_gpu/1e6:.0f} MB")
print(f" - full_attn: {full_attn_total/TP_SIZE/1e6:.0f} MB")
print(f" - linear state: {linear_total/TP_SIZE/1e6:.0f} MB")
print(f"CUDA overhead: ~{overhead:.0f} MB")
print()
# With TQ:
tq_kv_per_gpu = (new_num_blocks * new_block_cost) / TP_SIZE
tq_overhead = TQ_OVERHEAD_TOTAL / TP_SIZE
print(f"WITH TQ:")
print(f"KV cache pool: {tq_kv_per_gpu/1e6:.0f} MB (same budget, more tokens)")
print(f" - tq_full_attn: {new_num_blocks * tq_full_attn_per_block / TP_SIZE / 1e6:.0f} MB")
print(f" - linear state: {new_num_blocks * LINEAR_PER_BLOCK / TP_SIZE / 1e6:.0f} MB")
print(f"TQ overhead (Pi+S): {tq_overhead/1e6:.1f} MB")
print(f"VRAM freed: {(kv_cache_per_gpu - tq_kv_per_gpu)/1e6:.0f} MB/GPU (if not reused for more tokens)")
baseline_vs_tq_v2.py
-174
View File
@@ -1,174 +0,0 @@
#!/usr/bin/env python3
"""
Measure KV cache usage during CONCURRENT requests to see actual block allocation.
Also compute the definitive baseline vs TQ comparison.
"""
import json
import time
import subprocess
import tempfile
import os
import threading
BASE_URL = "http://localhost:8000/v1"
MODEL = "Qwen/Qwen3.5-35B-A3B"
NUM_GPU_BLOCKS = 5190
BLOCK_SIZE = 272
TP_SIZE = 8
# Per block costs
FULL_ATTN_PER_BLOCK = BLOCK_SIZE * 10 * 2 * 256 * 2 * 2 # 5.57 MB
LINEAR_PER_BLOCK = 278528 * 30 # 8.36 MB
TOTAL_PER_BLOCK = FULL_ATTN_PER_BLOCK + LINEAR_PER_BLOCK # 13.93 MB
# TQ per block
TQ_FULL_ATTN_PER_BLOCK = BLOCK_SIZE * 10 * 464 # 1.26 MB
TQ_TOTAL_PER_BLOCK = TQ_FULL_ATTN_PER_BLOCK + LINEAR_PER_BLOCK # 9.62 MB
def get_kv_usage():
result = subprocess.run(["curl", "-s", "http://localhost:8000/metrics"],
capture_output=True, text=True, timeout=10)
for line in result.stdout.splitlines():
if line.startswith("vllm:kv_cache_usage_perc{"):
return float(line.split()[-1])
return 0
def send_long_request(target_tokens, max_gen=500):
"""Send a request that generates many tokens to keep KV cache allocated."""
filler = "The quick brown fox jumps over the lazy dog. " * (target_tokens // 10)
filler = filler[:target_tokens * 4]
prompt = filler + "\n\nWrite a very long and detailed essay about the history of computing, at least 500 words."
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, dir='/tmp') as f:
json.dump({
"model": MODEL,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": max_gen,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
}, f)
tmppath = f.name
cmd = ["curl", "-s", "-X", "POST", f"{BASE_URL}/chat/completions",
"-H", "Content-Type: application/json", "-d", f"@{tmppath}",
"--max-time", "600"]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=610)
os.unlink(tmppath)
try:
return json.loads(result.stdout)
except:
return {"error": "failed"}
def monitor_kv_usage(duration, interval=0.2):
"""Monitor KV cache usage for a given duration."""
samples = []
t0 = time.time()
while time.time() - t0 < duration:
kv = get_kv_usage()
samples.append((time.time() - t0, kv))
time.sleep(interval)
return samples
print("=" * 80)
print(" ACTUAL KV CACHE MEASUREMENT DURING INFERENCE")
print("=" * 80)
print()
# Send a long-generation request and monitor KV usage during it
for target_ctx in [8000, 32000, 64000, 100000, 131000]:
print(f"--- Context ~{target_ctx:,} tokens + 500 generation tokens ---")
# Start monitoring in background
kv_samples = []
stop_monitor = threading.Event()
def monitor():
t0 = time.time()
while not stop_monitor.is_set():
kv = get_kv_usage()
kv_samples.append((time.time() - t0, kv))
time.sleep(0.1)
mon_thread = threading.Thread(target=monitor, daemon=True)
mon_thread.start()
# Send request
t0 = time.time()
resp = send_long_request(target_ctx, max_gen=500)
elapsed = time.time() - t0
# Let it settle
time.sleep(0.5)
stop_monitor.set()
mon_thread.join(timeout=2)
# Find peak KV usage
if kv_samples:
peak_kv = max(s[1] for s in kv_samples)
peak_time = [s[0] for s in kv_samples if s[1] == peak_kv][0]
else:
peak_kv = 0
peak_time = 0
# Compute actual bytes
blocks_used = peak_kv * NUM_GPU_BLOCKS
baseline_bytes = blocks_used * TOTAL_PER_BLOCK
baseline_full_attn = blocks_used * FULL_ATTN_PER_BLOCK
baseline_linear = blocks_used * LINEAR_PER_BLOCK
# TQ equivalent
tokens_in_cache = blocks_used * BLOCK_SIZE
tq_full_attn = tokens_in_cache * 10 * 464 # TQ compressed
tq_total = tq_full_attn + baseline_linear
savings = baseline_bytes - tq_total
usage = resp.get("usage", {}) if "error" not in resp else {}
prompt_tokens = usage.get("prompt_tokens", "?")
completion_tokens = usage.get("completion_tokens", "?")
print(f" Prompt: {prompt_tokens} tok, Generated: {completion_tokens} tok, Time: {elapsed:.1f}s")
print(f" Peak KV usage: {peak_kv:.4%} at t={peak_time:.1f}s")
print(f" Blocks used: {blocks_used:.0f} ({tokens_in_cache:.0f} tokens)")
print()
print(f" BASELINE (bf16 KV):")
print(f" full_attn KV: {baseline_full_attn/1e6:>10.1f} MB total ({baseline_full_attn/TP_SIZE/1e6:>8.1f} MB/GPU)")
print(f" linear state: {baseline_linear/1e6:>10.1f} MB total ({baseline_linear/TP_SIZE/1e6:>8.1f} MB/GPU)")
print(f" TOTAL: {baseline_bytes/1e6:>10.1f} MB total ({baseline_bytes/TP_SIZE/1e6:>8.1f} MB/GPU)")
print()
print(f" WITH TQ (3b key / 2b val):")
print(f" tq full_attn: {tq_full_attn/1e6:>10.1f} MB total ({tq_full_attn/TP_SIZE/1e6:>8.1f} MB/GPU)")
print(f" linear state: {baseline_linear/1e6:>10.1f} MB total ({baseline_linear/TP_SIZE/1e6:>8.1f} MB/GPU)")
print(f" TOTAL: {tq_total/1e6:>10.1f} MB total ({tq_total/TP_SIZE/1e6:>8.1f} MB/GPU)")
print()
print(f" SAVINGS: {savings/1e6:>10.1f} MB total ({savings/TP_SIZE/1e6:>8.1f} MB/GPU) ({savings/baseline_bytes*100:.1f}% reduction)" if baseline_bytes > 0 else " SAVINGS: N/A")
print()
# Final summary
print("=" * 80)
print(" SUMMARY: What TQ buys you on this model")
print("=" * 80)
print()
print("KV cache pool: 9,035 MB/GPU (72.3 GB total)")
print(f" 40% is full_attention (3,614 MB/GPU) -- compressible by TQ at 4.4x")
print(f" 60% is linear_attention (5,421 MB/GPU) -- NOT compressible")
print()
print("With TQ applied to full_attention layers:")
print(f" full_attn compressed from 3,614 to ~820 MB/GPU (4.4x)")
print(f" Total KV pool: 820 + 5,421 = 6,241 MB/GPU (vs 9,035 baseline)")
print(f" VRAM savings: 2,794 MB/GPU = 22.4 GB total")
print()
print("Context extension:")
print(f" Baseline: {NUM_GPU_BLOCKS * BLOCK_SIZE:,} tokens capacity")
new_blocks = int(NUM_GPU_BLOCKS * TOTAL_PER_BLOCK / TQ_TOTAL_PER_BLOCK)
print(f" With TQ: {new_blocks * BLOCK_SIZE:,} tokens capacity")
print(f" Extension: {new_blocks * BLOCK_SIZE / (NUM_GPU_BLOCKS * BLOCK_SIZE):.2f}x")
print()
print("OR: use freed VRAM for larger batch size / more concurrent requests")
print(f" Extra VRAM: 2,794 MB/GPU could serve ~{int(2794 / (TOTAL_PER_BLOCK/BLOCK_SIZE*131072/1e6/TP_SIZE))} additional 131k-context requests concurrently")
diagnose_quality.py
-60
View File
@@ -1,60 +0,0 @@
#!/usr/bin/env python3
"""Diagnose: is cos_sim=0.93 from key or value quantization?"""
import sys; sys.path.insert(0, "/tmp")
import torch, torch.nn.functional as F, math
from turboquant.quantizer import TurboQuantProd
from turboquant.kv_cache import quantize_values, dequantize_values
torch.manual_seed(42)
D=256; H=2; N=8192; SCALE=1.0/math.sqrt(D); dev="cuda:0"
keys = torch.randn(1, H, N, D, device=dev) * 0.02
values = torch.randn(1, H, N, D, device=dev) * 0.02
query = torch.randn(1, H, 1, D, device=dev) * 0.02
true_scores = torch.matmul(query, keys.transpose(-2, -1)) * SCALE
true_w = F.softmax(true_scores, dim=-1)
true_out = torch.matmul(true_w, values)
# Case 1: TQ keys, exact values
q = TurboQuantProd(dim=D, bits=3, device=dev, seed=42)
key_q = q.quantize(keys)
tq_scores = q.attention_score(query, key_q) * SCALE
tq_w = F.softmax(tq_scores, dim=-1)
tq_out_exact_v = torch.matmul(tq_w, values)
cos1 = F.cosine_similarity(true_out.reshape(-1, D), tq_out_exact_v.reshape(-1, D), dim=-1).mean().item()
# Case 2: Exact keys, 2-bit values
val_q = quantize_values(values, bits=2, group_size=32)
v_dequant = dequantize_values(val_q, group_size=32)
exact_out_tq_v = torch.matmul(true_w, v_dequant)
cos2 = F.cosine_similarity(true_out.reshape(-1, D), exact_out_tq_v.reshape(-1, D), dim=-1).mean().item()
# Case 3: TQ keys + 2-bit values
tq_out_both = torch.matmul(tq_w, v_dequant)
cos3 = F.cosine_similarity(true_out.reshape(-1, D), tq_out_both.reshape(-1, D), dim=-1).mean().item()
# Case 4: Exact keys, 4-bit values
val_q4 = quantize_values(values, bits=4, group_size=32)
v_dequant4 = dequantize_values(val_q4, group_size=32)
exact_out_4v = torch.matmul(true_w, v_dequant4)
cos4 = F.cosine_similarity(true_out.reshape(-1, D), exact_out_4v.reshape(-1, D), dim=-1).mean().item()
# Value reconstruction quality
v_cos = F.cosine_similarity(values.reshape(-1, D), v_dequant.reshape(-1, D), dim=-1).mean().item()
v4_cos = F.cosine_similarity(values.reshape(-1, D), v_dequant4.reshape(-1, D), dim=-1).mean().item()
print("DIAGNOSIS: What causes cos_sim drop to 0.93?")
print(f" TQ keys + exact values: cos={cos1:.6f} -- key quant only")
print(f" Exact keys + 2b values: cos={cos2:.6f} -- value quant only")
print(f" TQ keys + 2b values: cos={cos3:.6f} -- both")
print(f" Exact keys + 4b values: cos={cos4:.6f} -- value quant 4-bit")
print()
print("Value vector reconstruction:")
print(f" 2-bit cos_sim: {v_cos:.6f}")
print(f" 4-bit cos_sim: {v4_cos:.6f}")
print()
if cos2 < cos1:
print("VERDICT: The quality drop is from 2-bit VALUE quantization, not TQ key compression.")
else:
print("VERDICT: The quality drop is from TQ key compression.")
explainer.html
-298
View File
@@ -1,298 +0,0 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>TurboQuant — How It Works</title>
<style>
:root { --bg: #0f0f1a; --card: #1a1a2e; --accent: #4fc3f7; --green: #66bb6a; --red: #ef5350; --yellow: #ffd54f; --text: #e0e0e0; --dim: #888; }
* { margin: 0; padding: 0; box-sizing: border-box; }
body { font-family: 'Segoe UI', system-ui, -apple-system, sans-serif; background: var(--bg); color: var(--text); line-height: 1.6; }
.container { max-width: 960px; margin: 0 auto; padding: 40px 24px; }
h1 { font-size: 2.8rem; text-align: center; margin-bottom: 8px; }
h1 span { color: var(--accent); }
.subtitle { text-align: center; color: var(--dim); font-size: 1.1rem; margin-bottom: 48px; }
h2 { font-size: 1.6rem; margin: 48px 0 20px; padding-bottom: 8px; border-bottom: 2px solid #333; }
h3 { font-size: 1.15rem; color: var(--accent); margin: 16px 0 8px; }
p, li { font-size: 0.95rem; }
ul { padding-left: 20px; margin: 8px 0; }
li { margin: 4px 0; }
.hero-stats { display: grid; grid-template-columns: repeat(3, 1fr); gap: 16px; margin: 32px 0; }
.stat { background: var(--card); border-radius: 12px; padding: 24px; text-align: center; border: 1px solid #333; }
.stat .num { font-size: 2.4rem; font-weight: 800; }
.stat .label { font-size: 0.85rem; color: var(--dim); margin-top: 4px; }
.stat.green .num { color: var(--green); }
.stat.red .num { color: var(--red); }
.stat.blue .num { color: var(--accent); }
.flow { display: flex; align-items: stretch; gap: 0; margin: 24px 0; flex-wrap: wrap; justify-content: center; }
.flow-step { background: var(--card); border: 1px solid #333; border-radius: 10px; padding: 16px; flex: 1; min-width: 140px; max-width: 200px; text-align: center; position: relative; }
.flow-step .icon { font-size: 1.8rem; margin-bottom: 6px; }
.flow-step .title { font-weight: 700; font-size: 0.9rem; }
.flow-step .desc { font-size: 0.78rem; color: var(--dim); margin-top: 4px; }
.flow-arrow { display: flex; align-items: center; font-size: 1.4rem; color: var(--dim); padding: 0 6px; }
.two-col { display: grid; grid-template-columns: 1fr 1fr; gap: 24px; margin: 16px 0; }
@media (max-width: 700px) { .two-col { grid-template-columns: 1fr; } .hero-stats { grid-template-columns: 1fr; } }
.card { background: var(--card); border: 1px solid #333; border-radius: 10px; padding: 20px; }
.card h3 { margin-top: 0; }
.diagram { background: #111; border: 1px solid #333; border-radius: 10px; padding: 24px; margin: 20px 0; font-family: 'Courier New', monospace; font-size: 0.82rem; line-height: 1.8; overflow-x: auto; white-space: pre; color: #ccc; }
.diagram .hl { color: var(--accent); font-weight: 700; }
.diagram .gr { color: var(--green); }
.diagram .rd { color: var(--red); }
.diagram .yw { color: var(--yellow); }
table { width: 100%; border-collapse: collapse; margin: 16px 0; font-size: 0.9rem; }
th { background: #222; padding: 10px 12px; text-align: left; font-weight: 600; border-bottom: 2px solid #444; }
td { padding: 8px 12px; border-bottom: 1px solid #2a2a2a; }
tr:hover td { background: #1a1a2e; }
.g { color: var(--green); font-weight: 700; }
.r { color: var(--red); }
.math { font-family: 'Times New Roman', serif; font-style: italic; font-size: 1.05em; }
.code { background: #222; padding: 2px 6px; border-radius: 4px; font-family: monospace; font-size: 0.85em; }
.section-num { display: inline-block; background: var(--accent); color: #000; font-weight: 800; width: 28px; height: 28px; border-radius: 50%; text-align: center; line-height: 28px; font-size: 0.85rem; margin-right: 8px; vertical-align: middle; }
.tag { display: inline-block; padding: 2px 10px; border-radius: 12px; font-size: 0.75rem; font-weight: 600; }
.tag-green { background: #1b5e20; color: #a5d6a7; }
.tag-red { background: #b71c1c; color: #ef9a9a; }
.tag-blue { background: #0d47a1; color: #90caf9; }
</style>
</head>
<body>
<div class="container">
<h1>Turbo<span>Quant</span></h1>
<div class="subtitle">KV Cache Compression for vLLM &mdash; 2&times; context, 0% overhead, 30 GB freed</div>
<div class="hero-stats">
<div class="stat green"><div class="num">30 GB</div><div class="label">KV cache freed (4 GPUs)</div></div>
<div class="stat blue"><div class="num">2.0&times;</div><div class="label">context capacity</div></div>
<div class="stat red"><div class="num">0%</div><div class="label">throughput overhead</div></div>
</div>
<!-- ─── THE PROBLEM ─── -->
<h2><span class="section-num">1</span> The Problem</h2>
<div class="two-col">
<div class="card">
<h3>KV cache grows with every token</h3>
<p>During generation, every layer stores a <strong>Key</strong> and <strong>Value</strong> vector for every previous token. For long contexts this dominates GPU memory.</p>
<div class="diagram">Token 1 &rarr; K₁ V₁
Token 2 &rarr; K₁ V₁ K₂ V₂
Token 3 &rarr; K₁ V₁ K₂ V₂ K₃ V₃
...
<span class="rd">Token N &rarr; K₁..Kₙ V₁..Vₙ &larr; O(N) memory</span></div>
</div>
<div class="card">
<h3>Qwen3.5-27B on 4&times; RTX 3090</h3>
<ul>
<li>Model weights: ~13 GB / GPU</li>
<li><strong>KV cache: ~7.5 GB / GPU</strong> (30 GB total)</li>
<li>Max context: 457,072 tokens</li>
<li>Want longer? Need more GPUs &mdash; expensive</li>
</ul>
<p style="margin-top:12px; color: var(--yellow);">What if we could compress the KV cache to 39% of its size?</p>
</div>
</div>
<!-- ─── HIGH-LEVEL PIPELINE ─── -->
<h2><span class="section-num">2</span> The TurboQuant Pipeline</h2>
<p>Each KV pair goes through a compression pipeline that reduces <strong>512 bytes &rarr; 198 bytes</strong> per token (2.6&times; compression):</p>
<div class="flow">
<div class="flow-step"><div class="icon">🔑</div><div class="title">Raw Key</div><div class="desc">256-dim bf16<br>512 bytes</div></div>
<div class="flow-arrow">&rarr;</div>
<div class="flow-step" style="border-color: var(--accent);"><div class="icon">🔄</div><div class="title">Rotate</div><div class="desc">Random orthogonal<br>matrix &Pi;</div></div>
<div class="flow-arrow">&rarr;</div>
<div class="flow-step" style="border-color: var(--accent);"><div class="icon">📊</div><div class="title">Quantize</div><div class="desc">3-bit Lloyd-Max<br>8 centroids</div></div>
<div class="flow-arrow">&rarr;</div>
<div class="flow-step" style="border-color: var(--accent);"><div class="icon">±</div><div class="title">QJL Signs</div><div class="desc">Residual direction<br>sign bits</div></div>
<div class="flow-arrow">&rarr;</div>
<div class="flow-step" style="border-color: var(--green);"><div class="icon">✅</div><div class="title">Compressed</div><div class="desc">~110 bytes<br><strong>4.7&times; smaller</strong></div></div>
</div>
<div class="flow" style="margin-top: 16px;">
<div class="flow-step"><div class="icon">💎</div><div class="title">Raw Value</div><div class="desc">256-dim bf16<br>512 bytes</div></div>
<div class="flow-arrow">&rarr;</div>
<div class="flow-step" style="border-color: var(--yellow);"><div class="icon">📦</div><div class="title">Group Quantize</div><div class="desc">2-bit per value<br>groups of 32</div></div>
<div class="flow-arrow">&rarr;</div>
<div class="flow-step" style="border-color: var(--yellow);"><div class="icon">🗜️</div><div class="title">Bit-Pack</div><div class="desc">4 values per byte<br>+ scales & zeros</div></div>
<div class="flow-arrow">&rarr;</div>
<div class="flow-step" style="border-color: var(--green);"><div class="icon">✅</div><div class="title">Compressed</div><div class="desc">~88 bytes<br><strong>5.8&times; smaller</strong></div></div>
</div>
<!-- ─── KEY COMPRESSION DETAIL ─── -->
<h2><span class="section-num">3</span> Key Compression (3-bit)</h2>
<div class="two-col">
<div class="card">
<h3>Step A: Random Rotation</h3>
<p>Multiply keys by a random orthogonal matrix <span class="math">&Pi;</span>. This spreads information uniformly across dimensions so no single dimension carries disproportionate signal.</p>
<p style="margin-top:8px;"><span class="math">K<sub>rot</sub> = K &middot; &Pi;</span> &nbsp;&nbsp; (norm-preserving)</p>
<h3 style="margin-top:16px;">Step B: Lloyd-Max Quantization</h3>
<p>After rotation, each dimension follows a <strong>Beta distribution</strong>. Lloyd-Max finds the <strong>MSE-optimal</strong> 8 centroids for this distribution.</p>
<p style="margin-top:8px;">Codebooks are <strong>pre-computed offline</strong> for d=256 &mdash; zero runtime overhead.</p>
</div>
<div class="card">
<h3>Step C: QJL Residual Signs</h3>
<p>The quantization residual <span class="math">r = K<sub>rot</sub> - K&#x0302;</span> still contains directional information. QJL (Quantized Johnson-Lindenstrauss) projects it into a compact sign vector:</p>
<p style="margin-top:8px;"><span class="math">signs = sign(S &middot; r)</span></p>
<p>where <span class="math">S</span> is a random &plusmn;1 matrix. This preserves inner-product relationships with high probability.</p>
<h3 style="margin-top:16px;">Storage Breakdown</h3>
<table>
<tr><td>3-bit indices (256 dims)</td><td style="text-align:right;">96 B</td></tr>
<tr><td>Norm (fp32)</td><td style="text-align:right;">4 B</td></tr>
<tr><td>Residual norm (fp32)</td><td style="text-align:right;">4 B</td></tr>
<tr><td>QJL signs</td><td style="text-align:right;">~8 B</td></tr>
<tr><th>Total per token per head</th><th style="text-align:right;">~110 B</th></tr>
</table>
</div>
</div>
<!-- ─── VALUE COMPRESSION DETAIL ─── -->
<h2><span class="section-num">4</span> Value Compression (2-bit)</h2>
<div class="card">
<div class="two-col">
<div>
<h3>Group Quantization</h3>
<p>Values are split into groups of 32 dimensions. Each group gets its own scale and zero point, then each value is mapped to one of <strong>4 levels</strong> (2 bits).</p>
<div class="diagram"><span class="yw">For group [d₀ .. d₃₁]:</span>
scale = (max - min) / 3
zero = min
idx_i = round((val_i - zero) / scale) <span class="hl">&larr; 0,1,2,3</span>
<span class="yw">Bit-packing (4 values per byte):</span>
byte = idx₀ | (idx₁&lt;&lt;2) | (idx₂&lt;&lt;4) | (idx₃&lt;&lt;6)
<span class="gr">256 dims &rarr; 64 packed bytes + 16B scales + 16B zeros = 88 bytes</span></div>
</div>
<div>
<h3>Why 2-bit works for values</h3>
<ul>
<li>Values are <strong>summed with softmax weights</strong> during attention &mdash; small errors average out across many tokens</li>
<li>Per-group scales preserve the <strong>dynamic range</strong> of each dimension cluster</li>
<li>4&times; denser than 8-bit quantization</li>
</ul>
<h3 style="margin-top:16px;">Combined Compression</h3>
<table>
<tr><th>Component</th><th>Original</th><th>Compressed</th></tr>
<tr><td>Key</td><td class="r">512 B</td><td class="g">110 B</td></tr>
<tr><td>Value</td><td class="r">512 B</td><td class="g">88 B</td></tr>
<tr><th>Total</th><th class="r">1,024 B</th><th class="g">198 B (2.6&times;)</th></tr>
</table>
</div>
</div>
</div>
<!-- ─── VLLM INTEGRATION ─── -->
<h2><span class="section-num">5</span> vLLM Integration</h2>
<div class="diagram" style="font-size: 0.78rem; line-height: 2;">
<span class="yw">═══ PREFILL (processing the prompt) ═════════════════════════════════════════</span>
User prompt &rarr; vLLM tokenizer &rarr; <span class="hl">FlashAttention forward(Q, K, V)</span>
&darr;
Writes K,V to paged KV cache (normal)
&darr;
<span class="gr">TQ Hook intercepts K,V &rarr; Rotate &rarr; Quantize &rarr; Store compressed</span>
&darr;
Returns attention output (from flash)
<span class="yw">═══ DECODE (generating each new token) ══════════════════════════════════════</span>
New token &rarr; vLLM &rarr; <span class="gr">TQ ACTIVE mode intercepts forward()</span>
&darr;
<span class="gr">Reconstruct K,V from compressed store (Triton kernel)</span>
<span class="gr">Compute attention entirely in TQ</span>
<span class="gr">Flash attention SKIPPED &mdash; paged cache not read</span>
&darr;
Returns attention output
<span class="yw">═══ FREE KV CACHE ═══════════════════════════════════════════════════════════</span>
After prefill, TQ compressed store has all token data.
<span class="rd">Replace paged KV cache tensors with 1-byte dummies.</span>
<span class="gr">torch.cuda.empty_cache() &rarr; 30 GB VRAM returned to pool!</span>
Decode continues using only TQ compressed store.
</div>
<!-- ─── DECODE KERNEL ─── -->
<h2><span class="section-num">6</span> Fused Triton Decode Kernel</h2>
<div class="two-col">
<div class="card">
<h3>What happens every decode step</h3>
<p>For each new query token <span class="math">q</span>, against all <span class="math">N</span> cached tokens:</p>
<ol style="padding-left: 20px; margin-top: 8px;">
<li><strong>Dequantize keys:</strong> indices &rarr; centroids, undo rotation via <span class="math">&Pi;<sup>T</sup></span></li>
<li><strong>Attention scores:</strong> <span class="math">&alpha;<sub>i</sub> = q &middot; k<sub>i</sub> / &radic;d</span></li>
<li><strong>Dequantize values:</strong> unpack 2-bit &rarr; scale &middot; idx + zero</li>
<li><strong>Weighted sum:</strong> <span class="math">out = &sum; softmax(&alpha;<sub>i</sub>) &middot; v<sub>i</sub></span></li>
</ol>
<p style="margin-top: 12px;"><strong>All 4 steps fused into ONE Triton kernel</strong> &mdash; no intermediate tensors allocated.</p>
</div>
<div class="card">
<h3>Why zero overhead</h3>
<ul>
<li>Decode is <strong>compute-bound</strong>, not memory-bound</li>
<li>Reading compressed data is <strong>faster</strong> (less memory bandwidth)</li>
<li>Dequantization is cheap ALU ops fused into the attention loop</li>
<li>Works with <strong>GQA</strong> (Grouped-Query Attention)</li>
</ul>
<table style="margin-top: 16px;">
<tr><th>Metric</th><th>Baseline</th><th>TQ</th></tr>
<tr><td>Throughput</td><td>6.0 tok/s</td><td class="g">6.0 tok/s</td></tr>
<tr><td>Overhead</td><td>&mdash;</td><td class="g">0%</td></tr>
<tr><td>Output</td><td>reference</td><td class="g">identical</td></tr>
</table>
</div>
</div>
<!-- ─── RESULTS ─── -->
<h2><span class="section-num">7</span> Results</h2>
<p style="text-align:center; color: var(--dim); margin-bottom: 16px;">Qwen3.5-27B &middot; 4&times; RTX 3090 &middot; TP=4 &middot; bf16 &middot; gpu_mem=0.90</p>
<table>
<tr><th>Metric</th><th>Baseline vLLM</th><th>TurboQuant</th><th>Change</th></tr>
<tr><td>KV cache blocks</td><td>583</td><td>583 &rarr; freed</td><td class="g">&minus;30 GB</td></tr>
<tr><td>VRAM / GPU</td><td>21,539 MB</td><td>21,299 MB</td><td class="g">&minus;240 MB visible</td></tr>
<tr><td>Tensor bytes freed / GPU</td><td>&mdash;</td><td>7,489 MB</td><td class="g">available in CUDA pool</td></tr>
<tr><td>Total freed</td><td>&mdash;</td><td><strong>30.0 GB</strong></td><td></td></tr>
<tr><td>Max token capacity</td><td>457,072</td><td><strong>914,144</strong></td><td class="g"><strong>2.0&times;</strong></td></tr>
<tr><td>Throughput</td><td>6.0 tok/s</td><td>6.0 tok/s</td><td class="g">0% overhead</td></tr>
<tr><td>Output quality</td><td>reference</td><td>identical</td><td>&mdash;</td></tr>
</table>
<!-- ─── ARCHITECTURE ─── -->
<h2><span class="section-num">8</span> Code Architecture</h2>
<div class="diagram" style="line-height: 2.2; font-size: 0.8rem;">turboquant/
<span class="hl">codebook.py</span> &larr; Lloyd-Max optimal quantizer for Beta distribution
<span class="hl">codebooks/</span> &larr; Pre-generated codebook files (d=256, bits 2/3/4)
<span class="hl">rotation.py</span> &larr; Random orthogonal &Pi; + QJL projection S
<span class="hl">quantizer.py</span> &larr; TurboQuantMSE + TurboQuantProd pipeline
<span class="hl">kv_cache.py</span> &larr; KV cache manager, value bit-packing
<span class="gr">triton_kernels.py</span> &larr; 3 fused Triton kernels (decode attention)
<span class="yw">vllm_attn_backend.py</span> &larr; Monkey-patch hooks, free_kv_cache(), enable_no_alloc()
<span class="rd">proof.py</span> &larr; Definitive A/B benchmark (separate processes)
<span class="rd">benchmark.py</span> &larr; Throughput + quality + VRAM benchmark</div>
<div style="text-align: center; margin-top: 48px; padding: 32px; border-top: 1px solid #333;">
<div style="font-size: 1.4rem; font-weight: 700;">Same model. Same GPUs. <span style="color: var(--accent);">2&times; the context.</span></div>
<div style="margin-top: 12px;"><a href="https://github.com/0xSero/turboquant" style="color: var(--accent); text-decoration: none;">github.com/0xSero/turboquant</a></div>
</div>
</div>
</body>
</html>
profile_100k.py
-500
View File
@@ -1,500 +0,0 @@
#!/usr/bin/env python3
"""
Full GPU profiling for Qwen3.5-35B-A3B at 100k+ context.
Measures every metric that matters:
- Prefill speed (tok/s)
- Generation speed (tok/s)
- Time to First Token (TTFT)
- VRAM used per GPU
- KV cache size in VRAM
- Activation memory during prefill
- CPU usage during inference
- Context size tested
- Quality: needle retrieval, coherence, logprobs/perplexity
Run via: python3 /tmp/profile_100k.py
Requires: vLLM server running on localhost:8000 with sufficient max_model_len
"""
import json
import time
import subprocess
import math
import os
import threading
import re
BASE_URL = "http://localhost:8000/v1"
MODEL = "Qwen/Qwen3.5-35B-A3B"
def curl_post(endpoint, data, timeout=600):
cmd = [
"curl", "-s", "-X", "POST",
f"{BASE_URL}/{endpoint}",
"-H", "Content-Type: application/json",
"-d", json.dumps(data),
"--max-time", str(timeout),
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 30)
if result.returncode != 0:
return {"error": result.stderr}
try:
return json.loads(result.stdout)
except json.JSONDecodeError:
return {"error": f"Invalid JSON: {result.stdout[:500]}"}
def get_gpu_stats():
"""Get per-GPU memory and utilization."""
result = subprocess.run(
["nvidia-smi", "--query-gpu=index,memory.used,memory.free,memory.total,utilization.gpu,temperature.gpu,power.draw",
"--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=10,
)
gpus = []
for line in result.stdout.strip().splitlines():
parts = [x.strip() for x in line.split(",")]
gpus.append({
"idx": int(parts[0]),
"mem_used_mb": int(parts[1]),
"mem_free_mb": int(parts[2]),
"mem_total_mb": int(parts[3]),
"gpu_util_pct": int(parts[4]),
"temp_c": int(parts[5]),
"power_w": float(parts[6]),
})
return gpus
def get_vllm_metrics():
"""Get KV cache usage and request stats from vLLM."""
result = subprocess.run(
["curl", "-s", "http://localhost:8000/metrics"],
capture_output=True, text=True, timeout=10,
)
metrics = {}
for line in result.stdout.splitlines():
if line.startswith("vllm:kv_cache_usage_perc{"):
metrics["kv_usage_pct"] = float(line.split()[-1])
elif line.startswith("vllm:num_requests_running{"):
metrics["requests_running"] = float(line.split()[-1])
elif line.startswith("vllm:num_requests_waiting{"):
metrics["requests_waiting"] = float(line.split()[-1])
return metrics
def get_cpu_usage():
"""Get CPU usage percentage."""
result = subprocess.run(
["bash", "-c", "top -bn1 | head -3 | grep 'Cpu'"],
capture_output=True, text=True, timeout=5,
)
line = result.stdout.strip()
# Parse: %Cpu(s): 12.3 us, 2.1 sy, ...
match = re.search(r'(\d+\.?\d*)\s*us', line)
if match:
return float(match.group(1))
return None
def build_haystack_prompt(target_tokens, needles, question):
"""Build a prompt with needles placed at specific positions in filler text."""
filler_block = (
"In the vast digital landscape, information flows through networks of interconnected systems. "
"Each node processes data according to its designated protocols and algorithms. "
"The architecture of distributed computing enables parallel processing at scale. "
"Modern cloud infrastructure supports millions of concurrent operations across data centers. "
"Load balancers distribute traffic evenly among server clusters for optimal performance. "
"Database sharding partitions large datasets across multiple storage nodes for efficiency. "
"Container orchestration platforms manage the lifecycle of microservices deployments. "
"Network latency optimization involves routing traffic through geographically optimal paths. "
)
target_chars = target_tokens * 4
needle_positions = sorted(needles.keys())
parts = []
current_pos = 0
for needle_pos in needle_positions:
char_pos = int(target_chars * needle_pos)
filler_needed = char_pos - current_pos
if filler_needed > 0:
filler = ""
while len(filler) < filler_needed:
filler += filler_block
parts.append(filler[:filler_needed])
parts.append(f"\n\n{needles[needle_pos]}\n\n")
current_pos = char_pos + len(needles[needle_pos]) + 4
remaining = target_chars - current_pos - len(question) - 50
if remaining > 0:
filler = ""
while len(filler) < remaining:
filler += filler_block
parts.append(filler[:remaining])
parts.append(f"\n\n{question}")
return "".join(parts)
def profile_context_length(target_tokens, run_needle=True, run_generation=True):
"""Full profile at a given context length."""
print(f"\n{'='*70}")
print(f" PROFILING AT ~{target_tokens:,} TOKENS")
print(f"{'='*70}")
result = {
"target_tokens": target_tokens,
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
}
# --- Pre-request GPU state ---
gpu_before = get_gpu_stats()
vllm_before = get_vllm_metrics()
result["gpu_before"] = gpu_before
result["vllm_before"] = vllm_before
# --- CPU monitoring thread ---
cpu_samples = []
stop_cpu = threading.Event()
def sample_cpu():
while not stop_cpu.is_set():
usage = get_cpu_usage()
if usage is not None:
cpu_samples.append(usage)
time.sleep(0.5)
cpu_thread = threading.Thread(target=sample_cpu, daemon=True)
cpu_thread.start()
# --- GPU monitoring thread ---
gpu_samples = []
stop_gpu = threading.Event()
def sample_gpu():
while not stop_gpu.is_set():
stats = get_gpu_stats()
gpu_samples.append(stats)
time.sleep(0.5)
gpu_thread = threading.Thread(target=sample_gpu, daemon=True)
gpu_thread.start()
# ========= TEST 1: Needle-in-haystack =========
if run_needle:
needles = {
0.1: "The access code for Project Neptune is TRIDENT-5582.",
0.3: "Dr. Chen's laboratory is located on floor 47 of Building Sigma.",
0.5: "The backup server IP address is 10.42.88.201 port 9443.",
0.7: "The quarterly budget for Division Omega is exactly $4,271,093.",
0.9: "The launch window for satellite Helios-7 opens at 03:42 UTC.",
}
expected = {
"TRIDENT-5582": "Project Neptune access code",
"floor 47": "Dr. Chen's lab location",
"10.42.88.201": "Backup server IP",
"4,271,093": "Division Omega budget",
"03:42 UTC": "Helios-7 launch window",
}
question = (
"Answer these questions with ONLY the specific answer, one per line:\n"
"1. What is the access code for Project Neptune?\n"
"2. What floor is Dr. Chen's laboratory on?\n"
"3. What is the backup server IP address?\n"
"4. What is the quarterly budget for Division Omega?\n"
"5. When does the launch window for satellite Helios-7 open?"
)
prompt = build_haystack_prompt(target_tokens, needles, question)
print(f"\n [Needle Test] Sending ~{target_tokens:,} token prompt...")
t_start = time.time()
resp = curl_post("chat/completions", {
"model": MODEL,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 200,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
"stream": False,
}, timeout=max(600, target_tokens // 100))
t_end = time.time()
ttft_total = t_end - t_start
if "error" in resp:
print(f" ERROR: {resp['error'][:200]}")
result["needle_error"] = resp["error"][:200]
else:
usage = resp.get("usage", {})
output = resp["choices"][0]["message"]["content"] if resp.get("choices") else ""
prompt_tokens = usage.get("prompt_tokens", 0)
completion_tokens = usage.get("completion_tokens", 0)
# Calculate speeds
# vLLM doesn't expose TTFT separately in non-streaming mode
# Approximate: prefill_time ~ (elapsed - completion_tokens * decode_time_per_token)
# We'll get more accurate numbers with streaming later
prefill_speed = prompt_tokens / ttft_total if ttft_total > 0 else 0
gen_speed = completion_tokens / ttft_total if ttft_total > 0 else 0
found = 0
for key, desc in expected.items():
if key.lower() in output.lower():
found += 1
result["needle"] = {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"elapsed_s": round(ttft_total, 3),
"approx_prefill_toks": prefill_speed,
"approx_gen_toks": gen_speed,
"needles_found": found,
"needles_total": len(expected),
"output": output[:500],
}
print(f" Prompt tokens: {prompt_tokens:,}")
print(f" Completion tokens: {completion_tokens}")
print(f" Total elapsed: {ttft_total:.2f}s")
print(f" Approx throughput: {prefill_speed:.0f} tok/s (prompt+gen combined)")
print(f" Needles: {found}/{len(expected)}")
print(f" Output: {output[:200]}...")
# ========= TEST 2: Streaming for accurate TTFT =========
if run_generation:
filler_block = (
"The digital frontier expands as new technologies emerge. "
"Data centers process millions of requests every second. "
"Cloud computing has transformed how organizations manage infrastructure. "
)
filler_chars = target_tokens * 4 - 200
filler = ""
while len(filler) < filler_chars:
filler += filler_block
filler = filler[:filler_chars]
gen_prompt = filler + "\n\nWrite a haiku about the ocean."
print(f"\n [Generation Test] Streaming for TTFT measurement...")
t_gen_start = time.time()
# Use streaming to measure TTFT accurately
cmd = [
"curl", "-s", "-N", "-X", "POST",
f"{BASE_URL}/chat/completions",
"-H", "Content-Type: application/json",
"-d", json.dumps({
"model": MODEL,
"messages": [{"role": "user", "content": gen_prompt}],
"max_tokens": 100,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
"stream": True,
}),
"--max-time", str(max(600, target_tokens // 100)),
]
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
ttft = None
tokens_received = 0
full_output = ""
for line in proc.stdout:
line = line.strip()
if not line.startswith("data: "):
continue
data_str = line[6:]
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
delta = chunk.get("choices", [{}])[0].get("delta", {})
content = delta.get("content", "")
if content:
if ttft is None:
ttft = time.time() - t_gen_start
tokens_received += 1
full_output += content
except json.JSONDecodeError:
pass
proc.wait(timeout=10)
t_gen_end = time.time()
total_gen_time = t_gen_end - t_gen_start
if ttft is not None:
decode_time = total_gen_time - ttft
decode_speed = tokens_received / decode_time if decode_time > 0 else 0
prefill_speed = target_tokens / ttft if ttft > 0 else 0
result["generation"] = {
"ttft_s": round(ttft, 4),
"total_s": round(total_gen_time, 3),
"tokens_generated": tokens_received,
"prefill_tok_s": round(prefill_speed, 1),
"decode_tok_s": round(decode_speed, 1),
"output": full_output[:200],
}
print(f" TTFT: {ttft:.3f}s")
print(f" Prefill speed: {prefill_speed:.0f} tok/s")
print(f" Decode speed: {decode_speed:.1f} tok/s")
print(f" Total: {total_gen_time:.2f}s for {tokens_received} tokens")
print(f" Output: {full_output[:100]}...")
else:
print(f" ERROR: No tokens received")
result["generation"] = {"error": "no tokens received"}
# --- Stop monitoring ---
stop_cpu.set()
stop_gpu.set()
cpu_thread.join(timeout=2)
gpu_thread.join(timeout=2)
# --- Post-request GPU state ---
gpu_after = get_gpu_stats()
vllm_after = get_vllm_metrics()
result["gpu_after"] = gpu_after
result["vllm_after"] = vllm_after
# --- Compute deltas ---
print(f"\n [GPU Memory]")
for g in gpu_after:
gb = gpu_before[g["idx"]]
delta = g["mem_used_mb"] - gb["mem_used_mb"]
print(f" GPU {g['idx']}: {g['mem_used_mb']} MB used ({delta:+d} MB delta), "
f"{g['gpu_util_pct']}% util, {g['temp_c']}C, {g['power_w']:.0f}W")
if gpu_samples:
peak_mem = [0] * 8
peak_util = [0] * 8
for sample in gpu_samples:
for g in sample:
peak_mem[g["idx"]] = max(peak_mem[g["idx"]], g["mem_used_mb"])
peak_util[g["idx"]] = max(peak_util[g["idx"]], g["gpu_util_pct"])
result["peak_mem_mb"] = peak_mem
result["peak_util_pct"] = peak_util
print(f" Peak mem: {[f'{m}MB' for m in peak_mem]}")
if cpu_samples:
result["cpu_avg_pct"] = round(sum(cpu_samples) / len(cpu_samples), 1)
result["cpu_peak_pct"] = round(max(cpu_samples), 1)
print(f"\n [CPU] Avg: {result['cpu_avg_pct']}%, Peak: {result['cpu_peak_pct']}%")
kv_delta = None
if "kv_usage_pct" in vllm_before and "kv_usage_pct" in vllm_after:
kv_delta = vllm_after["kv_usage_pct"] - vllm_before["kv_usage_pct"]
result["kv_cache_delta_pct"] = round(kv_delta * 100, 4)
print(f"\n [KV Cache] Usage: {vllm_after['kv_usage_pct']:.4%} (delta: {kv_delta:+.4%})")
# --- Theoretical memory breakdown ---
# Full attention KV for this prompt
prompt_tokens_actual = result.get("needle", result.get("generation", {})).get("prompt_tokens", target_tokens)
kv_full_attn_bytes = prompt_tokens_actual * 10 * 2 * 256 * 2 * 2 # 10 layers, 2 kv_heads, 256 dim, K+V, bf16
print(f"\n [Memory Breakdown (theoretical)]")
print(f" KV cache (10 full_attn layers): {kv_full_attn_bytes / 1e6:.1f} MB total, {kv_full_attn_bytes / 8 / 1e6:.1f} MB/GPU")
# Activation during prefill (per layer, FlashAttention)
# Q: (B, H, S, D) = 1 * 16 * S * 256 * 2 bytes
act_q = 1 * 16 * prompt_tokens_actual * 256 * 2
act_kv = 1 * 2 * prompt_tokens_actual * 256 * 2 * 2 # K and V
print(f" Activation Q per layer: {act_q / 1e6:.1f} MB")
print(f" Activation KV per layer: {act_kv / 1e6:.1f} MB")
# With FlashAttention: no S*S matrix, just O(S) workspace
print(f" (FlashAttention avoids O(S^2) attention matrix)")
return result
# ============================================================
# MAIN
# ============================================================
if __name__ == "__main__":
print("=" * 70)
print(" FULL GPU PROFILING: Qwen3.5-35B-A3B on 8x RTX 3090")
print("=" * 70)
print(f" Time: {time.strftime('%Y-%m-%d %H:%M:%S')}")
# Check server is up and get max_model_len
result = subprocess.run(
["curl", "-s", "http://localhost:8000/v1/models"],
capture_output=True, text=True, timeout=10,
)
try:
models = json.loads(result.stdout)
max_len = models["data"][0]["max_model_len"]
print(f" Model: {MODEL}")
print(f" max_model_len: {max_len}")
except (json.JSONDecodeError, KeyError, IndexError):
print(f" WARNING: Could not get model info. Server might be down.")
max_len = 8192
# Profile at increasing context lengths
all_results = []
# Start with smaller sizes for comparison, then go big
test_sizes = [1000, 4000, 8000]
# Add larger sizes if server supports them
if max_len >= 16384:
test_sizes.append(16000)
if max_len >= 32768:
test_sizes.append(32000)
if max_len >= 65536:
test_sizes.append(64000)
if max_len >= 131072:
test_sizes.extend([100000, 131000])
for size in test_sizes:
try:
r = profile_context_length(size)
all_results.append(r)
except Exception as e:
print(f"\n FAILED at {size}: {e}")
all_results.append({"target_tokens": size, "error": str(e)})
# ============================================================
# SUMMARY TABLE
# ============================================================
print("\n\n" + "=" * 70)
print(" SUMMARY TABLE")
print("=" * 70)
print()
print(f"{'Context':>10} {'Prefill':>10} {'Decode':>10} {'TTFT':>8} "
f"{'VRAM/GPU':>10} {'KV Cache':>10} {'CPU':>6} {'Needles':>8}")
print(f"{'tokens':>10} {'tok/s':>10} {'tok/s':>10} {'(s)':>8} "
f"{'(MB)':>10} {'(MB)':>10} {'(%)':>6} {'found':>8}")
print("-" * 80)
for r in all_results:
ctx = r.get("target_tokens", "?")
gen = r.get("generation", {})
needle = r.get("needle", {})
prefill = gen.get("prefill_tok_s", "")
decode = gen.get("decode_tok_s", "")
ttft = gen.get("ttft_s", "")
peak_mem = r.get("peak_mem_mb", [])
avg_peak = sum(peak_mem) / len(peak_mem) if peak_mem else ""
cpu = r.get("cpu_avg_pct", "")
needles = f"{needle.get('needles_found', '?')}/{needle.get('needles_total', '?')}" if needle else ""
# Theoretical KV cache for this context
actual_tokens = needle.get("prompt_tokens", gen.get("prompt_tokens", ctx))
if isinstance(actual_tokens, (int, float)):
kv_mb = actual_tokens * 10 * 2 * 256 * 2 * 2 / 1e6 / 8 # per GPU
else:
kv_mb = ""
print(f"{ctx:>10,} {prefill:>10} {decode:>10} {ttft:>8} "
f"{avg_peak:>10} {kv_mb:>10} {cpu:>6} {needles:>8}")
# Save full results
out_path = "/tmp/profile_100k_results.json"
with open(out_path, "w") as f:
json.dump(all_results, f, indent=2, default=str)
print(f"\nFull results saved to {out_path}")
profile_large.py
-372
View File
@@ -1,372 +0,0 @@
#!/usr/bin/env python3
"""
Large context profiling (64k-131k) using file-based payloads.
Fixes the 'Argument list too long' error from curl.
"""
import json
import time
import subprocess
import math
import os
import threading
import re
import tempfile
BASE_URL = "http://localhost:8000/v1"
MODEL = "Qwen/Qwen3.5-35B-A3B"
def curl_post_file(endpoint, data, timeout=1200):
"""POST with payload written to a temp file to avoid arg length limits."""
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, dir='/tmp') as f:
json.dump(data, f)
tmppath = f.name
try:
cmd = [
"curl", "-s", "-X", "POST",
f"{BASE_URL}/{endpoint}",
"-H", "Content-Type: application/json",
"-d", f"@{tmppath}",
"--max-time", str(timeout),
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 30)
if result.returncode != 0:
return {"error": result.stderr[:500]}
try:
return json.loads(result.stdout)
except json.JSONDecodeError:
return {"error": f"Invalid JSON: {result.stdout[:500]}"}
finally:
os.unlink(tmppath)
def curl_stream_file(endpoint, data, timeout=1200):
"""Streaming POST for TTFT measurement."""
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, dir='/tmp') as f:
json.dump(data, f)
tmppath = f.name
try:
cmd = [
"curl", "-s", "-N", "-X", "POST",
f"{BASE_URL}/{endpoint}",
"-H", "Content-Type: application/json",
"-d", f"@{tmppath}",
"--max-time", str(timeout),
]
return subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
except Exception as e:
os.unlink(tmppath)
raise
def get_gpu_stats():
result = subprocess.run(
["nvidia-smi", "--query-gpu=index,memory.used,memory.free,memory.total,utilization.gpu,temperature.gpu,power.draw",
"--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=10,
)
gpus = []
for line in result.stdout.strip().splitlines():
parts = [x.strip() for x in line.split(",")]
gpus.append({
"idx": int(parts[0]), "mem_used_mb": int(parts[1]),
"mem_free_mb": int(parts[2]), "mem_total_mb": int(parts[3]),
"gpu_util_pct": int(parts[4]), "temp_c": int(parts[5]),
"power_w": float(parts[6]),
})
return gpus
def get_vllm_metrics():
result = subprocess.run(["curl", "-s", "http://localhost:8000/metrics"],
capture_output=True, text=True, timeout=10)
metrics = {}
for line in result.stdout.splitlines():
if line.startswith("vllm:kv_cache_usage_perc{"):
metrics["kv_usage_pct"] = float(line.split()[-1])
return metrics
def get_cpu_usage():
result = subprocess.run(["bash", "-c", "top -bn1 | head -3 | grep 'Cpu'"],
capture_output=True, text=True, timeout=5)
match = re.search(r'(\d+\.?\d*)\s*us', result.stdout)
return float(match.group(1)) if match else None
def build_haystack_prompt(target_tokens, needles, question):
filler_block = (
"In the vast digital landscape, information flows through networks of interconnected systems. "
"Each node processes data according to its designated protocols and algorithms. "
"The architecture of distributed computing enables parallel processing at scale. "
"Modern cloud infrastructure supports millions of concurrent operations across data centers. "
"Load balancers distribute traffic evenly among server clusters for optimal performance. "
"Database sharding partitions large datasets across multiple storage nodes for efficiency. "
"Container orchestration platforms manage the lifecycle of microservices deployments. "
"Network latency optimization involves routing traffic through geographically optimal paths. "
)
target_chars = target_tokens * 4
needle_positions = sorted(needles.keys())
parts = []
current_pos = 0
for needle_pos in needle_positions:
char_pos = int(target_chars * needle_pos)
filler_needed = char_pos - current_pos
if filler_needed > 0:
filler = ""
while len(filler) < filler_needed:
filler += filler_block
parts.append(filler[:filler_needed])
parts.append(f"\n\n{needles[needle_pos]}\n\n")
current_pos = char_pos + len(needles[needle_pos]) + 4
remaining = target_chars - current_pos - len(question) - 50
if remaining > 0:
filler = ""
while len(filler) < remaining:
filler += filler_block
parts.append(filler[:remaining])
parts.append(f"\n\n{question}")
return "".join(parts)
def profile_large(target_tokens):
print(f"\n{'='*70}")
print(f" PROFILING AT ~{target_tokens:,} TOKENS")
print(f"{'='*70}")
needles = {
0.1: "The access code for Project Neptune is TRIDENT-5582.",
0.3: "Dr. Chen's laboratory is located on floor 47 of Building Sigma.",
0.5: "The backup server IP address is 10.42.88.201 port 9443.",
0.7: "The quarterly budget for Division Omega is exactly $4,271,093.",
0.9: "The launch window for satellite Helios-7 opens at 03:42 UTC.",
}
expected = {
"TRIDENT-5582": "Project Neptune access code",
"floor 47": "Dr. Chen lab location",
"10.42.88.201": "Backup server IP",
"4,271,093": "Division Omega budget",
"03:42 UTC": "Helios-7 launch window",
}
question = (
"Answer these questions with ONLY the specific answer, one per line:\n"
"1. What is the access code for Project Neptune?\n"
"2. What floor is Dr. Chen's laboratory on?\n"
"3. What is the backup server IP address?\n"
"4. What is the quarterly budget for Division Omega?\n"
"5. When does the launch window for satellite Helios-7 open?"
)
prompt = build_haystack_prompt(target_tokens, needles, question)
print(f" Prompt chars: {len(prompt):,}")
# GPU monitoring
gpu_samples = []
stop_gpu = threading.Event()
def sample_gpu():
while not stop_gpu.is_set():
try:
gpu_samples.append(get_gpu_stats())
except:
pass
time.sleep(1.0)
gpu_thread = threading.Thread(target=sample_gpu, daemon=True)
gpu_thread.start()
# CPU monitoring
cpu_samples = []
stop_cpu = threading.Event()
def sample_cpu():
while not stop_cpu.is_set():
u = get_cpu_usage()
if u is not None:
cpu_samples.append(u)
time.sleep(1.0)
cpu_thread = threading.Thread(target=sample_cpu, daemon=True)
cpu_thread.start()
gpu_before = get_gpu_stats()
vllm_before = get_vllm_metrics()
# ---- NEEDLE TEST ----
print(f" [Needle Test] Sending request...")
t0 = time.time()
resp = curl_post_file("chat/completions", {
"model": MODEL,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 200,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
}, timeout=max(1200, target_tokens // 50))
t1 = time.time()
needle_elapsed = t1 - t0
if "error" in resp:
print(f" NEEDLE ERROR: {resp['error'][:300]}")
else:
usage = resp.get("usage", {})
output = resp["choices"][0]["message"]["content"] if resp.get("choices") else ""
prompt_tokens = usage.get("prompt_tokens", 0)
completion_tokens = usage.get("completion_tokens", 0)
found = sum(1 for k in expected if k.lower() in output.lower())
print(f" Prompt tokens: {prompt_tokens:,}")
print(f" Completion tokens: {completion_tokens}")
print(f" Total elapsed: {needle_elapsed:.2f}s")
print(f" Needles: {found}/{len(expected)}")
print(f" Output: {output[:300]}...")
# ---- STREAMING TEST (TTFT) ----
print(f"\n [Streaming Test] Measuring TTFT...")
filler_block = "The digital frontier expands as new technologies emerge. Data centers process millions of requests. "
filler_chars = target_tokens * 4 - 200
filler = ""
while len(filler) < filler_chars:
filler += filler_block
filler = filler[:filler_chars]
gen_prompt = filler + "\n\nWrite a haiku about the ocean."
stream_data = {
"model": MODEL,
"messages": [{"role": "user", "content": gen_prompt}],
"max_tokens": 100,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
"stream": True,
}
# Write to file for streaming
stream_path = "/tmp/stream_payload.json"
with open(stream_path, 'w') as f:
json.dump(stream_data, f)
t_gen_start = time.time()
proc = subprocess.Popen(
["curl", "-s", "-N", "-X", "POST", f"{BASE_URL}/chat/completions",
"-H", "Content-Type: application/json", "-d", f"@{stream_path}",
"--max-time", str(max(1200, target_tokens // 50))],
stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
)
ttft = None
tokens_received = 0
full_output = ""
for line in proc.stdout:
line = line.strip()
if not line.startswith("data: "):
continue
data_str = line[6:]
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
delta = chunk.get("choices", [{}])[0].get("delta", {})
content = delta.get("content", "")
if content:
if ttft is None:
ttft = time.time() - t_gen_start
tokens_received += 1
full_output += content
except json.JSONDecodeError:
pass
proc.wait(timeout=30)
t_gen_end = time.time()
total_gen_time = t_gen_end - t_gen_start
if ttft is not None:
decode_time = total_gen_time - ttft
decode_speed = tokens_received / decode_time if decode_time > 0 else 0
prefill_speed = target_tokens / ttft if ttft > 0 else 0
print(f" TTFT: {ttft:.3f}s")
print(f" Prefill speed: {prefill_speed:.0f} tok/s")
print(f" Decode speed: {decode_speed:.1f} tok/s")
print(f" Total: {total_gen_time:.2f}s for {tokens_received} tokens")
print(f" Output: {full_output[:100]}...")
else:
print(" ERROR: No tokens received from streaming")
prefill_speed = 0
decode_speed = 0
ttft = 0
# Stop monitoring
stop_gpu.set()
stop_cpu.set()
gpu_thread.join(timeout=2)
cpu_thread.join(timeout=2)
gpu_after = get_gpu_stats()
vllm_after = get_vllm_metrics()
print(f"\n [GPU Memory]")
for g in gpu_after:
gb = gpu_before[g["idx"]]
delta = g["mem_used_mb"] - gb["mem_used_mb"]
print(f" GPU {g['idx']}: {g['mem_used_mb']} MB ({delta:+d}), "
f"{g['gpu_util_pct']}% util, {g['temp_c']}C, {g['power_w']:.0f}W")
peak_mem = [0] * 8
for sample in gpu_samples:
for g in sample:
peak_mem[g["idx"]] = max(peak_mem[g["idx"]], g["mem_used_mb"])
print(f" Peak: {[f'{m}' for m in peak_mem]}")
if cpu_samples:
print(f" [CPU] Avg: {sum(cpu_samples)/len(cpu_samples):.1f}%, Peak: {max(cpu_samples):.1f}%")
kv_pct = vllm_after.get("kv_usage_pct", 0)
print(f" [KV Cache] Usage: {kv_pct:.4%}")
# Theoretical
prompt_tok = resp.get("usage", {}).get("prompt_tokens", target_tokens) if "error" not in resp else target_tokens
kv_mb = prompt_tok * 10 * 2 * 256 * 2 * 2 / 1e6
print(f" [Theoretical KV] {kv_mb:.1f} MB total, {kv_mb/8:.1f} MB/GPU")
os.unlink(stream_path)
return {
"target_tokens": target_tokens,
"prompt_tokens": resp.get("usage", {}).get("prompt_tokens", "?") if "error" not in resp else "?",
"prefill_tok_s": round(prefill_speed, 1),
"decode_tok_s": round(decode_speed, 1),
"ttft_s": round(ttft, 3) if ttft else "?",
"needles": f"{found}/{len(expected)}" if "error" not in resp else "ERR",
"peak_mem_avg_mb": round(sum(peak_mem) / len(peak_mem)),
"kv_cache_mb_total": round(kv_mb, 1),
"cpu_avg": round(sum(cpu_samples) / len(cpu_samples), 1) if cpu_samples else "?",
}
if __name__ == "__main__":
print("=" * 70)
print(" LARGE CONTEXT PROFILING (64k-131k)")
print("=" * 70)
results = []
for size in [64000, 100000, 131000]:
try:
r = profile_large(size)
results.append(r)
except Exception as e:
print(f"\n FAILED at {size}: {e}")
import traceback
traceback.print_exc()
results.append({"target_tokens": size, "error": str(e)})
print("\n\n" + "=" * 70)
print(" LARGE CONTEXT SUMMARY")
print("=" * 70)
print()
print(f"{'Context':>10} {'Prefill':>10} {'Decode':>10} {'TTFT':>8} {'PeakVRAM':>10} {'KV total':>10} {'CPU':>6} {'Needles':>8}")
print(f"{'tokens':>10} {'tok/s':>10} {'tok/s':>10} {'(s)':>8} {'(MB/GPU)':>10} {'(MB)':>10} {'(%)':>6} {'found':>8}")
print("-" * 80)
for r in results:
if "error" in r:
print(f"{r['target_tokens']:>10,} {'ERROR':>10}")
continue
print(f"{r['target_tokens']:>10,} {r['prefill_tok_s']:>10} {r['decode_tok_s']:>10} "
f"{r['ttft_s']:>8} {r['peak_mem_avg_mb']:>10} {r['kv_cache_mb_total']:>10} "
f"{r['cpu_avg']:>6} {r['needles']:>8}")
with open("/tmp/profile_large_results.json", "w") as f:
json.dump(results, f, indent=2, default=str)
print(f"\nSaved to /tmp/profile_large_results.json")
scripts/tq_autoresearch_dashboard.py
-300
View File
@@ -1,300 +0,0 @@
#!/usr/bin/env python3
"""Render a single-file HTML dashboard for the TurboQuant autoresearch control plane."""
from __future__ import annotations
import argparse
import html
import json
from pathlib import Path
from typing import Any
from tq_harness_lib import atomic_write_text, utc_now_iso
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--status-json", default="test-output/autoresearch-driver/status.json")
parser.add_argument("--watch-report", default="test-output/autoresearch-driver/watch-report.json")
parser.add_argument("--output", required=True)
return parser.parse_args()
def load_json(path: Path) -> dict[str, Any] | None:
if not path.exists():
return None
return json.loads(path.read_text())
def summarize_stage_rows(driver_report: dict[str, Any] | None) -> list[dict[str, Any]]:
if not driver_report:
return []
rows: list[dict[str, Any]] = []
for stage in driver_report.get("stages", []):
rows.append(
{
"stage_name": stage.get("stage_name"),
"context_len": stage.get("context_len"),
"cases": stage.get("cases"),
"phases": stage.get("phases"),
"status": stage.get("status"),
"probe_campaign_root": stage.get("probe_campaign_root"),
"recovery_campaign_root": stage.get("recovery_campaign_root"),
}
)
return rows
def build_dashboard_payload(status_payload: dict[str, Any] | None, watch_payload: dict[str, Any] | None) -> dict[str, Any]:
status_payload = status_payload or {}
watch_payload = watch_payload or {}
latest_watch = (watch_payload.get("history") or [{}])[-1]
driver_report = status_payload.get("driver_report")
return {
"generated_at": utc_now_iso(),
"blocker_code": status_payload.get("blocker_code"),
"blocker_summary": status_payload.get("blocker_summary"),
"campaign_plan": status_payload.get("campaign_plan"),
"latest_probe": status_payload.get("latest_probe") or {},
"next_command": status_payload.get("next_command"),
"campaign_failed": status_payload.get("campaign_failed") or [],
"campaign_missing": status_payload.get("campaign_missing") or [],
"driver_dry_run": driver_report.get("dry_run") if driver_report else None,
"stage_rows": summarize_stage_rows(driver_report),
"watch_iterations": watch_payload.get("iterations"),
"watch_driver_triggered": watch_payload.get("driver_triggered"),
"watch_latest_blocker": latest_watch.get("blocker_code"),
"watch_latest_summary": latest_watch.get("blocker_summary"),
}
def render_list(items: list[str]) -> str:
if not items:
return "<p class='muted'>None</p>"
return "<ul>" + "".join(f"<li>{html.escape(item)}</li>" for item in items) + "</ul>"
def render_stage_rows(stage_rows: list[dict[str, Any]]) -> str:
if not stage_rows:
return "<p class='muted'>No stage plan available.</p>"
header = (
"<tr><th>Stage</th><th>Context</th><th>Cases</th><th>Phases</th><th>Status</th></tr>"
)
rows = []
for row in stage_rows:
rows.append(
"<tr>"
f"<td>{html.escape(str(row.get('stage_name') or ''))}</td>"
f"<td>{html.escape(str(row.get('context_len') or ''))}</td>"
f"<td>{html.escape(str(row.get('cases') or ''))}</td>"
f"<td>{html.escape(str(row.get('phases') or ''))}</td>"
f"<td>{html.escape(str(row.get('status') or ''))}</td>"
"</tr>"
)
return f"<table>{header}{''.join(rows)}</table>"
def render_html(payload: dict[str, Any]) -> str:
failed_items = [
f"{item.get('context_len')} / {item.get('case')} / {item.get('phase')} -> {item.get('status')} (exit={item.get('exit_code')})"
for item in payload.get("campaign_failed", [])
]
missing_items = [
f"{item.get('context_len')} / {item.get('case')} / {item.get('phase')}"
for item in payload.get("campaign_missing", [])
]
next_command = html.escape(payload.get("next_command") or "No next command available")
latest_probe = payload.get("latest_probe") or {}
return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>TurboQuant Autoresearch Dashboard</title>
<style>
:root {{
--bg: #0f1419;
--panel: #151c23;
--panel-strong: #1b2530;
--text: #edf2f7;
--muted: #95a3b3;
--accent: #6ee7b7;
--warn: #f59e0b;
--danger: #f87171;
--border: #2a3642;
--mono: "SFMono-Regular", "SF Mono", ui-monospace, Menlo, monospace;
--sans: "IBM Plex Sans", "Segoe UI", sans-serif;
}}
* {{ box-sizing: border-box; }}
body {{
margin: 0;
background: radial-gradient(circle at top right, #163041 0%, var(--bg) 42%);
color: var(--text);
font-family: var(--sans);
}}
.wrap {{
max-width: 1200px;
margin: 0 auto;
padding: 32px 20px 56px;
}}
h1, h2 {{ margin: 0 0 12px; line-height: 1.1; }}
h1 {{ font-size: 34px; letter-spacing: -0.03em; }}
h2 {{ font-size: 18px; }}
p {{ margin: 0; }}
.hero {{
display: grid;
grid-template-columns: 2fr 1fr;
gap: 16px;
margin-bottom: 16px;
}}
.grid {{
display: grid;
grid-template-columns: repeat(2, minmax(0, 1fr));
gap: 16px;
}}
.card {{
background: linear-gradient(180deg, var(--panel-strong), var(--panel));
border: 1px solid var(--border);
padding: 18px;
border-radius: 14px;
box-shadow: 0 18px 60px rgba(0, 0, 0, 0.28);
}}
.badge {{
display: inline-block;
padding: 6px 10px;
border-radius: 999px;
font-size: 12px;
font-weight: 700;
letter-spacing: 0.04em;
text-transform: uppercase;
margin-bottom: 12px;
}}
.badge-danger {{ background: rgba(248, 113, 113, 0.14); color: var(--danger); }}
.badge-ok {{ background: rgba(110, 231, 183, 0.14); color: var(--accent); }}
.meta {{
display: grid;
grid-template-columns: repeat(2, minmax(0, 1fr));
gap: 12px;
margin-top: 14px;
}}
.meta div {{
padding: 10px 12px;
background: rgba(255, 255, 255, 0.02);
border: 1px solid rgba(255, 255, 255, 0.04);
border-radius: 10px;
}}
.label {{
display: block;
font-size: 11px;
color: var(--muted);
text-transform: uppercase;
letter-spacing: 0.06em;
margin-bottom: 5px;
}}
.value {{ font-size: 14px; line-height: 1.5; }}
.muted {{ color: var(--muted); }}
pre {{
margin: 0;
padding: 14px;
overflow-x: auto;
border-radius: 12px;
background: #0b1116;
border: 1px solid var(--border);
color: #d7e3ef;
font-size: 12px;
line-height: 1.45;
font-family: var(--mono);
white-space: pre-wrap;
word-break: break-word;
}}
table {{
width: 100%;
border-collapse: collapse;
font-size: 13px;
}}
th, td {{
text-align: left;
border-bottom: 1px solid var(--border);
padding: 10px 8px;
vertical-align: top;
}}
th {{ color: var(--muted); font-size: 11px; text-transform: uppercase; letter-spacing: 0.06em; }}
ul {{ margin: 0; padding-left: 18px; }}
li + li {{ margin-top: 6px; }}
@media (max-width: 900px) {{
.hero, .grid {{ grid-template-columns: 1fr; }}
}}
</style>
</head>
<body>
<div class="wrap">
<section class="hero">
<div class="card">
<span class="badge {'badge-ok' if payload.get('blocker_code') == 'ready' else 'badge-danger'}">{html.escape(str(payload.get('blocker_code') or 'unknown'))}</span>
<h1>TurboQuant Autoresearch Dashboard</h1>
<p class="muted">{html.escape(str(payload.get('blocker_summary') or 'No summary available'))}</p>
<div class="meta">
<div><span class="label">Generated</span><span class="value">{html.escape(str(payload.get('generated_at') or ''))}</span></div>
<div><span class="label">50k Plan</span><span class="value">{html.escape(str(payload.get('campaign_plan') or 'unknown'))}</span></div>
<div><span class="label">Watch Iterations</span><span class="value">{html.escape(str(payload.get('watch_iterations') or 0))}</span></div>
<div><span class="label">Driver Triggered</span><span class="value">{html.escape(str(payload.get('watch_driver_triggered')))}</span></div>
</div>
</div>
<div class="card">
<h2>Latest Probe</h2>
<div class="meta">
<div><span class="label">Status</span><span class="value">{html.escape(str(latest_probe.get('status') or 'unknown'))}</span></div>
<div><span class="label">Error Type</span><span class="value">{html.escape(str(latest_probe.get('error_type') or 'none'))}</span></div>
</div>
<div style="margin-top:14px">
<span class="label">stderr tail</span>
<pre>{html.escape(str(latest_probe.get('stderr_tail') or ''))}</pre>
</div>
</div>
</section>
<section class="grid">
<div class="card">
<h2>Next Live Command</h2>
<pre>{next_command}</pre>
</div>
<div class="card">
<h2>Watch Gate</h2>
<div class="meta">
<div><span class="label">Latest Blocker</span><span class="value">{html.escape(str(payload.get('watch_latest_blocker') or 'unknown'))}</span></div>
<div><span class="label">Latest Summary</span><span class="value">{html.escape(str(payload.get('watch_latest_summary') or ''))}</span></div>
<div><span class="label">Driver Dry Run</span><span class="value">{html.escape(str(payload.get('driver_dry_run')))}</span></div>
</div>
</div>
<div class="card">
<h2>Failed Phases</h2>
{render_list(failed_items)}
</div>
<div class="card">
<h2>Missing Phases</h2>
{render_list(missing_items)}
</div>
<div class="card" style="grid-column: 1 / -1;">
<h2>Stage Plan</h2>
{render_stage_rows(payload.get('stage_rows') or [])}
</div>
</section>
</div>
</body>
</html>
"""
def main() -> int:
args = parse_args()
status_payload = load_json(Path(args.status_json))
watch_payload = load_json(Path(args.watch_report))
payload = build_dashboard_payload(status_payload, watch_payload)
atomic_write_text(Path(args.output), render_html(payload))
print(args.output)
return 0
if __name__ == "__main__":
raise SystemExit(main())
scripts/tq_autoresearch_driver.py
-189
View File
@@ -1,189 +0,0 @@
#!/usr/bin/env python3
"""Compose and optionally execute staged TurboQuant autoresearch recovery runs."""
from __future__ import annotations
import argparse
import json
import shlex
import subprocess
import sys
from pathlib import Path
from tq_harness_lib import atomic_write_json, ensure_dir, utc_now_iso
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--contexts", default="50000,80000,120000,200000")
parser.add_argument("--base-remote-root", default="/5090-qwen3.5-27b/logs/campaigns")
parser.add_argument("--remote-model-path", default="/5090-qwen3.5-27b/models/QuantTrio-Qwen3.5-27B-AWQ")
parser.add_argument("--host", default="your-gpu-host.example.com")
parser.add_argument("--port", type=int, default=3003)
parser.add_argument("--user", default="root")
parser.add_argument("--identity-file", default="~/.ssh/id_rsa")
parser.add_argument("--remote-scripts-dir", default="/5090-qwen3.5-27b/scripts")
parser.add_argument("--probe-interval-s", type=int, default=30)
parser.add_argument("--max-wait-s", type=int, default=300)
parser.add_argument("--ssh-timeout-s", type=int, default=20)
parser.add_argument("--prompt-seed", type=int, default=5090)
parser.add_argument("--max-output-tokens", type=int, default=24)
parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
parser.add_argument("--tensor-parallel-size", type=int, default=1)
parser.add_argument("--local-report-dir", default="test-output/autoresearch-driver")
parser.add_argument("--output", required=True)
parser.add_argument("--dry-run", action="store_true")
return parser.parse_args()
def parse_contexts(value: str) -> list[int]:
return [int(part.strip()) for part in value.split(",") if part.strip()]
def stage_name_for(context_len: int) -> str:
if context_len == 50000:
return "smoke-50k-planb"
if context_len == 200000:
return "smoke-200k-tq"
return f"smoke-{context_len // 1000}k-planb"
def stage_cases_for(context_len: int) -> str:
return "tq" if context_len >= 200000 else "baseline,tq"
def stage_phases_for(context_len: int) -> str:
if context_len >= 200000:
return "init,prefill_only,decode_only,full"
return "init,prefill_only,decode_only,full"
def build_stage_specs(args: argparse.Namespace) -> list[dict]:
contexts = parse_contexts(args.contexts)
specs: list[dict] = []
probe_root = f"{args.base_remote_root}/smoke-50k-harness"
local_report_dir = Path(args.local_report_dir)
for context_len in contexts:
stage_name = stage_name_for(context_len)
recovery_root = f"{args.base_remote_root}/{stage_name}"
local_report_path = local_report_dir / f"{stage_name}-resume-report.json"
specs.append(
{
"context_len": context_len,
"stage_name": stage_name,
"probe_campaign_root": probe_root,
"recovery_campaign_root": recovery_root,
"cases": stage_cases_for(context_len),
"phases": stage_phases_for(context_len),
"local_report_path": str(local_report_path),
"notes": [
"200k stage is TQ-only by default" if context_len >= 200000 else "A/B stage uses baseline and TQ",
"Split phases isolate prefill/decode pressure",
],
}
)
probe_root = recovery_root
return specs
def build_resume_command(args: argparse.Namespace, stage: dict, dry_run: bool | None = None) -> list[str]:
if dry_run is None:
dry_run = args.dry_run
cmd = [
sys.executable,
str(Path(__file__).resolve().parent / "tq_resume_recovery.py"),
"--host",
args.host,
"--port",
str(args.port),
"--user",
args.user,
"--identity-file",
args.identity_file,
"--ssh-timeout-s",
str(args.ssh_timeout_s),
"--probe-interval-s",
str(args.probe_interval_s),
"--max-wait-s",
str(args.max_wait_s),
"--remote-scripts-dir",
args.remote_scripts_dir,
"--remote-model-path",
args.remote_model_path,
"--probe-campaign-root",
stage["probe_campaign_root"],
"--recovery-campaign-root",
stage["recovery_campaign_root"],
"--contexts",
str(stage["context_len"]),
"--cases",
stage["cases"],
"--phases",
stage["phases"],
"--prompt-seed",
str(args.prompt_seed),
"--max-output-tokens",
str(args.max_output_tokens),
"--gpu-memory-utilization",
str(args.gpu_memory_utilization),
"--tensor-parallel-size",
str(args.tensor_parallel_size),
"--output",
stage["local_report_path"],
]
if dry_run:
cmd.append("--dry-run")
return cmd
def command_to_shell(cmd: list[str]) -> str:
return " ".join(shlex.quote(part) for part in cmd)
def main() -> int:
args = parse_args()
output_path = Path(args.output)
ensure_dir(output_path.parent)
ensure_dir(Path(args.local_report_dir))
stages = build_stage_specs(args)
report = {
"started_at": utc_now_iso(),
"status": "ok",
"dry_run": args.dry_run,
"contexts": [stage["context_len"] for stage in stages],
"stages": [],
}
for stage in stages:
cmd = build_resume_command(args, stage)
stage_report = {
**stage,
"command": cmd,
"command_shell": command_to_shell(cmd),
"status": "planned",
}
if not args.dry_run:
proc = subprocess.run(cmd, capture_output=True, text=True)
stage_report["status"] = "ok" if proc.returncode == 0 else "error"
stage_report["returncode"] = proc.returncode
stage_report["stdout_tail"] = proc.stdout[-4000:] if proc.stdout else ""
stage_report["stderr_tail"] = proc.stderr[-4000:] if proc.stderr else ""
report["stages"].append(stage_report)
if proc.returncode != 0:
report["status"] = "error"
report["error_stage"] = stage["stage_name"]
break
else:
report["stages"].append(stage_report)
report["ended_at"] = utc_now_iso()
atomic_write_json(output_path, report)
print(output_path)
return 0 if report["status"] == "ok" else 1
if __name__ == "__main__":
raise SystemExit(main())
scripts/tq_autoresearch_status.py
-154
View File
@@ -1,154 +0,0 @@
#!/usr/bin/env python3
"""Aggregate local TurboQuant autoresearch artifacts into a status packet."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any
from tq_harness_lib import atomic_write_json, atomic_write_text, utc_now_iso
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--remote-smoke-dir", default="test-output/remote-smoke-50k")
parser.add_argument("--driver-report", default="test-output/autoresearch-driver/driver-report.json")
parser.add_argument("--output-json", required=True)
parser.add_argument("--output-md", required=True)
return parser.parse_args()
def load_json_if_exists(path: Path) -> dict[str, Any] | None:
if not path.exists():
return None
return json.loads(path.read_text())
def choose_latest_probe(remote_smoke_dir: Path) -> tuple[str | None, dict[str, Any] | None]:
candidates = [
("ping_latest", remote_smoke_dir / "ping-snapshot-latest.json"),
("ping", remote_smoke_dir / "ping-snapshot.json"),
("summary", remote_smoke_dir / "health-snapshot-summary.json"),
("retry", remote_smoke_dir / "retry-snapshot.json"),
]
for name, path in candidates:
payload = load_json_if_exists(path)
if payload is not None:
return name, payload
return None, None
def summarize_blocker(probe: dict[str, Any] | None, recovery: dict[str, Any] | None) -> tuple[str, str]:
if probe and probe.get("status") == "error":
error_type = probe.get("error_type", "unknown")
stderr_tail = probe.get("stderr_tail", "")
if "Operation timed out" in stderr_tail:
return "ssh_outage", "Remote SSH is timing out on your-gpu-host.example.com:3003"
return error_type, probe.get("error_message", "Remote probe failed")
if recovery and recovery.get("status") == "error":
return recovery.get("error_type", "recovery_error"), recovery.get("error_message", "Recovery driver failed")
return "ready", "No current local blocker recorded"
def extract_next_command(driver_report: dict[str, Any] | None) -> str | None:
if not driver_report:
return None
for stage in driver_report.get("stages", []):
if stage.get("status") in {"planned", "error"}:
command_shell = stage.get("command_shell")
if not command_shell:
return None
return command_shell.replace(" --dry-run", "")
return None
def build_status_payload(remote_smoke_dir: Path, driver_report_path: Path) -> dict[str, Any]:
local_manifest = load_json_if_exists(remote_smoke_dir / "local-manifest.json")
recovery_report = load_json_if_exists(remote_smoke_dir / "resume-recovery-report.json")
driver_report = load_json_if_exists(driver_report_path)
probe_name, probe_payload = choose_latest_probe(remote_smoke_dir)
blocker_code, blocker_summary = summarize_blocker(probe_payload, recovery_report)
next_command = extract_next_command(driver_report)
return {
"generated_at": utc_now_iso(),
"remote_smoke_dir": str(remote_smoke_dir),
"driver_report_path": str(driver_report_path),
"campaign_plan": local_manifest.get("recommended_plan") if local_manifest else None,
"campaign_failed": local_manifest.get("failed", []) if local_manifest else [],
"campaign_missing": local_manifest.get("missing", []) if local_manifest else [],
"latest_probe_name": probe_name,
"latest_probe": probe_payload,
"recovery_report": recovery_report,
"driver_report": driver_report,
"blocker_code": blocker_code,
"blocker_summary": blocker_summary,
"next_command": next_command,
}
def render_markdown(payload: dict[str, Any]) -> str:
lines = [
"# TurboQuant Autoresearch Status",
"",
f"- Generated at: `{payload['generated_at']}`",
f"- Blocker: `{payload['blocker_code']}`",
f"- Summary: {payload['blocker_summary']}",
f"- 50k campaign plan: `{payload.get('campaign_plan')}`",
f"- Latest probe: `{payload.get('latest_probe_name')}`",
"",
]
latest_probe = payload.get("latest_probe") or {}
if latest_probe:
lines.extend(
[
"## Latest Probe",
f"- Status: `{latest_probe.get('status', 'unknown')}`",
f"- Error type: `{latest_probe.get('error_type')}`",
f"- stderr tail: `{latest_probe.get('stderr_tail', '')}`",
"",
]
)
failed = payload.get("campaign_failed") or []
missing = payload.get("campaign_missing") or []
if failed:
lines.append("## Failed")
for item in failed:
lines.append(
f"- {item.get('context_len')} / {item.get('case')} / {item.get('phase')} -> {item.get('status')} (exit={item.get('exit_code')})"
)
lines.append("")
if missing:
lines.append("## Missing")
for item in missing:
lines.append(f"- {item.get('context_len')} / {item.get('case')} / {item.get('phase')}")
lines.append("")
lines.append("## Next Command")
if payload.get("next_command"):
lines.append(f"```bash\n{payload['next_command']}\n```")
else:
lines.append("- No next command available")
lines.append("")
return "\n".join(lines)
def main() -> int:
args = parse_args()
remote_smoke_dir = Path(args.remote_smoke_dir)
driver_report_path = Path(args.driver_report)
payload = build_status_payload(remote_smoke_dir, driver_report_path)
markdown = render_markdown(payload)
atomic_write_json(Path(args.output_json), payload)
atomic_write_text(Path(args.output_md), markdown)
print(args.output_json)
print(args.output_md)
return 0
if __name__ == "__main__":
raise SystemExit(main())
scripts/tq_autoresearch_watch.py
-142
View File
@@ -1,142 +0,0 @@
#!/usr/bin/env python3
"""Watch the local autoresearch state and trigger staged recovery when ready."""
from __future__ import annotations
import argparse
import json
import subprocess
import sys
import time
from pathlib import Path
from tq_harness_lib import atomic_write_json, ensure_dir, utc_now_iso
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--status-json", default="test-output/autoresearch-driver/status.json")
parser.add_argument("--status-md", default="test-output/autoresearch-driver/status.md")
parser.add_argument("--status-script", default="scripts/tq_autoresearch_status.py")
parser.add_argument("--dashboard-script", default="scripts/tq_autoresearch_dashboard.py")
parser.add_argument("--dashboard-output", default="test-output/autoresearch-driver/dashboard.html")
parser.add_argument("--driver-script", default="scripts/tq_autoresearch_driver.py")
parser.add_argument("--driver-output", default="test-output/autoresearch-driver/driver-report.json")
parser.add_argument("--watch-report", required=True)
parser.add_argument("--poll-interval-s", type=int, default=60)
parser.add_argument("--max-iterations", type=int, default=0)
parser.add_argument("--run-driver-when-ready", action="store_true")
parser.add_argument("--driver-dry-run", action="store_true")
return parser.parse_args()
def build_status_command(args: argparse.Namespace) -> list[str]:
return [
sys.executable,
args.status_script,
"--output-json",
args.status_json,
"--output-md",
args.status_md,
]
def build_driver_command(args: argparse.Namespace) -> list[str]:
cmd = [
sys.executable,
args.driver_script,
"--output",
args.driver_output,
]
if args.driver_dry_run:
cmd.append("--dry-run")
return cmd
def build_dashboard_command(args: argparse.Namespace) -> list[str]:
return [
sys.executable,
args.dashboard_script,
"--status-json",
args.status_json,
"--watch-report",
args.watch_report,
"--output",
args.dashboard_output,
]
def load_status(path: Path) -> dict | None:
if not path.exists():
return None
return json.loads(path.read_text())
def should_trigger_driver(status_payload: dict | None) -> bool:
if not status_payload:
return False
return status_payload.get("blocker_code") == "ready"
def main() -> int:
args = parse_args()
watch_report_path = Path(args.watch_report)
ensure_dir(watch_report_path.parent)
history: list[dict] = []
iteration = 0
driver_triggered = False
while True:
iteration += 1
step = {"iteration": iteration, "checked_at": utc_now_iso()}
status_cmd = build_status_command(args)
status_proc = subprocess.run(status_cmd, capture_output=True, text=True)
step["status_command"] = status_cmd
step["status_returncode"] = status_proc.returncode
step["status_stdout_tail"] = status_proc.stdout[-4000:] if status_proc.stdout else ""
step["status_stderr_tail"] = status_proc.stderr[-4000:] if status_proc.stderr else ""
status_payload = load_status(Path(args.status_json))
step["blocker_code"] = status_payload.get("blocker_code") if status_payload else None
step["blocker_summary"] = status_payload.get("blocker_summary") if status_payload else None
dashboard_cmd = build_dashboard_command(args)
dashboard_proc = subprocess.run(dashboard_cmd, capture_output=True, text=True)
step["dashboard_command"] = dashboard_cmd
step["dashboard_returncode"] = dashboard_proc.returncode
step["dashboard_stdout_tail"] = dashboard_proc.stdout[-4000:] if dashboard_proc.stdout else ""
step["dashboard_stderr_tail"] = dashboard_proc.stderr[-4000:] if dashboard_proc.stderr else ""
if args.run_driver_when_ready and should_trigger_driver(status_payload):
driver_cmd = build_driver_command(args)
driver_proc = subprocess.run(driver_cmd, capture_output=True, text=True)
step["driver_command"] = driver_cmd
step["driver_returncode"] = driver_proc.returncode
step["driver_stdout_tail"] = driver_proc.stdout[-4000:] if driver_proc.stdout else ""
step["driver_stderr_tail"] = driver_proc.stderr[-4000:] if driver_proc.stderr else ""
driver_triggered = True
history.append(step)
report = {
"started_at": history[0]["checked_at"],
"ended_at": utc_now_iso(),
"iterations": iteration,
"driver_triggered": driver_triggered,
"history": history,
}
atomic_write_json(watch_report_path, report)
if driver_triggered:
print(watch_report_path)
return 0
if args.max_iterations and iteration >= args.max_iterations:
print(watch_report_path)
return 0
time.sleep(args.poll_interval_s)
if __name__ == "__main__":
raise SystemExit(main())
scripts/tq_campaign.py
-270
View File
@@ -1,270 +0,0 @@
#!/usr/bin/env python3
"""Run a resumable TurboQuant telemetry campaign across contexts, cases, and phases."""
from __future__ import annotations
import argparse
import os
import subprocess
import sys
from pathlib import Path
from typing import Any
from tq_harness_lib import (
DEFAULT_CASES,
DEFAULT_CONTEXTS,
DEFAULT_PHASES,
atomic_write_json,
atomic_write_text,
build_python_launcher,
collect_campaign_summary,
ensure_dir,
generate_prompt_artifacts,
parse_csv_ints,
parse_csv_strings,
phase_timeout_for,
scrub_gpu_processes,
should_skip_phase,
timestamp_slug,
utc_now_iso,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--campaign-root", default=None)
parser.add_argument("--contexts", default=None)
parser.add_argument("--cases", default=None)
parser.add_argument("--phases", default=None)
parser.add_argument("--model-path", required=True)
parser.add_argument("--prompt-seed", type=int, default=5090)
parser.add_argument("--max-output-tokens", type=int, default=24)
parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
parser.add_argument("--tensor-parallel-size", type=int, default=1)
parser.add_argument("--skip-existing", action="store_true")
parser.add_argument("--resume", action="store_true")
parser.add_argument("--strict-timeouts", action="store_true")
parser.add_argument("--force", action="store_true")
parser.add_argument("--sync-remote", action="store_true")
parser.add_argument("--remote-host", default="your-gpu-host.example.com")
parser.add_argument("--remote-port", type=int, default=3003)
parser.add_argument("--remote-user", default="root")
parser.add_argument("--remote-key", default="~/.ssh/id_rsa")
parser.add_argument("--remote-scripts-dir", default="/5090-qwen3.5-27b/scripts")
parser.add_argument("--dry-run", action="store_true")
return parser.parse_args()
def phase_runner_path() -> Path:
return Path(__file__).resolve().parent / "tq_phase_runner.py"
def collector_path() -> Path:
return Path(__file__).resolve().parent / "tq_collect_report.py"
def scripts_to_sync() -> list[Path]:
scripts_dir = Path(__file__).resolve().parent
return [
scripts_dir / "tq_harness_lib.py",
scripts_dir / "tq_phase_runner.py",
scripts_dir / "tq_campaign.py",
scripts_dir / "tq_collect_report.py",
]
def sync_remote(args: argparse.Namespace) -> None:
ssh_base = [
"ssh",
"-F",
"/dev/null",
"-o",
"StrictHostKeyChecking=no",
"-o",
"ConnectTimeout=12",
"-o",
f"IdentityFile={args.remote_key}",
"-o",
"IdentitiesOnly=yes",
f"{args.remote_user}@{args.remote_host}",
"-p",
str(args.remote_port),
]
subprocess.run(
ssh_base + [f"mkdir -p {args.remote_scripts_dir}"],
check=True,
)
for script_path in scripts_to_sync():
subprocess.run(
[
"scp",
"-F",
"/dev/null",
"-o",
"StrictHostKeyChecking=no",
"-o",
f"IdentityFile={args.remote_key}",
"-o",
"IdentitiesOnly=yes",
"-P",
str(args.remote_port),
str(script_path),
f"{args.remote_user}@{args.remote_host}:{args.remote_scripts_dir}/{script_path.name}",
],
check=True,
)
def build_campaign_root(args: argparse.Namespace) -> Path:
if args.campaign_root:
return Path(args.campaign_root).resolve()
return Path("logs") / "campaigns" / timestamp_slug()
def run_phase(args: argparse.Namespace, campaign_root: Path, campaign_id: str, context_len: int, case: str, phase: str) -> int:
output_path = campaign_root / str(context_len) / case / f"{phase}.json"
if should_skip_phase(output_path, skip_existing=args.skip_existing or args.resume, force=args.force):
return 0
prompt_dir = campaign_root / str(context_len) / "prompt"
timeout_s = phase_timeout_for(case, phase, context_len, strict_timeouts=args.strict_timeouts)
cmd = build_python_launcher() + [
str(phase_runner_path()),
"--case",
case,
"--phase",
phase,
"--context-len",
str(context_len),
"--output",
str(output_path),
"--timeout-s",
str(timeout_s),
"--prompt-seed",
str(args.prompt_seed),
"--max-output-tokens",
str(args.max_output_tokens),
"--model-path",
args.model_path,
"--prompt-dir",
str(prompt_dir),
"--campaign-id",
campaign_id,
"--gpu-memory-utilization",
str(args.gpu_memory_utilization),
"--tensor-parallel-size",
str(args.tensor_parallel_size),
]
if args.dry_run:
cmd.append("--dry-run")
proc = subprocess.run(cmd, cwd=Path(__file__).resolve().parent.parent)
return proc.returncode
def main() -> int:
args = parse_args()
contexts = parse_csv_ints(args.contexts, DEFAULT_CONTEXTS)
cases = parse_csv_strings(args.cases, DEFAULT_CASES)
phases = parse_csv_strings(args.phases, DEFAULT_PHASES)
campaign_root = build_campaign_root(args)
campaign_root = ensure_dir(campaign_root)
campaign_id = campaign_root.name
if args.sync_remote:
sync_remote(args)
config = {
"campaign_id": campaign_id,
"campaign_root": str(campaign_root),
"created_at": utc_now_iso(),
"model_path": args.model_path,
"prompt_seed": args.prompt_seed,
"contexts": contexts,
"cases": cases,
"phases": phases,
"gpu_memory_utilization": args.gpu_memory_utilization,
"tensor_parallel_size": args.tensor_parallel_size,
"strict_timeouts": args.strict_timeouts,
"dry_run": args.dry_run,
}
atomic_write_json(campaign_root / "campaign_config.json", config)
gate_stop = False
gate_reason = None
for context_len in contexts:
prompt_dir = campaign_root / str(context_len) / "prompt"
generate_prompt_artifacts(
model_path=args.model_path,
context_len=context_len,
seed=args.prompt_seed,
prompt_dir=prompt_dir,
force=args.force,
dry_run=args.dry_run,
)
for case in cases:
attempts = 2 if case == "baseline" and context_len >= 200000 else 1
for phase in phases:
if gate_stop:
break
scrub_info = scrub_gpu_processes()
atomic_write_json(
campaign_root / str(context_len) / case / f"{phase}.prescrub.json",
{
"campaign_id": campaign_id,
"context_len": context_len,
"case": case,
"phase": phase,
"scrub": scrub_info,
"timestamp": utc_now_iso(),
},
)
rc = 1
for _ in range(attempts):
rc = run_phase(args, campaign_root, campaign_id, context_len, case, phase)
if rc == 0:
break
manifest, summary = collect_campaign_summary(campaign_root)
atomic_write_json(campaign_root / "manifest.json", manifest)
atomic_write_text(campaign_root / "summary.md", summary)
if rc != 0 and context_len <= 120000 and args.strict_timeouts:
gate_stop = True
gate_reason = (
f"Strict timeout gate tripped at context {context_len} "
f"case={case} phase={phase}; stopping higher-context execution."
)
break
if gate_stop:
break
if gate_stop:
break
manifest, summary = collect_campaign_summary(campaign_root)
if gate_reason:
manifest["gate_stop_reason"] = gate_reason
summary += f"\n## Gate Stop\n- {gate_reason}\n"
atomic_write_json(campaign_root / "manifest.json", manifest)
atomic_write_text(campaign_root / "summary.md", summary)
collector_cmd = build_python_launcher() + [
str(collector_path()),
"--campaign-root",
str(campaign_root),
]
if args.dry_run:
collector_cmd.append("--dry-run")
subprocess.run(collector_cmd, cwd=Path(__file__).resolve().parent.parent, check=False)
print(f"campaign_root={campaign_root}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
scripts/tq_collect_report.py
-37
View File
@@ -1,37 +0,0 @@
#!/usr/bin/env python3
"""Collect campaign JSON artifacts into a manifest and summary."""
from __future__ import annotations
import argparse
from pathlib import Path
from tq_harness_lib import atomic_write_json, atomic_write_text, collect_campaign_summary
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--campaign-root", required=True)
parser.add_argument("--manifest-path", default=None)
parser.add_argument("--summary-path", default=None)
parser.add_argument("--dry-run", action="store_true")
return parser.parse_args()
def main() -> int:
args = parse_args()
campaign_root = Path(args.campaign_root).resolve()
manifest, summary = collect_campaign_summary(campaign_root)
manifest_path = Path(args.manifest_path).resolve() if args.manifest_path else campaign_root / "manifest.json"
summary_path = Path(args.summary_path).resolve() if args.summary_path else campaign_root / "summary.md"
atomic_write_json(manifest_path, manifest)
atomic_write_text(summary_path, summary)
print(f"manifest={manifest_path}")
print(f"summary={summary_path}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
scripts/tq_harness_lib.py
-703
View File
@@ -1,703 +0,0 @@
#!/usr/bin/env python3
"""Shared helpers for the TurboQuant telemetry harness."""
from __future__ import annotations
import csv
import hashlib
import json
import os
import random
import shutil
import socket
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parent
DEFAULT_CONTEXTS = [30000, 50000, 80000, 120000, 200000]
DEFAULT_CASES = ["baseline", "tq"]
DEFAULT_PHASES = ["init", "ttft", "full", "quality"]
DEFAULT_QUALITY_PROMPT = (
"Answer precisely with one numbered line per item: "
"1) Capital of France? "
"2) 17*23? "
"3) Chemical formula for water? "
"4) Author of Romeo and Juliet? "
"5) What does KV cache store in transformer inference?"
)
_PROMPT_PARAGRAPHS = [
"TurboQuant reduces KV memory by compressing historical keys and values while preserving exact recent tokens for decode stability.",
"The research operator records prompt hashes, phase runtimes, GPU memory, temperature, and power so regressions can be audited later.",
"Every benchmark phase is isolated because long chained jobs hide the difference between initialization cost, prefill cost, and decode cost.",
"The remote 5090 lane is intentionally single-GPU and uses bounded output tokens to keep comparisons about context handling rather than long generations.",
"Baseline and TurboQuant runs must reuse the exact same prompt text, tokenizer path, and output token budget or the telemetry is not comparable.",
"Exit code 137 is treated as a first-class artifact because host kills are part of the observed system behavior, not an external footnote.",
"A stable research lane writes prompt metadata, stdout tails, stderr tails, and partial manifests so interrupted runs still leave evidence.",
"No-alloc mode shifts pressure away from paged KV allocation and toward the compressed TurboQuant store, so hook installation state matters.",
"The campaign stops pretending that symmetric 200k telemetry is mandatory if baseline repeatedly stalls while TurboQuant remains healthy.",
"A useful long-context report includes both successes and explicit failures, because missing artifacts destroy trust faster than bad numbers.",
]
@dataclass
class PromptBundle:
text: str
prompt_hash: str
prompt_tokens: int
token_ids: list[int]
seed: int
context_len: int
def ensure_repo_import_path() -> Path:
"""Make the TurboQuant repo importable from local or mirrored script layouts."""
candidates = [
REPO_ROOT,
SCRIPT_DIR.parent / "turboquant",
]
for candidate in candidates:
if (candidate / "setup.py").exists() and (candidate / "turboquant").exists():
candidate_str = str(candidate)
if candidate_str not in sys.path:
sys.path.insert(0, candidate_str)
return candidate
raise RuntimeError(
f"Unable to locate TurboQuant repo root from {SCRIPT_DIR}. "
"Expected setup.py + turboquant package in a parent or sibling checkout."
)
def utc_now_iso() -> str:
return datetime.now(timezone.utc).isoformat()
def timestamp_slug() -> str:
return datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
def ensure_dir(path: Path) -> Path:
path.mkdir(parents=True, exist_ok=True)
return path
def parse_csv_ints(value: str | None, default: list[int]) -> list[int]:
if not value:
return list(default)
return [int(part.strip()) for part in value.split(",") if part.strip()]
def parse_csv_strings(value: str | None, default: list[str]) -> list[str]:
if not value:
return list(default)
return [part.strip() for part in value.split(",") if part.strip()]
def atomic_write_json(path: Path, payload: dict[str, Any]) -> None:
ensure_dir(path.parent)
fd, tmp_name = tempfile.mkstemp(prefix=path.name + ".", suffix=".tmp", dir=path.parent)
tmp_path = Path(tmp_name)
try:
with os.fdopen(fd, "w", encoding="utf-8") as handle:
json.dump(payload, handle, indent=2, sort_keys=True)
handle.write("\n")
handle.flush()
os.fsync(handle.fileno())
tmp_path.replace(path)
finally:
if tmp_path.exists():
tmp_path.unlink()
def atomic_write_text(path: Path, text: str) -> None:
ensure_dir(path.parent)
fd, tmp_name = tempfile.mkstemp(prefix=path.name + ".", suffix=".tmp", dir=path.parent)
tmp_path = Path(tmp_name)
try:
with os.fdopen(fd, "w", encoding="utf-8") as handle:
handle.write(text)
handle.flush()
os.fsync(handle.fileno())
tmp_path.replace(path)
finally:
if tmp_path.exists():
tmp_path.unlink()
def tail_text(text: str | None, limit: int = 4000) -> str:
if not text:
return ""
text = text.strip()
if len(text) <= limit:
return text
return text[-limit:]
def sha256_jsonable(payload: Any) -> str:
data = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
return hashlib.sha256(data).hexdigest()
def command_available(name: str) -> bool:
return shutil.which(name) is not None
def build_python_launcher() -> list[str]:
if command_available("uv"):
return ["uv", "run", "python"]
return [sys.executable]
def _query_nvidia_csv(query: str) -> list[dict[str, str]]:
try:
proc = subprocess.run(
[
"nvidia-smi",
f"--query-gpu={query}",
"--format=csv,noheader,nounits",
],
capture_output=True,
text=True,
check=True,
)
except (FileNotFoundError, subprocess.CalledProcessError):
return []
columns = [col.strip() for col in query.split(",")]
rows: list[dict[str, str]] = []
for line in proc.stdout.strip().splitlines():
values = next(csv.reader([line]))
rows.append({columns[idx]: values[idx].strip() for idx in range(min(len(columns), len(values)))})
return rows
def probe_gpu_metrics() -> dict[str, Any]:
rows = _query_nvidia_csv(
"name,memory.used,memory.free,utilization.gpu,power.draw,temperature.gpu"
)
if not rows:
return {
"gpu_name": None,
"memory_used_mb": None,
"memory_free_mb": None,
"utilization_gpu_pct": None,
"power_w": None,
"temp_c": None,
"gpus": [],
}
def _num(row: dict[str, str], key: str) -> float | int | None:
raw = row.get(key)
if raw in (None, "", "[Not Supported]"):
return None
if "." in raw:
return float(raw)
return int(raw)
parsed = [
{
"gpu_name": row.get("name"),
"memory_used_mb": _num(row, "memory.used"),
"memory_free_mb": _num(row, "memory.free"),
"utilization_gpu_pct": _num(row, "utilization.gpu"),
"power_w": _num(row, "power.draw"),
"temp_c": _num(row, "temperature.gpu"),
}
for row in rows
]
primary = dict(parsed[0])
primary["gpus"] = parsed
return primary
def scrub_gpu_processes() -> dict[str, Any]:
killed: list[int] = []
try:
proc = subprocess.run(
[
"nvidia-smi",
"--query-compute-apps=pid",
"--format=csv,noheader,nounits",
],
capture_output=True,
text=True,
check=True,
)
for raw in proc.stdout.splitlines():
raw = raw.strip()
if not raw:
continue
pid = int(raw)
try:
os.kill(pid, 9)
killed.append(pid)
except ProcessLookupError:
continue
except (FileNotFoundError, subprocess.CalledProcessError):
pass
return {
"killed_pids": killed,
"post_scrub_gpu": probe_gpu_metrics(),
}
def _build_prompt_units(seed: int) -> list[str]:
rng = random.Random(seed)
paragraphs = list(_PROMPT_PARAGRAPHS)
rng.shuffle(paragraphs)
units = [
"System note: this prompt is synthetic long-context research traffic for deterministic telemetry collection.",
"Operator requirements: preserve exact wording, record prompt hash, and keep outputs short.",
]
for idx in range(4096):
paragraph = paragraphs[idx % len(paragraphs)]
units.append(
f"Section {idx:04d}: {paragraph} "
f"Evidence marker {rng.randint(1000, 9999)}. "
f"Token budget focus: context stability, KV behavior, and staged retries."
)
return units
def build_prompt_bundle_from_tokenizer(tokenizer: Any, context_len: int, seed: int) -> PromptBundle:
units = _build_prompt_units(seed)
token_ids: list[int] = []
for unit in units:
rendered = unit + "\n"
token_ids.extend(tokenizer.encode(rendered, add_special_tokens=False))
if len(token_ids) >= max(context_len + 512, context_len):
break
if not token_ids:
raise RuntimeError("Prompt builder produced no token ids.")
token_ids = token_ids[:context_len]
text = tokenizer.decode(
token_ids,
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)
actual_ids = tokenizer.encode(text, add_special_tokens=False)
while len(actual_ids) > context_len and len(actual_ids) > 1:
token_ids = token_ids[: context_len - (len(actual_ids) - context_len)]
text = tokenizer.decode(
token_ids,
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)
actual_ids = tokenizer.encode(text, add_special_tokens=False)
prompt_hash = sha256_jsonable(actual_ids)
return PromptBundle(
text=text,
prompt_hash=prompt_hash,
prompt_tokens=len(actual_ids),
token_ids=actual_ids,
seed=seed,
context_len=context_len,
)
def generate_prompt_artifacts(
model_path: str,
context_len: int,
seed: int,
prompt_dir: Path,
force: bool = False,
dry_run: bool = False,
) -> dict[str, Any]:
ensure_dir(prompt_dir)
prompt_path = prompt_dir / "prompt.txt"
meta_path = prompt_dir / "prompt_meta.json"
if not force and prompt_path.exists() and meta_path.exists():
return json.loads(meta_path.read_text(encoding="utf-8"))
if dry_run:
prompt_text = (
f"DRY RUN prompt for context {context_len} seed {seed}. "
"This stub bypasses tokenizer/model requirements and only exists for contract testing."
)
prompt_tokens = min(context_len, max(8, len(prompt_text.split())))
prompt_hash = sha256_jsonable([context_len, seed, prompt_tokens, prompt_text])
payload = {
"context_len": context_len,
"prompt_seed": seed,
"prompt_hash": prompt_hash,
"prompt_tokens": prompt_tokens,
"model_path": model_path,
"tokenizer_path": None,
"generated_at": utc_now_iso(),
}
atomic_write_text(prompt_path, prompt_text)
atomic_write_json(meta_path, payload)
return payload
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
bundle = build_prompt_bundle_from_tokenizer(tokenizer, context_len, seed)
payload = {
"context_len": context_len,
"prompt_seed": seed,
"prompt_hash": bundle.prompt_hash,
"prompt_tokens": bundle.prompt_tokens,
"model_path": model_path,
"tokenizer_path": model_path,
"generated_at": utc_now_iso(),
}
atomic_write_text(prompt_path, bundle.text)
atomic_write_json(meta_path, payload)
return payload
def load_prompt_artifacts(prompt_dir: Path) -> tuple[str, dict[str, Any]]:
prompt_path = prompt_dir / "prompt.txt"
meta_path = prompt_dir / "prompt_meta.json"
if not prompt_path.exists() or not meta_path.exists():
raise FileNotFoundError(f"Missing prompt artifacts in {prompt_dir}")
return (
prompt_path.read_text(encoding="utf-8"),
json.loads(meta_path.read_text(encoding="utf-8")),
)
def extract_executor(llm: Any) -> Any:
engine = llm.llm_engine
core = getattr(engine, "engine_core", engine)
inner = getattr(core, "engine_core", core)
return getattr(inner, "model_executor", None)
def inspect_runtime_stats(llm: Any) -> dict[str, Any]:
executor = extract_executor(llm)
if executor is None:
return {
"tq_hooked_layers": None,
"shared_kv_layers": None,
"kv_reserved_gb": None,
"tq_total_memory_bytes": None,
"tq_total_compressed_tokens": None,
"tq_mode": None,
}
def _probe(worker):
model_runner = worker.model_runner
static_ctx = model_runner.compilation_config.static_forward_context
tq_states = (
getattr(model_runner, "_tq_layer_states", None)
or getattr(model_runner, "_tq_states", None)
or {}
)
shared_layers = 0
for attn_module in static_ctx.values():
if getattr(attn_module, "kv_sharing_target_layer_name", None):
shared_layers += 1
seen: dict[int, int] = {}
def _capture_tensor_bytes(value):
if hasattr(value, "data_ptr") and hasattr(value, "nelement") and hasattr(value, "element_size"):
ptr = value.data_ptr()
if ptr:
seen[ptr] = value.nelement() * value.element_size()
elif isinstance(value, (list, tuple)):
for item in value:
_capture_tensor_bytes(item)
for entry in getattr(model_runner, "kv_caches", []):
_capture_tensor_bytes(entry)
tq_total_memory_bytes = None
tq_total_compressed_tokens = None
tq_mode = None
try:
from turboquant.integration.vllm import get_stats as _get_stats
tq_stats = _get_stats(model_runner)
tq_total_memory_bytes = tq_stats.get("total_memory_bytes")
tq_total_compressed_tokens = tq_stats.get("total_compressed_tokens")
tq_mode = tq_stats.get("mode")
except Exception:
pass
return {
"tq_hooked_layers": len(tq_states),
"shared_kv_layers": shared_layers,
"kv_reserved_gb": round(sum(seen.values()) / 1e9, 6),
"tq_total_memory_bytes": tq_total_memory_bytes,
"tq_total_compressed_tokens": tq_total_compressed_tokens,
"tq_mode": tq_mode,
}
results = executor.collective_rpc(_probe)
return results[0] if results else {}
def maybe_free_tq_kv(llm: Any) -> dict[str, Any]:
executor = extract_executor(llm)
if executor is None:
return {"freed_kv_bytes": None}
def _free(worker):
from turboquant.vllm_attn_backend import free_kv_cache
return free_kv_cache(worker.model_runner)
try:
freed = executor.collective_rpc(_free)
except Exception:
return {"freed_kv_bytes": None}
return {"freed_kv_bytes": int(freed[0]) if freed else None}
def phase_timeout_for(case: str, phase: str, context_len: int, strict_timeouts: bool = False) -> int:
base = {
"init": 420,
"ttft": 900,
"prefill_only": 900,
"full": 1200,
"decode_only": 1200,
"quality": 600,
}[phase]
multiplier = max(1.0, context_len / 30000.0)
timeout = int(base * multiplier)
if case == "baseline" and context_len >= 200000:
timeout = 1800 if strict_timeouts else 2400
elif case == "tq" and context_len >= 200000:
timeout = 2400 if strict_timeouts else 3000
elif strict_timeouts:
timeout = int(timeout * 0.85)
return max(timeout, base)
def load_json(path: Path) -> dict[str, Any]:
return json.loads(path.read_text(encoding="utf-8"))
def should_skip_phase(output_path: Path, skip_existing: bool, force: bool) -> bool:
if not output_path.exists() or force:
return False
if skip_existing:
return True
try:
payload = load_json(output_path)
except Exception:
return False
return payload.get("status") == "ok"
def status_rank(status: str) -> int:
order = {
"ok": 0,
"partial": 1,
"timeout": 2,
"killed": 3,
"error": 4,
"missing": 5,
}
return order.get(status, 99)
def collect_campaign_summary(campaign_root: Path) -> tuple[dict[str, Any], str]:
config_path = campaign_root / "campaign_config.json"
config = load_json(config_path) if config_path.exists() else {}
contexts = config.get("contexts", [])
cases = config.get("cases", [])
phases = config.get("phases", [])
rows: list[dict[str, Any]] = []
missing: list[dict[str, Any]] = []
failed: list[dict[str, Any]] = []
prompt_mismatches: list[dict[str, Any]] = []
for context_len in contexts:
prompt_hashes: dict[str, str | None] = {}
for case in cases:
case_dir = campaign_root / str(context_len) / case
prompt_meta_path = campaign_root / str(context_len) / "prompt" / "prompt_meta.json"
if prompt_meta_path.exists():
prompt_meta = load_json(prompt_meta_path)
prompt_hashes[case] = prompt_meta.get("prompt_hash")
phase_statuses = {}
for phase in phases:
path = case_dir / f"{phase}.json"
if not path.exists():
phase_statuses[phase] = "missing"
missing.append(
{
"context_len": context_len,
"case": case,
"phase": phase,
"path": str(path),
}
)
continue
payload = load_json(path)
status = payload.get("status", "error")
phase_statuses[phase] = status
if status != "ok":
failed.append(
{
"context_len": context_len,
"case": case,
"phase": phase,
"status": status,
"path": str(path),
"exit_code": payload.get("exit_code"),
}
)
if case not in prompt_hashes and payload.get("prompt_hash"):
prompt_hashes[case] = payload.get("prompt_hash")
rows.append(
{
"context_len": context_len,
"case": case,
"phase": phase,
"status": status,
"elapsed_s": payload.get("elapsed_s"),
"prompt_hash": payload.get("prompt_hash"),
"sample_text_present": bool(payload.get("sample_text")),
}
)
baseline_hash = prompt_hashes.get("baseline")
tq_hash = prompt_hashes.get("tq")
if baseline_hash and tq_hash and baseline_hash != tq_hash:
prompt_mismatches.append(
{
"context_len": context_len,
"baseline_prompt_hash": baseline_hash,
"tq_prompt_hash": tq_hash,
}
)
ok_rows = sum(1 for row in rows if row["status"] == "ok")
non_ok_rows = len(rows) - ok_rows
exit_137_count = sum(1 for item in failed if item.get("exit_code") == 137 or item.get("status") == "killed")
tq_200k_ok = any(
row["context_len"] == 200000 and row["case"] == "tq" and row["phase"] == "full" and row["status"] == "ok"
for row in rows
)
baseline_200k_ok = any(
row["context_len"] == 200000 and row["case"] == "baseline" and row["phase"] == "full" and row["status"] == "ok"
for row in rows
)
stable_through_120k = all(
any(
row["context_len"] == context_len
and row["case"] == case
and row["phase"] == phase
and row["status"] == "ok"
for row in rows
)
for context_len in [c for c in contexts if c <= 120000]
for case in cases
for phase in phases
)
requested_contexts = sorted(int(context) for context in contexts)
only_sub_200k = bool(requested_contexts) and all(context < 200000 for context in requested_contexts)
all_requested_complete = all(
any(
row["context_len"] == context_len
and row["case"] == case
and row["phase"] == phase
and row["status"] == "ok"
for row in rows
)
for context_len in requested_contexts
for case in cases
for phase in phases
)
if all_requested_complete and only_sub_200k and not prompt_mismatches:
recommended_plan = "A"
elif stable_through_120k and tq_200k_ok and not prompt_mismatches:
recommended_plan = "A"
elif tq_200k_ok and not baseline_200k_ok:
recommended_plan = "C"
elif exit_137_count > 0 or any(item["status"] in {"timeout", "killed"} for item in failed):
recommended_plan = "B"
else:
recommended_plan = "investigate"
manifest = {
"campaign_id": config.get("campaign_id", campaign_root.name),
"campaign_root": str(campaign_root),
"generated_at": utc_now_iso(),
"host": socket.gethostname(),
"contexts": contexts,
"cases": cases,
"phases": phases,
"rows": rows,
"missing": missing,
"failed": failed,
"prompt_mismatches": prompt_mismatches,
"ok_rows": ok_rows,
"non_ok_rows": non_ok_rows,
"recommended_plan": recommended_plan,
}
lines = [
f"# TurboQuant Campaign Summary: {manifest['campaign_id']}",
"",
f"- Campaign root: `{campaign_root}`",
f"- Generated at: `{manifest['generated_at']}`",
f"- Recommended plan: `Plan {recommended_plan}`",
f"- OK rows: `{ok_rows}`",
f"- Non-OK rows: `{non_ok_rows}`",
f"- Prompt mismatches: `{len(prompt_mismatches)}`",
"",
"| Context | Case | Phase | Status | Elapsed (s) | Sample Text |",
"| --- | --- | --- | --- | ---: | --- |",
]
for row in sorted(rows, key=lambda item: (item["context_len"], item["case"], item["phase"])):
elapsed = "" if row["elapsed_s"] is None else f"{row['elapsed_s']:.3f}"
lines.append(
f"| {row['context_len']} | {row['case']} | {row['phase']} | {row['status']} | {elapsed} | "
f"{'yes' if row['sample_text_present'] else 'no'} |"
)
if missing:
lines.extend(
[
"",
"## Missing",
]
)
for item in missing:
lines.append(f"- {item['context_len']} / {item['case']} / {item['phase']} -> `{item['path']}`")
if failed:
lines.extend(
[
"",
"## Failed",
]
)
for item in failed:
lines.append(
f"- {item['context_len']} / {item['case']} / {item['phase']} -> "
f"{item['status']} (exit={item.get('exit_code')})"
)
if prompt_mismatches:
lines.extend(
[
"",
"## Prompt Mismatches",
]
)
for item in prompt_mismatches:
lines.append(
f"- {item['context_len']}: baseline `{item['baseline_prompt_hash']}` vs tq `{item['tq_prompt_hash']}`"
)
return manifest, "\n".join(lines) + "\n"
scripts/tq_phase_runner.py
-399
View File
@@ -1,399 +0,0 @@
#!/usr/bin/env python3
"""Run exactly one TurboQuant telemetry phase and always emit a JSON artifact."""
from __future__ import annotations
import argparse
import json
import os
import signal
import socket
import subprocess
import sys
import time
from pathlib import Path
from typing import Any
from tq_harness_lib import (
DEFAULT_QUALITY_PROMPT,
atomic_write_json,
ensure_dir,
ensure_repo_import_path,
inspect_runtime_stats,
load_prompt_artifacts,
maybe_free_tq_kv,
probe_gpu_metrics,
scrub_gpu_processes,
tail_text,
utc_now_iso,
)
PHASES = ("init", "ttft", "prefill_only", "full", "decode_only", "quality")
CASES = ("baseline", "tq")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--case", choices=CASES, required=True)
parser.add_argument("--phase", choices=PHASES, required=True)
parser.add_argument("--context-len", type=int, required=True)
parser.add_argument("--output", required=True)
parser.add_argument("--timeout-s", type=int, required=True)
parser.add_argument("--prompt-seed", type=int, required=True)
parser.add_argument("--max-output-tokens", type=int, required=True)
parser.add_argument("--model-path", required=True)
parser.add_argument("--prompt-dir", required=True)
parser.add_argument("--campaign-id", default=None)
parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
parser.add_argument("--tensor-parallel-size", type=int, default=1)
parser.add_argument("--trust-remote-code", action="store_true", default=True)
parser.add_argument("--enforce-eager", action="store_true")
parser.add_argument("--max-model-len", type=int, default=None)
parser.add_argument("--heartbeat-interval-s", type=float, default=15.0)
parser.add_argument("--worker-mode", action="store_true", help=argparse.SUPPRESS)
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--dry-run-sleep-s", type=float, default=0.0, help=argparse.SUPPRESS)
return parser.parse_args()
def status_path_for(output_path: Path) -> Path:
return output_path.with_name(output_path.stem + ".status.json")
def worker_command(args: argparse.Namespace) -> list[str]:
cmd = [
sys.executable,
str(Path(__file__).resolve()),
"--worker-mode",
"--case",
args.case,
"--phase",
args.phase,
"--context-len",
str(args.context_len),
"--output",
args.output,
"--timeout-s",
str(args.timeout_s),
"--prompt-seed",
str(args.prompt_seed),
"--max-output-tokens",
str(args.max_output_tokens),
"--model-path",
args.model_path,
"--prompt-dir",
args.prompt_dir,
"--gpu-memory-utilization",
str(args.gpu_memory_utilization),
"--tensor-parallel-size",
str(args.tensor_parallel_size),
"--heartbeat-interval-s",
str(args.heartbeat_interval_s),
]
if args.campaign_id:
cmd.extend(["--campaign-id", args.campaign_id])
if args.enforce_eager:
cmd.append("--enforce-eager")
if args.dry_run:
cmd.append("--dry-run")
if args.dry_run_sleep_s:
cmd.extend(["--dry-run-sleep-s", str(args.dry_run_sleep_s)])
if args.max_model_len is not None:
cmd.extend(["--max-model-len", str(args.max_model_len)])
return cmd
def build_supervisor_payload(args: argparse.Namespace, output_path: Path) -> dict[str, Any]:
return {
"campaign_id": args.campaign_id,
"host": socket.gethostname(),
"case": args.case,
"phase": args.phase,
"context_len": args.context_len,
"prompt_seed": args.prompt_seed,
"max_output_tokens": args.max_output_tokens,
"model_path": args.model_path,
"prompt_dir": args.prompt_dir,
"cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
"output_path": str(output_path),
"status": "running",
"start_ts": utc_now_iso(),
}
def write_heartbeat(status_path: Path, payload: dict[str, Any], process: subprocess.Popen[str]) -> None:
heartbeat = dict(payload)
heartbeat.update(
{
"status": "running",
"supervisor_pid": os.getpid(),
"worker_pid": process.pid,
"heartbeat_ts": utc_now_iso(),
}
)
atomic_write_json(status_path, heartbeat)
def parse_last_json_line(text: str) -> dict[str, Any] | None:
for line in reversed(text.strip().splitlines()):
line = line.strip()
if not line:
continue
try:
return json.loads(line)
except json.JSONDecodeError:
continue
return None
def run_supervisor(args: argparse.Namespace) -> int:
output_path = Path(args.output).resolve()
status_path = status_path_for(output_path)
ensure_dir(output_path.parent)
payload = build_supervisor_payload(args, output_path)
process = subprocess.Popen(
worker_command(args),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
env=os.environ.copy(),
)
start_monotonic = time.monotonic()
next_heartbeat = start_monotonic
timed_out = False
while True:
now = time.monotonic()
if now >= next_heartbeat:
write_heartbeat(status_path, payload, process)
next_heartbeat = now + args.heartbeat_interval_s
rc = process.poll()
if rc is not None:
break
if now - start_monotonic > args.timeout_s:
timed_out = True
process.terminate()
try:
process.wait(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout=5)
break
time.sleep(0.5)
stdout, stderr = process.communicate()
worker_payload = parse_last_json_line(stdout)
end_ts = utc_now_iso()
elapsed_s = round(time.monotonic() - start_monotonic, 6)
if worker_payload:
payload.update(worker_payload)
payload["stdout_tail"] = tail_text(stdout)
payload["stderr_tail"] = tail_text(stderr)
payload["exit_code"] = process.returncode
payload["end_ts"] = end_ts
payload["elapsed_s"] = payload.get("elapsed_s", elapsed_s)
if timed_out:
payload["status"] = "timeout"
payload["error_type"] = "phase_timeout"
elif process.returncode == 137 or process.returncode == -signal.SIGKILL:
payload["status"] = "killed"
payload["error_type"] = "exit_137"
elif process.returncode and payload.get("status") == "ok":
payload["status"] = "error"
payload["error_type"] = payload.get("error_type", "worker_nonzero_exit")
elif not worker_payload:
payload["status"] = "error"
payload["error_type"] = "missing_worker_json"
atomic_write_json(output_path, payload)
atomic_write_json(status_path, payload)
return 0 if payload["status"] == "ok" else 1
def compute_max_model_len(args: argparse.Namespace) -> int:
if args.max_model_len is not None:
return args.max_model_len
reserve = max(512, args.max_output_tokens + 256)
return args.context_len + reserve
def worker_main(args: argparse.Namespace) -> int:
try:
ensure_repo_import_path()
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
prompt_text, prompt_meta = load_prompt_artifacts(Path(args.prompt_dir))
prompt_hash = prompt_meta.get("prompt_hash")
prompt_tokens = prompt_meta.get("prompt_tokens")
init_gpu = probe_gpu_metrics()
result: dict[str, Any] = {
"campaign_id": args.campaign_id,
"host": socket.gethostname(),
"case": args.case,
"phase": args.phase,
"context_len": args.context_len,
"prompt_seed": args.prompt_seed,
"prompt_hash": prompt_hash,
"prompt_tokens": prompt_tokens,
"cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
"gpu_name": init_gpu.get("gpu_name"),
"memory_used_mb": init_gpu.get("memory_used_mb"),
"memory_free_mb": init_gpu.get("memory_free_mb"),
"power_w": init_gpu.get("power_w"),
"temp_c": init_gpu.get("temp_c"),
"start_ts": utc_now_iso(),
"status": "ok",
"stdout_tail": "",
"stderr_tail": "",
"error_type": None,
}
if args.dry_run:
if args.dry_run_sleep_s:
time.sleep(args.dry_run_sleep_s)
result.update(
{
"init_s": 0.123,
"ttft_s": 0.456 if args.phase in {"ttft", "prefill_only"} else None,
"prefill_tok_s": float(prompt_tokens or 0) / 0.456 if args.phase in {"ttft", "prefill_only"} else None,
"gen_tok_s": float(args.max_output_tokens) / 1.5 if args.phase in {"full", "decode_only", "quality"} else None,
"sample_text": "dry run output" if args.phase in {"full", "decode_only", "quality"} else "",
"quality_prompt": DEFAULT_QUALITY_PROMPT if args.phase == "quality" else None,
"status": "ok",
"end_ts": utc_now_iso(),
"elapsed_s": 0.789,
}
)
print(json.dumps(result))
return 0
init_start = time.perf_counter()
if args.case == "tq":
from turboquant.vllm_attn_backend import enable_no_alloc
enable_no_alloc(key_bits=3, value_bits=2, buffer_size=128, initial_layers_count=4)
from vllm import LLM, SamplingParams
llm = LLM(
model=args.model_path,
trust_remote_code=args.trust_remote_code,
gpu_memory_utilization=args.gpu_memory_utilization,
max_model_len=compute_max_model_len(args),
tensor_parallel_size=args.tensor_parallel_size,
max_num_seqs=1,
enforce_eager=args.enforce_eager,
)
init_elapsed = time.perf_counter() - init_start
post_init_gpu = probe_gpu_metrics()
runtime_stats = inspect_runtime_stats(llm)
result.update(
{
"init_s": round(init_elapsed, 6),
"memory_used_mb": post_init_gpu.get("memory_used_mb"),
"memory_free_mb": post_init_gpu.get("memory_free_mb"),
"power_w": post_init_gpu.get("power_w"),
"temp_c": post_init_gpu.get("temp_c"),
"gpu_name": post_init_gpu.get("gpu_name"),
"tq_hooked_layers": runtime_stats.get("tq_hooked_layers"),
"shared_kv_layers": runtime_stats.get("shared_kv_layers"),
"kv_reserved_gb": runtime_stats.get("kv_reserved_gb"),
"tq_total_memory_bytes": runtime_stats.get("tq_total_memory_bytes"),
"tq_total_compressed_tokens": runtime_stats.get("tq_total_compressed_tokens"),
"tq_mode": runtime_stats.get("tq_mode"),
}
)
if args.phase == "init":
result["elapsed_s"] = round(init_elapsed, 6)
result["end_ts"] = utc_now_iso()
print(json.dumps(result))
return 0
run_prompt = prompt_text
max_tokens = args.max_output_tokens
if args.phase == "quality":
run_prompt = DEFAULT_QUALITY_PROMPT
max_tokens = max(args.max_output_tokens, 64)
result["quality_prompt"] = DEFAULT_QUALITY_PROMPT
elif args.phase in {"ttft", "prefill_only"}:
max_tokens = 1
if args.phase == "prefill_only":
result["phase_note"] = (
"Prefill-only probe requests a single output token so elapsed time is dominated by long-context prefill."
)
elif args.phase == "decode_only":
result["phase_note"] = (
"Decode-only is an isolated multi-token generate run. "
"Because vLLM does not expose reusable prefill state across processes here, this is decode-dominant rather than perfectly decode-exclusive."
)
sample_start = time.perf_counter()
outputs = llm.generate([run_prompt], SamplingParams(temperature=0, max_tokens=max_tokens))
sample_elapsed = time.perf_counter() - sample_start
response = outputs[0].outputs[0]
post_phase_gpu = probe_gpu_metrics()
result["sample_text"] = response.text
result["output_tokens"] = len(response.token_ids)
result["elapsed_s"] = round(sample_elapsed, 6)
result["end_ts"] = utc_now_iso()
result["activation_est_mb"] = (
post_phase_gpu.get("memory_used_mb") - post_init_gpu.get("memory_used_mb")
if post_phase_gpu.get("memory_used_mb") is not None
and post_init_gpu.get("memory_used_mb") is not None
else None
)
result["memory_used_mb"] = post_phase_gpu.get("memory_used_mb")
result["memory_free_mb"] = post_phase_gpu.get("memory_free_mb")
result["power_w"] = post_phase_gpu.get("power_w")
result["temp_c"] = post_phase_gpu.get("temp_c")
if args.phase in {"ttft", "prefill_only"}:
result["ttft_s"] = round(sample_elapsed, 6)
result["prefill_tok_s"] = round((prompt_tokens or 0) / sample_elapsed, 6) if sample_elapsed else None
elif args.phase in {"full", "decode_only", "quality"}:
result["gen_tok_s"] = round(len(response.token_ids) / sample_elapsed, 6) if sample_elapsed else None
if args.case == "tq":
result.update(maybe_free_tq_kv(llm))
post_free_gpu = probe_gpu_metrics()
result["memory_after_free_mb"] = post_free_gpu.get("memory_used_mb")
result["memory_free_after_free_mb"] = post_free_gpu.get("memory_free_mb")
print(json.dumps(result))
return 0
except Exception as exc:
error_payload = {
"campaign_id": args.campaign_id,
"case": args.case,
"phase": args.phase,
"context_len": args.context_len,
"prompt_seed": args.prompt_seed,
"status": "error",
"error_type": exc.__class__.__name__,
"error_message": str(exc),
"end_ts": utc_now_iso(),
}
print(json.dumps(error_payload))
return 1
def main() -> int:
args = parse_args()
if args.worker_mode:
return worker_main(args)
return run_supervisor(args)
if __name__ == "__main__":
raise SystemExit(main())
scripts/tq_remote_snapshot.py
-236
View File
@@ -1,236 +0,0 @@
#!/usr/bin/env python3
"""Capture a remote TurboQuant campaign snapshot over SSH into local JSON."""
from __future__ import annotations
import argparse
import json
import subprocess
from pathlib import Path
from tq_harness_lib import atomic_write_json, atomic_write_text, ensure_dir, tail_text, utc_now_iso
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--host", default="your-gpu-host.example.com")
parser.add_argument("--port", type=int, default=3003)
parser.add_argument("--user", default="root")
parser.add_argument("--identity-file", default="~/.ssh/id_rsa")
parser.add_argument("--campaign-root", required=True)
parser.add_argument("--output", required=True)
parser.add_argument("--mirror-root", default=None)
parser.add_argument("--mode", choices=("ping", "summary", "full"), default="full")
parser.add_argument("--ssh-timeout-s", type=int, default=30)
return parser.parse_args()
def run_ssh(args: argparse.Namespace, remote_cmd: str) -> str:
proc = subprocess.run(
[
"ssh",
"-F",
"/dev/null",
"-o",
"StrictHostKeyChecking=no",
"-o",
"ConnectTimeout=12",
"-o",
f"IdentityFile={args.identity_file}",
"-o",
"IdentitiesOnly=yes",
f"{args.user}@{args.host}",
"-p",
str(args.port),
remote_cmd,
],
capture_output=True,
text=True,
check=True,
timeout=args.ssh_timeout_s,
)
return proc.stdout
def write_mirror_files(payload: dict, mirror_root: Path) -> None:
ensure_dir(mirror_root)
manifest = payload.get("manifest")
if manifest is not None:
atomic_write_json(mirror_root / "manifest.json", manifest)
summary_text = payload.get("summary_text")
if summary_text:
atomic_write_text(mirror_root / "summary.md", summary_text)
for entry in payload.get("files", []):
rel_path = entry.get("relative_path")
if not rel_path:
continue
target = mirror_root / rel_path
text = entry.get("text")
if text is None:
continue
if entry.get("type") == "json":
atomic_write_json(target, json.loads(text))
else:
atomic_write_text(target, text)
def build_remote_python(root: str, mode: str) -> str:
if mode == "ping":
return f"""
import json
import socket
from pathlib import Path
root = Path({root!r})
payload = {{
"campaign_root": str(root),
"generated_at": None,
"manifest": None,
"summary_head": None,
"summary_text": None,
"phase_files": [],
"files": [],
"mode": "ping",
"remote_host": socket.gethostname(),
"campaign_root_exists": root.exists(),
}}
print(json.dumps(payload))
"""
include_files = mode == "full"
return f"""
import json
from pathlib import Path
root = Path({root!r})
include_files = {include_files!r}
payload = {{
"campaign_root": str(root),
"generated_at": None,
"manifest": None,
"summary_head": None,
"summary_text": None,
"phase_files": [],
"files": [],
"mode": {mode!r},
}}
manifest_path = root / "manifest.json"
summary_path = root / "summary.md"
if manifest_path.exists():
manifest_text = manifest_path.read_text()
payload["manifest"] = json.loads(manifest_text)
if include_files:
payload["files"].append({{
"path": str(manifest_path),
"relative_path": str(manifest_path.relative_to(root)),
"type": "json",
"text": manifest_text,
}})
if summary_path.exists():
summary_text = summary_path.read_text()
payload["summary_text"] = summary_text
payload["summary_head"] = summary_text.splitlines()[:40]
if include_files:
payload["files"].append({{
"path": str(summary_path),
"relative_path": str(summary_path.relative_to(root)),
"type": "text",
"text": summary_text,
}})
for path in sorted(root.glob("*/*/*.json")):
if path.name.endswith(".prescrub.json"):
continue
entry = {{
"path": str(path),
"kind": "status" if path.name.endswith(".status.json") else "phase",
}}
try:
text = path.read_text()
data = json.loads(text)
for key in ("status", "elapsed_s", "ttft_s", "prefill_tok_s", "gen_tok_s", "error_type", "exit_code", "worker_pid", "heartbeat_ts"):
if key in data:
entry[key] = data.get(key)
if include_files:
payload["files"].append({{
"path": str(path),
"relative_path": str(path.relative_to(root)),
"type": "json",
"text": text,
}})
except Exception as exc:
entry["error"] = str(exc)
payload["phase_files"].append(entry)
if include_files:
for path in sorted(root.glob("*/*/*")):
if not path.is_file():
continue
if path.suffix not in {{".txt", ".md"}}:
continue
payload["files"].append({{
"path": str(path),
"relative_path": str(path.relative_to(root)),
"type": "text",
"text": path.read_text(),
}})
print(json.dumps(payload))
"""
def main() -> int:
args = parse_args()
root = args.campaign_root
output_path = Path(args.output)
fallback_payload = {
"campaign_root": root,
"generated_at": utc_now_iso(),
"manifest": None,
"summary_head": None,
"summary_text": None,
"phase_files": [],
"files": [],
}
remote_python = build_remote_python(root, args.mode)
try:
stdout = run_ssh(args, f"python3 - <<'PY'\n{remote_python}\nPY")
payload = json.loads(stdout)
payload["generated_at"] = utc_now_iso()
atomic_write_json(output_path, payload)
except subprocess.TimeoutExpired as exc:
fallback_payload.update(
{
"status": "error",
"error_type": "ssh_timeout",
"error_message": str(exc),
"ssh_timeout_s": args.ssh_timeout_s,
}
)
atomic_write_json(output_path, fallback_payload)
print(args.output)
return 1
except subprocess.CalledProcessError as exc:
fallback_payload.update(
{
"status": "error",
"error_type": "ssh_failed",
"error_message": str(exc),
"stdout_tail": tail_text(exc.stdout),
"stderr_tail": tail_text(exc.stderr),
}
)
atomic_write_json(output_path, fallback_payload)
print(args.output)
return 1
if args.mirror_root:
write_mirror_files(payload, Path(args.mirror_root))
print(args.output)
return 0
if __name__ == "__main__":
raise SystemExit(main())
scripts/tq_resume_recovery.py
-237
View File
@@ -1,237 +0,0 @@
#!/usr/bin/env python3
"""Wait for remote recovery, sync the harness, and launch the Plan B campaign."""
from __future__ import annotations
import argparse
import json
import math
import subprocess
import time
from pathlib import Path
from types import SimpleNamespace
from tq_campaign import sync_remote
from tq_harness_lib import atomic_write_json, tail_text, utc_now_iso
from tq_remote_snapshot import build_remote_python, run_ssh
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--host", default="your-gpu-host.example.com")
parser.add_argument("--port", type=int, default=3003)
parser.add_argument("--user", default="root")
parser.add_argument("--identity-file", default="~/.ssh/id_rsa")
parser.add_argument("--ssh-timeout-s", type=int, default=20)
parser.add_argument("--probe-interval-s", type=int, default=30)
parser.add_argument("--max-wait-s", type=int, default=300)
parser.add_argument("--remote-scripts-dir", default="/5090-qwen3.5-27b/scripts")
parser.add_argument("--remote-model-path", default="/5090-qwen3.5-27b/models/QuantTrio-Qwen3.5-27B-AWQ")
parser.add_argument("--probe-campaign-root", default="/5090-qwen3.5-27b/logs/campaigns/smoke-50k-harness")
parser.add_argument("--recovery-campaign-root", default="/5090-qwen3.5-27b/logs/campaigns/smoke-50k-planb")
parser.add_argument("--contexts", default="50000")
parser.add_argument("--cases", default="baseline,tq")
parser.add_argument("--phases", default="init,prefill_only,decode_only,full")
parser.add_argument("--prompt-seed", type=int, default=5090)
parser.add_argument("--max-output-tokens", type=int, default=24)
parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
parser.add_argument("--tensor-parallel-size", type=int, default=1)
parser.add_argument("--output", required=True)
parser.add_argument("--dry-run", action="store_true")
return parser.parse_args()
def probe_remote(args: argparse.Namespace, mode: str, campaign_root: str) -> dict:
payload = {
"mode": mode,
"campaign_root": campaign_root,
"ok": False,
"checked_at": utc_now_iso(),
}
remote_python = build_remote_python(campaign_root, mode)
try:
stdout = run_ssh(args, f"python3 - <<'PY'\n{remote_python}\nPY")
data = json.loads(stdout)
payload["ok"] = True
payload["payload"] = data
except subprocess.TimeoutExpired as exc:
payload["error_type"] = "ssh_timeout"
payload["error_message"] = str(exc)
except subprocess.CalledProcessError as exc:
payload["error_type"] = "ssh_failed"
payload["error_message"] = str(exc)
payload["stdout_tail"] = tail_text(exc.stdout)
payload["stderr_tail"] = tail_text(exc.stderr)
return payload
def build_remote_launch_command(args: argparse.Namespace) -> str:
launch_log = f"{args.recovery_campaign_root}/launch.log"
return (
f"mkdir -p {args.recovery_campaign_root} && "
f"cd {args.remote_scripts_dir} && "
"nohup env "
"CUDA_VISIBLE_DEVICES=0 "
"VLLM_ENABLE_V1_MULTIPROCESSING=0 "
"TOKENIZERS_PARALLELISM=false "
f"/5090-qwen3.5-27b/.venv/bin/python tq_campaign.py "
f"--model-path {args.remote_model_path} "
f"--contexts {args.contexts} "
f"--cases {args.cases} "
f"--phases {args.phases} "
f"--campaign-root {args.recovery_campaign_root} "
f"--prompt-seed {args.prompt_seed} "
f"--max-output-tokens {args.max_output_tokens} "
f"--gpu-memory-utilization {args.gpu_memory_utilization} "
f"--tensor-parallel-size {args.tensor_parallel_size} "
"--strict-timeouts --force "
f"> {launch_log} 2>&1 < /dev/null & echo $!"
)
def run_remote_command(args: argparse.Namespace, remote_cmd: str) -> subprocess.CompletedProcess[str]:
return subprocess.run(
[
"ssh",
"-F",
"/dev/null",
"-o",
"StrictHostKeyChecking=no",
"-o",
"ConnectTimeout=12",
"-o",
f"IdentityFile={args.identity_file}",
"-o",
"IdentitiesOnly=yes",
f"{args.user}@{args.host}",
"-p",
str(args.port),
remote_cmd,
],
capture_output=True,
text=True,
timeout=args.ssh_timeout_s,
check=True,
)
def main() -> int:
args = parse_args()
output_path = Path(args.output)
report = {
"started_at": utc_now_iso(),
"status": "error",
"host": args.host,
"port": args.port,
"probe_campaign_root": args.probe_campaign_root,
"recovery_campaign_root": args.recovery_campaign_root,
"phases": args.phases.split(","),
"checks": [],
"sync": None,
"launch": None,
}
attempts = max(1, math.ceil(args.max_wait_s / max(args.probe_interval_s, 1)))
for attempt in range(1, attempts + 1):
probe = probe_remote(args, "ping", args.probe_campaign_root)
probe["attempt"] = attempt
report["checks"].append(probe)
if probe["ok"]:
break
if attempt < attempts and not args.dry_run:
time.sleep(args.probe_interval_s)
if not report["checks"][-1]["ok"]:
report["error_type"] = report["checks"][-1].get("error_type", "probe_failed")
report["error_message"] = report["checks"][-1].get("error_message", "Ping probe did not recover")
report["ended_at"] = utc_now_iso()
atomic_write_json(output_path, report)
print(output_path)
return 1
summary_probe = probe_remote(args, "summary", args.probe_campaign_root)
report["checks"].append(summary_probe)
if not summary_probe["ok"]:
report["error_type"] = summary_probe.get("error_type", "summary_probe_failed")
report["error_message"] = summary_probe.get("error_message", "Summary probe failed")
report["ended_at"] = utc_now_iso()
atomic_write_json(output_path, report)
print(output_path)
return 1
if args.dry_run:
report["sync"] = {"status": "skipped", "reason": "dry_run"}
report["launch"] = {
"status": "skipped",
"reason": "dry_run",
"remote_command": build_remote_launch_command(args),
}
report["status"] = "ok"
report["ended_at"] = utc_now_iso()
atomic_write_json(output_path, report)
print(output_path)
return 0
sync_args = SimpleNamespace(
remote_host=args.host,
remote_port=args.port,
remote_user=args.user,
remote_key=args.identity_file,
remote_scripts_dir=args.remote_scripts_dir,
)
try:
sync_remote(sync_args)
report["sync"] = {"status": "ok", "synced_at": utc_now_iso()}
except Exception as exc: # noqa: BLE001
report["sync"] = {"status": "error", "error_type": exc.__class__.__name__, "error_message": str(exc)}
report["error_type"] = "sync_failed"
report["error_message"] = str(exc)
report["ended_at"] = utc_now_iso()
atomic_write_json(output_path, report)
print(output_path)
return 1
launch_cmd = build_remote_launch_command(args)
try:
proc = run_remote_command(args, launch_cmd)
report["launch"] = {
"status": "ok",
"launched_at": utc_now_iso(),
"remote_command": launch_cmd,
"stdout_tail": tail_text(proc.stdout),
"stderr_tail": tail_text(proc.stderr),
"worker_pid": proc.stdout.strip().splitlines()[-1] if proc.stdout.strip() else None,
"launch_log": f"{args.recovery_campaign_root}/launch.log",
}
except subprocess.TimeoutExpired as exc:
report["launch"] = {"status": "error", "error_type": "ssh_timeout", "error_message": str(exc)}
report["error_type"] = "launch_timeout"
report["error_message"] = str(exc)
report["ended_at"] = utc_now_iso()
atomic_write_json(output_path, report)
print(output_path)
return 1
except subprocess.CalledProcessError as exc:
report["launch"] = {
"status": "error",
"error_type": "ssh_failed",
"error_message": str(exc),
"stdout_tail": tail_text(exc.stdout),
"stderr_tail": tail_text(exc.stderr),
}
report["error_type"] = "launch_failed"
report["error_message"] = str(exc)
report["ended_at"] = utc_now_iso()
atomic_write_json(output_path, report)
print(output_path)
return 1
report["status"] = "ok"
report["ended_at"] = utc_now_iso()
atomic_write_json(output_path, report)
print(output_path)
return 0
if __name__ == "__main__":
raise SystemExit(main())
setup.py
+1 -1
View File
@@ -25,7 +25,7 @@ setup(
"Development Status :: 3 - Alpha",
"Intended Audience :: Science/Research",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"License :: OSI Approved :: MIT License",
"License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
"Programming Language :: Python :: 3",
],
)
test_modular.py
-566
View File
@@ -1,566 +0,0 @@
#!/usr/bin/env python3
"""
TurboQuant modular architecture tests.
Tests:
1. RingBuffer — write, overflow, drain, reset
2. CompressedKVStore — append_chunk, flat cache, chunked growth
3. KVCaptureEngine — prefill, decode, flush orchestration
4. compute_hybrid_attention — compressed + exact merge
5. Attention recall agreement — TQ vs exact top-k recall
6. Retrieval accuracy — needle-in-haystack style
"""
import math
import torch
import numpy as np
torch.manual_seed(42)
np.random.seed(42)
PASS = 0
FAIL = 0
def run_test(name, fn, **kwargs):
global PASS, FAIL
try:
fn(**kwargs)
print(f" PASS {name}")
PASS += 1
except Exception as e:
import traceback
print(f" FAIL {name}")
traceback.print_exc()
FAIL += 1
# ── Test 1: RingBuffer ───────────────────────────────────────────────
def test_ring_buffer_basic():
from turboquant.capture import RingBuffer
buf = RingBuffer(capacity=8, num_kv_heads=2, head_dim=4, device="cpu")
assert buf.size == 0
assert not buf.is_full
k = torch.randn(3, 2, 4)
v = torch.randn(3, 2, 4)
overflow = buf.write(k, v, 3)
assert overflow is None
assert buf.size == 3
data = buf.peek()
assert data is not None
assert data[0].shape == (3, 2, 4)
def test_ring_buffer_overflow():
from turboquant.capture import RingBuffer
buf = RingBuffer(capacity=4, num_kv_heads=2, head_dim=4, device="cpu")
k = torch.randn(6, 2, 4)
v = torch.randn(6, 2, 4)
overflow = buf.write(k, v, 6)
assert overflow is not None
ok, ov = overflow
assert ok.shape[0] == 4 # one full buffer drained
assert buf.size == 2 # 2 remaining in buffer
def test_ring_buffer_drain():
from turboquant.capture import RingBuffer
buf = RingBuffer(capacity=8, num_kv_heads=2, head_dim=4, device="cpu")
k = torch.randn(5, 2, 4)
v = torch.randn(5, 2, 4)
buf.write(k, v, 5)
data = buf.drain()
assert data is not None
assert data[0].shape == (5, 2, 4)
assert buf.size == 0
def test_ring_buffer_large_overflow():
"""Write more than 2x capacity in one call."""
from turboquant.capture import RingBuffer
buf = RingBuffer(capacity=4, num_kv_heads=1, head_dim=2, device="cpu")
k = torch.randn(11, 1, 2)
v = torch.randn(11, 1, 2)
overflow = buf.write(k, v, 11)
assert overflow is not None
ok, ov = overflow
# Two full buffers drained (4+4=8), 3 remain in buffer
assert ok.shape[0] == 8
assert buf.size == 3
# ── Test 2: CompressedKVStore ─────────────────────────────────────────
def test_store_basic():
from turboquant.store import CompressedKVStore
store = CompressedKVStore(
head_dim=64, num_kv_heads=4, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"), layer_idx=0,
)
assert store.num_tokens == 0
assert store.num_chunks == 0
k = torch.randn(32, 4, 64)
v = torch.randn(32, 4, 64)
store.append_chunk(k, v)
assert store.num_tokens == 32
assert store.num_chunks == 1
flat = store.get_flat_cache()
assert flat is not None
assert flat.num_tokens == 32
def test_store_multi_chunk():
from turboquant.store import CompressedKVStore
store = CompressedKVStore(
head_dim=64, num_kv_heads=2, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
for _ in range(3):
k = torch.randn(16, 2, 64)
v = torch.randn(16, 2, 64)
store.append_chunk(k, v)
assert store.num_tokens == 48
assert store.num_chunks == 3
flat = store.get_flat_cache()
assert flat is not None
assert flat.num_tokens == 48
# Second call should use cache
flat2 = store.get_flat_cache()
assert flat2 is flat
def test_store_invalidation():
from turboquant.store import CompressedKVStore
store = CompressedKVStore(
head_dim=64, num_kv_heads=2, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
store.append_chunk(torch.randn(16, 2, 64), torch.randn(16, 2, 64))
flat1 = store.get_flat_cache()
store.append_chunk(torch.randn(16, 2, 64), torch.randn(16, 2, 64))
flat2 = store.get_flat_cache()
assert flat2 is not flat1
assert flat2.num_tokens == 32
def test_store_memory():
from turboquant.store import CompressedKVStore
store = CompressedKVStore(
head_dim=128, num_kv_heads=8, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
k = torch.randn(512, 8, 128)
v = torch.randn(512, 8, 128)
store.append_chunk(k, v)
mem = store.memory_bytes()
fp16_mem = 512 * 8 * 128 * 2 * 2 # K+V in FP16
assert mem < fp16_mem, f"TQ mem {mem} should be < FP16 mem {fp16_mem}"
ratio = fp16_mem / mem
assert ratio > 2.0, f"Compression ratio {ratio:.1f}x should be > 2x"
# ── Test 3: KVCaptureEngine ──────────────────────────────────────────
def test_capture_engine_prefill():
from turboquant.capture import KVCaptureEngine
from turboquant.store import CompressedKVStore
store = CompressedKVStore(
head_dim=64, num_kv_heads=2, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
engine = KVCaptureEngine(store, ring_capacity=32, device=torch.device("cpu"))
k = torch.randn(100, 2, 64)
v = torch.randn(100, 2, 64)
engine.ingest_prefill(k, v, 100)
# 100 - 32 = 68 compressed, 32 buffered
assert engine.total_compressed_tokens == 68
assert engine.total_buffered_tokens == 32
assert engine.total_tokens == 100
def test_capture_engine_decode():
from turboquant.capture import KVCaptureEngine
from turboquant.store import CompressedKVStore
store = CompressedKVStore(
head_dim=64, num_kv_heads=2, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
engine = KVCaptureEngine(store, ring_capacity=8, device=torch.device("cpu"))
# Fill ring
for i in range(10):
k = torch.randn(1, 2, 64)
v = torch.randn(1, 2, 64)
engine.ingest_decode(k, v, 1)
# After 10 decode tokens with ring=8: 8 overflowed -> compressed, 2 in buffer
assert engine.total_compressed_tokens == 8
assert engine.total_buffered_tokens == 2
def test_capture_engine_flush():
from turboquant.capture import KVCaptureEngine
from turboquant.store import CompressedKVStore
store = CompressedKVStore(
head_dim=64, num_kv_heads=2, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
engine = KVCaptureEngine(store, ring_capacity=16, device=torch.device("cpu"))
k = torch.randn(5, 2, 64)
v = torch.randn(5, 2, 64)
engine.ingest_decode(k, v, 5)
assert engine.total_buffered_tokens == 5
assert engine.total_compressed_tokens == 0
engine.flush()
assert engine.total_buffered_tokens == 0
assert engine.total_compressed_tokens == 5
# ── Test 4: Hybrid attention ─────────────────────────────────────────
def test_hybrid_attention_compressed_only():
from turboquant.store import CompressedKVStore
from turboquant.score import compute_hybrid_attention
d = 64
H_kv = 2
Q = 4 # num_query_heads
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
k = torch.randn(64, H_kv, d)
v = torch.randn(64, H_kv, d)
store.append_chunk(k, v)
query = torch.randn(1, Q, d)
out = compute_hybrid_attention(
query=query, store=store, recent_k=None, recent_v=None,
num_query_heads=Q,
)
assert out.shape == (1, Q, d)
def test_hybrid_attention_exact_only():
from turboquant.store import CompressedKVStore
from turboquant.score import compute_hybrid_attention
d = 64
H_kv = 2
Q = 4
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
# No compressed data, just recent buffer
recent_k = torch.randn(16, H_kv, d)
recent_v = torch.randn(16, H_kv, d)
query = torch.randn(1, Q, d)
out = compute_hybrid_attention(
query=query, store=store, recent_k=recent_k, recent_v=recent_v,
num_query_heads=Q,
)
assert out.shape == (1, Q, d)
def test_hybrid_attention_both():
from turboquant.store import CompressedKVStore
from turboquant.score import compute_hybrid_attention
d = 64
H_kv = 2
Q = 4
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
k_hist = torch.randn(64, H_kv, d)
v_hist = torch.randn(64, H_kv, d)
store.append_chunk(k_hist, v_hist)
recent_k = torch.randn(8, H_kv, d)
recent_v = torch.randn(8, H_kv, d)
query = torch.randn(1, Q, d)
out = compute_hybrid_attention(
query=query, store=store, recent_k=recent_k, recent_v=recent_v,
num_query_heads=Q,
)
assert out.shape == (1, Q, d)
# ── Test 5: Attention recall agreement ────────────────────────────────
def test_attention_recall():
"""Verify that TQ-compressed keys produce similar attention ranking to exact keys.
Measures recall@k: what fraction of the true top-k keys by attention score
are also in the TQ top-k. Target: >= 0.90 for 3-bit at k=8.
"""
from turboquant.quantizer import TurboQuantProd
d = 128
N = 256
n_queries = 16
k = 8
bits = 3
device = torch.device("cpu")
quantizer = TurboQuantProd(dim=d, bits=bits, device=device, seed=42)
keys = torch.randn(1, 1, N, d, device=device) * 0.1
queries = torch.randn(1, 1, n_queries, d, device=device) * 0.1
# True scores
true_scores = torch.matmul(queries, keys.transpose(-2, -1)).squeeze(0).squeeze(0) # (n_q, N)
true_topk = true_scores.topk(k, dim=-1).indices # (n_q, k)
# TQ scores
key_q = quantizer.quantize(keys)
tq_scores = quantizer.attention_score(queries, key_q).squeeze(0).squeeze(0)
tq_topk = tq_scores.topk(k, dim=-1).indices
# Compute recall@k
recalls = []
for q_idx in range(n_queries):
true_set = set(true_topk[q_idx].tolist())
tq_set = set(tq_topk[q_idx].tolist())
recall = len(true_set & tq_set) / k
recalls.append(recall)
mean_recall = sum(recalls) / len(recalls)
assert mean_recall >= 0.50, f"Mean recall@{k} = {mean_recall:.3f} < 0.50"
def test_attention_recall_4bit():
"""4-bit should achieve very high recall."""
from turboquant.quantizer import TurboQuantProd
d = 128
N = 256
n_queries = 16
k = 8
bits = 4
device = torch.device("cpu")
quantizer = TurboQuantProd(dim=d, bits=bits, device=device, seed=42)
keys = torch.randn(1, 1, N, d, device=device) * 0.1
queries = torch.randn(1, 1, n_queries, d, device=device) * 0.1
true_scores = torch.matmul(queries, keys.transpose(-2, -1)).squeeze(0).squeeze(0)
true_topk = true_scores.topk(k, dim=-1).indices
key_q = quantizer.quantize(keys)
tq_scores = quantizer.attention_score(queries, key_q).squeeze(0).squeeze(0)
tq_topk = tq_scores.topk(k, dim=-1).indices
recalls = []
for q_idx in range(n_queries):
true_set = set(true_topk[q_idx].tolist())
tq_set = set(tq_topk[q_idx].tolist())
recalls.append(len(true_set & tq_set) / k)
mean_recall = sum(recalls) / len(recalls)
assert mean_recall >= 0.65, f"4-bit recall@{k} = {mean_recall:.3f} < 0.65"
# ── Test 6: Needle retrieval (compressed store) ──────────────────────
def test_needle_retrieval():
"""Place a distinctive key in a sea of noise. Verify TQ compressed store
still ranks the needle's position highest for a matching query."""
from turboquant.store import CompressedKVStore
d = 64
H_kv = 1
N = 200
needle_pos = 137
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
keys = torch.randn(N, H_kv, d) * 0.05
values = torch.randn(N, H_kv, d) * 0.05
# Plant a strong needle signal
needle_key = torch.randn(1, H_kv, d) * 2.0
keys[needle_pos] = needle_key.squeeze(0)
values[needle_pos] = torch.ones(H_kv, d)
store.append_chunk(keys, values)
flat = store.get_flat_cache()
assert flat is not None
# Query that matches the needle
query_vec = needle_key.squeeze(0).unsqueeze(0) # (1, H_kv, d)
# Dequantize and compute scores
k_dequant = store.quantizer.dequantize(flat.prod_q) # (H_kv, N, d)
scores = torch.bmm(
query_vec.float().transpose(0, 1), # (H_kv, 1, d)
k_dequant.float().transpose(1, 2), # (H_kv, d, N)
).squeeze(1) # (H_kv, N)
top_idx = scores.argmax(dim=-1).item()
assert top_idx == needle_pos, f"Needle at {needle_pos} but top score at {top_idx}"
def test_needle_retrieval_multi_chunk():
"""Needle retrieval across multiple chunks."""
from turboquant.store import CompressedKVStore
d = 64
H_kv = 2
needle_pos_in_chunk2 = 15
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=4, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
# Chunk 1: noise
store.append_chunk(torch.randn(32, H_kv, d) * 0.05, torch.randn(32, H_kv, d))
# Chunk 2: needle at position 15
keys2 = torch.randn(32, H_kv, d) * 0.05
needle_key = torch.randn(1, H_kv, d) * 3.0
keys2[needle_pos_in_chunk2] = needle_key.squeeze(0)
store.append_chunk(keys2, torch.randn(32, H_kv, d))
# Chunk 3: noise
store.append_chunk(torch.randn(32, H_kv, d) * 0.05, torch.randn(32, H_kv, d))
flat = store.get_flat_cache()
k_dequant = store.quantizer.dequantize(flat.prod_q)
query_vec = needle_key.squeeze(0).unsqueeze(0)
scores = torch.bmm(
query_vec.float().transpose(0, 1),
k_dequant.float().transpose(1, 2),
).squeeze(1)
# Needle is at global position 32 + 15 = 47
expected_pos = 32 + needle_pos_in_chunk2
# Check per head
for h in range(H_kv):
top_idx = scores[h].argmax().item()
assert top_idx == expected_pos, \
f"Head {h}: needle at {expected_pos} but top at {top_idx}"
# ── Test 7: Score rank correlation ────────────────────────────────────
def test_score_rank_correlation():
"""Spearman rank correlation between exact and TQ scores should be high."""
from turboquant.quantizer import TurboQuantProd
d = 128
N = 128
bits = 3
device = torch.device("cpu")
quantizer = TurboQuantProd(dim=d, bits=bits, device=device, seed=42)
keys = torch.randn(1, 1, N, d, device=device) * 0.1
query = torch.randn(1, 1, 1, d, device=device) * 0.1
true_scores = torch.matmul(query, keys.transpose(-2, -1)).squeeze()
key_q = quantizer.quantize(keys)
tq_scores = quantizer.attention_score(query, key_q).squeeze()
# Spearman rank correlation
true_ranks = true_scores.argsort().argsort().float()
tq_ranks = tq_scores.argsort().argsort().float()
rank_corr = torch.corrcoef(torch.stack([true_ranks, tq_ranks]))[0, 1].item()
assert rank_corr > 0.80, f"Rank correlation {rank_corr:.3f} < 0.80"
# ── Main ──────────────────────────────────────────────────────────────
if __name__ == "__main__":
print()
print("=" * 60)
print("TurboQuant Modular Architecture Tests")
print("=" * 60)
print()
print("-- RingBuffer --")
run_test("basic write/peek", test_ring_buffer_basic)
run_test("overflow", test_ring_buffer_overflow)
run_test("drain", test_ring_buffer_drain)
run_test("large overflow (>2x capacity)", test_ring_buffer_large_overflow)
print()
print("-- CompressedKVStore --")
run_test("single chunk", test_store_basic)
run_test("multi chunk + caching", test_store_multi_chunk)
run_test("flat cache invalidation", test_store_invalidation)
run_test("memory savings vs FP16", test_store_memory)
print()
print("-- KVCaptureEngine --")
run_test("prefill", test_capture_engine_prefill)
run_test("decode with overflow", test_capture_engine_decode)
run_test("explicit flush", test_capture_engine_flush)
print()
print("-- Hybrid Attention --")
run_test("compressed only", test_hybrid_attention_compressed_only)
run_test("exact only", test_hybrid_attention_exact_only)
run_test("compressed + exact merge", test_hybrid_attention_both)
print()
print("-- Retrieval Quality --")
run_test("attention recall@8 (3-bit)", test_attention_recall)
run_test("attention recall@8 (4-bit)", test_attention_recall_4bit)
run_test("needle retrieval (single chunk)", test_needle_retrieval)
run_test("needle retrieval (multi chunk)", test_needle_retrieval_multi_chunk)
run_test("score rank correlation", test_score_rank_correlation)
print()
print("=" * 60)
print(f"Results: {PASS} passed, {FAIL} failed (total {PASS + FAIL})")
if FAIL == 0:
print("All tests passed!")
else:
print(f"{FAIL} test(s) failed")
exit(1)
print("=" * 60)
test_triton_kernels.py
-296
View File
@@ -1,296 +0,0 @@
#!/usr/bin/env python3
"""
Test TurboQuant Triton kernels against PyTorch reference implementation.
Tests:
1. MSE score kernel vs PyTorch dequantize-then-matmul
2. QJL score kernel vs PyTorch unpack-sketch-dot
3. Combined attention score vs full PyTorch path
4. Fused decode (scores + softmax + value aggr) vs PyTorch
5. Performance benchmark: Triton vs PyTorch
"""
import math
import os
import time
import torch
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from turboquant.quantizer import TurboQuantProd, TurboQuantMSE, MSEQuantized, ProdQuantized
from turboquant.kv_cache import quantize_values, dequantize_values
from turboquant.triton_kernels import (
turboquant_mse_score,
turboquant_qjl_score,
turboquant_attention_score,
turboquant_fused_decode,
_get_packing_params,
)
torch.manual_seed(42)
device = torch.device("cuda")
PASS = 0
FAIL = 0
def run_test(name, fn, **kwargs):
global PASS, FAIL
try:
fn(**kwargs)
print(f" ✓ {name}")
PASS += 1
except Exception as e:
import traceback
print(f" ✗ {name}")
traceback.print_exc()
FAIL += 1
# ─── Test 1: MSE Score Kernel ─────────────────────────────────────────
def test_mse_score(B=2, H=4, N=128, D=64, bits=3):
"""Compare Triton MSE score vs PyTorch reference."""
BH = B * H
# Create quantizer and random keys
quantizer = TurboQuantMSE(dim=D, bits=bits - 1, device=device, seed=42) # bits-1 for MSE stage
keys = torch.randn(BH, N, D, device=device, dtype=torch.float16)
# Quantize
mse_q = quantizer.quantize(keys)
# --- PyTorch reference ---
keys_dequant = quantizer.dequantize(mse_q)
query = torch.randn(BH, D, device=device, dtype=torch.float16)
ref_scores = torch.matmul(query.float().unsqueeze(1), keys_dequant.float().transpose(-2, -1)).squeeze(1)
# --- Triton kernel ---
# Rotate query: q_rot = q @ Pi^T
q_rot = torch.matmul(query.float(), quantizer.Pi.T)
triton_scores = turboquant_mse_score(
q_rot, mse_q.indices, mse_q.norms, quantizer.centroids, mse_q.bits
)
# Compare
torch.testing.assert_close(triton_scores, ref_scores.float(), atol=1e-2, rtol=1e-2)
# ─── Test 2: QJL Score Kernel ─────────────────────────────────────────
def test_qjl_score(B=2, H=4, N=128, D=64, bits=3):
"""Compare Triton QJL score vs PyTorch reference."""
BH = B * H
# Create full quantizer
quantizer = TurboQuantProd(dim=D, bits=bits, device=device, seed=42)
keys = torch.randn(BH, N, D, device=device, dtype=torch.float16)
query = torch.randn(BH, D, device=device, dtype=torch.float16)
# Quantize
prod_q = quantizer.quantize(keys)
# --- PyTorch reference: QJL part only ---
signs = quantizer._unpack_qjl_signs(prod_q.qjl_signs)
q_sketched = torch.matmul(query.float(), quantizer.S.T)
ref_qjl = torch.matmul(q_sketched.unsqueeze(1), signs.transpose(-2, -1)).squeeze(1)
ref_qjl = ref_qjl * (quantizer.qjl_scale * prod_q.residual_norms)
# --- Triton kernel ---
q_sketch = torch.matmul(query.float(), quantizer.S.T)
triton_qjl = turboquant_qjl_score(
q_sketch, prod_q.qjl_signs, prod_q.residual_norms, quantizer.qjl_scale
)
torch.testing.assert_close(triton_qjl, ref_qjl.float(), atol=1e-2, rtol=1e-2)
# ─── Test 3: Combined Score ───────────────────────────────────────────
def test_combined_score(B=2, H=2, N=256, D=64, bits=3):
"""Compare Triton combined score vs PyTorch TurboQuantProd.attention_score."""
BH = B * H
quantizer = TurboQuantProd(dim=D, bits=bits, device=device, seed=42)
keys = torch.randn(BH, N, D, device=device, dtype=torch.float16)
query = torch.randn(BH, 1, D, device=device, dtype=torch.float16)
prod_q = quantizer.quantize(keys)
# --- PyTorch reference ---
ref_scores = quantizer.attention_score(query, prod_q).squeeze(1) # (BH, N)
# --- Triton ---
triton_scores = turboquant_attention_score(
query, prod_q,
quantizer.mse_quantizer.Pi,
quantizer.S,
quantizer.mse_quantizer.centroids,
prod_q.mse_bits,
quantizer.qjl_scale,
)
torch.testing.assert_close(triton_scores, ref_scores.float(), atol=0.1, rtol=0.05)
# ─── Test 4: Fused Decode ─────────────────────────────────────────────
def test_fused_decode(B=1, H=2, N=64, D=64, bits=3, group_size=32):
"""Compare Triton fused decode vs PyTorch score+softmax+value."""
BH = B * H
quantizer = TurboQuantProd(dim=D, bits=bits, device=device, seed=42)
keys = torch.randn(BH, N, D, device=device, dtype=torch.float16)
values = torch.randn(BH, N, D, device=device, dtype=torch.float16)
query = torch.randn(BH, 1, D, device=device, dtype=torch.float16)
sm_scale = 1.0 / math.sqrt(D)
# Quantize
prod_q = quantizer.quantize(keys)
val_q = quantize_values(values, bits=2, group_size=group_size)
# --- PyTorch reference ---
ref_scores = quantizer.attention_score(query, prod_q) # (BH, 1, N)
ref_scores = ref_scores * sm_scale
ref_weights = torch.softmax(ref_scores.float(), dim=-1)
ref_values = dequantize_values(val_q, group_size) # (BH, N, D)
ref_output = torch.matmul(ref_weights, ref_values.float()).squeeze(1) # (BH, D)
# --- Triton fused ---
triton_output = turboquant_fused_decode(
query, prod_q, val_q,
quantizer.mse_quantizer.Pi,
quantizer.S,
quantizer.mse_quantizer.centroids,
prod_q.mse_bits,
quantizer.qjl_scale,
sm_scale,
group_size,
)
# Attention output should be close (not exact due to float accumulation order)
cos_sim = torch.nn.functional.cosine_similarity(
triton_output.flatten().float(),
ref_output.flatten().float(),
dim=0,
)
print(f" cosine_sim={cos_sim.item():.6f}, "
f"mse={((triton_output.float() - ref_output.float()) ** 2).mean().item():.2e}")
assert cos_sim > 0.95, f"Cosine similarity too low: {cos_sim.item()}"
# ─── Test 5: Various configs ──────────────────────────────────────────
def test_various_configs():
"""Test with different bit widths and dimensions."""
for D in [64, 128]:
for bits in [2, 3, 4]:
BH, N = 4, 64
quantizer = TurboQuantProd(dim=D, bits=bits, device=device, seed=42)
keys = torch.randn(BH, N, D, device=device, dtype=torch.float16)
query = torch.randn(BH, 1, D, device=device, dtype=torch.float16)
prod_q = quantizer.quantize(keys)
ref = quantizer.attention_score(query, prod_q).squeeze(1)
triton_out = turboquant_attention_score(
query, prod_q,
quantizer.mse_quantizer.Pi,
quantizer.S,
quantizer.mse_quantizer.centroids,
prod_q.mse_bits,
quantizer.qjl_scale,
)
max_err = (triton_out - ref.float()).abs().max().item()
print(f" D={D}, bits={bits}: max_err={max_err:.4f}")
assert max_err < 1.0, f"Max error too high: {max_err}"
# ─── Test 6: Benchmark ───────────────────────────────────────────────
def benchmark():
"""Benchmark Triton vs PyTorch for score computation."""
print()
print(" ── Performance Benchmark ──")
B, H, D, bits = 1, 32, 64, 3
BH = B * H
quantizer = TurboQuantProd(dim=D, bits=bits, device=device, seed=42)
for N in [256, 512, 1024, 2048, 4096]:
keys = torch.randn(BH, N, D, device=device, dtype=torch.float16)
query = torch.randn(BH, 1, D, device=device, dtype=torch.float16)
prod_q = quantizer.quantize(keys)
# Warmup
for _ in range(5):
_ = quantizer.attention_score(query, prod_q)
_ = turboquant_attention_score(
query, prod_q, quantizer.mse_quantizer.Pi, quantizer.S,
quantizer.mse_quantizer.centroids, prod_q.mse_bits, quantizer.qjl_scale,
)
torch.cuda.synchronize()
# PyTorch
N_ITER = 50
torch.cuda.synchronize()
t0 = time.time()
for _ in range(N_ITER):
_ = quantizer.attention_score(query, prod_q)
torch.cuda.synchronize()
pytorch_ms = (time.time() - t0) / N_ITER * 1000
# Triton
torch.cuda.synchronize()
t0 = time.time()
for _ in range(N_ITER):
_ = turboquant_attention_score(
query, prod_q, quantizer.mse_quantizer.Pi, quantizer.S,
quantizer.mse_quantizer.centroids, prod_q.mse_bits, quantizer.qjl_scale,
)
torch.cuda.synchronize()
triton_ms = (time.time() - t0) / N_ITER * 1000
speedup = pytorch_ms / triton_ms if triton_ms > 0 else float('inf')
print(f" N={N:5d} PyTorch: {pytorch_ms:6.2f}ms Triton: {triton_ms:6.2f}ms Speedup: {speedup:.2f}x")
# ─── Run ──────────────────────────────────────────────────────────────
print(f"GPU: {torch.cuda.get_device_name()}")
print()
print("═" * 60)
print("TurboQuant Triton Kernel Tests")
print("═" * 60)
run_test("MSE score (B=2 H=4 N=128 D=64 bits=3)", test_mse_score)
run_test("MSE score (B=1 H=8 N=512 D=64 bits=4)", test_mse_score, B=1, H=8, N=512, D=64, bits=4)
run_test("MSE score (B=1 H=4 N=256 D=128 bits=3)", test_mse_score, B=1, H=4, N=256, D=128, bits=3)
run_test("QJL score (B=2 H=4 N=128 D=64 bits=3)", test_qjl_score)
run_test("QJL score (B=1 H=8 N=512 D=64 bits=4)", test_qjl_score, B=1, H=8, N=512, D=64, bits=4)
run_test("QJL score (B=1 H=4 N=256 D=128 bits=3)", test_qjl_score, B=1, H=4, N=256, D=128, bits=3)
run_test("Combined score (B=2 H=2 N=256 D=64 bits=3)", test_combined_score)
run_test("Combined score (B=1 H=4 N=512 D=128 bits=4)", test_combined_score, B=1, H=4, N=512, D=128, bits=4)
run_test("Fused decode (B=1 H=2 N=64 D=64 bits=3)", test_fused_decode)
run_test("Fused decode (B=1 H=4 N=128 D=64 bits=3)", test_fused_decode, B=1, H=4, N=128)
run_test("Various configs (D=64/128, bits=2/3/4)", test_various_configs)
run_test("Benchmark", benchmark)
print()
print("═" * 60)
print(f"Results: {PASS} passed, {FAIL} failed (total {PASS + FAIL})")
if FAIL == 0:
print("🎉 All tests passed!")
else:
print(f"⚠️ {FAIL} test(s) failed")
print("═" * 60)
test_turboquant.py
-387
View File
@@ -1,387 +0,0 @@
#!/usr/bin/env python3
"""
TurboQuant test suite — validates the core algorithms against paper's theorems.
Run: python test_turboquant.py
"""
import time
import torch
import numpy as np
import math
torch.manual_seed(42)
np.random.seed(42)
def test_codebook():
"""Test Lloyd-Max codebook computation."""
print("=" * 60)
print("TEST: Lloyd-Max codebook computation")
print("=" * 60)
from turboquant.codebook import compute_lloyd_max_codebook
# Paper Table: for b=1,2,3,4, MSE ≈ 0.36, 0.117, 0.03, 0.009
expected_mse = {1: 0.36, 2: 0.117, 3: 0.03, 4: 0.009}
for bits in [1, 2, 3, 4]:
t0 = time.time()
cb = compute_lloyd_max_codebook(d=128, bits=bits)
elapsed = time.time() - t0
mse = cb["mse_total"]
exp = expected_mse[bits]
print(f" bits={bits}: MSE={mse:.4f} (expected ≈{exp:.3f}), "
f"n_centroids={len(cb['centroids'])}, time={elapsed:.2f}s")
# Allow 30% tolerance from paper values (paper uses d→∞ approximation)
assert abs(mse - exp) / exp < 0.30, f"MSE {mse} too far from expected {exp}"
print(f" ✓ Within tolerance of paper values")
print()
return True
def test_rotation():
"""Test that rotation preserves norms and produces Beta-distributed coords."""
print("=" * 60)
print("TEST: Random rotation properties")
print("=" * 60)
from turboquant.rotation import generate_rotation_matrix, rotate_forward
d = 128
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
Pi = generate_rotation_matrix(d, device)
# Check orthogonality
I = Pi @ Pi.T
orth_error = (I - torch.eye(d, device=device)).abs().max().item()
print(f" Orthogonality error: {orth_error:.2e}")
assert orth_error < 1e-5, f"Rotation matrix not orthogonal: {orth_error}"
print(f" ✓ Matrix is orthogonal")
# Check norm preservation
x = torch.randn(1000, d, device=device)
x = x / x.norm(dim=-1, keepdim=True)
y = rotate_forward(x, Pi)
norm_diff = (y.norm(dim=-1) - 1.0).abs().max().item()
print(f" Max norm deviation after rotation: {norm_diff:.2e}")
assert norm_diff < 1e-5, f"Norms not preserved: {norm_diff}"
print(f" ✓ Norms preserved")
# Check that rotated coords follow Beta distribution (approximately Gaussian for d=128)
coords = y[:, 0].cpu().numpy()
expected_std = 1.0 / math.sqrt(d)
actual_std = coords.std()
print(f" Expected coord std: {expected_std:.4f}, actual: {actual_std:.4f}")
assert abs(actual_std - expected_std) / expected_std < 0.15
print(f" ✓ Coordinate distribution matches theory")
print()
return True
def test_mse_quantizer():
"""Test TurboQuant MSE quantizer distortion bounds."""
print("=" * 60)
print("TEST: TurboQuant MSE quantizer")
print("=" * 60)
from turboquant.quantizer import TurboQuantMSE
d = 128
n_vectors = 1000
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Paper Theorem 1 exact values for b=1,2,3,4
expected_mse_total = {1: 0.36, 2: 0.117, 3: 0.03, 4: 0.009}
for bits in [1, 2, 3, 4]:
quantizer = TurboQuantMSE(dim=d, bits=bits, device=device)
# Generate random unit vectors
x = torch.randn(n_vectors, d, device=device)
x = x / x.norm(dim=-1, keepdim=True)
# Quantize and dequantize
q = quantizer.quantize(x)
x_hat = quantizer.dequantize(q)
# Compute MSE
mse = ((x - x_hat) ** 2).sum(dim=-1).mean().item()
exp = expected_mse_total[bits]
ratio = mse / exp
status = "✓" if 0.7 <= ratio <= 1.5 else "✗"
print(f" bits={bits}: MSE={mse:.4f} (expected ≈{exp:.3f}, ratio={ratio:.2f}) {status}")
print()
return True
def test_prod_quantizer():
"""Test TurboQuant inner product quantizer — unbiasedness and distortion."""
print("=" * 60)
print("TEST: TurboQuant inner product quantizer")
print("=" * 60)
from turboquant.quantizer import TurboQuantProd
d = 128
n_vectors = 500
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
for bits in [2, 3, 4]:
quantizer = TurboQuantProd(dim=d, bits=bits, device=device)
# Generate random vectors
x = torch.randn(n_vectors, d, device=device)
x = x / x.norm(dim=-1, keepdim=True)
y = torch.randn(n_vectors, d, device=device)
# Quantize x, compute <y, dequant(x)>
q = quantizer.quantize(x)
x_hat = quantizer.dequantize(q)
# True inner products
true_ip = (x * y).sum(dim=-1)
# Estimated inner products
est_ip = (x_hat * y).sum(dim=-1)
# Check unbiasedness: E[est] ≈ true
bias = (est_ip - true_ip).mean().item()
relative_bias = abs(bias) / true_ip.abs().mean().item()
# Check distortion
distortion = ((est_ip - true_ip) ** 2).mean().item()
# Paper Theorem 2: Dprod ≤ sqrt(3)π²/(2d) · ||y||² · 1/4^b
y_norm_sq = (y ** 2).sum(dim=-1).mean().item()
theoretical_bound = math.sqrt(3) * math.pi**2 / (2 * d) * y_norm_sq / 4**bits
print(f" bits={bits}:")
print(f" Bias: {bias:.6f} (relative: {relative_bias:.4f})")
print(f" Distortion: {distortion:.6f} (theoretical bound: {theoretical_bound:.6f})")
print(f" {'✓' if relative_bias < 0.1 else '✗'} Approximately unbiased")
print()
return True
def test_attention_score():
"""Test TurboQuant attention score computation."""
print("=" * 60)
print("TEST: Attention score computation")
print("=" * 60)
from turboquant.quantizer import TurboQuantProd
d = 128
n_keys = 200
n_queries = 4
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
min_corr = {2: 0.7, 3: 0.85, 4: 0.95}
for bits in [2, 3, 4]:
quantizer = TurboQuantProd(dim=d, bits=bits, device=device)
# Simulate attention: queries (n_q, d) × keys (n_k, d)
queries = torch.randn(1, 1, n_queries, d, device=device) * 0.1
keys = torch.randn(1, 1, n_keys, d, device=device) * 0.1
# True attention scores
true_scores = torch.matmul(queries, keys.transpose(-2, -1))
# Quantized attention scores
q_keys = quantizer.quantize(keys)
est_scores = quantizer.attention_score(queries, q_keys)
# Check correlation
true_flat = true_scores.flatten().cpu()
est_flat = est_scores.flatten().cpu()
correlation = torch.corrcoef(torch.stack([true_flat, est_flat]))[0, 1].item()
# MSE
score_mse = ((true_scores - est_scores) ** 2).mean().item()
threshold = min_corr[bits]
print(f" bits={bits}: correlation={correlation:.4f}, score_MSE={score_mse:.6f}")
assert correlation > threshold, f"Correlation {correlation:.4f} < {threshold} for {bits}-bit"
print(f" ✓ Correlation > {threshold} for {bits}-bit quantization")
print()
return True
def test_kv_cache():
"""Test the full KV cache with prefill and decode."""
print("=" * 60)
print("TEST: TurboQuant KV cache")
print("=" * 60)
from turboquant.kv_cache import TurboQuantKVCache
d = 128
n_heads = 8
batch_size = 1
prefill_len = 512
decode_steps = 64
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cache = TurboQuantKVCache(
head_dim=d,
key_bits=3,
value_bits=2,
buffer_size=128,
device=device,
layer_idx=0,
)
# Generate random KV states
keys = torch.randn(batch_size, n_heads, prefill_len, d, device=device) * 0.1
values = torch.randn(batch_size, n_heads, prefill_len, d, device=device) * 0.1
# Prefill
t0 = time.time()
cache.prefill(keys, values)
prefill_time = time.time() - t0
print(f" Prefill ({prefill_len} tokens): {prefill_time*1000:.1f}ms")
# Decode
t0 = time.time()
for i in range(decode_steps):
new_k = torch.randn(batch_size, n_heads, 1, d, device=device) * 0.1
new_v = torch.randn(batch_size, n_heads, 1, d, device=device) * 0.1
cache.append(new_k, new_v)
decode_time = time.time() - t0
print(f" Decode ({decode_steps} tokens): {decode_time*1000:.1f}ms "
f"({decode_time/decode_steps*1000:.2f}ms/token)")
# Attention score
query = torch.randn(batch_size, n_heads, 1, d, device=device) * 0.1
t0 = time.time()
scores = cache.attention_scores(query)
score_time = time.time() - t0
print(f" Attention score: {score_time*1000:.2f}ms, shape={scores.shape}")
# Attention output
attn_weights = torch.softmax(scores, dim=-1)
t0 = time.time()
output = cache.attend(attn_weights)
attend_time = time.time() - t0
print(f" Attend: {attend_time*1000:.2f}ms, shape={output.shape}")
# Memory
mem = cache.memory_bytes()
total_fp16 = batch_size * n_heads * cache.seq_len * d * 2 * 2 # K+V in fp16
compression = total_fp16 / max(mem["total"], 1)
print(f" Memory: {mem['total']/1024:.1f}KB (vs {total_fp16/1024:.1f}KB FP16)")
print(f" Compression ratio: {compression:.1f}x")
assert scores.shape == (batch_size, n_heads, 1, prefill_len + decode_steps)
assert output.shape == (batch_size, n_heads, 1, d)
print(f" ✓ KV cache functioning correctly")
print()
return True
def test_memory_savings():
"""Compare memory usage vs FP16 and FP8 baselines."""
print("=" * 60)
print("TEST: Memory comparison")
print("=" * 60)
from turboquant.kv_cache import TurboQuantKVCache
d = 128
n_heads = 32 # Realistic for 7B model
n_kv_heads = 8 # GQA
seq_len = 8192
batch_size = 1
n_layers = 32
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# FP16 baseline
fp16_per_layer = batch_size * n_kv_heads * seq_len * d * 2 * 2 # K+V
fp16_total = fp16_per_layer * n_layers
# FP8 baseline
fp8_per_layer = batch_size * n_kv_heads * seq_len * d * 1 * 2 # K+V in FP8
fp8_total = fp8_per_layer * n_layers
# TurboQuant 3.5-bit keys, 2-bit values
tq_total = 0
for layer_idx in range(n_layers):
cache = TurboQuantKVCache(
head_dim=d,
key_bits=4 if layer_idx < 4 else 3,
value_bits=2,
buffer_size=128,
device=device,
layer_idx=layer_idx,
)
keys = torch.randn(batch_size, n_kv_heads, seq_len, d, device=device) * 0.1
values = torch.randn(batch_size, n_kv_heads, seq_len, d, device=device) * 0.1
cache.prefill(keys, values)
tq_total += cache.memory_bytes()["total"]
print(f" Config: {n_layers} layers, {n_kv_heads} KV heads, head_dim={d}, seq_len={seq_len}")
print(f" FP16: {fp16_total/1024/1024:.1f} MB")
print(f" FP8: {fp8_total/1024/1024:.1f} MB ({fp16_total/fp8_total:.1f}x vs FP16)")
print(f" TurboQuant: {tq_total/1024/1024:.1f} MB ({fp16_total/tq_total:.1f}x vs FP16)")
assert tq_total < fp8_total, "TurboQuant should use less memory than FP8"
print(f" ✓ TurboQuant uses {fp8_total/tq_total:.1f}x less memory than FP8")
print()
return True
if __name__ == "__main__":
print("\n🚀 TurboQuant Test Suite\n")
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")
if device == "cuda":
print(f"GPU: {torch.cuda.get_device_name()}\n")
else:
print()
results = {}
tests = [
("Codebook", test_codebook),
("Rotation", test_rotation),
("MSE Quantizer", test_mse_quantizer),
("Prod Quantizer", test_prod_quantizer),
("Attention Score", test_attention_score),
("KV Cache", test_kv_cache),
("Memory Savings", test_memory_savings),
]
for name, test_fn in tests:
try:
results[name] = test_fn()
except Exception as e:
print(f" ✗ FAILED: {e}")
import traceback
traceback.print_exc()
results[name] = False
print("=" * 60)
print("SUMMARY")
print("=" * 60)
for name, passed in results.items():
print(f" {'✓' if passed else '✗'} {name}")
print()
all_passed = all(results.values())
if all_passed:
print("🎉 All tests passed!")
else:
print("❌ Some tests failed")
exit(1)
tests/test_tq_harness.py
-485
View File
@@ -1,485 +0,0 @@
#!/usr/bin/env python3
"""Local contract tests for the TurboQuant telemetry harness."""
from __future__ import annotations
import json
import subprocess
import tempfile
import unittest
import unittest.mock
from pathlib import Path
from types import SimpleNamespace
ROOT = Path(__file__).resolve().parents[1]
SCRIPTS = ROOT / "scripts"
import sys
if str(SCRIPTS) not in sys.path:
sys.path.insert(0, str(SCRIPTS))
from tq_harness_lib import (
atomic_write_json,
build_prompt_bundle_from_tokenizer,
collect_campaign_summary,
probe_gpu_metrics,
should_skip_phase,
)
from tq_remote_snapshot import build_remote_python, run_ssh, write_mirror_files
from tq_autoresearch_driver import build_resume_command, build_stage_specs
from tq_autoresearch_dashboard import build_dashboard_payload, render_html
from tq_autoresearch_status import build_status_payload, render_markdown
from tq_autoresearch_watch import build_dashboard_command, build_driver_command, build_status_command, should_trigger_driver
from tq_resume_recovery import build_remote_launch_command
class FakeTokenizer:
def encode(self, text, add_special_tokens=False):
return [ord(ch) for ch in text]
def decode(self, token_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False):
return "".join(chr(token) for token in token_ids)
class HarnessTests(unittest.TestCase):
def test_prompt_bundle_deterministic(self):
tokenizer = FakeTokenizer()
first = build_prompt_bundle_from_tokenizer(tokenizer, context_len=128, seed=42)
second = build_prompt_bundle_from_tokenizer(tokenizer, context_len=128, seed=42)
self.assertEqual(first.prompt_hash, second.prompt_hash)
self.assertEqual(first.prompt_tokens, second.prompt_tokens)
self.assertEqual(first.text, second.text)
def test_atomic_write_json(self):
with tempfile.TemporaryDirectory() as tmpdir:
path = Path(tmpdir) / "artifact.json"
atomic_write_json(path, {"status": "ok", "value": 7})
self.assertTrue(path.exists())
self.assertEqual(json.loads(path.read_text()), {"status": "ok", "value": 7})
leftovers = list(path.parent.glob("*.tmp"))
self.assertEqual(leftovers, [])
def test_collect_campaign_summary(self):
with tempfile.TemporaryDirectory() as tmpdir:
root = Path(tmpdir)
atomic_write_json(
root / "campaign_config.json",
{
"campaign_id": "demo",
"contexts": [30000],
"cases": ["baseline", "tq"],
"phases": ["init", "full"],
},
)
atomic_write_json(
root / "30000" / "prompt" / "prompt_meta.json",
{"prompt_hash": "abc", "prompt_tokens": 30000},
)
atomic_write_json(
root / "30000" / "baseline" / "init.json",
{"status": "ok", "prompt_hash": "abc", "elapsed_s": 1.2},
)
atomic_write_json(
root / "30000" / "baseline" / "full.json",
{"status": "ok", "prompt_hash": "abc", "elapsed_s": 3.4, "sample_text": "hi"},
)
atomic_write_json(
root / "30000" / "tq" / "init.json",
{"status": "timeout", "prompt_hash": "abc", "elapsed_s": 4.5, "exit_code": 137},
)
manifest, summary = collect_campaign_summary(root)
self.assertEqual(manifest["recommended_plan"], "B")
self.assertEqual(len(manifest["missing"]), 1)
self.assertIn("30000", summary)
self.assertIn("timeout", summary)
def test_collect_campaign_summary_complete_sub_200k_is_plan_a(self):
with tempfile.TemporaryDirectory() as tmpdir:
root = Path(tmpdir)
atomic_write_json(
root / "campaign_config.json",
{
"campaign_id": "demo-complete",
"contexts": [30000],
"cases": ["baseline", "tq"],
"phases": ["init", "ttft", "full"],
},
)
atomic_write_json(
root / "30000" / "prompt" / "prompt_meta.json",
{"prompt_hash": "abc", "prompt_tokens": 30000},
)
for case in ("baseline", "tq"):
for phase in ("init", "ttft", "full"):
atomic_write_json(
root / "30000" / case / f"{phase}.json",
{
"status": "ok",
"prompt_hash": "abc",
"elapsed_s": 1.0,
"sample_text": "done" if phase == "full" else "",
},
)
manifest, _summary = collect_campaign_summary(root)
self.assertEqual(manifest["recommended_plan"], "A")
def test_should_skip_phase(self):
with tempfile.TemporaryDirectory() as tmpdir:
path = Path(tmpdir) / "full.json"
atomic_write_json(path, {"status": "ok"})
self.assertTrue(should_skip_phase(path, skip_existing=False, force=False))
self.assertTrue(should_skip_phase(path, skip_existing=True, force=False))
self.assertFalse(should_skip_phase(path, skip_existing=False, force=True))
def test_should_not_skip_non_ok_phase_without_skip_existing(self):
with tempfile.TemporaryDirectory() as tmpdir:
path = Path(tmpdir) / "full.json"
atomic_write_json(path, {"status": "error"})
self.assertFalse(should_skip_phase(path, skip_existing=False, force=False))
self.assertTrue(should_skip_phase(path, skip_existing=True, force=False))
def test_remote_snapshot_ssh_wrapper(self):
with unittest.mock.patch("subprocess.run") as mocked:
mocked.return_value = SimpleNamespace(stdout='{"ok":true}\n')
args = SimpleNamespace(
host="example.com",
port=2222,
user="root",
identity_file="/tmp/key",
ssh_timeout_s=30,
)
out = run_ssh(args, "echo hi")
self.assertEqual(out, '{"ok":true}\n')
def test_remote_snapshot_writes_local_mirror(self):
with tempfile.TemporaryDirectory() as tmpdir:
mirror_root = Path(tmpdir) / "mirror"
payload = {
"manifest": {"campaign_id": "demo"},
"summary_text": "# Demo\n",
"files": [
{
"relative_path": "30000/baseline/init.json",
"type": "json",
"text": json.dumps({"status": "ok", "elapsed_s": 1.23}),
},
{
"relative_path": "30000/prompt/prompt.txt",
"type": "text",
"text": "prompt body",
},
],
}
write_mirror_files(payload, mirror_root)
self.assertEqual(json.loads((mirror_root / "manifest.json").read_text()), {"campaign_id": "demo"})
self.assertEqual((mirror_root / "summary.md").read_text(), "# Demo\n")
self.assertEqual(
json.loads((mirror_root / "30000" / "baseline" / "init.json").read_text()),
{"status": "ok", "elapsed_s": 1.23},
)
self.assertEqual((mirror_root / "30000" / "prompt" / "prompt.txt").read_text(), "prompt body")
def test_remote_snapshot_summary_mode_skips_file_payloads(self):
remote_python = build_remote_python("/tmp/demo-campaign", "summary")
self.assertIn("include_files = False", remote_python)
self.assertIn('"mode": \'summary\'', remote_python)
def test_remote_snapshot_ping_mode_is_minimal(self):
remote_python = build_remote_python("/tmp/demo-campaign", "ping")
self.assertIn('"mode": "ping"', remote_python)
self.assertIn('"campaign_root_exists": root.exists()', remote_python)
self.assertNotIn("manifest_path = root / \"manifest.json\"", remote_python)
def test_resume_recovery_builds_plan_b_launch_command(self):
args = SimpleNamespace(
recovery_campaign_root="/remote/logs/smoke-50k-planb",
remote_scripts_dir="/remote/scripts",
remote_model_path="/remote/models/model",
contexts="50000",
cases="baseline,tq",
phases="init,prefill_only,decode_only,full",
prompt_seed=5090,
max_output_tokens=24,
gpu_memory_utilization=0.9,
tensor_parallel_size=1,
)
cmd = build_remote_launch_command(args)
self.assertIn("nohup env CUDA_VISIBLE_DEVICES=0", cmd)
self.assertIn("--phases init,prefill_only,decode_only,full", cmd)
self.assertIn("--campaign-root /remote/logs/smoke-50k-planb", cmd)
self.assertIn("> /remote/logs/smoke-50k-planb/launch.log 2>&1 < /dev/null & echo $!", cmd)
def test_autoresearch_driver_builds_stage_specs(self):
args = SimpleNamespace(
contexts="50000,80000,120000,200000",
base_remote_root="/remote/logs",
local_report_dir="test-output/autoresearch-driver",
)
stages = build_stage_specs(args)
self.assertEqual([stage["context_len"] for stage in stages], [50000, 80000, 120000, 200000])
self.assertEqual(stages[0]["stage_name"], "smoke-50k-planb")
self.assertEqual(stages[0]["probe_campaign_root"], "/remote/logs/smoke-50k-harness")
self.assertEqual(stages[1]["probe_campaign_root"], "/remote/logs/smoke-50k-planb")
self.assertEqual(stages[-1]["cases"], "tq")
def test_autoresearch_driver_builds_resume_command(self):
args = SimpleNamespace(
dry_run=True,
host="your-gpu-host.example.com",
port=3003,
user="root",
identity_file="/tmp/key",
ssh_timeout_s=20,
probe_interval_s=30,
max_wait_s=300,
remote_scripts_dir="/remote/scripts",
remote_model_path="/remote/models/model",
prompt_seed=5090,
max_output_tokens=24,
gpu_memory_utilization=0.9,
tensor_parallel_size=1,
)
stage = {
"context_len": 80000,
"probe_campaign_root": "/remote/logs/smoke-50k-planb",
"recovery_campaign_root": "/remote/logs/smoke-80k-planb",
"cases": "baseline,tq",
"phases": "init,prefill_only,decode_only,full",
"local_report_path": "test-output/autoresearch-driver/smoke-80k-planb-resume-report.json",
}
cmd = build_resume_command(args, stage)
self.assertIn("--contexts", cmd)
self.assertIn("80000", cmd)
self.assertIn("--recovery-campaign-root", cmd)
self.assertIn("/remote/logs/smoke-80k-planb", cmd)
self.assertEqual(cmd[-1], "--dry-run")
def test_autoresearch_status_builds_blocked_payload(self):
with tempfile.TemporaryDirectory() as tmpdir:
root = Path(tmpdir)
remote = root / "remote-smoke-50k"
driver_dir = root / "autoresearch-driver"
remote.mkdir()
driver_dir.mkdir()
atomic_write_json(
remote / "local-manifest.json",
{
"recommended_plan": "B",
"failed": [{"context_len": 50000, "case": "baseline", "phase": "ttft", "status": "killed", "exit_code": -9}],
"missing": [{"context_len": 50000, "case": "tq", "phase": "init"}],
},
)
atomic_write_json(
remote / "ping-snapshot-latest.json",
{
"status": "error",
"error_type": "ssh_failed",
"stderr_tail": "ssh: connect to host your-gpu-host.example.com port 3003: Operation timed out",
},
)
atomic_write_json(
remote / "resume-recovery-report.json",
{
"status": "error",
"error_type": "ssh_failed",
"error_message": "connect timed out",
},
)
atomic_write_json(
driver_dir / "driver-report.json",
{
"stages": [
{
"status": "planned",
"command_shell": "uv run python scripts/tq_resume_recovery.py --output report.json --dry-run",
}
]
},
)
payload = build_status_payload(remote, driver_dir / "driver-report.json")
self.assertEqual(payload["blocker_code"], "ssh_outage")
self.assertEqual(payload["campaign_plan"], "B")
self.assertIn("uv run python", payload["next_command"])
self.assertNotIn("--dry-run", payload["next_command"])
def test_autoresearch_status_renders_markdown(self):
payload = {
"generated_at": "2026-03-27T00:00:00+00:00",
"blocker_code": "ssh_outage",
"blocker_summary": "Remote SSH is timing out on your-gpu-host.example.com:3003",
"campaign_plan": "B",
"latest_probe_name": "ping_latest",
"latest_probe": {"status": "error", "error_type": "ssh_failed", "stderr_tail": "timeout"},
"campaign_failed": [{"context_len": 50000, "case": "baseline", "phase": "ttft", "status": "killed", "exit_code": -9}],
"campaign_missing": [{"context_len": 50000, "case": "tq", "phase": "init"}],
"next_command": "uv run python scripts/tq_autoresearch_driver.py --output report.json",
}
md = render_markdown(payload)
self.assertIn("# TurboQuant Autoresearch Status", md)
self.assertIn("ssh_outage", md)
self.assertIn("50000 / baseline / ttft", md)
self.assertIn("```bash", md)
def test_autoresearch_watch_builds_commands(self):
args = SimpleNamespace(
status_script="scripts/tq_autoresearch_status.py",
status_json="test-output/autoresearch-driver/status.json",
status_md="test-output/autoresearch-driver/status.md",
dashboard_script="scripts/tq_autoresearch_dashboard.py",
dashboard_output="test-output/autoresearch-driver/dashboard.html",
watch_report="test-output/autoresearch-driver/watch-report.json",
driver_script="scripts/tq_autoresearch_driver.py",
driver_output="test-output/autoresearch-driver/driver-report.json",
driver_dry_run=True,
)
status_cmd = build_status_command(args)
dashboard_cmd = build_dashboard_command(args)
driver_cmd = build_driver_command(args)
self.assertEqual(status_cmd[1], "scripts/tq_autoresearch_status.py")
self.assertIn("--output-json", status_cmd)
self.assertEqual(dashboard_cmd[1], "scripts/tq_autoresearch_dashboard.py")
self.assertIn("--watch-report", dashboard_cmd)
self.assertEqual(driver_cmd[1], "scripts/tq_autoresearch_driver.py")
self.assertEqual(driver_cmd[-1], "--dry-run")
def test_autoresearch_watch_trigger_gate(self):
self.assertFalse(should_trigger_driver(None))
self.assertFalse(should_trigger_driver({"blocker_code": "ssh_outage"}))
self.assertTrue(should_trigger_driver({"blocker_code": "ready"}))
def test_autoresearch_dashboard_payload(self):
status_payload = {
"blocker_code": "ssh_outage",
"blocker_summary": "Remote SSH is timing out on your-gpu-host.example.com:3003",
"campaign_plan": "B",
"latest_probe": {"status": "error", "error_type": "ssh_failed", "stderr_tail": "timeout"},
"next_command": "uv run python scripts/tq_resume_recovery.py --output report.json",
"campaign_failed": [{"context_len": 50000, "case": "baseline", "phase": "ttft", "status": "killed", "exit_code": -9}],
"campaign_missing": [{"context_len": 50000, "case": "tq", "phase": "init"}],
"driver_report": {
"dry_run": True,
"stages": [{"stage_name": "smoke-50k-planb", "context_len": 50000, "cases": "baseline,tq", "phases": "init,prefill_only", "status": "planned"}],
},
}
watch_payload = {
"iterations": 1,
"driver_triggered": False,
"history": [{"blocker_code": "ssh_outage", "blocker_summary": "Remote SSH is timing out on your-gpu-host.example.com:3003"}],
}
payload = build_dashboard_payload(status_payload, watch_payload)
self.assertEqual(payload["blocker_code"], "ssh_outage")
self.assertEqual(payload["watch_iterations"], 1)
self.assertEqual(len(payload["stage_rows"]), 1)
def test_autoresearch_dashboard_renders_html(self):
payload = {
"generated_at": "2026-03-27T00:00:00+00:00",
"blocker_code": "ssh_outage",
"blocker_summary": "Remote SSH is timing out on your-gpu-host.example.com:3003",
"campaign_plan": "B",
"latest_probe": {"status": "error", "error_type": "ssh_failed", "stderr_tail": "timeout"},
"next_command": "uv run python scripts/tq_resume_recovery.py --output report.json",
"campaign_failed": [{"context_len": 50000, "case": "baseline", "phase": "ttft", "status": "killed", "exit_code": -9}],
"campaign_missing": [{"context_len": 50000, "case": "tq", "phase": "init"}],
"driver_dry_run": True,
"stage_rows": [{"stage_name": "smoke-50k-planb", "context_len": 50000, "cases": "baseline,tq", "phases": "init,prefill_only", "status": "planned"}],
"watch_iterations": 1,
"watch_driver_triggered": False,
"watch_latest_blocker": "ssh_outage",
"watch_latest_summary": "Remote SSH is timing out on your-gpu-host.example.com:3003",
}
html_text = render_html(payload)
self.assertIn("<!DOCTYPE html>", html_text)
self.assertIn("TurboQuant Autoresearch Dashboard", html_text)
self.assertIn("ssh_outage", html_text)
self.assertIn("smoke-50k-planb", html_text)
def test_phase_runner_timeout_writes_artifact(self):
with tempfile.TemporaryDirectory() as tmpdir:
prompt_dir = Path(tmpdir) / "prompt"
atomic_write_json(
prompt_dir / "prompt_meta.json",
{"prompt_hash": "dry", "prompt_tokens": 12},
)
(prompt_dir / "prompt.txt").write_text("dry run prompt", encoding="utf-8")
output = Path(tmpdir) / "phase.json"
cmd = [
sys.executable,
str(SCRIPTS / "tq_phase_runner.py"),
"--case",
"baseline",
"--phase",
"full",
"--context-len",
"30000",
"--output",
str(output),
"--timeout-s",
"1",
"--prompt-seed",
"7",
"--max-output-tokens",
"24",
"--model-path",
"unused",
"--prompt-dir",
str(prompt_dir),
"--dry-run",
"--dry-run-sleep-s",
"2",
]
proc = subprocess.run(cmd, capture_output=True, text=True)
self.assertNotEqual(proc.returncode, 0)
payload = json.loads(output.read_text())
self.assertEqual(payload["status"], "timeout")
self.assertEqual(payload["error_type"], "phase_timeout")
def test_phase_runner_prefill_only_dry_run(self):
with tempfile.TemporaryDirectory() as tmpdir:
prompt_dir = Path(tmpdir) / "prompt"
atomic_write_json(
prompt_dir / "prompt_meta.json",
{"prompt_hash": "dry", "prompt_tokens": 12},
)
(prompt_dir / "prompt.txt").write_text("dry run prompt", encoding="utf-8")
output = Path(tmpdir) / "phase.json"
cmd = [
sys.executable,
str(SCRIPTS / "tq_phase_runner.py"),
"--case",
"baseline",
"--phase",
"prefill_only",
"--context-len",
"30000",
"--output",
str(output),
"--timeout-s",
"30",
"--prompt-seed",
"7",
"--max-output-tokens",
"24",
"--model-path",
"unused",
"--prompt-dir",
str(prompt_dir),
"--dry-run",
]
proc = subprocess.run(cmd, capture_output=True, text=True)
self.assertEqual(proc.returncode, 0)
payload = json.loads(output.read_text())
self.assertEqual(payload["status"], "ok")
self.assertIsNotNone(payload["ttft_s"])
self.assertIsNotNone(payload["prefill_tok_s"])
def test_gpu_probe_payload_is_json_serializable(self):
payload = probe_gpu_metrics()
json.dumps(payload)
if __name__ == "__main__":
unittest.main()
tq_5090_results.html
-69
View File
@@ -1,69 +0,0 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>TurboQuant 5090 Results</title>
<style>
body { font-family: Inter, Arial, sans-serif; margin: 24px; background:#0b1020; color:#e8ecf3; }
h1,h2 { margin: 0 0 12px; }
h1 { font-size: 28px; }
h2 { margin-top: 28px; font-size: 20px; color:#9cc3ff; }
.meta { color:#b7c2d8; margin-bottom: 18px; }
table { border-collapse: collapse; width: 100%; margin: 12px 0 24px; background:#111831; }
th, td { border: 1px solid #2b355c; padding: 10px 12px; text-align: left; font-size: 14px; }
th { background:#182244; color:#cfe1ff; }
tr:nth-child(even) td { background:#0f1630; }
.good { color:#7ee787; font-weight: 600; }
.note { color:#b7c2d8; font-size: 14px; }
code { background:#17203d; padding:2px 6px; border-radius:6px; }
</style>
</head>
<body>
<h1>TurboQuant on RTX 5090</h1>
<div class="meta">Model: <code>QuantTrio/Qwen3.5-27B-AWQ</code> · Runtime: <code>vllm==0.18.0</code> · GPU: <code>RTX 5090 32GB</code></div>
<h2>Summary</h2>
<table>
<tr><th>Metric</th><th>Baseline</th><th>TurboQuant</th></tr>
<tr><td>Status</td><td>init_ok</td><td>init_ok</td></tr>
<tr><td>KV cache blocks</td><td>57</td><td>57</td></tr>
<tr><td>Measured baseline context</td><td>44,688</td><td>-</td></tr>
<tr><td>Verified max working context</td><td>44,688</td><td class="good">81,920</td></tr>
<tr><td>Improvement</td><td>-</td><td class="good">1.83x</td></tr>
<tr><td>Hooks installed</td><td>-</td><td>16</td></tr>
<tr><td>Freed bytes</td><td>-</td><td>2,928,672,768</td></tr>
<tr><td>Freed GiB</td><td>-</td><td>2.73</td></tr>
<tr><td>Elapsed time (s)</td><td>126.07</td><td>57.65</td></tr>
</table>
<h2>GPU Memory</h2>
<table>
<tr><th>Phase</th><th>Used MiB</th><th>Free MiB</th></tr>
<tr><td>Baseline after init</td><td>27895</td><td>4218</td></tr>
<tr><td>Baseline after generation</td><td>28155</td><td>3958</td></tr>
<tr><td>TurboQuant after init</td><td>27895</td><td>4218</td></tr>
<tr><td>TurboQuant after generation</td><td>28173</td><td>3940</td></tr>
<tr><td>TurboQuant after free_kv_cache()</td><td>27915</td><td>4198</td></tr>
</table>
<h2>Capacity Search Results</h2>
<table>
<tr><th>max_model_len</th><th>Status</th><th>Hooks</th><th>Elapsed (s)</th><th>Sample Output</th></tr>
<tr><td>44,688</td><td class="good">ok</td><td>16</td><td>57.65</td><td>KV cache compression</td></tr>
<tr><td>49,152</td><td class="good">ok</td><td>16</td><td>69.66</td><td>KV cache compression</td></tr>
<tr><td>53,248</td><td class="good">ok</td><td>16</td><td>53.76</td><td>KV cache compression</td></tr>
<tr><td>57,344</td><td class="good">ok</td><td>16</td><td>63.03</td><td>KV cache compression</td></tr>
<tr><td>61,440</td><td class="good">ok</td><td>16</td><td>54.65</td><td>KV cache compression</td></tr>
<tr><td>65,536</td><td class="good">ok</td><td>16</td><td>54.02</td><td>KV cache compression</td></tr>
<tr><td>73,728</td><td class="good">ok</td><td>16</td><td>72.70</td><td>KV cache compression</td></tr>
<tr><td>81,920</td><td class="good">ok</td><td>16</td><td>62.59</td><td>KV cache compression</td></tr>
</table>
<h2>Notes</h2>
<div class="note">
<p>These values are from completed remote experiments only.</p>
<p>The TurboQuant repository code itself was not modified.</p>
</div>
</body>
</html>
turboquant/integration/vllm.py
+48 -17
View File
@@ -180,8 +180,21 @@ def _no_alloc_prefill_attention(
return out.squeeze(0).transpose(0, 1)
def _make_patched_forward(orig_fn, state: LayerState, no_alloc: bool = False):
"""Intercept forward to optionally use TQ decode."""
def _make_patched_forward(orig_fn, state: LayerState, no_alloc: bool = False,
capture_in_forward: bool = False):
"""Intercept forward to optionally use TQ decode.
If capture_in_forward=True, also capture K/V from forward args
(needed when the backend has no separate do_kv_cache_update method).
"""
def _capture_kv(key, value, attn_metadata):
"""Capture K/V tensors into TQ store."""
num_tokens = getattr(attn_metadata, 'num_actual_tokens', key.shape[0])
if num_tokens <= 1:
state.engine.ingest_decode(key[:num_tokens], value[:num_tokens], num_tokens)
else:
state.engine.ingest_prefill(key[:num_tokens], value[:num_tokens], num_tokens)
def patched(
self_impl,
@@ -197,6 +210,10 @@ def _make_patched_forward(orig_fn, state: LayerState, no_alloc: bool = False):
):
mode = _GLOBAL_MODE
# Capture K/V when no separate kv_update hook exists
if capture_in_forward and mode not in (MODE_OFF,) and attn_metadata is not None:
_capture_kv(key, value, attn_metadata)
# Off or capture-only: always use flash
if mode in (MODE_OFF, MODE_CAPTURE_ONLY):
return orig_fn(
@@ -384,27 +401,41 @@ def install_hooks(
layer_states[layer_name] = state
if backend_kind == "flash":
patched_update = _make_patched_kv_update(
impl.do_kv_cache_update.__func__, state, no_alloc=no_alloc
)
has_separate_kv_update = hasattr(impl, "do_kv_cache_update")
needs_forward_capture = not has_separate_kv_update
if has_separate_kv_update:
patched_update = _make_patched_kv_update(
impl.do_kv_cache_update.__func__, state, no_alloc=no_alloc
)
impl.do_kv_cache_update = types.MethodType(
lambda self, *a, _p=patched_update, **kw: _p(self, *a, **kw), impl
)
patched_forward = _make_patched_forward(
impl.forward.__func__, state, no_alloc=no_alloc
)
impl.do_kv_cache_update = types.MethodType(
lambda self, *a, _p=patched_update, **kw: _p(self, *a, **kw), impl
impl.forward.__func__, state, no_alloc=no_alloc,
capture_in_forward=needs_forward_capture,
)
impl.forward = types.MethodType(
lambda self, *a, _p=patched_forward, **kw: _p(self, *a, **kw), impl
)
if needs_forward_capture and layer_idx == 0:
logger.info(
"[TurboQuant] No do_kv_cache_update found (vLLM 0.16 FlashInfer); "
"capturing K/V in forward()"
)
else:
patched_update = _make_patched_mla_update(impl.do_kv_cache_update.__func__, state)
patched_fwd = _make_patched_mla_forward(impl.forward_mqa.__func__, state)
impl.do_kv_cache_update = types.MethodType(
lambda self, *a, _p=patched_update, **kw: _p(self, *a, **kw), impl
)
impl.forward_mqa = types.MethodType(
lambda self, *a, _p=patched_fwd, **kw: _p(self, *a, **kw), impl
)
if hasattr(impl, "do_kv_cache_update"):
patched_update = _make_patched_mla_update(impl.do_kv_cache_update.__func__, state)
impl.do_kv_cache_update = types.MethodType(
lambda self, *a, _p=patched_update, **kw: _p(self, *a, **kw), impl
)
if hasattr(impl, "forward_mqa"):
patched_fwd = _make_patched_mla_forward(impl.forward_mqa.__func__, state)
impl.forward_mqa = types.MethodType(
lambda self, *a, _p=patched_fwd, **kw: _p(self, *a, **kw), impl
)
impl._tq_layer_state = state
layer_idx += 1
turboquant/vllm_attn_backend.py
+52 -2
View File
@@ -114,6 +114,9 @@ def enable_no_alloc(
def patched_get_kv_cache_specs(self):
cfg = _TQ_NO_ALLOC_CONFIG
with open("/tmp/tq_debug.log", "a") as f:
f.write(f"patched_get_kv_cache_specs called pid={os.getpid()} cfg={cfg is not None}\n")
f.flush()
if cfg is None:
return orig_get_specs(self)
@@ -155,12 +158,59 @@ def enable_no_alloc(
"shared_layers": shared_layers,
}
hooks = self.collective_rpc(_worker_install_tq)
logger.info(f"[TurboQuant] Installed no_alloc hooks: {hooks}")
try:
hooks = self.collective_rpc(_worker_install_tq)
print(f"[TurboQuant] Installed no_alloc hooks: {hooks}", flush=True)
except Exception as e:
import traceback
print(f"[TurboQuant] collective_rpc FAILED: {e}", flush=True)
traceback.print_exc()
return orig_get_specs(self)
Executor.get_kv_cache_specs = patched_get_kv_cache_specs
Executor._tq_patched = True
# Patch the worker's load_model (NOT decorated, so our patch won't be bypassed)
try:
from vllm.v1.worker.gpu_worker import GPUWorker as WorkerCls
except ImportError:
try:
from vllm.v1.worker.gpu_worker import Worker as WorkerCls
except ImportError:
WorkerCls = None
if WorkerCls is not None:
orig_worker_load = WorkerCls.load_model
def patched_worker_load(self_worker):
orig_worker_load(self_worker)
cfg = _TQ_NO_ALLOC_CONFIG
if cfg:
try:
import sys
sys.path.insert(0, '/tmp')
from turboquant.vllm_attn_backend import install_turboquant_hooks, MODE_ACCUMULATE
tq = install_turboquant_hooks(
self_worker.model_runner,
key_bits=cfg["key_bits"],
value_bits=cfg["value_bits"],
buffer_size=cfg["buffer_size"],
initial_layers_count=cfg["initial_layers_count"],
mode=MODE_ACCUMULATE,
no_alloc=False,
)
with open("/tmp/tq_debug.log", "a") as f:
f.write(f"TQ hooks: {len(tq)} layers pid={os.getpid()}\n")
f.flush()
except Exception as e:
with open("/tmp/tq_debug.log", "a") as f:
import traceback
f.write(f"TQ FAIL pid={os.getpid()}: {e}\n")
traceback.print_exc(file=f)
f.flush()
WorkerCls.load_model = patched_worker_load
logger.info("[TurboQuant] Patched Executor for auto TQ hook installation")
validate_moe.py
-409
View File
@@ -1,409 +0,0 @@
#!/usr/bin/env python3
"""
TurboQuant MoE Validation on Qwen3.5-35B-A3B (8x RTX 3090)
Phase 1: Baseline measurements via vLLM OpenAI API
- Max context length the server can handle (configured at 8192)
- KV cache usage at different context lengths
- TTFT and generation speed
- Output coherence: needle-in-haystack with REAL model output
Phase 2: TurboQuant integration
- Offline: capture KV states from the model, compress with TQ
- Measure compression ratio on REAL activations (not synthetic)
- Measure attention output quality degradation
Architecture: Qwen3.5-35B-A3B (pruned MoE)
- 40 layers: 30 linear_attention + 10 full_attention (every 4th)
- full_attention: head_dim=256, num_attention_heads=16, num_kv_heads=2
- linear_attention: linear_key_head_dim=128, linear_num_key_heads=16
- 205 experts, 8 active per token
- TQ only applies to the 10 full_attention layers
"""
import json
import time
import sys
import os
import subprocess
BASE_URL = "http://localhost:8000/v1"
MODEL = "Qwen/Qwen3.5-35B-A3B"
def curl_post(endpoint, data, timeout=120):
"""Make a POST request via curl (available on the remote machine)."""
cmd = [
"curl", "-s", "-X", "POST",
f"{BASE_URL}/{endpoint}",
"-H", "Content-Type: application/json",
"-d", json.dumps(data),
"--max-time", str(timeout),
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 10)
if result.returncode != 0:
return {"error": result.stderr}
try:
return json.loads(result.stdout)
except json.JSONDecodeError:
return {"error": f"Invalid JSON: {result.stdout[:500]}"}
def curl_get(endpoint):
cmd = ["curl", "-s", f"http://localhost:8000/{endpoint}"]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
return result.stdout
def get_kv_cache_usage():
metrics = curl_get("metrics")
for line in metrics.splitlines():
if line.startswith("vllm:kv_cache_usage_perc{"):
return float(line.split()[-1])
return None
def get_gpu_memory():
result = subprocess.run(
["nvidia-smi", "--query-gpu=memory.used,memory.free", "--format=csv,noheader,nounits"],
capture_output=True, text=True,
)
gpus = []
for line in result.stdout.strip().splitlines():
used, free = [int(x.strip()) for x in line.split(",")]
gpus.append({"used_mb": used, "free_mb": free})
return gpus
# ============================================================
# PHASE 1A: Baseline — generation at different prompt lengths
# ============================================================
def build_needle_prompt(context_len, needle_pos_frac=0.5):
"""Build a needle-in-haystack prompt at the given context length.
Places a specific fact in a sea of filler text, then asks about it.
Uses token-efficient filler (repeated simple sentences).
"""
needle = "The secret code for project Zephyr is AURORA-7742."
question = "What is the secret code for project Zephyr? Answer with just the code."
# Rough estimate: 1 token ~ 4 chars for English
target_chars = context_len * 4
filler_block = (
"The quick brown fox jumps over the lazy dog. "
"Lorem ipsum dolor sit amet, consectetur adipiscing elit. "
"In a distant galaxy, stars are born and die in cycles of cosmic renewal. "
"The weather today is partly cloudy with a chance of rain in the afternoon. "
"Scientists recently discovered a new species of deep-sea fish near hydrothermal vents. "
)
needle_char_pos = int(target_chars * needle_pos_frac)
# Build filler before needle
filler_before = ""
while len(filler_before) < needle_char_pos:
filler_before += filler_block
filler_before = filler_before[:needle_char_pos]
# Build filler after needle
remaining = target_chars - len(filler_before) - len(needle) - len(question) - 100
filler_after = ""
while len(filler_after) < remaining:
filler_after += filler_block
filler_after = filler_after[:max(0, remaining)]
prompt = f"{filler_before}\n\n{needle}\n\n{filler_after}\n\n{question}"
return prompt, "AURORA-7742"
def run_baseline_tests():
print("=" * 60)
print(" PHASE 1: Baseline Measurements (vLLM API)")
print("=" * 60)
print()
# Check current state
kv_usage = get_kv_cache_usage()
gpu_mem = get_gpu_memory()
print(f"Initial KV cache usage: {kv_usage:.4%}")
print(f"GPU memory (avg): {sum(g['used_mb'] for g in gpu_mem)/len(gpu_mem):.0f} MB used, "
f"{sum(g['free_mb'] for g in gpu_mem)/len(gpu_mem):.0f} MB free")
print()
# Test at increasing context lengths
results = []
for ctx_len in [512, 1024, 2048, 4096, 6144, 7680]:
print(f"--- Context length ~{ctx_len} tokens ---")
prompt, expected = build_needle_prompt(ctx_len, needle_pos_frac=0.5)
kv_before = get_kv_cache_usage()
t0 = time.time()
resp = curl_post("chat/completions", {
"model": MODEL,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 50,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
}, timeout=180)
elapsed = time.time() - t0
kv_after = get_kv_cache_usage()
if "error" in resp:
print(f" ERROR: {resp['error'][:200]}")
results.append({"ctx_len": ctx_len, "error": True})
continue
usage = resp.get("usage", {})
prompt_tokens = usage.get("prompt_tokens", 0)
completion_tokens = usage.get("completion_tokens", 0)
output = ""
if resp.get("choices"):
output = resp["choices"][0].get("message", {}).get("content", "")
needle_found = expected in output
ttft = elapsed # approximate (includes network + generation)
tok_per_sec = completion_tokens / elapsed if elapsed > 0 else 0
result = {
"ctx_len": ctx_len,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"elapsed_s": round(elapsed, 2),
"tok_per_sec": round(tok_per_sec, 1),
"kv_delta": round(kv_after - kv_before, 6) if kv_before and kv_after else None,
"needle_found": needle_found,
"output_preview": output[:100],
}
results.append(result)
print(f" Prompt tokens: {prompt_tokens}")
print(f" Completion tokens: {completion_tokens}")
print(f" Elapsed: {elapsed:.2f}s ({tok_per_sec:.1f} tok/s)")
print(f" KV cache delta: {result['kv_delta']}")
print(f" Needle found: {needle_found}")
print(f" Output: {output[:80]}...")
print()
return results
# ============================================================
# PHASE 1B: KV cache memory analysis
# ============================================================
def analyze_kv_memory():
print("=" * 60)
print(" KV Cache Memory Analysis")
print("=" * 60)
print()
# From config:
# full_attention layers: 10 (every 4th of 40)
# head_dim=256, num_kv_heads=2
# linear_attention layers: 30
# linear_key_head_dim=128, linear_num_key_heads=16, linear_num_value_heads=32
#
# vLLM block_size=272
# num_gpu_blocks=4553
n_full_attn_layers = 10
full_kv_heads = 2
full_head_dim = 256
n_linear_layers = 30
# Linear attention uses a recurrent state, not a standard KV cache.
# For Qwen3.5 MoE, linear attention layers use a "linear attention"
# mechanism where the state size is fixed (not proportional to seq_len).
# KV cache per token per full_attention layer:
# K: kv_heads * head_dim * dtype_size = 2 * 256 * 2 = 1024 bytes (bf16)
# V: same = 1024 bytes
# Total per token per layer: 2048 bytes
kv_per_token_per_layer = full_kv_heads * full_head_dim * 2 * 2 # 2 for K+V, 2 for bf16
kv_per_token_total = kv_per_token_per_layer * n_full_attn_layers
print(f"Full attention layers: {n_full_attn_layers} (of 40 total)")
print(f"KV heads: {full_kv_heads}, head_dim: {full_head_dim}")
print(f"KV cache per token per full_attn layer: {kv_per_token_per_layer} bytes")
print(f"KV cache per token (all full_attn): {kv_per_token_total} bytes ({kv_per_token_total/1024:.1f} KB)")
print()
# vLLM block info
block_size = 272 # tokens per block
num_blocks = 4553
total_kv_bytes = num_blocks * block_size * kv_per_token_total
print(f"vLLM blocks: {num_blocks} x {block_size} tokens")
print(f"Total KV capacity: {num_blocks * block_size:,} tokens")
print(f"Total KV memory: {total_kv_bytes / 1e9:.2f} GB")
print()
# With TurboQuant at 3-bit keys + 2-bit values on the 10 full_attn layers:
# Per token per layer:
# Key: 3 bits * head_dim = 3 * 256 = 768 bits = 96 bytes (per head)
# * 2 kv_heads = 192 bytes
# Plus norms: 2 * 4 bytes * 2 kv_heads = 16 bytes
# Value: 2 bits * head_dim = 2 * 256 = 512 bits = 64 bytes (per head)
# * 2 kv_heads = 128 bytes
# Plus scales/zeros: 2 * 2 * (256/32) * 2 kv_heads = 64 bytes
# Total TQ: ~400 bytes per token per layer
#
# vs FP16: 2048 bytes per token per layer
# Ratio: 2048 / 400 = ~5.1x on the full_attn layers only
tq_key_bytes = full_kv_heads * (full_head_dim * 3 // 8 + 4 + full_head_dim // 8 + 4) # mse_indices + norms + qjl_signs + residual_norms
tq_val_bytes = full_kv_heads * (full_head_dim * 2 // 8 + (full_head_dim // 32) * 4) # packed values + scales/zeros
tq_per_token_per_layer = tq_key_bytes + tq_val_bytes
tq_per_token_total = tq_per_token_per_layer * n_full_attn_layers
ratio = kv_per_token_total / tq_per_token_total
print(f"TQ KV per token per full_attn layer: {tq_per_token_per_layer} bytes")
print(f"TQ KV per token (all full_attn): {tq_per_token_total} bytes ({tq_per_token_total/1024:.1f} KB)")
print(f"Compression ratio (full_attn only): {ratio:.2f}x")
print()
# But linear attention layers also use KV cache in vLLM (even if it's recurrent)
# The actual savings depend on what fraction of total KV memory is full_attention
# Let's estimate: vLLM treats linear attention layers as having a fixed-size state
# So the KV cache memory is dominated by full_attention layers for long sequences.
for seq_len in [1024, 2048, 4096, 8192]:
fp16_full_attn = seq_len * kv_per_token_total
tq_full_attn = seq_len * tq_per_token_total
savings_mb = (fp16_full_attn - tq_full_attn) / 1e6
# Per GPU (TP=8): each GPU holds 1/8 of KV heads...
# Actually with 2 KV heads and TP=8, some GPUs may have 0 KV heads for full_attn
# Let's assume KV is distributed evenly
savings_per_gpu = savings_mb / 8
print(f" seq_len={seq_len:>5}: FP16={fp16_full_attn/1e6:.1f}MB, TQ={tq_full_attn/1e6:.1f}MB, "
f"savings={savings_mb:.1f}MB total ({savings_per_gpu:.1f}MB/GPU)")
print()
print("NOTE: With only 2 KV heads and TP=8, each GPU handles 0.25 KV heads on average.")
print("The KV cache is already tiny relative to model weights.")
print("TQ savings are real but small in absolute terms for this model.")
# ============================================================
# PHASE 1C: Coherence test — multi-needle at different depths
# ============================================================
def run_coherence_test():
print()
print("=" * 60)
print(" PHASE 1C: Multi-Needle Coherence Test")
print("=" * 60)
print()
# Place 3 facts at different positions, then ask about all of them
facts = [
("The capital of the fictional country Zephyria is Windholm.", "Windholm"),
("Agent Cooper's badge number is 4491.", "4491"),
("The password to the vault is CRIMSON-DELTA-9.", "CRIMSON-DELTA-9"),
]
ctx_len = 6144 # Use most of the 8192 context
filler_block = (
"The quick brown fox jumps over the lazy dog. "
"Lorem ipsum dolor sit amet, consectetur adipiscing elit. "
"In a distant galaxy, stars are born and die in cycles. "
"The weather is partly cloudy with rain expected later. "
"Scientists discovered a new species near deep-sea vents. "
)
target_chars = ctx_len * 4
segment_len = target_chars // 4 # 4 segments: filler, fact1, filler, fact2, filler, fact3, filler
parts = []
for i, (fact, _) in enumerate(facts):
filler = ""
while len(filler) < segment_len:
filler += filler_block
filler = filler[:segment_len]
parts.append(filler)
parts.append(f"\n{fact}\n")
# Final filler
filler = ""
while len(filler) < segment_len:
filler += filler_block
parts.append(filler[:segment_len])
question = (
"\n\nAnswer these three questions:\n"
"1. What is the capital of Zephyria?\n"
"2. What is Agent Cooper's badge number?\n"
"3. What is the password to the vault?\n"
"Give just the answers, one per line."
)
parts.append(question)
prompt = "".join(parts)
print(f"Prompt length: ~{len(prompt)//4} tokens (estimated)")
t0 = time.time()
resp = curl_post("chat/completions", {
"model": MODEL,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 100,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
}, timeout=180)
elapsed = time.time() - t0
if "error" in resp:
print(f"ERROR: {resp['error'][:200]}")
return
output = ""
if resp.get("choices"):
output = resp["choices"][0].get("message", {}).get("content", "")
usage = resp.get("usage", {})
print(f"Prompt tokens: {usage.get('prompt_tokens', '?')}")
print(f"Elapsed: {elapsed:.2f}s")
print(f"Output:\n{output}")
print()
found = 0
for fact_text, expected in facts:
if expected.lower() in output.lower():
found += 1
print(f" FOUND: {expected}")
else:
print(f" MISSED: {expected}")
print(f"\nCoherence score: {found}/{len(facts)} facts retrieved")
# ============================================================
# MAIN
# ============================================================
if __name__ == "__main__":
print(f"TurboQuant MoE Validation")
print(f"Model: {MODEL}")
print(f"Time: {time.strftime('%Y-%m-%d %H:%M:%S')}")
print()
# Phase 1A: Baseline generation at different context lengths
results = run_baseline_tests()
# KV memory analysis
analyze_kv_memory()
# Phase 1C: Coherence test
run_coherence_test()
# Save results
with open("/tmp/tq_moe_baseline.json", "w") as f:
json.dump(results, f, indent=2)
print(f"\nResults saved to /tmp/tq_moe_baseline.json")
validate_moe_phase2.py
-435
View File
@@ -1,435 +0,0 @@
#!/usr/bin/env python3
"""
Phase 2: TurboQuant validation on REAL KV activations from Qwen3.5-35B-A3B MoE.
Strategy:
We can't monkey-patch the running vLLM server. Instead we:
1. Load the model in a SEPARATE process on a single GPU (--device_map auto won't work
with 8 GPUs at 98% utilization).
Actually -- we only have ~2.8 GB free per GPU. Can't load the model again.
Better approach: Use vLLM's logprobs API to get token-level probabilities, then
compare baseline (full KV) vs simulated TQ (offline compression) quality.
BEST approach given constraints:
Capture KV activations from the model's attention layers via a hook, then
offline-quantize with TQ and measure the attention output degradation.
BUT: can't hook into the running server without restarting it.
PRACTICAL approach:
1. Generate a corpus of prompts at different lengths
2. Get vLLM completions WITH logprobs (this is our ground truth)
3. Measure: at what context length does quality degrade?
4. Theoretical analysis: compute TQ compression on head_dim=256, kv_heads=2
5. Synthetic validation: create realistic KV cache tensors matching
the model's statistics (head_dim=256) and measure TQ quality
This script does the synthetic validation with model-realistic parameters.
"""
import sys
sys.path.insert(0, "/tmp")
import math
import time
import json
import torch
import torch.nn.functional as F
import numpy as np
from turboquant.quantizer import TurboQuantProd
from turboquant.kv_cache import TurboQuantKVCache, quantize_values, dequantize_values
torch.manual_seed(42)
np.random.seed(42)
# Qwen3.5-35B-A3B full_attention layer config
HEAD_DIM = 256
NUM_KV_HEADS = 2
NUM_Q_HEADS = 16 # GQA ratio 8:1
GQA_RATIO = NUM_Q_HEADS // NUM_KV_HEADS # 8
NUM_FULL_ATTN_LAYERS = 10
SCALE = 1.0 / math.sqrt(HEAD_DIM)
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
def section(title):
print(f"\n{'='*60}")
print(f" {title}")
print(f"{'='*60}\n")
# ============================================================
# TEST 1: TQ quality at head_dim=256 (this model's actual dim)
# ============================================================
section("TEST 1: TQ quantization quality at head_dim=256")
print(f"Model parameters: head_dim={HEAD_DIM}, kv_heads={NUM_KV_HEADS}, "
f"q_heads={NUM_Q_HEADS}, GQA={GQA_RATIO}:1")
print(f"Device: {DEVICE}")
print()
for bits in [3, 4]:
q = TurboQuantProd(dim=HEAD_DIM, bits=bits, device=DEVICE, seed=42)
# Generate realistic KV cache: N tokens, kv_heads, head_dim
for N in [512, 1024, 4096, 8192]:
keys = torch.randn(1, NUM_KV_HEADS, N, HEAD_DIM, device=DEVICE) * 0.02
queries = torch.randn(1, NUM_KV_HEADS, 1, HEAD_DIM, device=DEVICE) * 0.02
# Exact attention scores
true_scores = torch.matmul(queries, keys.transpose(-2, -1)) * SCALE # (1, H, 1, N)
true_weights = F.softmax(true_scores, dim=-1)
# TQ quantized scores
keys_4d = keys # already (1, H, N, D)
key_q = q.quantize(keys_4d)
# Dequantize for attention score computation
keys_dequant = q.dequantize(key_q)
tq_scores = torch.matmul(queries, keys_dequant.transpose(-2, -1)) * SCALE
tq_weights = F.softmax(tq_scores, dim=-1)
# Also test the direct attention_score method
tq_scores_direct = q.attention_score(queries, key_q) * SCALE
tq_weights_direct = F.softmax(tq_scores_direct, dim=-1)
# Generate random values for output comparison
values = torch.randn(1, NUM_KV_HEADS, N, HEAD_DIM, device=DEVICE) * 0.02
true_out = torch.matmul(true_weights, values)
tq_out = torch.matmul(tq_weights_direct, values)
# Metrics
cos = F.cosine_similarity(
true_out.reshape(-1, HEAD_DIM),
tq_out.reshape(-1, HEAD_DIM),
dim=-1,
).mean().item()
kl = F.kl_div(
tq_weights_direct.log().clamp(min=-20),
true_weights,
reduction="batchmean",
).item()
out_mse = ((true_out - tq_out) ** 2).sum(dim=-1).mean().item()
out_mag = (true_out ** 2).sum(dim=-1).mean().item()
rel_mse = out_mse / max(out_mag, 1e-10)
# Recall@1 and Recall@8
true_top1 = true_scores.squeeze().argmax(dim=-1)
tq_top1 = tq_scores_direct.squeeze().argmax(dim=-1)
recall1 = (true_top1 == tq_top1).float().mean().item()
# Recall@8 per head
recall8_total = 0
for h in range(NUM_KV_HEADS):
true_top8 = set(true_scores[0, h, 0].topk(8).indices.tolist())
tq_top8 = set(tq_scores_direct[0, h, 0].topk(8).indices.tolist())
recall8_total += len(true_top8 & tq_top8) / 8
recall8 = recall8_total / NUM_KV_HEADS
if N == 512:
print(f"{bits}-bit, head_dim={HEAD_DIM}:")
print(f" {'N':>6} {'cos_sim':>10} {'KL_div':>10} {'rel_MSE':>12} {'recall@1':>10} {'recall@8':>10}")
print(f" {'-'*60}")
print(f" {N:>6} {cos:>10.6f} {kl:>10.6f} {rel_mse:>12.8f} {recall1:>10.1%} {recall8:>10.1%}")
print()
# ============================================================
# TEST 2: Memory savings analysis (real numbers for this model)
# ============================================================
section("TEST 2: Memory savings for Qwen3.5-35B-A3B")
# Per-token KV cache for full_attention layers
fp16_per_token_per_layer = NUM_KV_HEADS * HEAD_DIM * 2 * 2 # K+V, bf16
# TQ at 3-bit keys, 2-bit values
for key_bits, val_bits in [(3, 2), (4, 2), (4, 4)]:
q = TurboQuantProd(dim=HEAD_DIM, bits=key_bits, device=DEVICE, seed=42)
# Measure actual byte sizes from quantized output
test_keys = torch.randn(1, NUM_KV_HEADS, 1000, HEAD_DIM, device=DEVICE)
test_vals = torch.randn(1, NUM_KV_HEADS, 1000, HEAD_DIM, device=DEVICE)
key_q = q.quantize(test_keys)
val_q = quantize_values(test_vals, bits=val_bits, group_size=32)
# Count actual bytes
key_bytes = (
key_q.mse_indices.nelement() * key_q.mse_indices.element_size() +
key_q.qjl_signs.nelement() * key_q.qjl_signs.element_size() +
key_q.residual_norms.nelement() * 4 + # float32
key_q.norms.nelement() * 4 # float32
)
val_bytes = (
val_q.data.nelement() * val_q.data.element_size() +
val_q.scales.nelement() * 4 + # float32
val_q.zeros.nelement() * 4 # float32
)
total_tq = key_bytes + val_bytes
total_fp16 = test_keys.nelement() * 2 + test_vals.nelement() * 2
ratio = total_fp16 / total_tq
# Per-token overhead (Pi and S matrices, amortized)
pi_bytes = HEAD_DIM * HEAD_DIM * 4 # 256*256*4 = 262KB
s_bytes = HEAD_DIM * HEAD_DIM * 4 # 262KB
overhead_per_layer = pi_bytes + s_bytes # 524KB per layer
print(f"Key {key_bits}-bit / Value {val_bits}-bit:")
print(f" FP16 KV for 1000 tokens: {total_fp16/1e3:.1f} KB")
print(f" TQ KV for 1000 tokens: {total_tq/1e3:.1f} KB")
print(f" Compression ratio: {ratio:.2f}x")
print(f" Per-layer overhead (Pi+S): {overhead_per_layer/1e3:.0f} KB ({overhead_per_layer/1e6:.2f} MB)")
print(f" Overhead for {NUM_FULL_ATTN_LAYERS} layers: {overhead_per_layer * NUM_FULL_ATTN_LAYERS / 1e6:.1f} MB")
print()
# At what N does the overhead become < 10% of the savings?
savings_per_token = (total_fp16 - total_tq) / 1000 * NUM_FULL_ATTN_LAYERS
total_overhead = overhead_per_layer * NUM_FULL_ATTN_LAYERS
breakeven_tokens = total_overhead / max(savings_per_token, 1)
print(f" Break-even: overhead < 10% of savings at N > {int(breakeven_tokens * 10)}")
print()
# ============================================================
# TEST 3: Context extension potential
# ============================================================
section("TEST 3: Context extension potential")
# vLLM allocated 4553 blocks of 272 tokens = 1,238,416 token capacity
# But the model is served with max_model_len=8192
# With TP=8, each GPU handles a shard of the KV cache
# Current: ~2.8 GB free per GPU
# Model weights use ~21.3 GB of the 24 GB per GPU
# KV cache allocated from the remaining ~2.7 GB per GPU * 8 = 21.6 GB total
TOTAL_KV_GB = 25.36 # from vLLM: 4553 * 272 * 20480 bytes
FREE_PER_GPU_MB = 2850 # approximate
TOTAL_FREE_MB = FREE_PER_GPU_MB * 8
print(f"Current vLLM KV cache: {TOTAL_KV_GB:.1f} GB ({4553} blocks of 272 tokens)")
print(f"Max model len: 8192 tokens")
print(f"Free GPU memory: ~{FREE_PER_GPU_MB} MB/GPU, {TOTAL_FREE_MB/1000:.1f} GB total")
print()
# KV per token (full attention only, the linear attention layers have O(1) state)
kv_per_token_bytes = NUM_FULL_ATTN_LAYERS * NUM_KV_HEADS * HEAD_DIM * 2 * 2 # 20 KB
# But linear attention layers also have KV cache in vLLM's block allocator
# Let's check the actual block size from vLLM config: block_size=272
# This includes ALL layers, not just full_attention
# With TQ on the 10 full_attention layers:
# Savings per token = (fp16 - tq) for full_attn layers only
# Linear attention layers are unchanged
for key_bits, val_bits, label in [(3, 2, "aggressive"), (4, 2, "balanced"), (4, 4, "conservative")]:
q = TurboQuantProd(dim=HEAD_DIM, bits=key_bits, device=DEVICE, seed=42)
test_keys = torch.randn(1, 1, 100, HEAD_DIM, device=DEVICE)
test_vals = torch.randn(1, 1, 100, HEAD_DIM, device=DEVICE)
key_q = q.quantize(test_keys)
val_q = quantize_values(test_vals, bits=val_bits, group_size=32)
tq_bytes_per_100 = (
key_q.mse_indices.nelement() + key_q.qjl_signs.nelement() +
key_q.residual_norms.nelement() * 4 + key_q.norms.nelement() * 4 +
val_q.data.nelement() + val_q.scales.nelement() * 4 + val_q.zeros.nelement() * 4
)
tq_per_token_per_head = tq_bytes_per_100 / 100
tq_per_token_full_attn = tq_per_token_per_head * NUM_KV_HEADS * NUM_FULL_ATTN_LAYERS
fp16_per_token_full_attn = NUM_FULL_ATTN_LAYERS * NUM_KV_HEADS * HEAD_DIM * 2 * 2
savings_per_token = fp16_per_token_full_attn - tq_per_token_full_attn
savings_fraction = savings_per_token / fp16_per_token_full_attn
# How many MORE tokens can we fit in the freed memory?
# At max_model_len=8192, total KV for full_attn = 8192 * 20KB = 163.8 MB
current_kv_8192 = 8192 * fp16_per_token_full_attn
tq_kv_8192 = 8192 * tq_per_token_full_attn
freed_mb = (current_kv_8192 - tq_kv_8192) / 1e6
# Extra tokens that fit in freed memory
extra_tokens = int(freed_mb * 1e6 / tq_per_token_full_attn)
print(f"{label.upper()} ({key_bits}b keys / {val_bits}b vals):")
print(f" Savings per token: {savings_per_token:.0f} bytes ({savings_fraction:.0%} of full_attn KV)")
print(f" At 8192 tokens: free {freed_mb:.1f} MB -> fit {extra_tokens:,} more tokens")
print(f" New max context: ~{8192 + extra_tokens:,} tokens ({(8192+extra_tokens)/8192:.1f}x)")
print()
print("IMPORTANT CAVEATS:")
print(" 1. These savings are ONLY for the 10 full_attention layers")
print(" 2. The 30 linear_attention layers also use KV cache blocks in vLLM")
print(" 3. vLLM's block allocator may not efficiently reuse freed blocks")
print(" 4. The model already has max_position_embeddings=262144")
print(" 5. The bottleneck for context length is vLLM's max_model_len setting,")
print(" not memory — increasing it requires restarting the server")
# ============================================================
# TEST 4: Quality at head_dim=256 with spiky attention patterns
# ============================================================
section("TEST 4: Quality with realistic (spiky) attention patterns")
print("Testing with attention patterns where top-5 tokens carry >80% weight:")
print()
n_queries = 200
for bits in [3, 4]:
q = TurboQuantProd(dim=HEAD_DIM, bits=bits, device=DEVICE, seed=42)
print(f"{bits}-bit:")
print(f" {'N':>6} {'cos_sim':>10} {'KL_div':>10} {'rel_MSE':>12} {'top5_agree':>12}")
print(f" {'-'*52}")
for N in [512, 1024, 4096, 8192]:
cos_sims = []; kl_divs = []; rel_mses = []; top5_agrees = []
keys = torch.randn(1, NUM_KV_HEADS, N, HEAD_DIM, device=DEVICE) * 0.02
values = torch.randn(1, NUM_KV_HEADS, N, HEAD_DIM, device=DEVICE) * 0.02
key_q = q.quantize(keys)
for qi in range(n_queries):
# Create spiky query: strongly correlated with a few random keys
target_indices = torch.randint(0, N, (5,))
query = keys[:, :, target_indices[0]:target_indices[0]+1, :] * 3.0 # amplify
query = query + torch.randn_like(query) * 0.02
true_scores = torch.matmul(query, keys.transpose(-2, -1)) * SCALE
true_weights = F.softmax(true_scores, dim=-1)
true_out = torch.matmul(true_weights, values)
tq_scores = q.attention_score(query, key_q) * SCALE
tq_weights = F.softmax(tq_scores, dim=-1)
tq_out = torch.matmul(tq_weights, values)
cos = F.cosine_similarity(
true_out.reshape(-1, HEAD_DIM), tq_out.reshape(-1, HEAD_DIM), dim=-1
).mean().item()
cos_sims.append(cos)
kl = F.kl_div(
tq_weights.log().clamp(min=-20), true_weights, reduction="batchmean"
).item()
kl_divs.append(abs(kl))
mse = ((true_out - tq_out)**2).sum(dim=-1).mean().item()
mag = (true_out**2).sum(dim=-1).mean().item()
rel_mses.append(mse / max(mag, 1e-10))
# Top-5 agreement per head
for h in range(NUM_KV_HEADS):
t5 = set(true_weights[0, h, 0].topk(5).indices.tolist())
tq5 = set(tq_weights[0, h, 0].topk(5).indices.tolist())
top5_agrees.append(len(t5 & tq5) / 5)
print(f" {N:>6} {np.mean(cos_sims):>10.6f} {np.mean(kl_divs):>10.6f} "
f"{np.mean(rel_mses):>12.8f} {np.mean(top5_agrees):>12.1%}")
print()
# ============================================================
# TEST 5: Full KV cache simulation (all 10 full_attention layers)
# ============================================================
section("TEST 5: Full cache simulation — 10 layers, 8192 tokens")
N = 8192
print(f"Simulating full KV cache: {NUM_FULL_ATTN_LAYERS} layers, {N} tokens, "
f"{NUM_KV_HEADS} KV heads, head_dim={HEAD_DIM}")
print()
for bits in [3, 4]:
print(f"--- {bits}-bit TQ ---")
total_fp16_mb = 0
total_tq_mb = 0
layer_cos_sims = []
t0 = time.time()
for layer in range(NUM_FULL_ATTN_LAYERS):
cache = TurboQuantKVCache(
head_dim=HEAD_DIM,
key_bits=bits,
value_bits=2,
value_group_size=32,
buffer_size=128,
device=DEVICE,
layer_idx=layer,
)
# Simulate prefill
keys = torch.randn(1, NUM_KV_HEADS, N, HEAD_DIM, device=DEVICE) * 0.02
values = torch.randn(1, NUM_KV_HEADS, N, HEAD_DIM, device=DEVICE) * 0.02
cache.prefill(keys, values)
# Memory
mem = cache.memory_bytes()
fp16_bytes = N * NUM_KV_HEADS * HEAD_DIM * 2 * 2
total_fp16_mb += fp16_bytes / 1e6
total_tq_mb += mem["total"] / 1e6
# Quality: run a few test queries
query = torch.randn(1, NUM_KV_HEADS, 1, HEAD_DIM, device=DEVICE) * 0.02
tq_scores = cache.attention_scores(query)
true_scores = torch.matmul(query, keys.transpose(-2, -1)) * SCALE
tq_w = F.softmax(tq_scores, dim=-1)
true_w = F.softmax(true_scores, dim=-1)
tq_out = cache.attend(tq_w)
true_out = torch.matmul(true_w, values)
cos = F.cosine_similarity(
true_out.reshape(-1, HEAD_DIM), tq_out.reshape(-1, HEAD_DIM), dim=-1
).mean().item()
layer_cos_sims.append(cos)
# Free memory
del cache, keys, values, query
torch.cuda.empty_cache()
elapsed = time.time() - t0
print(f" FP16 total: {total_fp16_mb:.1f} MB")
print(f" TQ total: {total_tq_mb:.1f} MB")
print(f" Ratio: {total_fp16_mb / total_tq_mb:.2f}x")
print(f" Per-layer cos_sim: {[f'{c:.6f}' for c in layer_cos_sims]}")
print(f" Mean cos_sim: {np.mean(layer_cos_sims):.8f}")
print(f" Time: {elapsed:.1f}s")
print()
# ============================================================
# SUMMARY
# ============================================================
section("SUMMARY")
print("Model: Qwen3.5-35B-A3B (pruned MoE, 205 experts)")
print(f"Architecture: {NUM_FULL_ATTN_LAYERS} full_attention + 30 linear_attention layers")
print(f"Full attention: head_dim={HEAD_DIM}, {NUM_KV_HEADS} KV heads (GQA {GQA_RATIO}:1)")
print()
print("Baseline (no TQ):")
print(" - Single needle: PASS at all context lengths (512-7680)")
print(" - Multi-needle: 3/3 PASS at 5040 tokens")
print(" - Generation speed: 8.2-46 tok/s depending on context")
print(" - Max context: 8192 (vLLM configured)")
print()
print("TQ Impact on full_attention layers:")
print(" - See test results above for quality metrics")
print(" - Memory savings are modest due to only 10/40 layers having standard attention")
print(" - The linear_attention layers (30/40) are NOT compressible by TQ")
print()
print("Key takeaway: This MoE model is NOT the ideal target for TQ because")
print("75% of its layers use linear attention (O(1) state). TQ shines on")
print("dense transformers with 100% standard attention layers.")
validate_moe_phase3.py
-269
View File
@@ -1,269 +0,0 @@
#!/usr/bin/env python3
"""
Phase 3: Measure actual context extension capability.
Since vLLM is running with max_model_len=8192, we measure:
1. How much KV cache memory the model actually uses at different lengths
2. Perplexity-like quality metric via logprobs at different context lengths
3. Whether the model maintains coherence at 8192 (its max)
"""
import json
import time
import subprocess
import math
BASE_URL = "http://localhost:8000/v1"
MODEL = "Qwen/Qwen3.5-35B-A3B"
def curl_post(endpoint, data, timeout=180):
cmd = [
"curl", "-s", "-X", "POST",
f"{BASE_URL}/{endpoint}",
"-H", "Content-Type: application/json",
"-d", json.dumps(data),
"--max-time", str(timeout),
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 10)
if result.returncode != 0:
return {"error": result.stderr}
try:
return json.loads(result.stdout)
except json.JSONDecodeError:
return {"error": f"Invalid JSON: {result.stdout[:500]}"}
def get_metrics():
cmd = ["curl", "-s", "http://localhost:8000/metrics"]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
metrics = {}
for line in result.stdout.splitlines():
if line.startswith("vllm:kv_cache_usage_perc{"):
metrics["kv_usage"] = float(line.split()[-1])
return metrics
# ============================================================
# Test 1: Perplexity-like quality via logprobs
# ============================================================
print("=" * 60)
print(" Logprobs Quality Test at Different Context Lengths")
print("=" * 60)
print()
# Use a structured passage that requires understanding context
test_passages = {
"short": (
"The Fibonacci sequence starts with 0 and 1. Each subsequent number is the "
"sum of the two preceding ones. So the sequence goes: 0, 1, 1, 2, 3, 5, 8, 13, 21, 34. "
"The ratio of consecutive Fibonacci numbers approaches the golden ratio, which is approximately"
),
"medium": None, # will be padded version
"long": None, # will be padded version
}
# Filler that's semantically unrelated
filler = (
"The quick brown fox jumps over the lazy dog near the riverbank. "
"Lorem ipsum dolor sit amet, consectetur adipiscing elit. "
"Stars twinkle in the night sky as the moon rises above the hills. "
"The ocean waves crash against the rocky shore in steady rhythm. "
)
core_passage = test_passages["short"]
print("Testing: Can the model complete 'golden ratio is approximately...' after varying filler")
print()
for target_tokens, label in [(200, "~200 tok"), (1000, "~1k tok"), (3000, "~3k tok"),
(5000, "~5k tok"), (7000, "~7k tok"), (8000, "~8k tok")]:
# Build prompt with filler BEFORE the passage
filler_needed = max(0, target_tokens * 4 - len(core_passage))
padded_filler = ""
while len(padded_filler) < filler_needed:
padded_filler += filler
padded_filler = padded_filler[:filler_needed]
prompt = padded_filler + "\n\n" + core_passage
t0 = time.time()
resp = curl_post("chat/completions", {
"model": MODEL,
"messages": [{"role": "user", "content": prompt + "\n\nComplete the number:"}],
"max_tokens": 20,
"temperature": 0.0,
"logprobs": True,
"top_logprobs": 5,
"chat_template_kwargs": {"enable_thinking": False},
}, timeout=120)
elapsed = time.time() - t0
if "error" in resp:
print(f" {label}: ERROR - {resp['error'][:100]}")
continue
usage = resp.get("usage", {})
output = resp["choices"][0]["message"]["content"] if resp.get("choices") else ""
# Check if model knows the golden ratio
correct = any(x in output for x in ["1.618", "1.62", "phi", "1.6180"])
# Get logprobs from response
logprobs_data = resp["choices"][0].get("logprobs", {})
token_logprobs = []
if logprobs_data and logprobs_data.get("content"):
for tok_info in logprobs_data["content"]:
if tok_info.get("logprob") is not None:
token_logprobs.append(tok_info["logprob"])
avg_logprob = sum(token_logprobs) / len(token_logprobs) if token_logprobs else 0
perplexity = math.exp(-avg_logprob) if avg_logprob < 0 else float("inf")
print(f" {label:>10} | tokens={usage.get('prompt_tokens', '?'):>5} | "
f"ppl={perplexity:>6.2f} | correct={correct} | "
f"time={elapsed:.2f}s | output: {output[:50]}")
# ============================================================
# Test 2: Multi-needle at maximum context
# ============================================================
print()
print("=" * 60)
print(" Multi-Needle at Maximum Context (8192)")
print("=" * 60)
print()
facts = [
("The password for vault Alpha is PHOENIX-2891.", "PHOENIX-2891", 0.1),
("The activation code for system Beta is 7743-OMEGA.", "7743-OMEGA", 0.3),
("Agent Rivera's emergency contact number is 555-0173.", "555-0173", 0.5),
("The launch sequence for Project Gamma is DELTA-ECHO-9.", "DELTA-ECHO-9", 0.7),
("The encryption key for Channel Sigma is XJ7-KAPPA-412.", "XJ7-KAPPA-412", 0.9),
]
# Build prompt that fills nearly all of 8192 tokens
target_chars = 7500 * 4 # leave room for question
filler_block = (
"In the vast expanse of the digital frontier, data streams flow like rivers. "
"Each byte carries information across networks spanning the globe. "
"Servers hum in climate-controlled rooms, processing millions of requests. "
"The architecture of modern computing relies on layers of abstraction. "
"From silicon transistors to high-level programming languages, each layer "
"builds upon the previous one to create increasingly complex systems. "
)
segments = []
chars_per_segment = target_chars // (len(facts) + 1)
for i, (fact, _, pos) in enumerate(facts):
fill = ""
while len(fill) < chars_per_segment:
fill += filler_block
fill = fill[:chars_per_segment]
segments.append(fill)
segments.append(f"\n\n{fact}\n\n")
# Final filler
fill = ""
while len(fill) < chars_per_segment:
fill += filler_block
segments.append(fill[:chars_per_segment])
question = "\n\nAnswer these questions with just the answers, one per line:\n"
question += "1. What is the password for vault Alpha?\n"
question += "2. What is the activation code for system Beta?\n"
question += "3. What is Agent Rivera's emergency contact number?\n"
question += "4. What is the launch sequence for Project Gamma?\n"
question += "5. What is the encryption key for Channel Sigma?\n"
prompt = "".join(segments) + question
t0 = time.time()
resp = curl_post("chat/completions", {
"model": MODEL,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 200,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
}, timeout=180)
elapsed = time.time() - t0
if "error" in resp:
print(f"ERROR: {resp['error'][:200]}")
else:
usage = resp.get("usage", {})
output = resp["choices"][0]["message"]["content"] if resp.get("choices") else ""
print(f"Prompt tokens: {usage.get('prompt_tokens', '?')}")
print(f"Elapsed: {elapsed:.2f}s")
print(f"Output:\n{output}")
print()
found = 0
for fact_text, expected, _ in facts:
if expected.lower() in output.lower():
found += 1
print(f" FOUND: {expected}")
else:
print(f" MISSED: {expected}")
print(f"\nCoherence score: {found}/{len(facts)} facts retrieved at near-max context")
# ============================================================
# Test 3: Generation quality — does output make sense at 8k context?
# ============================================================
print()
print("=" * 60)
print(" Generation Quality at Max Context")
print("=" * 60)
print()
# Fill context then ask a complex question requiring reasoning
filler_chars = 7000 * 4
padded = ""
while len(padded) < filler_chars:
padded += filler_block
padded = padded[:filler_chars]
complex_question = (
"\n\nIgnore all the text above. Now answer this question:\n"
"A farmer has 3 fields. Field A produces 120 bushels of wheat per acre. "
"Field B produces 85 bushels per acre. Field C produces 150 bushels per acre. "
"If the farmer plants 10 acres of A, 15 acres of B, and 8 acres of C, "
"how many total bushels does he produce? Show your calculation."
)
prompt = padded + complex_question
t0 = time.time()
resp = curl_post("chat/completions", {
"model": MODEL,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 300,
"temperature": 0.0,
"chat_template_kwargs": {"enable_thinking": False},
}, timeout=180)
elapsed = time.time() - t0
if "error" in resp:
print(f"ERROR: {resp['error'][:200]}")
else:
usage = resp.get("usage", {})
output = resp["choices"][0]["message"]["content"] if resp.get("choices") else ""
correct_answer = 10*120 + 15*85 + 8*150 # = 1200 + 1275 + 1200 = 3675
print(f"Prompt tokens: {usage.get('prompt_tokens', '?')}")
print(f"Elapsed: {elapsed:.2f}s")
print(f"Correct answer: {correct_answer}")
print(f"Output:\n{output[:500]}")
print(f"\nAnswer correct: {'3675' in output}")
print()
print("=" * 60)
print(" DONE")
print("=" * 60)
validate_paper.py
-353
View File
@@ -1,353 +0,0 @@
#!/usr/bin/env python3
"""
TurboQuant validation against paper claims (arXiv:2504.19874).
Validates:
1. MSE distortion matches paper's Theorem 1 bounds
2. Inner-product estimator is unbiased (Theorem 2)
3. Inner-product distortion within paper's Theorem 3 bounds
4. Attention recall@k at realistic scale (d=128, N=4096)
5. Needle-in-haystack retrieval at scale (d=128, N=8192, multiple needle depths)
6. Compression ratio matches paper's claimed 2.6x per layer
7. Codebook MSE matches paper's Table 1 values exactly
"""
import math
import sys
import torch
import numpy as np
torch.manual_seed(42)
np.random.seed(42)
PASS = 0
FAIL = 0
def run(name, fn):
global PASS, FAIL
try:
fn()
print(f" PASS {name}")
PASS += 1
except Exception as e:
import traceback
print(f" FAIL {name}")
traceback.print_exc()
FAIL += 1
# ---------- 1. MSE distortion (Theorem 1) ----------
def test_mse_distortion_bounds():
"""Paper Theorem 1: MSE <= sqrt(3)*pi/2 * 1/4^b per coordinate.
Bounds: b=1: 0.360, b=2: 0.117, b=3: 0.030, b=4: 0.009"""
from turboquant.quantizer import TurboQuantMSE
d = 128
N = 10000
bounds = {1: 0.360, 2: 0.117, 3: 0.030, 4: 0.009}
for bits, expected in bounds.items():
q = TurboQuantMSE(dim=d, bits=bits, device="cpu", seed=42)
x = torch.randn(N, d)
x = x / x.norm(dim=-1, keepdim=True) # unit norm
x_hat = q(x)
mse_per_coord = ((x - x_hat) ** 2).mean().item()
# Allow 15% tolerance (paper bound is an upper bound, empirical should be at or below)
assert mse_per_coord <= expected * 1.15, \
f"bits={bits}: MSE/coord={mse_per_coord:.4f} > bound*1.15={expected*1.15:.4f}"
def test_mse_codebook_table1():
"""Validate codebook total MSE values match paper Table 1."""
from turboquant.codebook import get_codebook
# Paper Table 1 values are total MSE for d-dimensional unit vector
paper_values = {1: 0.360, 2: 0.117, 3: 0.030, 4: 0.009}
for bits, expected in paper_values.items():
cb = get_codebook(128, bits)
actual = cb["mse_total"]
ratio = actual / expected
assert 0.85 <= ratio <= 1.20, \
f"bits={bits}: total MSE={actual:.4f}, expected~{expected:.3f}, ratio={ratio:.3f}"
# ---------- 2. Unbiasedness (Theorem 2) ----------
def test_prod_unbiased():
"""Paper Theorem 2: E[<y, x_tilde>] = <y, x>"""
from turboquant.quantizer import TurboQuantProd
d = 128
N = 5000
n_trials = 20
for bits in [2, 3, 4]:
biases = []
for trial in range(n_trials):
q = TurboQuantProd(dim=d, bits=bits, device="cpu", seed=trial * 100)
x = torch.randn(1, 1, N, d)
y = torch.randn(1, 1, 1, d)
true_ip = (y * x).sum(dim=-1) # (1, 1, N)
key_q = q.quantize(x)
est_ip = q.attention_score(y, key_q) # (1, 1, 1, N)
bias = (est_ip.squeeze() - true_ip.squeeze()).mean().item()
biases.append(bias)
mean_bias = np.mean(biases)
# Unbiased means mean bias should be near zero relative to signal magnitude
assert abs(mean_bias) < 0.05, \
f"bits={bits}: mean bias={mean_bias:.4f} (should be ~0)"
# ---------- 3. Inner-product distortion (Theorem 3) ----------
def test_prod_distortion_scaling():
"""Paper Theorem 3: D_prod <= sqrt(3)*pi^2*||y||^2/d * 1/4^b.
Distortion should decrease ~4x when adding 1 bit."""
from turboquant.quantizer import TurboQuantProd
d = 128
N = 2000
distortions = {}
for bits in [2, 3, 4]:
q = TurboQuantProd(dim=d, bits=bits, device="cpu", seed=42)
x = torch.randn(1, 1, N, d)
y = torch.randn(1, 1, 1, d)
true_ip = (y * x).sum(dim=-1).squeeze()
key_q = q.quantize(x)
est_ip = q.attention_score(y, key_q).squeeze()
mse = ((est_ip - true_ip) ** 2).mean().item()
distortions[bits] = mse
# Each extra bit should reduce distortion by roughly 4x (1/4^b scaling)
ratio_2_to_3 = distortions[2] / distortions[3]
ratio_3_to_4 = distortions[3] / distortions[4]
assert ratio_2_to_3 > 2.0, \
f"2->3 bit distortion ratio={ratio_2_to_3:.2f} (expected ~4x, at least >2x)"
assert ratio_3_to_4 > 2.0, \
f"3->4 bit distortion ratio={ratio_3_to_4:.2f} (expected ~4x, at least >2x)"
# ---------- 4. Attention recall@k at scale ----------
def test_recall_at_scale():
"""Recall@8 with d=128, N=4096 (realistic LLM KV size)."""
from turboquant.quantizer import TurboQuantProd
d = 128
N = 4096
n_queries = 32
k = 8
results = {}
for bits in [3, 4]:
q = TurboQuantProd(dim=d, bits=bits, device="cpu", seed=42)
keys = torch.randn(1, 1, N, d) * 0.1
queries = torch.randn(1, 1, n_queries, d) * 0.1
true_scores = torch.matmul(queries, keys.transpose(-2, -1)).squeeze(0).squeeze(0)
true_topk = true_scores.topk(k, dim=-1).indices
key_q = q.quantize(keys)
tq_scores = q.attention_score(queries, key_q).squeeze(0).squeeze(0)
tq_topk = tq_scores.topk(k, dim=-1).indices
recalls = []
for qi in range(n_queries):
true_set = set(true_topk[qi].tolist())
tq_set = set(tq_topk[qi].tolist())
recalls.append(len(true_set & tq_set) / k)
results[bits] = np.mean(recalls)
assert results[3] >= 0.40, f"3-bit recall@8={results[3]:.3f} < 0.40"
assert results[4] >= 0.55, f"4-bit recall@8={results[4]:.3f} < 0.55"
assert results[4] > results[3], "4-bit should have better recall than 3-bit"
# ---------- 5. Needle-in-haystack at multiple depths ----------
def test_needle_retrieval_depths():
"""Needle retrieval at different positions in a 4096-token context."""
from turboquant.store import CompressedKVStore
d = 128
H_kv = 4
N = 4096
depths = [0.1, 0.25, 0.5, 0.75, 0.9] # fraction into context
for bits in [3, 4]:
for depth in depths:
needle_pos = int(N * depth)
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=bits, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
keys = torch.randn(N, H_kv, d) * 0.02
values = torch.randn(N, H_kv, d)
needle_key = torch.randn(1, H_kv, d) * 3.0
keys[needle_pos] = needle_key.squeeze(0)
store.append_chunk(keys, values)
flat = store.get_flat_cache()
k_dequant = store.quantizer.dequantize(flat.prod_q)
query_vec = needle_key.squeeze(0).unsqueeze(0)
scores = torch.bmm(
query_vec.float().transpose(0, 1),
k_dequant.float().transpose(1, 2),
).squeeze(1)
for h in range(H_kv):
top_idx = scores[h].argmax().item()
assert top_idx == needle_pos, \
f"bits={bits} depth={depth} head={h}: needle@{needle_pos} top@{top_idx}"
def test_needle_chunked_8192():
"""Needle retrieval in 8192 tokens split across multiple chunks."""
from turboquant.store import CompressedKVStore
d = 128
H_kv = 2
total = 8192
chunk_size = 1024
needle_pos = 5555
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
all_keys = torch.randn(total, H_kv, d) * 0.02
all_values = torch.randn(total, H_kv, d)
needle_key = torch.randn(1, H_kv, d) * 3.0
all_keys[needle_pos] = needle_key.squeeze(0)
for i in range(0, total, chunk_size):
store.append_chunk(all_keys[i:i+chunk_size], all_values[i:i+chunk_size])
flat = store.get_flat_cache()
k_dequant = store.quantizer.dequantize(flat.prod_q)
query_vec = needle_key.squeeze(0).unsqueeze(0)
scores = torch.bmm(
query_vec.float().transpose(0, 1),
k_dequant.float().transpose(1, 2),
).squeeze(1)
for h in range(H_kv):
top_idx = scores[h].argmax().item()
assert top_idx == needle_pos, \
f"head={h}: needle@{needle_pos} top@{top_idx}"
# ---------- 6. Compression ratio ----------
def test_compression_ratio():
"""Verify compression matches claimed 2.6x for 3-bit keys + 2-bit values."""
from turboquant.store import CompressedKVStore
d = 128
H_kv = 8
N = 4096
store = CompressedKVStore(
head_dim=d, num_kv_heads=H_kv, key_bits=3, value_bits=2,
value_group_size=32, device=torch.device("cpu"),
)
k = torch.randn(N, H_kv, d)
v = torch.randn(N, H_kv, d)
store.append_chunk(k, v)
tq_bytes = store.memory_bytes()
fp16_bytes = N * H_kv * d * 2 * 2 # K+V in FP16
ratio = fp16_bytes / tq_bytes
assert ratio > 2.0, f"Compression ratio {ratio:.2f}x < 2.0x"
# Implementation uses 3-bit keys + 2-bit values + overhead, so ratio should be 2-6x
assert ratio < 8.0, f"Compression ratio {ratio:.2f}x > 8.0x (suspiciously high)"
# ---------- 7. Rank correlation at scale ----------
def test_rank_correlation_scale():
"""Spearman rank correlation at d=128, N=2048."""
from turboquant.quantizer import TurboQuantProd
d = 128
N = 2048
for bits in [3, 4]:
q = TurboQuantProd(dim=d, bits=bits, device="cpu", seed=42)
keys = torch.randn(1, 1, N, d) * 0.1
query = torch.randn(1, 1, 1, d) * 0.1
true_scores = torch.matmul(query, keys.transpose(-2, -1)).squeeze()
key_q = q.quantize(keys)
tq_scores = q.attention_score(query, key_q).squeeze()
true_ranks = true_scores.argsort().argsort().float()
tq_ranks = tq_scores.argsort().argsort().float()
corr = torch.corrcoef(torch.stack([true_ranks, tq_ranks]))[0, 1].item()
if bits == 3:
assert corr > 0.75, f"3-bit rank corr={corr:.3f} < 0.75"
elif bits == 4:
assert corr > 0.90, f"4-bit rank corr={corr:.3f} < 0.90"
# ---------- Main ----------
if __name__ == "__main__":
print()
print("=" * 60)
print("TurboQuant Paper Validation (arXiv:2504.19874)")
print("=" * 60)
print()
print("-- Theorem 1: MSE Distortion Bounds --")
run("MSE distortion <= paper bound (b=1..4)", test_mse_distortion_bounds)
run("Codebook MSE matches Table 1", test_mse_codebook_table1)
print()
print("-- Theorem 2: Unbiasedness --")
run("E[<y, x~>] = <y, x> (bits=2,3,4)", test_prod_unbiased)
print()
print("-- Theorem 3: Inner-Product Distortion --")
run("Distortion scales as 1/4^b", test_prod_distortion_scaling)
print()
print("-- Attention Quality --")
run("Recall@8 at d=128, N=4096 (bits=3,4)", test_recall_at_scale)
run("Rank correlation at d=128, N=2048", test_rank_correlation_scale)
print()
print("-- Needle-in-Haystack --")
run("Needle at 5 depths in 4096 tokens (bits=3,4)", test_needle_retrieval_depths)
run("Needle in 8192 tokens, chunked (3-bit)", test_needle_chunked_8192)
print()
print("-- Compression --")
run("Compression ratio > 2x", test_compression_ratio)
print()
print("=" * 60)
print(f"Results: {PASS} passed, {FAIL} failed (total {PASS + FAIL})")
if FAIL == 0:
print("All validations passed against paper claims.")
else:
print(f"{FAIL} validation(s) failed")
print("=" * 60)
sys.exit(1 if FAIL > 0 else 0)