ran into some obscure compilation error in true triton fashion. get around it by separating gathering of gradients for diagonal block causal from the sparse kv blocks, parallelize later

Author: lucidrains