[Feature] Support fp8 e5m2 kv cache with flashinfer

[ispobock]

Motivation

Support fp8 e5m2 kv cache with flashinfer.

Usage

Add --kv-cache-dtype fp8_e5m2 to enable this feature. Currently it only works when flashinfer is not disabled.

Performance & Accuracy

Tested with llama2-13b-chat on A100, the throughput increased by 17.8% without accuracy degradation.

Enable fp8_e5m2 kv cache	Throughput	MMLU (nsub=10) Avg Accuracy	gsm8k Accuracy
N	7.04 req/s	0.487	0.340
Y	8.30 req/s	0.488	0.334

Reproduce

# w/o fp8_e5m2 kv cache
python3 -m sglang.launch_server --model-path meta-llama/Llama-2-13b-chat-hf --port 30000 --trust-remote-code --disable-radix-cache --tp=1
# w/ fp8_e5m2 kv cache
python3 -m sglang.launch_server --model-path meta-llama/Llama-2-13b-chat-hf --port 30000 --trust-remote-code --disable-radix-cache --tp=1 --kv-cache-dtype fp8_e5m2

# benchmark
python3 -m sglang.bench_serving --backend sglang --tokenizer meta-llama/Llama-2-13b-chat-hf  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json --num-prompts 1000 --request-rate 128

# evaluation
python3 benchmark/mmlu/bench_sglang.py --nsub 10
python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319

The performance boost is model dependent. llama3-8b was also tested, but the performance was not improved.

[zhyncs]

Nice work! I'll review it asap. May we also support FP8 E4M3?

[ispobock]

May we also support FP8 E4M3?

FP8 E4M3 needs scale factor and calibration. We may add it in the future.

[zhyncs]

Currently, only FlashInfer is supported and not Triton, due to the issue of insufficient smem. This needs to be fixed in another PR.

[zhyncs]

workaround for float8_e5m2

Store as torch.uint8 because Tensor index_put is not implemented for torch.float8_e5m2

[qeternity]

Sorry to dig this up but - are we suggesting that fp8 kv cache increased accuracy in both mmlu and gsm8k? Are we sure we don't have those values in the table reversed?

[ispobock]

@qeternity In the previous evaluation, I tested gsm8k with 200 questions (default setting in the benchmark script), so the result may not be reliable enough. I just test all the datasets and update the result in the table.