Fix AttnRes silently dropping blocks past 32

[adityasingh2400]

_get_max_blocks capped MAX_BLOCKS at 32 for any n_blocks:

for mb in [4, 8, 16, 32]:
    if n_blocks <= mb:
        return mb
return 32

MAX_BLOCKS is the static trip count of the tl.static_range(0, MAX_BLOCKS) loops in both AttnRes kernels, and the width of the register-held scores and alpha vectors. So with more than 32 blocks the forward normalized the depth softmax over only the first 32 blocks and dropped the remainder from the weighted sum, and the backward left dV[32:] as the uninitialized torch.empty_like buffer. Both are silent, there is no shape error and no warning.

The fix rounds up to the next power of two instead. MAX_BLOCKS still has to be a power of two because of tl.arange(0, MAX_BLOCKS), and it still has a floor of 4.

This is behavior-preserving for every n_blocks at or below 32, since rounding to the next power of two reproduces the old ladder exactly: 1, 2, 3 and 4 give 4, then 5 to 8 give 8, 9 to 16 give 16, and 17 to 32 give 32. Only the previously broken range above 32 changes.

Added N=40 to the forward and backward correctness parametrizations, which is the smallest case that exceeds the old ceiling.

Verification note, stated plainly: these tests need CUDA and I do not have a GPU, so I did not run them. The argument here is by reading, plus the arithmetic above showing the mapping is unchanged at or below 32. Please run the AttnRes suite before merging.

[adityasingh2400]

Updated the branch onto main, so this is no longer behind and is mergeable again.

Worth flagging that no CI has run on this yet. All four workflows have been sitting at action_required since the PR opened:

• Checkstyle
• NVIDIA GPU
• AMD GPU
• Intel GPU

Approving the workflow run is all it needs to produce a signal.

[Tcc0403]

Great catch, lgtm