LLM batch inference latency predictor and profiler CLI tool
Project description
LLM Execution Time Predictor
A small utility to help train a regression model given to predict prefill/decode times. By using the batch size and input, the prefill/decode execution times are very predictable.
This can be plugged into a simulator for faster experiments.
A more complicated version is done by https://github.com/microsoft/vidur but it trains every component of the model forwarding. This utility instead just profiles the full model forwarding as a unit to simplify research.
The tool https://modal.com/llm-almanac/advisor is nice visualizer but it doesn't let you train a local version and specify an exact bs/input
Installation
Option 1: Install from PyPI (Recommended)
pip install llm_execution_time_predictor
Option 2: Install from Source
git clone <this-repo>
cd <repo-directory>
pip install -r requirements.txt
Using Prefill/Decode execution time for predictors
A very small set of features are used to train the predictor. Num new tokens: total tokens processed/generated:
- for decode, it's the batch size. for prefill, it's the full input chunk Product ext cost: Represents the cost of attention
- For prefill, it's O(seq_len^2) so we do bs * input^2
- For decode, it's just O(seq_len) Total context tokens:
- Total tokens processed across batch * input representing the cache usage Time of kernel
Tested on both prefill/decode the decode time
Usage
Using the PyPI Package
# Profile a model and generate benchmark data
llm-execution-time-predictor profile <model_name> --tp_size <tp_size>
# Train models from benchmark data
llm-execution-time-predictor train_models <config_name> <benchmark_file> [--predictor-file <output_file>]
# Make predictions using trained models
llm-execution-time-predictor predict <predictor_file> <config_name> --mode <prefill/decode> --bs <batch_size> --input-len <input_length>
# View trained models and make interactive predictions (CLI)
llm-execution-time-predictor view [--predictor-file <predictor_file>]
# Launch web-based viewer with interactive plots
llm-execution-time-predictor webview [--predictor-file <predictor_file>] [--host <host>] [--port <port>]
Using from Source
# Profile a model and generate benchmark data
python llm_execution_time_predictor/llm_forward_predictor_cli.py profile <model_name> --tp_size <tp_size>
# Train models from benchmark data
python llm_execution_time_predictor/llm_forward_predictor_cli.py train_models <config_name> <benchmark_file> [--predictor-file <output_file>]
# Make predictions using trained models
python llm_execution_time_predictor/llm_forward_predictor_cli.py predict <predictor_file> <config_name> --mode <prefill/decode> --bs <batch_size> --input-len <input_length>
# View trained models and make interactive predictions (CLI)
python llm_execution_time_predictor/llm_forward_predictor_cli.py view [--predictor-file <predictor_file>]
# Launch web-based viewer with interactive plots
python llm_execution_time_predictor/llm_forward_predictor_cli.py webview [--predictor-file <predictor_file>] [--host <host>] [--port <port>]
The trained predictor file format:
{
"config_name": {
"prefill": {
"weights": [0.1234, 0.5678, 0.9012, 0.3456],
"bias": 0.0123,
"model_type": "linear"
},
"decode": {
"weights": [0.2345, 0.6789, 0.0123, 0.4567],
"bias": 0.0456,
"model_type": "linear"
}
}
}
Feature order: [num_new_tokens, prod_ext_ctx, num_context_tokens, batch_size]
Webviewer
Quickstart workflow
Using PyPI Package
llm-execution-time-predictor profile Qwen/Qwen3-4B --tp_size 1
llm-execution-time-predictor train_models tp1_config benchmark_data_Qwen_Qwen3-4B_TP_1_PP_1.json --predictor-file trained_predictors.json
llm-execution-time-predictor predict trained_predictors.json tp1_config --mode decode --bs 8 --input-len 1024
llm-execution-time-predictor webview --predictor-file trained_predictors.json
Using from Source
python llm_execution_time_predictor/llm_forward_predictor_cli.py profile Qwen/Qwen3-4B --tp_size 1
python llm_execution_time_predictor/llm_forward_predictor_cli.py train_models tp1_config benchmark_data_Qwen_Qwen3-4B_TP_1_PP_1.json --predictor-file trained_predictors.json
python llm_execution_time_predictor/llm_forward_predictor_cli.py predict trained_predictors.json tp1_config --mode decode --bs 8 --input-len 1024
python llm_execution_time_predictor/llm_forward_predictor_cli.py webview --predictor-file trained_predictors.json
TODO
- Fix vLLM force one batch with vllm backend, currently vLLM might run more than 1 batch making some of the profiling innacurate skewing the model. Currently no good solution for this.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file llm_execution_time_predictor-0.1.0.tar.gz.
File metadata
- Download URL: llm_execution_time_predictor-0.1.0.tar.gz
- Upload date:
- Size: 26.5 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.1.0 CPython/3.13.5
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
f59eacfb4316f51d14209958370d3a800b5c4160c41f63b55b6496af4eafa8ce
|
|
| MD5 |
fbf3ac34d088cec1c80746f0c3dafbd1
|
|
| BLAKE2b-256 |
e795a0f34a768c8d7c3e44872c1466f49940570a47533fc75aec5be0cb71d549
|
File details
Details for the file llm_execution_time_predictor-0.1.0-py3-none-any.whl.
File metadata
- Download URL: llm_execution_time_predictor-0.1.0-py3-none-any.whl
- Upload date:
- Size: 29.4 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.1.0 CPython/3.13.5
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
fb8e1e71c997c81a56972ab8963fba0fdcbf0e3f47da2b4f5e00828101d61930
|
|
| MD5 |
dc8042233d8ec4e7992b28270b8e8af7
|
|
| BLAKE2b-256 |
e6959d1ab41c364b070c698961235d83aad57cb6419e274bc2e5ad3fee2f4527
|