OneRuby.devAN ENGINEERING NOTEBOOK

AI · 5 min read

Timing LLM inference in Python without inventing a CUDA speedup

Build a tested inference timer around an offline GPT-2 fixture, count generated tokens correctly and state exactly what a CUDA measurement would include.

Put a clock around a call to generate(), divide output tokens by elapsed seconds, and you have a number. Whether that number measures completed GPU work depends on where the clock stops. Whether it says anything useful about a deployed model depends on considerably more.

This note builds a small timing harness around an actual Hugging Face GPT-2 implementation. The model has random weights and only 16,864 parameters. It runs offline, without a tokenizer or checkpoint download. That makes it suitable for testing token counts and timing boundaries; it makes it completely unsuitable for predicting the throughput of a trained 7B model.

The recorded run used CPU. CUDA is an explicit, unexecuted path. The aim is to have a measurement method worth carrying to a GPU, before putting a speedup table in an article.

Decide what the timer includes

The fixture creates a four-token input and asks for eight new tokens. Model construction and tensor placement happen before warmup. Two warmup requests happen before five measured requests. Each measured request includes the complete generate() call, which processes the prompt and generates the continuation.

It excludes model loading, tokenization, transfer of a new request, text decoding, HTTP handling and queueing. This is generation-call latency for one already prepared request. Calling it end-to-end application latency would discard most of the application.

For the measurement itself, the helper accepts a synchronization function:

Python
def timed_call(generate, synchronize, clock=time.perf_counter):
synchronize()
start = clock()
output = generate()
synchronize()
elapsed = clock() - start
if elapsed <= 0:
raise ValueError("nonpositive elapsed time")
return output, elapsed

The first synchronization clears outstanding work before the start timestamp. The second waits for completion before taking the end timestamp. On CPU the supplied function does nothing. On CUDA it calls torch.cuda.synchronize(device) for the selected device.

PyTorch documents that GPU operations are normally asynchronous and that unsynchronized host timing can be inaccurate. Its CUDA semantics guide also describes CUDA events as another timing mechanism. Choose the mechanism for the boundary you want to measure; a host timer and a device event do not automatically include the same work.

Count the continuation, not the prompt

The fixture uses greedy generation, caching enabled and an explicit attention mask. It sets the end-of-sequence token to None so this artificial model always produces the requested continuation length. That is a test convenience, not a recommendation for a real language model.

After generation, the harness calculates output.shape[1] - inputs.shape[1]. The four prompt tokens are excluded from the new-token count. A model that stops early must use its actual continuation length, not the requested maximum.

Hugging Face's versioned generation reference distinguishes max_new_tokens from a length including the prompt and documents greedy decoding and the cache option. Those definitions matter more than the attractive precision of a tokens-per-second figure.

For this single-sequence fixture, the shape difference is sufficient. Padded batches, per-sequence end tokens and streaming require more careful accounting. A batch-level output width is not automatically the number of useful tokens generated for every member.

What actually ran

Download benchmark.py, test_benchmark.py and the environment notes. The local run used Python 3.11.5, PyTorch 2.9.1 and Transformers 4.32.1. This is a pinned compatibility fixture, not a claim that those are the latest library versions. It uses random initialization, with no remote model code or pretrained weights.

Run python3 -B benchmark.py --device cpu. The JSON records the runtime versions, device, dtype, parameter count, prompt length, warmup count and all five samples. Each sample contains elapsed seconds and the actual new-token count. The fixture produced eight new tokens per measured request.

The CPU entry point selects one PyTorch thread so that thread count is not an implicit variable in this particular run. The median is included as a compact description of five samples. Five samples are not enough to characterize a production latency tail, and a median says nothing about request failures or concurrent load.

Run python3 -B -m unittest -v test_benchmark.py to check the harness. Five tests passed. They verify the synchronization/clock ordering, reject invalid sample counts and lengths, and ensure CPU execution never calls CUDA synchronization. An actual generation test requests three new tokens and checks both the resulting count and the rate arithmetic. A separate test forces CUDA to be unavailable and verifies a clear failure.

Those checks can catch a broken timer or count. They cannot validate GPU kernels on a machine without a GPU. The CUDA option is retained so the boundary is visible and can be tested on suitable hardware; it does not produce a substitute GPU result on this computer.

Carry a controlled experiment to the real model

Before comparing CPU and CUDA, retain the same model revision, tokenizer, prompt fixtures, output limits, decoding settings and batch structure. Record hardware, library builds, dtype and device placement. A run that changes several of these at once may still be useful, but its speedup belongs to the whole configuration change.

Measure prompt processing and token generation separately if that is the question driving the optimization. Long prompts and long continuations stress different parts of the request. This harness combines them, so its tokens-per-second value should not be labelled steady-state decoding throughput.

If you evaluate a lower precision, measure output quality as well as latency. If you change caching or compilation, record warmup and memory behavior. Current allocation is not peak memory, and weight bytes alone do not describe the request's complete memory cost; the QLoRA memory note develops that distinction for training.

The useful result of the local exercise is a harness with explicit boundaries and tested arithmetic. A CUDA speedup still needs a CUDA run, with the workload and evidence attached.

Found a mistake or tried a different approach?

Send Alex a note ↗