misc: add device guard for kernels

flashinfer-ai / flashinfer

FlashInfer: Kernel Library for LLM Serving

https://flashinfer.ai

Apache License 2.0

1.46k stars 143 forks source link

misc: add device guard for kernels #611

Closed jeejeelee closed 1 week ago

jeejeelee commented 1 week ago

plan

[x] Check all kernels and add device guard
[x] Complete the tests

FIX: https://github.com/flashinfer-ai/flashinfer/issues/452

jeejeelee commented 1 week ago

Hi @jeejeelee , thanks for bringing this up, would you mind adding device guard for JIT templates as well?

https://github.com/flashinfer-ai/flashinfer/blob/a3360ff9c85e7d0bae8bb4ca6dbbca69cfadea37/python/flashinfer/jit/batch_decode_mla_templ.py

https://github.com/flashinfer-ai/flashinfer/blob/a3360ff9c85e7d0bae8bb4ca6dbbca69cfadea37/python/flashinfer/jit/batch_decode_templ.py

https://github.com/flashinfer-ai/flashinfer/blob/a3360ff9c85e7d0bae8bb4ca6dbbca69cfadea37/python/flashinfer/jit/batch_prefill_templ.py

https://github.com/flashinfer-ai/flashinfer/blob/a3360ff9c85e7d0bae8bb4ca6dbbca69cfadea37/python/flashinfer/jit/single_decode_templ.py

https://github.com/flashinfer-ai/flashinfer/blob/a3360ff9c85e7d0bae8bb4ca6dbbca69cfadea37/python/flashinfer/jit/single_prefill_templ.py

Okay, I will

jeejeelee commented 1 week ago

@yzh119 I have checked all kernels and added the device guards. Additionally, I tested the ROPE and NORM kernel in a CUDA 11.8 environment to verify they pass the tests. Please let me know if any additional testing is needed.