perf(store): Cache output stride parameters in registers to reduce global loads

[garethpaul]

Summary

This PR optimizes the store() function in the FlashMLA kernel by caching frequently used output stride parameters (i.e., o_batch_stride, o_row_stride, and o_head_stride) into local registers. By using the __ldg intrinsic to read these values once per thread, we reduce repetitive global memory accesses and potentially lower memory latency. This improvement should help improve kernel performance without affecting functionality.

Key Changes

• Added local caching for params.o_batch_stride, params.o_row_stride, and params.o_head_stride in the store() function.
• Utilizes the __ldg intrinsic to hint to the compiler that these values are read-only.
• Maintains functional consistency, ensuring that the remainder of the function remains unchanged.

Impact

This change is a targeted performance optimization that minimizes redundant memory loads from global memory, improving efficiency without altering correctness.

[beginlner]

I don't think the changes make any difference on performance. If so, we should deal with other attributes of params in the same way. Could you please update a benchmark comparison?