The tensor-cores accelerated GQA in our blog post was not enabled by default (user need to use Prefill kernels/wrappers for decode to get such acceleration).
In this PR we add an option use_tensor_cores to decode operators/wrappers, and user can select whether to use tensor_cores for acceleration depending on use cases.
Not that our prefill kernels are compiled for all possible group sizes (#301 ), but decode kernels are not. So if user wants to use general group size, it's encouraged to set use_tensor_cores=True.
The tensor-cores accelerated GQA in our blog post was not enabled by default (user need to use Prefill kernels/wrappers for decode to get such acceleration).
In this PR we add an option
use_tensor_coresto decode operators/wrappers, and user can select whether to usetensor_coresfor acceleration depending on use cases.Not that our prefill kernels are compiled for all possible group sizes (#301 ), but decode kernels are not. So if user wants to use general group size, it's encouraged to set
use_tensor_cores=True.