export via wrapper

app.py

@@ -516,13 +516,12 @@ def aoti_diag() -> str:
 
             with torch.no_grad():
                 exported = torch.export.export(
-                    aoti_attention.sparse_attention_functional,
+                    aoti_attention.sparse_module(geometry.topk, geometry.num_prefix_tiles),
                     args=(
                         q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
                         gate.transpose(1, 2), geometry.untile_index,
                         geometry.variable_block_sizes, geometry.tile_divisor,
                     ),
-                    kwargs={"topk": geometry.topk, "num_prefix_tiles": geometry.num_prefix_tiles},
                 )
             package = "/tmp/vsa_sparse_aoti.pt2"
             aoti_compile_and_package(exported, package_path=package)