[Tentative] Adding 192 head dim (step_size = 12)

[Narsil]

Not sure if this PR actually works correctly (I'm going to check it).

Deepseek models use head_dim=192, and this cannot be compiled because of this static assert.
This modification works by jumping like step_size=4 + jumping by 8 every iteration.

Let me know if this is interesting in here.

[zhyncs]

@Narsil May you write unit test for it? And you can ref https://github.com/zhyncs/dl/blob/master/flashinfer_build.sh to compile from source.

[Narsil]

Which tests would you like me to add, batch_prefill_kernels ? others ?

Wdym ref to build from source ? I am building already.

[zhyncs]

Which tests would you like me to add, batch_prefill_kernels ? others ?

https://github.com/flashinfer-ai/flashinfer/blob/main/python/tests/test_batch_prefill_kernels.py
https://github.com/flashinfer-ai/flashinfer/blob/main/python/tests/test_batch_decode_kernels.py

I am building already.

ok

[Narsil]

Are the tests ran anywhere ?

[zhyncs]

Are the tests ran anywhere ?

There is currently no CI configured, you can use pytest in the local development environment to run.

[yzh119]

Hi @Narsil , thanks for your contribution!
You can add 192 to

flashinfer/python/setup.py

Line 67 in 0d61871

head_dims = os.environ.get("FLASHINFER_HEAD_DIMS", "64,128,256").split(",")

if compilation successes, you can run unittests such as https://github.com/flashinfer-ai/flashinfer/blob/0d618712faff20a84bbd513d02ac01e16be19306/python/tests/test_batch_prefill_kernels.py and see how does it work.

[zhyncs]

Hi @Narsil Any update?