jax

f4448517 - Move epilogue SMEM allocation to the kernel function

Move epilogue SMEM allocation to the kernel function This lets us deduplicate the logic between the matmul and collective matmul wrappers PiperOrigin-RevId: 817591533