misc: add device guard for kernels

[jeejeelee]

plan

• Check all kernels and add device guard
• Complete the tests

FIX: #452

[yzh119]

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

[jeejeelee]

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]

@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.

[yzh119]

LGTM in general, thanks for the fix @jeejeelee !