[Pallas:MGPU] Allow specifying layout for TMEM refs
We can no longer get by with just packing, since e.g. scales for block-scaled MMAs
have their own special layouts. This will be used in the next CL that exposes
block-scaed MMAs, but it's a big change that was reasonably easy to split.
PiperOrigin-RevId: 781538447
