Support for per-token latency tracking in `generate()` (suggested options: using callback, profiler class, or using a config flag)
Author: spsagar13
Pull request: no
Labels: name: Feature request

### Feature request

**Summary**

I would like to propose a feature to enable **per-token latency tracking** during generation in Hugging Face’s `generate()` (via _beam_search, sample etc) loop. This is extremely useful for benchmarking models across hardware (e.g. CPU vs GPU, Arm vs x86), identifying bottlenecks, and understanding real-world inference performance at a granular level.

I’m suggesting **three high level designs** to add this feature support.

**Option 1: `per_token_latency_callback` (minimal, flexible)**
_This approach offers maximum flexibility with minimal intrusion - meaning - no config changes, just a simple hook for custom tracking._
By adding a new argument to `generate()`:
```python
def generate(..., per_token_latency_callback: Optional[Callable[[int, float], None]] = None):
    for token_idx in range(max_length):
        start = time.perf_counter()
        # generation step
        end = time.perf_counter()
        if per_token_latency_callback:
            per_token_latency_callback(token_idx, end - start)
```

Example Usage (for benchmarking)

```python
# Simple function to print per-token latency.  
# Can be extended to log, store, or analyze latencies as needed
def log_token_latency(token_idx, latency):
    print(f"Token {token_idx}: {latency * 1000:.2f} ms")

output = model.generate(
    input_ids,
    max_new_tokens=5,
    per_token_latency_callback=log_token_latency  # Usage with new proposed argument
)
```
This would allow users to easily plug in their own logic to log, store, or analyze latency **without touching the internal logic.**



**Option 2:` Profiler Class` along with `per_token_latency_callback` (structured, user-friendly)**
_This approach introduces a clean, reusable structure that can be extended to track other metrics like memory or FLOPs - ideal for advanced profiling_
Passing a reusable profiler object, for example:

```python
class TokenLatencyProfiler:
    def __init__(self):
        self.token_latencies = []

    def track(self, token_idx, latency):
        self.token_latencies.append(latency)

profiler = TokenLatencyProfiler()
output = model.generate(input_ids, per_token_latency_callback=profiler.track)

print(profiler.token_latencies)
```

**This structure makes it easy to support other metrics in the future** (e.g. memory, FLOPs, peak RAM/GPU usage).



