[XLA:Mosaic] Add internal scratch VMEM
- Make internal scratch size configurable.
- Pass the number of max sublanes allowed in scratch to apply-vector-layout pass.
- Create a helper function to fetch internal scratch VMEM address.
PiperOrigin-RevId: 644184896
