[Mosaic GPU] Fix two subtle issues with kernel lowering
1. The MLIR context is created by the user and its lifetime is not
in our control. To avoid depending on it, we serialize the module.
2. The operand and result layout requirements were missing from the custom call.
PiperOrigin-RevId: 657164985