payload["constants_map"] = constants

optimization_utils.py

@@ -86,6 +86,9 @@ class ZeroGPUCompiledModelFromDict:
         if payload.get("format") != "zerogpu_aoti_v1":
             raise ValueError(f"Unsupported payload format: {payload.get('format')}")
         self.archive_file = BytesIO(payload["archive_bytes"])
+        constants = payload["constants_map"]
+        constants = {k: v.to(device=device, dtype=torch.bfloat16).contiguous() for k, v in constants.items()}
+        payload["constants_map"] = constants
         self.constants_map_cpu: dict[str, torch.Tensor] = payload["constants_map"]
         self.device = device
         self.compiled_model: ContextVar[AOTICompiledModel | None] = ContextVar("compiled_model", default=None)
@@ -98,7 +101,7 @@ class ZeroGPUCompiledModelFromDict:
 
         if not self._loaded_constants:
             # Move constants to target device (cuda) and keep dtype as-is (bf16)
-            constants_map = {k: v.to(self.device, non_blocking=True) for k, v in self.constants_map_cpu.items()}
+            constants_map = {k: v.to(device="cuda", dtype=torch.bfloat16).contiguous() for k, v in self.constants_map_cpu.items()}
             compiled_model.load_constants(constants_map, check_full_update=True, user_managed=True)
             self._loaded_constants = True