rigtune 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- rigtune-0.2.0/LICENSE +21 -0
- rigtune-0.2.0/PKG-INFO +107 -0
- rigtune-0.2.0/README.md +69 -0
- rigtune-0.2.0/pyproject.toml +67 -0
- rigtune-0.2.0/rigtune.egg-info/PKG-INFO +107 -0
- rigtune-0.2.0/rigtune.egg-info/SOURCES.txt +69 -0
- rigtune-0.2.0/rigtune.egg-info/dependency_links.txt +1 -0
- rigtune-0.2.0/rigtune.egg-info/entry_points.txt +2 -0
- rigtune-0.2.0/rigtune.egg-info/requires.txt +16 -0
- rigtune-0.2.0/rigtune.egg-info/top_level.txt +1 -0
- rigtune-0.2.0/setup.cfg +4 -0
- rigtune-0.2.0/tests/test_ablation.py +115 -0
- rigtune-0.2.0/tests/test_ablation_shrink.py +22 -0
- rigtune-0.2.0/tests/test_all_levels.py +31 -0
- rigtune-0.2.0/tests/test_attention_check.py +31 -0
- rigtune-0.2.0/tests/test_baselines.py +81 -0
- rigtune-0.2.0/tests/test_break_even.py +54 -0
- rigtune-0.2.0/tests/test_budget.py +65 -0
- rigtune-0.2.0/tests/test_cache_and_pipeline.py +117 -0
- rigtune-0.2.0/tests/test_cli.py +71 -0
- rigtune-0.2.0/tests/test_combinations.py +157 -0
- rigtune-0.2.0/tests/test_compare.py +117 -0
- rigtune-0.2.0/tests/test_compare_heldout.py +26 -0
- rigtune-0.2.0/tests/test_compare_repeats.py +76 -0
- rigtune-0.2.0/tests/test_confirm.py +70 -0
- rigtune-0.2.0/tests/test_contenders.py +114 -0
- rigtune-0.2.0/tests/test_cpu_quota.py +32 -0
- rigtune-0.2.0/tests/test_dataset_workloads.py +85 -0
- rigtune-0.2.0/tests/test_drift.py +237 -0
- rigtune-0.2.0/tests/test_early_stop_check.py +40 -0
- rigtune-0.2.0/tests/test_explore.py +79 -0
- rigtune-0.2.0/tests/test_fit_per_machine.py +62 -0
- rigtune-0.2.0/tests/test_gguf_and_backends.py +129 -0
- rigtune-0.2.0/tests/test_hardware_and_config.py +75 -0
- rigtune-0.2.0/tests/test_hub_search.py +41 -0
- rigtune-0.2.0/tests/test_int8_kv.py +53 -0
- rigtune-0.2.0/tests/test_load_pick_order.py +27 -0
- rigtune-0.2.0/tests/test_make_dolly_prompts.py +29 -0
- rigtune-0.2.0/tests/test_make_oasst_prompts.py +28 -0
- rigtune-0.2.0/tests/test_measure_and_proxy.py +252 -0
- rigtune-0.2.0/tests/test_measure_levels.py +102 -0
- rigtune-0.2.0/tests/test_memcal.py +174 -0
- rigtune-0.2.0/tests/test_memory.py +88 -0
- rigtune-0.2.0/tests/test_nvlink.py +67 -0
- rigtune-0.2.0/tests/test_objectives_and_search.py +219 -0
- rigtune-0.2.0/tests/test_oom_retry.py +60 -0
- rigtune-0.2.0/tests/test_phases.py +418 -0
- rigtune-0.2.0/tests/test_power.py +428 -0
- rigtune-0.2.0/tests/test_predict.py +150 -0
- rigtune-0.2.0/tests/test_process_group.py +33 -0
- rigtune-0.2.0/tests/test_progress_line.py +42 -0
- rigtune-0.2.0/tests/test_quality.py +196 -0
- rigtune-0.2.0/tests/test_quant_default.py +58 -0
- rigtune-0.2.0/tests/test_report.py +112 -0
- rigtune-0.2.0/tests/test_rescore_ttft.py +43 -0
- rigtune-0.2.0/tests/test_search_budget.py +83 -0
- rigtune-0.2.0/tests/test_search_view.py +115 -0
- rigtune-0.2.0/tests/test_selector.py +48 -0
- rigtune-0.2.0/tests/test_sglang_python.py +31 -0
- rigtune-0.2.0/tests/test_spec_first.py +88 -0
- rigtune-0.2.0/tests/test_spec_methods.py +50 -0
- rigtune-0.2.0/tests/test_strategies.py +640 -0
- rigtune-0.2.0/tests/test_strategy_report.py +98 -0
- rigtune-0.2.0/tests/test_supervisor.py +123 -0
- rigtune-0.2.0/tests/test_tail.py +83 -0
- rigtune-0.2.0/tests/test_task_quality.py +29 -0
- rigtune-0.2.0/tests/test_ttft_percentile.py +52 -0
- rigtune-0.2.0/tests/test_vllm_cli.py +39 -0
- rigtune-0.2.0/tests/test_w8a8.py +87 -0
- rigtune-0.2.0/tests/test_workload_file.py +70 -0
- rigtune-0.2.0/tests/test_workloads.py +127 -0
rigtune-0.2.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Aagam Sachin Bothara
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
rigtune-0.2.0/PKG-INFO
ADDED
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: rigtune
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Autotuner for LLM serving: measures vLLM, SGLang and llama.cpp settings on your GPU and your traffic, then serves the fastest
|
|
5
|
+
Author-email: Aagam Sachin Bothara <mweeb19@gmail.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/Aagam-Bothara/rigtune
|
|
8
|
+
Project-URL: Documentation, https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md
|
|
9
|
+
Project-URL: Benchmarks, https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md
|
|
10
|
+
Project-URL: Issues, https://github.com/Aagam-Bothara/rigtune/issues
|
|
11
|
+
Keywords: llm,inference,vllm,llama.cpp,serving,autotuning
|
|
12
|
+
Classifier: Development Status :: 3 - Alpha
|
|
13
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Requires-Python: <3.14,>=3.10
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Requires-Dist: pydantic>=2.0
|
|
24
|
+
Requires-Dist: typer>=0.12
|
|
25
|
+
Requires-Dist: rich>=13.0
|
|
26
|
+
Requires-Dist: httpx>=0.27
|
|
27
|
+
Requires-Dist: fastapi>=0.110
|
|
28
|
+
Requires-Dist: uvicorn>=0.29
|
|
29
|
+
Requires-Dist: psutil>=5.9
|
|
30
|
+
Requires-Dist: huggingface_hub>=0.23
|
|
31
|
+
Provides-Extra: nvml
|
|
32
|
+
Requires-Dist: nvidia-ml-py>=12.535.133; extra == "nvml"
|
|
33
|
+
Provides-Extra: dev
|
|
34
|
+
Requires-Dist: pytest>=7.0; extra == "dev"
|
|
35
|
+
Requires-Dist: pytest-asyncio>=0.23; extra == "dev"
|
|
36
|
+
Requires-Dist: ruff==0.15.0; extra == "dev"
|
|
37
|
+
Dynamic: license-file
|
|
38
|
+
|
|
39
|
+
# rigtune
|
|
40
|
+
|
|
41
|
+
[](https://github.com/Aagam-Bothara/rigtune/actions/workflows/tests.yml)
|
|
42
|
+
|
|
43
|
+
*Formerly PolyServe, renamed because a separate [2025 paper](https://arxiv.org/abs/2507.17769) already has that name. Profiles, results and settings from before the rename keep working ([details](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md#renamed-from-polyserve)).*
|
|
44
|
+
|
|
45
|
+
**rigtune finds the fastest way to serve an LLM on your GPU and your traffic, then serves it behind an OpenAI-compatible API.** It measures real configurations — engine, weight precision, batch size, KV-cache type, speculative decoding — instead of guessing, because the settings that won changed on every card tried.
|
|
46
|
+
|
|
47
|
+
**It beats a hand-tuned configuration by 56–97%**, which is the comparison that matters — not the one against untouched defaults.
|
|
48
|
+
|
|
49
|
+
| Llama 3.1 8B, held-out prompts | rigtune | a fixed rule of thumb † | (stock defaults) |
|
|
50
|
+
|---|---|---|---|
|
|
51
|
+
| RTX 4090 | **1041 tok/s** | 667 — **+56%** | SGLang 448 |
|
|
52
|
+
| A100 | **1155 tok/s** | 674 — **+71%** | vLLM 726 |
|
|
53
|
+
| A40 | **506 tok/s** | 257 — **+97%** | vLLM 262 |
|
|
54
|
+
|
|
55
|
+
† fp8 weights and KV cache, a short context, batch 256: what an informed engineer sets without measuring. That rule was **slower than changing nothing** on both Ampere cards, because an fp8 KV cache costs throughput there — which is the point: the settings that win move from card to card, so a rule that is right somewhere is wrong elsewhere. The stock column is context rather than a claim; nobody serious ships untouched defaults. Each row was measured on prompts held out of calibration ([full sessions](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#llama-31-8b-on-four-gpus)).
|
|
56
|
+
|
|
57
|
+
## Results
|
|
58
|
+
|
|
59
|
+
- **Your traffic decides what pays, not the card alone.** On one A40 with Qwen2.5-3B the winning setting changed with the prompts: suffix decoding on Dolly prompts (+122% over stock), a draft model on news-article extraction (+66.5%), an fp8 KV cache on real chat (+10%). A tuner that ignores your prompts cannot find these.
|
|
60
|
+
- **The pick holds on prompts it never saw.** Every number here was measured on prompts held out of calibration: 300 Dolly-15k prompts to tune on, 300 others to score on. An H100 not in the table gained +24% over stock vLLM with fp8 weights.
|
|
61
|
+
- **Calibration pays for itself within hours.** 22–56 minutes per model and workload, repaid by 0.5–2.6 hours of busy serving. Measuring the busiest load first cuts 29–45% off that without changing a pick.
|
|
62
|
+
- **What the search order is worth: one win, one tie.** Given the same time, random sampling of the same space matched this staged search on all three 8B cards. On a 14B in fp8 on a 24 GB card the staged search won by 4.4% (385 against 369 tok/s); on the same card in the 7× larger space that includes 4-bit checkpoints, the two tied (1412 against 1416) even though random search was given less time ([the evidence](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#qwen25-14b-where-memory-binds)). What staging reliably gives is the same answer every run, every strategy tried at least once, and a profile saying what each setting was worth — the gains come from measuring, not from the order.
|
|
63
|
+
|
|
64
|
+
Every result, with methods, ablations, quality checks and limits: [docs/benchmarks.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md).
|
|
65
|
+
|
|
66
|
+
## How this differs from vLLM's auto-tune and AIConfigurator
|
|
67
|
+
|
|
68
|
+
Not benchmarked against either — the comparison below is one of *kind*, not of measured results, and saying otherwise would be the overclaiming this project tries to avoid.
|
|
69
|
+
|
|
70
|
+
- **vLLM's `auto_tune`** sweeps one engine's throughput knobs — batch size and batched-token budget — to find the best setting that still meets a latency target. Same spirit, narrower question: it tunes vLLM once you have already chosen vLLM, your weight precision and your decoding strategy. rigtune treats the engine itself as a variable, alongside precision, KV-cache type and speculative decoding — the dimensions that carried almost all the gain here. On an RTX 4090 SGLang led the early trials and vLLM only overtook it once speculative decoding and a quantized cache were tried — tuning the early leader alone would have eliminated the eventual winner, and a single-engine sweep never sees the crossover at all.
|
|
71
|
+
- **AIConfigurator** predicts good deployments analytically from a performance model, aimed at TensorRT-LLM, and answers sizing questions — how many GPUs, which parallelism — very cheaply because it does not launch anything. rigtune measures instead, which costs 20–50 minutes per workload and is why its answers are specific to your prompts. We tried the analytical shortcut and kept it small for a reason: our own fitted predictor reaches a rank correlation of 0.88–0.94, enough to prune quantizations that cannot win, but it ranks the true winner at a median of 15th out of 25 measured configurations — so prediction is a filter here, never the decision ([why](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/decisions.md)).
|
|
72
|
+
- **What none of the three does** is prove the others wrong. The honest gap is that no one outside this project has run rigtune, and we have not run their tools on our hardware.
|
|
73
|
+
|
|
74
|
+
## Quickstart
|
|
75
|
+
|
|
76
|
+
```bash
|
|
77
|
+
pip install rigtune
|
|
78
|
+
pip install "vllm==0.29.0" # driver 580+; older drivers: "vllm==0.11.0" "transformers>=4.56,<5"
|
|
79
|
+
rigtune serve Qwen/Qwen2.5-3B-Instruct --workload chat
|
|
80
|
+
curl localhost:8000/v1/chat/completions -H 'content-type: application/json' \
|
|
81
|
+
-d '{"model":"Qwen/Qwen2.5-3B-Instruct","messages":[{"role":"user","content":"hi"}]}'
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
The first launch calibrates and caches the result per machine, model, objective and workload; later launches serve at once. To tune for your own traffic, pass a sample of it: `--workload-file prompts.jsonl`, one prompt per line. Optional: `pip install arctic-inference==0.1.1` adds suffix decoding; SGLang can live in its own environment (set `SGLANG_PYTHON` to its `python`); llama.cpp needs `llama-server` built with CUDA, on `PATH` or in `LLAMA_SERVER`.
|
|
85
|
+
|
|
86
|
+
With Docker (built on the official vLLM image):
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
docker build -t rigtune .
|
|
90
|
+
docker run --gpus all --ipc=host -p 8000:8000 -v ~/.cache/huggingface:/root/.cache/huggingface \
|
|
91
|
+
-v ~/.rigtune:/root/.rigtune rigtune serve Qwen/Qwen2.5-3B-Instruct
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
### Supported hardware
|
|
95
|
+
|
|
96
|
+
Linux, Python 3.10–3.13. NVIDIA GPUs of compute capability 7.5 or newer run vLLM, SGLang and llama.cpp (measured on an A40, A100, H100 NVL, L4 and RTX 4090); older NVIDIA GPUs and x86 CPUs run llama.cpp. vLLM-CPU and pre-Turing GPUs have never been benchmarked ([details](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#hardware-and-engines-measured)).
|
|
97
|
+
|
|
98
|
+
## Learn more
|
|
99
|
+
|
|
100
|
+
- [docs/benchmarks.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md): every result, at a glance and in full, with methods, ablations, quality, limits and what is still unmeasured.
|
|
101
|
+
- [docs/usage.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md): how calibration works, workloads, objectives, every search option and the CLI.
|
|
102
|
+
- [docs/writeup.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/writeup.md): the design of the memory planner, calibration and the predictor, what the evidence does and does not support, and the roadmap.
|
|
103
|
+
- [docs/decisions.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/decisions.md): why each choice is the way it is — including why not just random search — with the measurement that settled it.
|
|
104
|
+
- [benchmarks/strategies/SUMMARY.md](https://github.com/Aagam-Bothara/rigtune/blob/main/benchmarks/strategies/SUMMARY.md): every table, regenerated from the raw JSON.
|
|
105
|
+
- [CONTRIBUTING.md](https://github.com/Aagam-Bothara/rigtune/blob/main/CONTRIBUTING.md): development setup, tests and adding a backend.
|
|
106
|
+
|
|
107
|
+
MIT licensed.
|
rigtune-0.2.0/README.md
ADDED
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
# rigtune
|
|
2
|
+
|
|
3
|
+
[](https://github.com/Aagam-Bothara/rigtune/actions/workflows/tests.yml)
|
|
4
|
+
|
|
5
|
+
*Formerly PolyServe, renamed because a separate [2025 paper](https://arxiv.org/abs/2507.17769) already has that name. Profiles, results and settings from before the rename keep working ([details](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md#renamed-from-polyserve)).*
|
|
6
|
+
|
|
7
|
+
**rigtune finds the fastest way to serve an LLM on your GPU and your traffic, then serves it behind an OpenAI-compatible API.** It measures real configurations — engine, weight precision, batch size, KV-cache type, speculative decoding — instead of guessing, because the settings that won changed on every card tried.
|
|
8
|
+
|
|
9
|
+
**It beats a hand-tuned configuration by 56–97%**, which is the comparison that matters — not the one against untouched defaults.
|
|
10
|
+
|
|
11
|
+
| Llama 3.1 8B, held-out prompts | rigtune | a fixed rule of thumb † | (stock defaults) |
|
|
12
|
+
|---|---|---|---|
|
|
13
|
+
| RTX 4090 | **1041 tok/s** | 667 — **+56%** | SGLang 448 |
|
|
14
|
+
| A100 | **1155 tok/s** | 674 — **+71%** | vLLM 726 |
|
|
15
|
+
| A40 | **506 tok/s** | 257 — **+97%** | vLLM 262 |
|
|
16
|
+
|
|
17
|
+
† fp8 weights and KV cache, a short context, batch 256: what an informed engineer sets without measuring. That rule was **slower than changing nothing** on both Ampere cards, because an fp8 KV cache costs throughput there — which is the point: the settings that win move from card to card, so a rule that is right somewhere is wrong elsewhere. The stock column is context rather than a claim; nobody serious ships untouched defaults. Each row was measured on prompts held out of calibration ([full sessions](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#llama-31-8b-on-four-gpus)).
|
|
18
|
+
|
|
19
|
+
## Results
|
|
20
|
+
|
|
21
|
+
- **Your traffic decides what pays, not the card alone.** On one A40 with Qwen2.5-3B the winning setting changed with the prompts: suffix decoding on Dolly prompts (+122% over stock), a draft model on news-article extraction (+66.5%), an fp8 KV cache on real chat (+10%). A tuner that ignores your prompts cannot find these.
|
|
22
|
+
- **The pick holds on prompts it never saw.** Every number here was measured on prompts held out of calibration: 300 Dolly-15k prompts to tune on, 300 others to score on. An H100 not in the table gained +24% over stock vLLM with fp8 weights.
|
|
23
|
+
- **Calibration pays for itself within hours.** 22–56 minutes per model and workload, repaid by 0.5–2.6 hours of busy serving. Measuring the busiest load first cuts 29–45% off that without changing a pick.
|
|
24
|
+
- **What the search order is worth: one win, one tie.** Given the same time, random sampling of the same space matched this staged search on all three 8B cards. On a 14B in fp8 on a 24 GB card the staged search won by 4.4% (385 against 369 tok/s); on the same card in the 7× larger space that includes 4-bit checkpoints, the two tied (1412 against 1416) even though random search was given less time ([the evidence](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#qwen25-14b-where-memory-binds)). What staging reliably gives is the same answer every run, every strategy tried at least once, and a profile saying what each setting was worth — the gains come from measuring, not from the order.
|
|
25
|
+
|
|
26
|
+
Every result, with methods, ablations, quality checks and limits: [docs/benchmarks.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md).
|
|
27
|
+
|
|
28
|
+
## How this differs from vLLM's auto-tune and AIConfigurator
|
|
29
|
+
|
|
30
|
+
Not benchmarked against either — the comparison below is one of *kind*, not of measured results, and saying otherwise would be the overclaiming this project tries to avoid.
|
|
31
|
+
|
|
32
|
+
- **vLLM's `auto_tune`** sweeps one engine's throughput knobs — batch size and batched-token budget — to find the best setting that still meets a latency target. Same spirit, narrower question: it tunes vLLM once you have already chosen vLLM, your weight precision and your decoding strategy. rigtune treats the engine itself as a variable, alongside precision, KV-cache type and speculative decoding — the dimensions that carried almost all the gain here. On an RTX 4090 SGLang led the early trials and vLLM only overtook it once speculative decoding and a quantized cache were tried — tuning the early leader alone would have eliminated the eventual winner, and a single-engine sweep never sees the crossover at all.
|
|
33
|
+
- **AIConfigurator** predicts good deployments analytically from a performance model, aimed at TensorRT-LLM, and answers sizing questions — how many GPUs, which parallelism — very cheaply because it does not launch anything. rigtune measures instead, which costs 20–50 minutes per workload and is why its answers are specific to your prompts. We tried the analytical shortcut and kept it small for a reason: our own fitted predictor reaches a rank correlation of 0.88–0.94, enough to prune quantizations that cannot win, but it ranks the true winner at a median of 15th out of 25 measured configurations — so prediction is a filter here, never the decision ([why](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/decisions.md)).
|
|
34
|
+
- **What none of the three does** is prove the others wrong. The honest gap is that no one outside this project has run rigtune, and we have not run their tools on our hardware.
|
|
35
|
+
|
|
36
|
+
## Quickstart
|
|
37
|
+
|
|
38
|
+
```bash
|
|
39
|
+
pip install rigtune
|
|
40
|
+
pip install "vllm==0.29.0" # driver 580+; older drivers: "vllm==0.11.0" "transformers>=4.56,<5"
|
|
41
|
+
rigtune serve Qwen/Qwen2.5-3B-Instruct --workload chat
|
|
42
|
+
curl localhost:8000/v1/chat/completions -H 'content-type: application/json' \
|
|
43
|
+
-d '{"model":"Qwen/Qwen2.5-3B-Instruct","messages":[{"role":"user","content":"hi"}]}'
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
The first launch calibrates and caches the result per machine, model, objective and workload; later launches serve at once. To tune for your own traffic, pass a sample of it: `--workload-file prompts.jsonl`, one prompt per line. Optional: `pip install arctic-inference==0.1.1` adds suffix decoding; SGLang can live in its own environment (set `SGLANG_PYTHON` to its `python`); llama.cpp needs `llama-server` built with CUDA, on `PATH` or in `LLAMA_SERVER`.
|
|
47
|
+
|
|
48
|
+
With Docker (built on the official vLLM image):
|
|
49
|
+
|
|
50
|
+
```bash
|
|
51
|
+
docker build -t rigtune .
|
|
52
|
+
docker run --gpus all --ipc=host -p 8000:8000 -v ~/.cache/huggingface:/root/.cache/huggingface \
|
|
53
|
+
-v ~/.rigtune:/root/.rigtune rigtune serve Qwen/Qwen2.5-3B-Instruct
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
### Supported hardware
|
|
57
|
+
|
|
58
|
+
Linux, Python 3.10–3.13. NVIDIA GPUs of compute capability 7.5 or newer run vLLM, SGLang and llama.cpp (measured on an A40, A100, H100 NVL, L4 and RTX 4090); older NVIDIA GPUs and x86 CPUs run llama.cpp. vLLM-CPU and pre-Turing GPUs have never been benchmarked ([details](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#hardware-and-engines-measured)).
|
|
59
|
+
|
|
60
|
+
## Learn more
|
|
61
|
+
|
|
62
|
+
- [docs/benchmarks.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md): every result, at a glance and in full, with methods, ablations, quality, limits and what is still unmeasured.
|
|
63
|
+
- [docs/usage.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md): how calibration works, workloads, objectives, every search option and the CLI.
|
|
64
|
+
- [docs/writeup.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/writeup.md): the design of the memory planner, calibration and the predictor, what the evidence does and does not support, and the roadmap.
|
|
65
|
+
- [docs/decisions.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/decisions.md): why each choice is the way it is — including why not just random search — with the measurement that settled it.
|
|
66
|
+
- [benchmarks/strategies/SUMMARY.md](https://github.com/Aagam-Bothara/rigtune/blob/main/benchmarks/strategies/SUMMARY.md): every table, regenerated from the raw JSON.
|
|
67
|
+
- [CONTRIBUTING.md](https://github.com/Aagam-Bothara/rigtune/blob/main/CONTRIBUTING.md): development setup, tests and adding a backend.
|
|
68
|
+
|
|
69
|
+
MIT licensed.
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61.0", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "rigtune"
|
|
7
|
+
version = "0.2.0"
|
|
8
|
+
description = "Autotuner for LLM serving: measures vLLM, SGLang and llama.cpp settings on your GPU and your traffic, then serves the fastest"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10,<3.14"
|
|
11
|
+
license = {text = "MIT"}
|
|
12
|
+
authors = [{name = "Aagam Sachin Bothara", email = "mweeb19@gmail.com"}]
|
|
13
|
+
keywords = ["llm", "inference", "vllm", "llama.cpp", "serving", "autotuning"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 3 - Alpha",
|
|
16
|
+
"Operating System :: POSIX :: Linux",
|
|
17
|
+
"Programming Language :: Python :: 3",
|
|
18
|
+
"Programming Language :: Python :: 3.10",
|
|
19
|
+
"Programming Language :: Python :: 3.11",
|
|
20
|
+
"Programming Language :: Python :: 3.12",
|
|
21
|
+
"Programming Language :: Python :: 3.13",
|
|
22
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
23
|
+
]
|
|
24
|
+
dependencies = [
|
|
25
|
+
"pydantic>=2.0",
|
|
26
|
+
"typer>=0.12",
|
|
27
|
+
"rich>=13.0",
|
|
28
|
+
"httpx>=0.27",
|
|
29
|
+
"fastapi>=0.110",
|
|
30
|
+
"uvicorn>=0.29",
|
|
31
|
+
"psutil>=5.9",
|
|
32
|
+
"huggingface_hub>=0.23",
|
|
33
|
+
]
|
|
34
|
+
|
|
35
|
+
[project.optional-dependencies]
|
|
36
|
+
nvml = ["nvidia-ml-py>=12.535.133"]
|
|
37
|
+
# llmtrace, the GPU tracer rigtune prefers for telemetry (it falls back to pynvml, then psutil), is not an
|
|
38
|
+
# extra: PyPI refuses packages with git dependencies, and "llmtrace" on PyPI is an unrelated project. Install
|
|
39
|
+
# it from git: pip install "llmtrace @ git+https://github.com/Aagam-Bothara/llmtrace.git"
|
|
40
|
+
# ruff is pinned: its default rules change between releases (0.16 added pyupgrade), so an unpinned
|
|
41
|
+
# linter passed locally and failed in CI on the same code.
|
|
42
|
+
dev = ["pytest>=7.0", "pytest-asyncio>=0.23", "ruff==0.15.0"]
|
|
43
|
+
|
|
44
|
+
[project.urls]
|
|
45
|
+
Homepage = "https://github.com/Aagam-Bothara/rigtune"
|
|
46
|
+
Documentation = "https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md"
|
|
47
|
+
Benchmarks = "https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md"
|
|
48
|
+
Issues = "https://github.com/Aagam-Bothara/rigtune/issues"
|
|
49
|
+
|
|
50
|
+
[project.scripts]
|
|
51
|
+
rigtune = "rigtune.cli:app"
|
|
52
|
+
|
|
53
|
+
[tool.setuptools.packages.find]
|
|
54
|
+
where = ["."]
|
|
55
|
+
include = ["rigtune*"]
|
|
56
|
+
|
|
57
|
+
[tool.pytest.ini_options]
|
|
58
|
+
testpaths = ["tests"]
|
|
59
|
+
addopts = "-q"
|
|
60
|
+
asyncio_mode = "auto"
|
|
61
|
+
|
|
62
|
+
[tool.ruff]
|
|
63
|
+
line-length = 110
|
|
64
|
+
target-version = "py310"
|
|
65
|
+
|
|
66
|
+
[tool.ruff.lint]
|
|
67
|
+
select = ["E4", "E7", "E9", "F"] # explicit, so the rule set does not move with ruff's defaults
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: rigtune
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Autotuner for LLM serving: measures vLLM, SGLang and llama.cpp settings on your GPU and your traffic, then serves the fastest
|
|
5
|
+
Author-email: Aagam Sachin Bothara <mweeb19@gmail.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/Aagam-Bothara/rigtune
|
|
8
|
+
Project-URL: Documentation, https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md
|
|
9
|
+
Project-URL: Benchmarks, https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md
|
|
10
|
+
Project-URL: Issues, https://github.com/Aagam-Bothara/rigtune/issues
|
|
11
|
+
Keywords: llm,inference,vllm,llama.cpp,serving,autotuning
|
|
12
|
+
Classifier: Development Status :: 3 - Alpha
|
|
13
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Requires-Python: <3.14,>=3.10
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Requires-Dist: pydantic>=2.0
|
|
24
|
+
Requires-Dist: typer>=0.12
|
|
25
|
+
Requires-Dist: rich>=13.0
|
|
26
|
+
Requires-Dist: httpx>=0.27
|
|
27
|
+
Requires-Dist: fastapi>=0.110
|
|
28
|
+
Requires-Dist: uvicorn>=0.29
|
|
29
|
+
Requires-Dist: psutil>=5.9
|
|
30
|
+
Requires-Dist: huggingface_hub>=0.23
|
|
31
|
+
Provides-Extra: nvml
|
|
32
|
+
Requires-Dist: nvidia-ml-py>=12.535.133; extra == "nvml"
|
|
33
|
+
Provides-Extra: dev
|
|
34
|
+
Requires-Dist: pytest>=7.0; extra == "dev"
|
|
35
|
+
Requires-Dist: pytest-asyncio>=0.23; extra == "dev"
|
|
36
|
+
Requires-Dist: ruff==0.15.0; extra == "dev"
|
|
37
|
+
Dynamic: license-file
|
|
38
|
+
|
|
39
|
+
# rigtune
|
|
40
|
+
|
|
41
|
+
[](https://github.com/Aagam-Bothara/rigtune/actions/workflows/tests.yml)
|
|
42
|
+
|
|
43
|
+
*Formerly PolyServe, renamed because a separate [2025 paper](https://arxiv.org/abs/2507.17769) already has that name. Profiles, results and settings from before the rename keep working ([details](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md#renamed-from-polyserve)).*
|
|
44
|
+
|
|
45
|
+
**rigtune finds the fastest way to serve an LLM on your GPU and your traffic, then serves it behind an OpenAI-compatible API.** It measures real configurations — engine, weight precision, batch size, KV-cache type, speculative decoding — instead of guessing, because the settings that won changed on every card tried.
|
|
46
|
+
|
|
47
|
+
**It beats a hand-tuned configuration by 56–97%**, which is the comparison that matters — not the one against untouched defaults.
|
|
48
|
+
|
|
49
|
+
| Llama 3.1 8B, held-out prompts | rigtune | a fixed rule of thumb † | (stock defaults) |
|
|
50
|
+
|---|---|---|---|
|
|
51
|
+
| RTX 4090 | **1041 tok/s** | 667 — **+56%** | SGLang 448 |
|
|
52
|
+
| A100 | **1155 tok/s** | 674 — **+71%** | vLLM 726 |
|
|
53
|
+
| A40 | **506 tok/s** | 257 — **+97%** | vLLM 262 |
|
|
54
|
+
|
|
55
|
+
† fp8 weights and KV cache, a short context, batch 256: what an informed engineer sets without measuring. That rule was **slower than changing nothing** on both Ampere cards, because an fp8 KV cache costs throughput there — which is the point: the settings that win move from card to card, so a rule that is right somewhere is wrong elsewhere. The stock column is context rather than a claim; nobody serious ships untouched defaults. Each row was measured on prompts held out of calibration ([full sessions](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#llama-31-8b-on-four-gpus)).
|
|
56
|
+
|
|
57
|
+
## Results
|
|
58
|
+
|
|
59
|
+
- **Your traffic decides what pays, not the card alone.** On one A40 with Qwen2.5-3B the winning setting changed with the prompts: suffix decoding on Dolly prompts (+122% over stock), a draft model on news-article extraction (+66.5%), an fp8 KV cache on real chat (+10%). A tuner that ignores your prompts cannot find these.
|
|
60
|
+
- **The pick holds on prompts it never saw.** Every number here was measured on prompts held out of calibration: 300 Dolly-15k prompts to tune on, 300 others to score on. An H100 not in the table gained +24% over stock vLLM with fp8 weights.
|
|
61
|
+
- **Calibration pays for itself within hours.** 22–56 minutes per model and workload, repaid by 0.5–2.6 hours of busy serving. Measuring the busiest load first cuts 29–45% off that without changing a pick.
|
|
62
|
+
- **What the search order is worth: one win, one tie.** Given the same time, random sampling of the same space matched this staged search on all three 8B cards. On a 14B in fp8 on a 24 GB card the staged search won by 4.4% (385 against 369 tok/s); on the same card in the 7× larger space that includes 4-bit checkpoints, the two tied (1412 against 1416) even though random search was given less time ([the evidence](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#qwen25-14b-where-memory-binds)). What staging reliably gives is the same answer every run, every strategy tried at least once, and a profile saying what each setting was worth — the gains come from measuring, not from the order.
|
|
63
|
+
|
|
64
|
+
Every result, with methods, ablations, quality checks and limits: [docs/benchmarks.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md).
|
|
65
|
+
|
|
66
|
+
## How this differs from vLLM's auto-tune and AIConfigurator
|
|
67
|
+
|
|
68
|
+
Not benchmarked against either — the comparison below is one of *kind*, not of measured results, and saying otherwise would be the overclaiming this project tries to avoid.
|
|
69
|
+
|
|
70
|
+
- **vLLM's `auto_tune`** sweeps one engine's throughput knobs — batch size and batched-token budget — to find the best setting that still meets a latency target. Same spirit, narrower question: it tunes vLLM once you have already chosen vLLM, your weight precision and your decoding strategy. rigtune treats the engine itself as a variable, alongside precision, KV-cache type and speculative decoding — the dimensions that carried almost all the gain here. On an RTX 4090 SGLang led the early trials and vLLM only overtook it once speculative decoding and a quantized cache were tried — tuning the early leader alone would have eliminated the eventual winner, and a single-engine sweep never sees the crossover at all.
|
|
71
|
+
- **AIConfigurator** predicts good deployments analytically from a performance model, aimed at TensorRT-LLM, and answers sizing questions — how many GPUs, which parallelism — very cheaply because it does not launch anything. rigtune measures instead, which costs 20–50 minutes per workload and is why its answers are specific to your prompts. We tried the analytical shortcut and kept it small for a reason: our own fitted predictor reaches a rank correlation of 0.88–0.94, enough to prune quantizations that cannot win, but it ranks the true winner at a median of 15th out of 25 measured configurations — so prediction is a filter here, never the decision ([why](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/decisions.md)).
|
|
72
|
+
- **What none of the three does** is prove the others wrong. The honest gap is that no one outside this project has run rigtune, and we have not run their tools on our hardware.
|
|
73
|
+
|
|
74
|
+
## Quickstart
|
|
75
|
+
|
|
76
|
+
```bash
|
|
77
|
+
pip install rigtune
|
|
78
|
+
pip install "vllm==0.29.0" # driver 580+; older drivers: "vllm==0.11.0" "transformers>=4.56,<5"
|
|
79
|
+
rigtune serve Qwen/Qwen2.5-3B-Instruct --workload chat
|
|
80
|
+
curl localhost:8000/v1/chat/completions -H 'content-type: application/json' \
|
|
81
|
+
-d '{"model":"Qwen/Qwen2.5-3B-Instruct","messages":[{"role":"user","content":"hi"}]}'
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
The first launch calibrates and caches the result per machine, model, objective and workload; later launches serve at once. To tune for your own traffic, pass a sample of it: `--workload-file prompts.jsonl`, one prompt per line. Optional: `pip install arctic-inference==0.1.1` adds suffix decoding; SGLang can live in its own environment (set `SGLANG_PYTHON` to its `python`); llama.cpp needs `llama-server` built with CUDA, on `PATH` or in `LLAMA_SERVER`.
|
|
85
|
+
|
|
86
|
+
With Docker (built on the official vLLM image):
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
docker build -t rigtune .
|
|
90
|
+
docker run --gpus all --ipc=host -p 8000:8000 -v ~/.cache/huggingface:/root/.cache/huggingface \
|
|
91
|
+
-v ~/.rigtune:/root/.rigtune rigtune serve Qwen/Qwen2.5-3B-Instruct
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
### Supported hardware
|
|
95
|
+
|
|
96
|
+
Linux, Python 3.10–3.13. NVIDIA GPUs of compute capability 7.5 or newer run vLLM, SGLang and llama.cpp (measured on an A40, A100, H100 NVL, L4 and RTX 4090); older NVIDIA GPUs and x86 CPUs run llama.cpp. vLLM-CPU and pre-Turing GPUs have never been benchmarked ([details](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md#hardware-and-engines-measured)).
|
|
97
|
+
|
|
98
|
+
## Learn more
|
|
99
|
+
|
|
100
|
+
- [docs/benchmarks.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/benchmarks.md): every result, at a glance and in full, with methods, ablations, quality, limits and what is still unmeasured.
|
|
101
|
+
- [docs/usage.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/usage.md): how calibration works, workloads, objectives, every search option and the CLI.
|
|
102
|
+
- [docs/writeup.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/writeup.md): the design of the memory planner, calibration and the predictor, what the evidence does and does not support, and the roadmap.
|
|
103
|
+
- [docs/decisions.md](https://github.com/Aagam-Bothara/rigtune/blob/main/docs/decisions.md): why each choice is the way it is — including why not just random search — with the measurement that settled it.
|
|
104
|
+
- [benchmarks/strategies/SUMMARY.md](https://github.com/Aagam-Bothara/rigtune/blob/main/benchmarks/strategies/SUMMARY.md): every table, regenerated from the raw JSON.
|
|
105
|
+
- [CONTRIBUTING.md](https://github.com/Aagam-Bothara/rigtune/blob/main/CONTRIBUTING.md): development setup, tests and adding a backend.
|
|
106
|
+
|
|
107
|
+
MIT licensed.
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
rigtune.egg-info/PKG-INFO
|
|
5
|
+
rigtune.egg-info/SOURCES.txt
|
|
6
|
+
rigtune.egg-info/dependency_links.txt
|
|
7
|
+
rigtune.egg-info/entry_points.txt
|
|
8
|
+
rigtune.egg-info/requires.txt
|
|
9
|
+
rigtune.egg-info/top_level.txt
|
|
10
|
+
tests/test_ablation.py
|
|
11
|
+
tests/test_ablation_shrink.py
|
|
12
|
+
tests/test_all_levels.py
|
|
13
|
+
tests/test_attention_check.py
|
|
14
|
+
tests/test_baselines.py
|
|
15
|
+
tests/test_break_even.py
|
|
16
|
+
tests/test_budget.py
|
|
17
|
+
tests/test_cache_and_pipeline.py
|
|
18
|
+
tests/test_cli.py
|
|
19
|
+
tests/test_combinations.py
|
|
20
|
+
tests/test_compare.py
|
|
21
|
+
tests/test_compare_heldout.py
|
|
22
|
+
tests/test_compare_repeats.py
|
|
23
|
+
tests/test_confirm.py
|
|
24
|
+
tests/test_contenders.py
|
|
25
|
+
tests/test_cpu_quota.py
|
|
26
|
+
tests/test_dataset_workloads.py
|
|
27
|
+
tests/test_drift.py
|
|
28
|
+
tests/test_early_stop_check.py
|
|
29
|
+
tests/test_explore.py
|
|
30
|
+
tests/test_fit_per_machine.py
|
|
31
|
+
tests/test_gguf_and_backends.py
|
|
32
|
+
tests/test_hardware_and_config.py
|
|
33
|
+
tests/test_hub_search.py
|
|
34
|
+
tests/test_int8_kv.py
|
|
35
|
+
tests/test_load_pick_order.py
|
|
36
|
+
tests/test_make_dolly_prompts.py
|
|
37
|
+
tests/test_make_oasst_prompts.py
|
|
38
|
+
tests/test_measure_and_proxy.py
|
|
39
|
+
tests/test_measure_levels.py
|
|
40
|
+
tests/test_memcal.py
|
|
41
|
+
tests/test_memory.py
|
|
42
|
+
tests/test_nvlink.py
|
|
43
|
+
tests/test_objectives_and_search.py
|
|
44
|
+
tests/test_oom_retry.py
|
|
45
|
+
tests/test_phases.py
|
|
46
|
+
tests/test_power.py
|
|
47
|
+
tests/test_predict.py
|
|
48
|
+
tests/test_process_group.py
|
|
49
|
+
tests/test_progress_line.py
|
|
50
|
+
tests/test_quality.py
|
|
51
|
+
tests/test_quant_default.py
|
|
52
|
+
tests/test_report.py
|
|
53
|
+
tests/test_rescore_ttft.py
|
|
54
|
+
tests/test_search_budget.py
|
|
55
|
+
tests/test_search_view.py
|
|
56
|
+
tests/test_selector.py
|
|
57
|
+
tests/test_sglang_python.py
|
|
58
|
+
tests/test_spec_first.py
|
|
59
|
+
tests/test_spec_methods.py
|
|
60
|
+
tests/test_strategies.py
|
|
61
|
+
tests/test_strategy_report.py
|
|
62
|
+
tests/test_supervisor.py
|
|
63
|
+
tests/test_tail.py
|
|
64
|
+
tests/test_task_quality.py
|
|
65
|
+
tests/test_ttft_percentile.py
|
|
66
|
+
tests/test_vllm_cli.py
|
|
67
|
+
tests/test_w8a8.py
|
|
68
|
+
tests/test_workload_file.py
|
|
69
|
+
tests/test_workloads.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
rigtune-0.2.0/setup.cfg
ADDED
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
"""The strategy ablation: variants of the pick, result loading, sweeps and crossover."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
|
|
7
|
+
import pytest
|
|
8
|
+
|
|
9
|
+
from rigtune.backends import get_backend
|
|
10
|
+
from rigtune.bench.ablation import (
|
|
11
|
+
crossover, load_pick, markdown, score_row, strategy_variants, sweep_workload,
|
|
12
|
+
)
|
|
13
|
+
from rigtune.bench.compare import results_path
|
|
14
|
+
from rigtune.calibrate.objectives import Constraints
|
|
15
|
+
from rigtune.calibrate.workload import get_workload
|
|
16
|
+
from rigtune.models import Config, ModelSpec, TrialMetrics, TrialResult
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@pytest.fixture
|
|
20
|
+
def drafts(monkeypatch):
|
|
21
|
+
import rigtune.backends.vllm as vl
|
|
22
|
+
import rigtune.speculative as S
|
|
23
|
+
|
|
24
|
+
monkeypatch.setattr(S, "vllm_supports_draft", lambda version=None: True)
|
|
25
|
+
real = vl.importlib.util.find_spec
|
|
26
|
+
monkeypatch.setattr(vl.importlib.util, "find_spec", lambda n: object() if n == "flashinfer" else real(n))
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
@pytest.mark.usefixtures("drafts")
|
|
30
|
+
def test_variants_flip_each_strategy(hw_a100, prepared_vllm):
|
|
31
|
+
be = get_backend("vllm")
|
|
32
|
+
model = prepared_vllm.model_copy(update={"weights_bytes": {**prepared_vllm.weights_bytes, "awq": 2_000_000_000},
|
|
33
|
+
"hf_paths": {"awq": "org/Llama-3.2-3B-Instruct-AWQ"}})
|
|
34
|
+
pick = Config(backend="vllm", quant="fp8", ctx=4096, batch=64, gpu_memory_utilization=0.9)
|
|
35
|
+
labels = [v.label for v in strategy_variants(pick, be, hw_a100, model, get_workload("chat-system"))]
|
|
36
|
+
assert labels == ["+awq", "+kv:fp8_e5m2", "-prefix", "+spec:ngram:4",
|
|
37
|
+
"+spec:draft:meta-llama/Llama-3.2-1B-Instruct:4"]
|
|
38
|
+
no_prefix = [v.label for v in strategy_variants(pick, be, hw_a100, model, get_workload("chat"))]
|
|
39
|
+
assert "-prefix" not in no_prefix # nothing shared, nothing for a prefix cache to do
|
|
40
|
+
|
|
41
|
+
used = pick.model_copy(update={"quant": "awq", "kv_dtype": "fp8", "spec_decode": "ngram:4"})
|
|
42
|
+
vs = {v.label: v.config for v in strategy_variants(used, be, hw_a100, model, get_workload("chat"))}
|
|
43
|
+
assert set(vs) == {"-int4 (fp8)", "-kv", "-spec"}
|
|
44
|
+
assert vs["-int4 (fp8)"].quant == "fp8" and vs["-kv"].kv_dtype == "auto" and vs["-spec"].spec_decode is None
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def test_removing_kv_quant_steps_the_batch_down_to_fit(hw_gtx1080, prepared_llamacpp):
|
|
48
|
+
pick = Config(backend="llamacpp-cuda", quant="Q4_K_M", ctx=4096, batch=16, n_gpu_layers=29, n_batch=512,
|
|
49
|
+
kv_dtype="q4_0")
|
|
50
|
+
(kv,) = [v for v in strategy_variants(pick, get_backend("llamacpp-cuda"), hw_gtx1080, prepared_llamacpp,
|
|
51
|
+
get_workload("chat")) if v.strategy == "kv"]
|
|
52
|
+
assert kv.label == "-kv" and kv.config.kv_dtype == "auto" and kv.config.batch == 8
|
|
53
|
+
assert "does not fit at batch 16" in kv.note
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def test_llamacpp_prefix_flags_are_added_or_removed(hw_gtx1080, prepared_llamacpp):
|
|
57
|
+
be = get_backend("llamacpp-cuda")
|
|
58
|
+
pick = Config(backend="llamacpp-cuda", quant="Q4_K_M", ctx=4096, batch=4, n_gpu_layers=29, n_batch=512)
|
|
59
|
+
shared = get_workload("chat-system")
|
|
60
|
+
added = [v.label for v in strategy_variants(pick, be, hw_gtx1080, prepared_llamacpp, shared)
|
|
61
|
+
if v.strategy == "prefix"]
|
|
62
|
+
assert added == ["+prefix:cache_reuse", "+prefix:kv_unified"]
|
|
63
|
+
tuned = pick.model_copy(update={"extra": {"kv_unified": True}})
|
|
64
|
+
(removed,) = [v for v in strategy_variants(tuned, be, hw_gtx1080, prepared_llamacpp, shared)
|
|
65
|
+
if v.strategy == "prefix"]
|
|
66
|
+
assert removed.label == "-prefix" and removed.config.extra == {}
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def test_load_pick_prefers_this_machine(tmp_path):
|
|
70
|
+
spec = ModelSpec(hf_id="org/M")
|
|
71
|
+
|
|
72
|
+
def write(hw_hash: str, quant: str) -> None:
|
|
73
|
+
rows = [{"label": "vllm-default", "config": Config(backend="vllm", quant="bf16").model_dump()},
|
|
74
|
+
{"label": "rigtune", "config": Config(backend="vllm", quant=quant).model_dump()}]
|
|
75
|
+
results_path(hw_hash, spec, "chat", "balanced", tmp_path).write_text(json.dumps({"rows": rows}))
|
|
76
|
+
|
|
77
|
+
write("aaaa", "bf16")
|
|
78
|
+
write("zzzz", "fp8")
|
|
79
|
+
assert load_pick(tmp_path, "org/M", "chat", "balanced", hw_hash="zzzz").quant == "fp8"
|
|
80
|
+
assert load_pick(tmp_path, "org/M", "chat", "balanced").quant == "bf16"
|
|
81
|
+
assert load_pick(tmp_path, "org/M", "rag", "balanced") is None
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def test_sweep_workload_and_crossover():
|
|
85
|
+
base = get_workload("rag-shared")
|
|
86
|
+
wl = sweep_workload(base, [8, 1, 64, 4])
|
|
87
|
+
assert wl.concurrencies == (1, 4, 8, 64) and wl.n_prompts == 128 and len(wl.prompts) == 128
|
|
88
|
+
assert wl.prefix_text == base.prefix_text and all(p.startswith(wl.prefix_text) for p in wl.prompts)
|
|
89
|
+
|
|
90
|
+
off = {"1": {"ok": True, "tok_s": 60, "e2e_ms": 2000}, "8": {"ok": True, "tok_s": 400, "e2e_ms": 2500},
|
|
91
|
+
"32": {"ok": True, "tok_s": 1200, "e2e_ms": 3500}}
|
|
92
|
+
on = {"1": {"ok": True, "tok_s": 110, "e2e_ms": 1200}, "8": {"ok": True, "tok_s": 390, "e2e_ms": 2400},
|
|
93
|
+
"32": {"ok": True, "tok_s": 900, "e2e_ms": 4800}}
|
|
94
|
+
assert crossover(off, on, "tok_s") == 8
|
|
95
|
+
assert crossover(off, on, "e2e_ms") == 32
|
|
96
|
+
assert crossover(off, {"1": {"ok": True, "tok_s": 999, "e2e_ms": 1}}) is None
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def test_score_row_and_markdown():
|
|
100
|
+
def result(tok_s: float, cfg: Config) -> TrialResult:
|
|
101
|
+
level = TrialMetrics(tok_s=tok_s, ttft_ms=30, tpot_ms=10, requests=16, output_tokens=2048, concurrency=8)
|
|
102
|
+
return TrialResult(config=cfg, stage="t", metrics=level.model_copy(update={"by_concurrency": {"8": level}}))
|
|
103
|
+
|
|
104
|
+
cons = Constraints(ttft_ceiling_ms=500)
|
|
105
|
+
base = Config(backend="vllm", quant="fp8")
|
|
106
|
+
rows = [{"label": label, "config_key": cfg.key(), **score_row(result(tok, cfg), "balanced", cons)}
|
|
107
|
+
for label, tok, cfg in (("rigtune", 500, base), ("+kv:fp8", 550, base.model_copy(update={"kv_dtype": "fp8"})))]
|
|
108
|
+
rows.append({"label": "-kv", "config_key": "-", "ok": False, "note": "does not fit in memory"})
|
|
109
|
+
top = rows[0]
|
|
110
|
+
assert top["ok"] and top["concurrency"] == 8 and top["meets_slo"]
|
|
111
|
+
assert top["e2e_ms"] == pytest.approx(30 + 10 * 127) and top["levels"]["8"]["tok_s"] == 500
|
|
112
|
+
md = markdown(rows)
|
|
113
|
+
assert "+10.0%" in md and "does not fit in memory" in md
|
|
114
|
+
failed = score_row(TrialResult(config=base, stage="t", metrics=TrialMetrics(), error="boom"), "balanced", cons)
|
|
115
|
+
assert failed == {"ok": False, "error": "boom"}
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
"""Removing a memory-saving strategy steps the batch down to what would have fitted without it."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from types import SimpleNamespace
|
|
6
|
+
|
|
7
|
+
import rigtune.bench.ablation as AB
|
|
8
|
+
from rigtune.backends import get_backend
|
|
9
|
+
from rigtune.calibrate.workload import get_workload
|
|
10
|
+
from rigtune.models import Config
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def test_int4_removal_steps_the_batch_down_to_fit(monkeypatch, hw_a100, prepared_vllm):
|
|
14
|
+
# Measured on an A40 (chat-system): fp8 weights at the 4-bit pick's batch 256 did not fit.
|
|
15
|
+
monkeypatch.setattr(AB, "estimate",
|
|
16
|
+
lambda hw, model, c, mm: SimpleNamespace(feasible=c.quant == "awq" or c.batch <= 64))
|
|
17
|
+
model = prepared_vllm.model_copy(update={"weights_bytes": {**prepared_vllm.weights_bytes, "awq": 2_000_000_000}})
|
|
18
|
+
pick = Config(backend="vllm", quant="awq", ctx=16384, batch=256, gpu_memory_utilization=0.95)
|
|
19
|
+
(w,) = [v for v in AB.strategy_variants(pick, get_backend("vllm"), hw_a100, model, get_workload("chat"))
|
|
20
|
+
if v.strategy == "weights"]
|
|
21
|
+
assert w.label == "-int4 (fp8)" and w.config.quant == "fp8" and w.config.batch == 64
|
|
22
|
+
assert w.note == "batch 64: fp8 weights does not fit at batch 256"
|