XQA kernel works slower with fp8 kv than with fp16 kv on H100

System Info

Hi!

I'm running speculative execution TRT-LLM engine with 4 or 5 generation length, and I noticed that fp8 kv cache attention works slower than fp16 kv cache attention. Would be great to improve fp8 kv cache performance.

I run it on 1x H100 SXM, Llama 3.1 8b model, fp8 weights/activations and either fp16 or fp8 key value cache. I run it with batch size 64. For fp16 kv cache I observe close enough performance for no drafts/1 draft/3 drafts & 4 drafts (~60/60/70/80 microseconds per kernel execution). For fp8 I see significant drop from no drafts to 1 draft, and then from 3 to 4 (40/60/70/110 microseconds per kernel execution). I.e. you can see that for 4 drafts fp16 kv cache works faster than fp8 kv cache (and comparable for 1&2 drafts).

Thank you.

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

• Build Llama 3.1 8B model with speculative decoding and different quantizations

Expected behavior

• FP8 kv cache is faster than FP16 kv cache
• Adding drafts from 3 to 4 doesn't increase attention runtime significantly

actual behavior

• FP8 kv cache is slower than FP16
• Adding drafts from 3 to 4 increases attention runtime significantly

additional notes

I've tested it with other batch sizes (16 & 32) and it shows similar behavior