fix: compile forward instead of module in load_checkpoint

self = torch.compile(self) in the inner module's load_checkpoint was a
no-op — reassigning self in a method doesn't affect the caller's
reference. Move compilation to the outer wrapper and target
self.model.forward directly, which has no KV cache writes and compiles
cleanly, giving 1.48x end-to-end speedup on H100.