A critical bug in attention kernel after refactoring

It seems there's a critical bug introduced by #53 
Running the `single_query_cached_kv_attention` kernel with certain configurations leads to CUDA illegal memory access errors. I found the bug in the unit tests.