Make chunks always a power of 2

[GMNGeoffrey]

Summary
For some reason, the largest chunk size in chunk size tuning is incremented by 4. This was added in a9a12890d without explanation. We can hypothesize that it was related to adding 4 to an input dimension in trace_utils.py (trying to get a test case to fit in one chunk?), but that file was long ago deleted. This just looks like a bug and makes us hit unhappy paths all over the place.

Changes

• Delete the addition of 4 to the largest chunk size

Related Issues

• Fixes [BUG/QUESTION] Chunk size incremented by 4 off largest power of 2 #203

Testing

• Confirmed that attention kernel calling pattern moves from batches of 129 to 128 (since it is divided by 4 for attention).
• I didn't find a good way to unit test this that was actually a useful test (e.g. I can assert that chunk_size is now 512 instead of 516, but I don't think that test carries its weight).

Other Notes

• Split out from Apply power-of-two chunking consistently #207 as a separate PR since that one had ongoing discussion about the best approach (e.g. Avoid flattening and expanding broadcast dimensions when chunking #213 instead)

[GMNGeoffrey]

Hey @christinaflo, I broke this out of #207 since we were exploring more complicated approaches for that and I think this change is good on its own.

[christinaflo]

lgtm!