Skip to main content

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

Web Viewer

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

  1. 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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

llm_execution_time_predictor-0.1.0.tar.gz (26.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

llm_execution_time_predictor-0.1.0-py3-none-any.whl (29.4 kB view details)

Uploaded Python 3

File details

Details for the file llm_execution_time_predictor-0.1.0.tar.gz.

File metadata

File hashes

Hashes for llm_execution_time_predictor-0.1.0.tar.gz
Algorithm Hash digest
SHA256 f59eacfb4316f51d14209958370d3a800b5c4160c41f63b55b6496af4eafa8ce
MD5 fbf3ac34d088cec1c80746f0c3dafbd1
BLAKE2b-256 e795a0f34a768c8d7c3e44872c1466f49940570a47533fc75aec5be0cb71d549

See more details on using hashes here.

File details

Details for the file llm_execution_time_predictor-0.1.0-py3-none-any.whl.

File metadata

File hashes

Hashes for llm_execution_time_predictor-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 fb8e1e71c997c81a56972ab8963fba0fdcbf0e3f47da2b4f5e00828101d61930
MD5 dc8042233d8ec4e7992b28270b8e8af7
BLAKE2b-256 e6959d1ab41c364b070c698961235d83aad57cb6419e274bc2e5ad3fee2f4527

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page