Google Gemma running error with half dtype

[hnyls2002]

Using flashinfer in sglang with google/gemma-7b-it

  File "/home/ubuntu/sglang-venv/lib/python3.11/site-packages/flashinfer/prefill.py", line 462, in forward
    return self._wrapper.forward(
           ^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: BatchPrefillWithPagedKVCache failed to dispatch with dtype Half

I don't know if this is caused by Gemma's bfloat16 dtype or my inappropriate usage.

[yzh119]

Sorry the error message was confusing.

It's because flashinfer 0.0.2 do not support head dim 256 (enabled in #132 ), I'll trigger v0.0.3 build and release the new version tonight.

[hnyls2002]

@yzh119 Thanks, looking forward to the new version!