module wrapper

aoti_attention.py

@@ -64,6 +64,24 @@ def _ns() -> torch.library.Library:
     return lib
 
 
+def sparse_module(topk: int, num_prefix_tiles: int) -> torch.nn.Module:
+    """`torch.export.export` requires an `nn.Module`; wrap the functional with baked scalars."""
+
+    class _SparseModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.topk = topk
+            self.num_prefix_tiles = num_prefix_tiles
+
+        def forward(self, q, k, v, gate, untile_index, variable_block_sizes, tile_divisor):
+            return sparse_attention_functional(
+                q, k, v, gate, untile_index, variable_block_sizes, tile_divisor,
+                self.topk, self.num_prefix_tiles,
+            )
+
+    return _SparseModule()
+
+
 def sparse_attention_functional(
     query: torch.Tensor,
     key: torch.Tensor,