Layernorm performance optimization and partition size pybind

[mawong-amd]

This PR does two things:

1. Provide performance optimizations of the fused_add_rms_norm kernels (used in some layernorms, e.g. input and post_attention per Llama decode layer). The performance benefits largely originate from two optimizations: use of shared memory to cache intermediate computation results (as opposed to the existing use of global memory), and use of packed operations for FP16 inputs. Unrolling was attempted but this was not found to affect performance. Another optimization attempted was the specialization of blockReduceSum/warpReduceSum to AMD wavesizes of 64 (as opposed to the existing CUDA-compatible warp size of 32), which should in theory reduce the number of shuffles by (1024/32 * 5 + 5) / (1024/64 * 6 + 4) = 1.65 times, but this was also not found to measurably affect performance.

Typical performance improvements are on the order of 10% based on average runtime of the kernels as tested on Llama2-7B and Llama2-70B models on MI300X. Performance can probably be improved further, but we are running into diminishing returns as the total runtime of layernorm is not large. One interesting observation is that the post-attention layernorm on LL2-70B consistently takes 4-5 us longer than the input layernorm to complete (16 us vs 20 us), while this behavior is not observed for LL2-7B. It is unclear why this is the case, could be cache-related.

2. Add platform-specific paged attention v2 partition sizes and expose these to Python.

[mawong-amd]

Pending work on the prefill side of things