Add support for the Gemma 2 model

[EricLBuehler]

Refs #79

[EricLBuehler]

@guoqingbao the model runs but seems to give garbage output. Do you see anything which stands out?

[guoqingbao]

@guoqingbao the model runs but seems to give garbage output. Do you see anything which stands out?

Perhaps there is a need for 'attn_logit_softcapping' if 'final_logit_softcapping' is used.

[guoqingbao]

I found the recent update #83 overwritten previous bug fixes including #80 @EricLBuehler

[guoqingbao]

@guoqingbao the model runs but seems to give garbage output. Do you see anything which stands out?

Perhaps there is a need for 'attn_logit_softcapping' if 'final_logit_softcapping' is used.

I have submitted a PR #86 to resolve the issue where both 'attn_logit_softcapping' and 'final_logit_softcapping' are necessary for gemma-2 inference. The corresponding PA kernel has also been revised. I'm curious why this isn't supported in vLLM. Google only mentioned that softcapping is beneficial for training.

[EricLBuehler]

@guoqingbao sounds good, I'll close this PR!