PyTorch flex_gemm adds support for chunked GLU epilogue outputs

A new update to PyTorch's flex_gemm component introduces support for the chunked form of interleaved weights. This change enables concatenated gate/up gradients and allows for saving unchanged accumulators from chunked GLU forward epilogues. The implementation utilizes existing tile stores without requiring new CUDA kernels.

Technical details of the flex_gemm update

• Supports concatenated gate/up gradients.
• Enables saving the unchanged accumulator from a chunked GLU forward epilogue.
• Existing tile stores write into aligned halves of a single allocation.
• Does not require new CUDA kernels or changes to vendored QuACK.
• Preserves existing interleaved paths and scalar/dtype suffixes.

Why it matters

This optimization improves the efficiency of handling interleaved weights in transformer-based models within the PyTorch ecosystem.