Gemma Support

[hnyls2002]

No description provided.

[hnyls2002]

Due to Gemma's head_dim=256, when launching the Triton kernel on CUDA devices with capability>=8, come across the Triton hardware limit error

triton.runtime.autotuner.OutOfResources: out of resource: shared memory, Required: 163840, Hardware limit: 101376. Reducing block sizes or `num_stages` may help.

[hnyls2002]

When launching gemma-7b-it without an instruction, e.g. <start_of_turn>user, the dot result of q and k can be very large. This may be due to the inappropriate use of the instructed weight or the head_dim=256.

To fix this, we can use flashinfer's attention kernel, which converts the intermediate result into float32. However, this feature is not yet supported in v0.0.2, but will be available soon. You can check its status at flashinfer-ai/flashinfer#157.

Alternatively, we can convert the intermediate result in our Triton's token attention kernel into float32. But be aware, this could result in a global slowdown and it has not yet been fully tested.

[hnyls2002]

If there is a nan problem when multinomial sampling, add --attention-reduce-in-fp32 to solve it.