[Usage]: trying to use generation_tokens_total and prompt_tokens_total to get total tokens in the current batch

[annapendleton]

Your current environment

I'm testing this with a single A100 40GB GPU using meta-llama/Meta-Llama-3-8B model weights. The server prints out the following info around expected memory reserved for the kv cache on startup:

INFO 01-13 11:35:21 model_runner.py:1099] Loading model weights took 14.9595 GB
INFO 01-13 11:35:23 worker.py:241] Memory profiling takes 1.32 seconds
INFO 01-13 11:35:23 worker.py:241] the current vLLM instance can use total_gpu_memory (39.50GiB) x gpu_memory_utilization (0.90) = 35.55GiB
INFO 01-13 11:35:23 worker.py:241] model weights take 14.96GiB; non_torch_memory takes 0.10GiB; PyTorch activation peak memory takes 9.37GiB; the rest of the memory reserved for KV Cache is 11.12GiB.
INFO 01-13 11:35:23 gpu_executor.py:76] # GPU blocks: 5691, # CPU blocks: 2048
INFO 01-13 11:35:23 gpu_executor.py:80] Maximum concurrency for 8192 tokens per request: 11.12x

How would you like to use vllm

I'm currently trying to use the sum of these two metrics, prompt_tokens_total and generation_tokens_total, to track how many tokens are in the current kv cache (aka how many tokens in the current batch per iteration).

Based on the above environment information, I would expect the max total tokens in the kv cache to be around 8192 * 11.12 = ~91k tokens. Note I've set the max concurrent requests allowed high enough to enable the full kv cache to be utilized.

However, in my load test where I max out the kv cache utilization, I am seeing both prompt_tokens_total and generation_tokens_total max out at around ~4k when the batch is full, summing to ~8k tokens. That's ~10x less than the expected 90k tokens summed.

I'm curious if the sum of these metrics represent the total tokens in the batch? If so, why is there a mismatch between the expected tokens for a full kv cache and the actual metrics emitted? If not, is there any metric that exists today that can give me the total tokens in the batch?

cc @robertgshaw2-redhat @DarkLight1337 who have worked with the total tokens metric that relies on these two metrics

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[DarkLight1337]

cc @WoosukKwon @comaniac as well

[comaniac]

IIRC, prompt_tokens_total is only the number of prompt tokens being prefilled in the current batch, so for example if you have max_num_batched_tokens=4096, then prompt_tokens_total is 4k at most.

Meanwhile, although your log indeed shows that you could have at most 5691 *16=91,056 tokens, this number not only includes prompt tokens being prefilled, but also the kv-cache tokens being used for decoding. So it could be like the following (assuming each request has 7k prompt tokens and decoding length is 1k):

Step 1 (total tokens: 4k, total reqs: 1):
KV-cache: 0
Prefill: 4k (req1)
Decode: 0

Step 2 (total tokens: 8k, total reqs: 2):
KV-cache: 4k (req1)
Prefill: 3k (req1) + 1k (req2)
Decode: 0

Step 3 (total tokens: ~12k, total reqs: 2):
KV-cache: 7k (req1) + 1k (req2)
Prefill: ~4k (req2)
Decode: 1 (req1)

Step 4 (total tokens: ~16k, total reqs: 3):
KV-cache: 7k (req1) + ~5k (req2)
Prefill: ~2k (req2) + ~2k (req3)
Decode: 1 (req1)

...

Step N (total tokens: ~91k, total reqs: 11):
KV-cache: ~8k*11
Prefill: 0
Decode: 11

[annapendleton]

@comaniac thanks for confirming that prompt_tokens_total represents the prefill tokens, the example is very useful! One follow up question, I thought that the generation_tokens_total metric would represent the kv-cache tokens. Is that intuition correct or is there a separate metric I could use to represent the tokens in the kv cache?

Currently generation_tokens_total is only showing 4k tokens as well for the metric emitted.

[comaniac]

It's again the decode tokens in a batch:

num_generation_tokens_iter = (
    actual_num_batched_tokens - num_prompt_tokens_iter +
    num_generation_tokens_from_prefill_groups
)

[annapendleton]

Okay, so what I'm understanding:

1. generation_tokens_total represents decode tokens in the kv cache
2. prompt_tokens_total is the prompt tokens being prefilled in the kv cache

Adding these two metrics up should give the total tokens in the kv cache.

So the one thing I'm not 100% clear on, but it seems like the information is pointing to, just because we could get up to ~90k tokens in the batch does not mean we will have ~90k tokens in the batch. So we might be seeing generation_tokens_total = 4k and prompt_tokens_total = 4k because we're having an even split of prefill and decode going on per iteration.

So my question is, why do we see 100% kv cache utilization with ~8k total tokens in the kv cache, knowing that we could get up to ~90k tokens? I was working under the assumption that kv cache utilization would have a 1:1 relationship to the total number of tokens in the kv cache / max tokens that could fit in the kv cache.

[robertgshaw2-redhat]

Your understanding is incorrect. These are not the number of tokens in the cache these are the number of tokens in the batch.

Generally, there are way more tokens in the cache than there are in a batch.

[annapendleton]

Ah okay, I understand better, thank for clarifying! The batch is only the tokens being actively worked on in the iteration cycle, not necessarily all the tokens that are stored in the kv cache.

@robertgshaw2-redhat I'm assuming there is no metric emitted today that gives the numbers of the tokens in the cache. Let me know if that's the case :)

[robertgshaw2-redhat]

We report the kv cache utilization, which is the % of KV cache blocks that are being utilized. But not the number itself.