[Kernel][Model] logits_soft_cap for Gemma2 with flashinfer

[LiuXiaoxuanPKU]

Add logits_soft_cap for flashinfer, which is needed by Gemma2 model, also add a simple gemma2 test.

[Yard1]

we should check the flashinfer version and raise if it's too old

[Yard1]

we should just raise an exception IMO.

[yongxb]

Could I check if logits_soft_cap is supposed to be the attn_logit_softcapping value instead? The two values are different in the Gemma2 config.

"attn_logit_softcapping": 50.0,
"final_logit_softcapping": 30.0,

[WoosukKwon]

@yongxb Nice catch! final_logit_softcapping is used to cap the final logits before sampling. @LiuXiaoxuanPKU Could you please fix this?

[zifeitong]

I think this warning can be removed to avoid confusion:

vllm/vllm/model_executor/models/gemma2.py

Line 140 in 7cd2ebb

if self.config.attn_logit_softcapping is not None:

[zifeitong]

I am able to run and reproduce the reported MMLU scores for both 9b and 27b models 👍

However, if I don't disable CUDA graph, vLLM will crash with this error:

ERROR 07-04 03:34:04 async_llm_engine.py:53]   File "vllm/vllm/worker/model_runner.py", line 1202, in execute_model
ERROR 07-04 03:34:04 async_llm_engine.py:53]     model_input.attn_metadata.decode_wrapper = self.graph_runners[
ERROR 07-04 03:34:04 async_llm_engine.py:53] IndexError: list index out of range

[LiuXiaoxuanPKU]

I am able to run and reproduce the reported MMLU scores for both 9b and 27b models 👍

However, if I don't disable CUDA graph, vLLM will crash with this error:

ERROR 07-04 03:34:04 async_llm_engine.py:53]   File "vllm/vllm/worker/model_runner.py", line 1202, in execute_model
ERROR 07-04 03:34:04 async_llm_engine.py:53]     model_input.attn_metadata.decode_wrapper = self.graph_runners[
ERROR 07-04 03:34:04 async_llm_engine.py:53] IndexError: list index out of range

Thanks for reporting! Could you give me an minimal reproducible example since I can run gemma-2 with flashinfer cudagraph on my end. Thanks!

[zifeitong]

I am using the run_batch script:

python -m vllm.entrypoints.openai.run_batch -i requests.jsonl -o /dev/null --model google/gemma-2-9b-it --disable-log-request

requests.jsonl.zip

[LiuXiaoxuanPKU]

I am using the run_batch script:

python -m vllm.entrypoints.openai.run_batch -i requests.jsonl -o /dev/null --model google/gemma-2-9b-it --disable-log-request

requests.jsonl.zip

I tried the script and data on H100, it seems work. Could you report your environment? Flashinfer only supports GPU with compute capability greater than 8.0 (https://developer.nvidia.com/cuda-gpus). Not sure if that might be the problem.

[zifeitong]

I am using the run_batch script:

python -m vllm.entrypoints.openai.run_batch -i requests.jsonl -o /dev/null --model google/gemma-2-9b-it --disable-log-request

requests.jsonl.zip

I tried the script and data on H100, it seems work. Could you report your environment? Flashinfer only supports GPU with compute capability greater than 8.0 (https://developer.nvidia.com/cuda-gpus). Not sure if that might be the problem.

I am using H100 with CUDA 12.5. Can you try sync you branch to the latest? #4412 might be related (it refactors graph_runners).

[LiuXiaoxuanPKU]

I am using the run_batch script:

python -m vllm.entrypoints.openai.run_batch -i requests.jsonl -o /dev/null --model google/gemma-2-9b-it --disable-log-request

requests.jsonl.zip

I tried the script and data on H100, it seems work. Could you report your environment? Flashinfer only supports GPU with compute capability greater than 8.0 (https://developer.nvidia.com/cuda-gpus). Not sure if that might be the problem.

I am using H100 with CUDA 12.5. Can you try sync you branch to the latest? #4412 might be related (it refactors graph_runners).

yes, it's a merge conflict. Just fixed, please try again. Thanks!

[zifeitong]

I am using the run_batch script:

python -m vllm.entrypoints.openai.run_batch -i requests.jsonl -o /dev/null --model google/gemma-2-9b-it --disable-log-request

requests.jsonl.zip

I tried the script and data on H100, it seems work. Could you report your environment? Flashinfer only supports GPU with compute capability greater than 8.0 (https://developer.nvidia.com/cuda-gpus). Not sure if that might be the problem.

I am using H100 with CUDA 12.5. Can you try sync you branch to the latest? #4412 might be related (it refactors graph_runners).

yes, it's a merge conflict. Just fixed, please try again. Thanks!

Thanks for the fix. It works now, w/ or w/o CUDA graph.