Get real-time GPU consumption for each generation from a model

[AI-jitsu]

Hello

I am struggling to find how to get how much GPU I consume each time from each inference. Using torch in this case does not really work since vllm allocates 90% memory for the KV cache by default but I would like to know, for each inference from a LLM, how much GPU have been consumed to generate the certain amount of tokens in output.

Thanks