[KVCache][Test] Fix TIR attn kernels for uncommon group size

apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators

https://tvm.apache.org/

Apache License 2.0

11.42k stars 3.4k forks source link

Closed MasterJH5574 closed 1 month ago

MasterJH5574 commented 1 month ago

This PR fixes the TIR attention kernels in PagedKVCache tests, which had issues when handling uncommon GQA group size (e.g., 6).