Tweak work group size and fix bugs