Cuda level metrics collection!

[iAdtya]

Feature request

Hey Folks,

I am using TEI for embedding creation, and I am constantly getting (OOM)
I need to profile at the CUDA level and understand what is happening under the hood.

• vLLM quickly saturates GPU memory but preempts gracefully.
• TEI grows linearly with batch size and crashes on OOM, without recovery.

This one!

Motivation

I am using TEI for embedding creation, and I am constantly getting (OOM)

Your contribution

I cannot contribute, as I dont know Rust!

[kozistr]

@iAdtya Hi! maybe this issue is related to #670 (comment). Currently, we need to manually set the max-batch-tokens to fit in your environment!

[iAdtya]

Hi @kozistr it's not entirely related to the batch tokens even with 16k tokens lower than the given limit of 32k

I am still getting the OOM issue on L40S 48 GB to run Qwen3-Embedding-4B

It is because there is no mechanism like VLLM has of Preemption and retry, drops sequences if space fills to avoid OOM

Can we still get THE CUDA LEVEL METRICS? feature? I want to know what's causing this to happen!

[kozistr]

Hi @kozistr it's not entirely related to the batch tokens even with 16k tokens lower than the given limit of 32k

I am still getting the OOM issue on L40S 48 GB to run Qwen3-Embedding-4B

It is because there is no mechanism like VLLM has of Preemption and retry, drops sequences if space fills to avoid OOM

Can we still get THE CUDA LEVEL METRICS? feature? I want to know what's causing this to happen!

Oh, sorry for the misunderstanding. I thought the OOM occurred when TEI started.

I might be wrong, but, afaik, it could be hard to collect CUDA-level metrics because there's no native, pure Rust API available for GPU metrics gathering. (or we need to update the FA source code to add the logging part, which seems not a good approach, I guess)

Maybe we could manually utilize Nsight or nvprof to profile instead!

[iAdtya]

Nsight would be of great help

Oh, sorry for the misunderstanding. I thought the OOM occurred when TEI started.

Hi @kozistr

OOM is an issue occurring in the middle of my embedding pipeline after embedding quite a few chunks of documents

I suspect that the VRAM is being linearly occupied, and because there is no fail-safe technique, it simply crashes.

We cannot put this in Production, and I cannot use VLLM as they have a 5-minute Cuda graph creation time, which TEI doesn't.

---

As mentioned above, I think this is the issue!

Takeaway:

• vLLM quickly saturates GPU memory but preempts gracefully.
• TEI grows linearly with batch size and crashes on OOM, without recovery.