**Option 3: `track_token_latency=True` in GenerationConfig (inspired by Intel's IPEX)**
_This approach adds convenience for quick profiling but introduces some coupling between the generation config and the return object structure._

As a higher-level alternative - conceptually inspired by Intel’s IPEX benchmarking approach - this method introduces a dedicated configuration flag within GenerationConfig to enable built-in per-token latency tracking.
_Unlike IPEX, which sets this behavior via setattr() on the model config at runtime, this proposal integrates the flag directly into GenerationConfig_

- Update GenerationConfig

> ```python
> track_token_latency: bool = field(
>     default=False,
>     metadata={"help": "If True, enables per-token latency tracking during generation."}
> )
> ```

- Use the flag internally

> ```python
> def generate(self, input_ids, generation_config=None, max_length=20, ...):
>     generation_config = generation_config or self.generation_config
>     token_latencies = [] if generation_config.track_token_latency else None
> 
>     for token_idx in range(max_length):
>         start = time.perf_counter()
>         # generation step (e.g., logits -> sampling -> input_ids update)
>         end = time.perf_counter()
> 
>         if token_latencies is not None:
>             token_latencies.append(end - start)
> 
>     output = GenerationOutput(...)  # or SampleDecoderOnlyOutput, etc.
>     if token_latencies is not None:
>         output.token_latencies = token_latencies
> 
>     return output
> ```

- Example Usage

> ```python
> from transformers import GenerationConfig
> 
> gen_config = GenerationConfig(track_token_latency=True)
> output = model.generate(input_ids, generation_config=gen_config)
> 
> print(output.token_latencies)  # → [0.012, 0.010, 0.011, ...]
> ```


### Motivation

As a performance engineer working with LLMs on hardware like Arm and x86 CPUs, I often need to measure **token-by-token latency** during generation (e.g., to analyze startup cost, cache reuse, SMT impact, or I/O stalls). 

However:
- There is currently **no easy way** to track **per-token latency** via Hugging Face’s `generate()` API.
- The existing `output.scores` provides logits but **not timing or performance hooks**.
- I discovered that Intel’s IPEX benchmark internally **tracks token latencies**, which was helpful for comparison - but it’s not easily accessible or extendable outside Intel’s stack.
- So I had to patch generate logic locally to do something like this:
  
```python
# inside _beam_search under transformers/src/transformers/generation/utils.py

for token_idx in range(max_length):
    start = time.perf_counter()
    # generation
    end = time.perf_counter()
    latencies.append(end - start)
```

_But this (my patch) could break with updates to the Transformers library or models that override generate() with custom loops._

Having official support would benefit researchers, profiling tools, downstream libraries, and anyone benchmarking LLMs on novel hardware.

This also helps with:

- Debugging regressions in low-level compute kernels (e.g., matmul, softmax) that impact token generation performance
- Optimizing latency-critical use cases (e.g., serverless inference, streaming chat) by isolating startup cost, monitoring per-token response time, and diagnosing real-time slowdowns
- Analyzing time-to-first-token vs. steady-state generation performance

### Your contribution

Yes, I’m happy to submit a PR.

I can start by adding **Option 1** (`per_token_latency_callback`) as a lightweight, non-breaking change to `GenerationMixin`. Optionally, I can follow up with Option 2 (profiler class) and/or Option 3 (`track_token_latency=True` in `GenerationConfig`) depending on maintainers’ preferences and/or with additional recommendations/improvements. 

Please let me know if these changes are fine and if so which direction is preferred, and I’ll start implementing.

Comments:
cc @gante @zucchini-nlp 

Hi @gante, @zucchini-nlp, hope you're well. Please let me know if you have any queries. I'd be happy to set up a call and discuss further. 

@spsagar13 👋  Thank you for opening the discussion!

Some context: We're currently trying to lower the complexity of the codebase and improve its readability, aka debloating. `generate` is one of the targets: its body is bloated, and it has way too many arguments (especially if we consider the `generation_config`).

Profiling is powerful, and essential to design good HW/SW. Would it be possible to explore some alternative that doesn't conflict with the debloating goal? For instance, if we isolate the logic of each iteration of the decoding loop into a function, could we achieve the goals you propose with e.g. a decorator?

e.g.
```py
model = AutoModelForCausalLM.from_pretrained(...)

# trigger profiling; `model._sample_iteration` contains one iteration of the decoding loop for `_sample`
profiling_data = []
model._sample_iteration = some_profiling_fn(model._sample_iteration, profiling_data)

# generate as usual -> `profiling_data` gets populated
model.generate(...)
```

I really like the proposal! I am benchmarking models and face the same problem. I am wondering if we have existing tool to measure the prefill and decode time (not yet to the level of per token latency). 

Hi @gante, Thank you for the response - I do like the direction you're suggesting and agree to debloating `generate()` API. The`_sample_iteration` seems like an extensible solution that aligns with both our goals. 

Just to confirm if my understanding is correct - please correct if I'm off:

You're suggesting modularizing the per token generation step (inside _sample, _beam_search etc), into the function `model._sample_itearation(....)`, which would encapsulate the logic for generation of 1 token`( forward pass -> sample/greedy decode -> appending new token)`.

So, I could use something like below to get token latency or other profiling data,

```python 
def time_each_token_step(fn, latencies):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = fn(*args, **kwargs)
        end = time.perf_counter()
        latencies.append(end - start)
        return result
    return wrapper

latencies = []
model._sample_iteration = time_each_token_step(model._sample_iteration, latencies)

# Normal generation call
model.generate(input_ids) 

print(latencies) # here i will get latencies of all tokens generated
```

If I’m understanding correctly, then yes - this approach aligns well with my goals.

I also agree with @siyuanchai1999 : if we could modularize the `prefill phase` as well, that would complement with my suggestion, which primarily covers the `decode phase` (via _sample_iteration).

I also will be happy to contribute to this effort - _sample_iteration refactoring for the greedy and sampling decoding parts.

Please let me know.

@spsagar13 yes, precisely: the token generation step becomes modularized and accessible as a model object's method, which can then be wrapped. If more granularity is needed, `model.forward`, `model.prepare_inputs_for_generation` and a few others are already accessible, so the same strategy can be applied there.

Action points:
1. for me -- execute this modularization, and isolate prefil into a separate function
2. for you -- prepare some benchmarking tools / demos based on this structure

Does this sound good to you? 🤗 

Sounds great - I’m aligned with the action plan! And if there’s any part of the modularization, you'd like me join, feel free to let me know - happy to contribute wherever it’s helpful! 🤗

Hi @gante and @spsagar13, thank you for implementing this feature. While you are working on it, I did a custom way to benchmark prefill and decode time. Would you mind taking a look if I am doing anything off the track? Thanks!

```python
prefill_start = time.time()
past_key_values = DynamicCache()
# Generate only the first token (prefill)
with torch.no_grad():
    prefill_outputs = model.generate(
        **inputs, 
        max_new_tokens=1,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
        use_cache=True,
        past_key_values=past_key_values
    )

torch.cuda.synchronize()
prefill_time = time.time() - prefill_start

# Decode phase - generate remaining tokens one by one
decode_start = time.time()

# Start with prefill output for decode phase
current_ids = prefill_outputs

with torch.no_grad():
    current_ids = model.generate(
        current_ids,
        max_new_tokens= args.gen_len - 1,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
        use_cache=True,
        past_key_values=past_key_values
    )

torch.cuda.synchronize()
decode_time = time.time() - decode_start

print("Total: " + str(prefill_time + decode_time) + " Prefill: " + str(prefill_time) + " Decode: " + str(decode_time))
```

Hi @siyuanchai1999,

I think the `time to first token` when using `model.generate()` includes both:

- Prefill time - the time to process the input prompt and to compute the key/value (KV) cache.
- Decode time - the time to generate the first new token.

My understanding is, we currently treat the `time to first token` as `prefill time`, since the decoding time for a single token is relatively small (?). 

So, your implementation seems to be correct. @gante, please correct if my understanding is wrong. 

Once this PR/feature lands, we should be able to separately measure the actual prefill time by subtracting the first token decode time. (even if the prefill time part is not abstracted using a modular function like `_sample_iteration`)

@siyuanchai1999 it's a good approximation! I think the only issue with the script you shared is that some of the overhead regarding preparing inputs in `generate` (things like building the `cache_position`) gets accounted for twice.

[Some time estimates: working towards this task won't be a priority for me in August, as I'm focusing on `transformers serve`. Perhaps in September]

Hi! Thanks for double checking this! I am very looking forward to the patch that enables more detailed latency!

> My understanding is, we currently treat the time to first token as prefill time, since the decoding time for a single token is relatively small (?).

I guess I may have use a less precise term. I think the standard is to measure Time to First Token (TTFT) and Intern Token Latency (ITL). https://docs.nvidia.com/nim/benchmarking/llm/latest/metrics.html. Yes, TTFT includes prefill processing time, and I believe it could be challenging to fully isolate the prefill processing time. I will change the name in my script to reflect it more precisely.

> I think the only issue with the script you shared is that some of the overhead regarding preparing inputs in generate (things like building the cache_position) gets accounted for twice.

How large do you think the overhead could be? I benchmarked the total end-to-end latency vs. `prefill time` (TTFT more precisely) and decode time. It runs OPT-13B model on A100-80G GPU with prompt length of 1920 and generation length of 128. I found little difference in total time.

```
Total end-to-end latency: 6.7192 seconds
Prefill (TTFT): 1.0132386684417725 Decode: 5.695080041885376 Sum: 6.708318710327148
```


cc @manueldeprada 

Hi @gante, hope you're doing well. If there’s a branch or draft I can pull, I’m happy to test locally and share results before you can push. Please let me know if that would be helpful. 

👋 

The refactors needed to enable this are ongoing. @manueldeprada is finalizing: 
1. a separate prefil function 
2. generic optional output manipulation

After these, we'll need to isolate the logic of an iteration of each decoding method into their own function -- help would be appreciated at this point 🤗  

Hi @manueldeprada, @gante, 👋 hope you're doing well. 

Just wanted to check whether you’re at the stage where my help may be useful. If so, I’d be happy to assist. 

Thank you! 🤗