Update aoti APIs

Author: angelayi

torchchat/cli/builder.py

@@ -571,9 +571,8 @@ def do_nothing(max_batch_size, max_seq_length):
             # attributes will NOT be seen on by AOTI-compiled forward
             # function, e.g. calling model.setup_cache will NOT touch
             # AOTI compiled and maintained model buffers such as kv_cache.
-            from torch._inductor.package import load_package
 
-            aoti_compiled_model = load_package(
+            aoti_compiled_model = torch._inductor.aoti_load_package(
                 str(builder_args.aoti_package_path.absolute())
             )
 

torchchat/export.py

@@ -72,16 +72,17 @@ def export_for_server(
         if not package:
             options = {"aot_inductor.output_path": output_path}
 
-        path = torch._export.aot_compile(
+        ep = torch.export.export(
             model,
             example_inputs,
             dynamic_shapes=dynamic_shapes,
-            options=options,
         )
+        options=options,
 
         if package:
-            from torch._inductor.package import package_aoti
-            path = package_aoti(output_path, path)
+            path = torch._inductor.aoti_copmile_and_package(ep, package_path=output_path, inductor_configs=options)
+        else:
+            path = torch._inductor.aot_compile(ep.module(), example_inputs, options=options)
 
     print(f"The generated packaged model can be found at: {path}")
     return path